实习08-Mamba 和 SSM

第一部分Mamba 基础概念先补地基1.1 什么是 State Space Model (SSM)[公式] - SSM 思想SSM 源自控制理论核心是一个连续时间系统# 连续形式控制理论 h(t) A·h(t) B·x(t) # 状态更新 y(t) C·h(t) D·x(t) # 输出t连续时间可以理解为词在句子里的位置x(t)当前输入当前 token 的向量h(t)Mamba 里一直在更新的那个 “记忆口袋”它会随着阅读不断变化h ′(t)记忆随时间的变化率记忆怎么更新A控制 “过去的记忆” 以什么速度保留或遗忘用来控制遗忘强度B当前输入 x (t) 对记忆产生多大影响C从记忆 h (t) 里提取信息变成输出 y (t)D让输入 x 直接加到输出y(t)最终输出[例子] - 把 SSM 当成 一个人在读书x(t) 当前看到的字h(t) 你脑子里的记忆A 你忘性多大遗忘率B 新内容记进去多少C 从记忆里提取内容回答问题y(t) 你理解后输出的意思[记忆更新]当前记忆 保留一部分旧记忆 写入一部分新输入对应公式 hₜ Ā hₜ₋₁ B̄ xₜ[输出]输出 从记忆里提取信息 直接参考当前字[和 Self-Attention 的区别]① Self-Attention 怎么做每个词查一遍整个历史看一遍所有前面的词算相似度QK加权求和V训练计算量一定是 O(n2d)推理为 O(nd)所以说SSM的目的主要是优化训练方面的速度② SSM / Mamba 怎么做把历史压缩成一个记忆 h不断更新来一个新词把信息揉进记忆 h输出时从 h 里读信息Self-AttentionSSM / Mamba怎么获取历史每次都重新看全部历史把历史存在 h 里只看 h计算方式矩阵乘法 O(n²)递推更新 O(n)核心相似度加权状态记忆结果上下文表示上下文表示最终目标一模一样理解上下文一模一样理解上下文[结论]Self-Attention 和 SSMMamba做的是完全一样的事情让当前 token 能看到前面所有 token 的信息也就是建模上下文依赖context dependency1.2 Mamba 的核心创新选择性 (Selection)[区别] Self-Attention vs MambaTransformer靠 “查历史” 建模上下文-Mamba靠 “记记忆” 建模上下文Mamba的核心突破 让记忆会看内容选择性 计算永远 O (n)不随长度爆炸**早期 SSM 的问题是 **A,B,C是固定参数无法根据输入内容动态调整无法感知输入类似 RAG 的优化 → 无法做内容感知推理不管输入是 “重要的话” 还是 “废话”记忆的写入、遗忘、输出规则全都一样。无法根据输入进行动态选择[改动] Mamba 做了什么逆天改动改动Mamba让B、C、Δ时间步不再是固定值而是由输入x算出来的B现在输入决定 “要记多少进记忆”C现在输入决定 “从记忆里读多少出来”Δ现在输入决定 “记忆更新得多快”好处1对于输入遇到重要内容 → 多记一点相反遇到废话 → 少记或不记2对于时间步遇到长期依赖 → 慢慢记、慢慢忘当遇到短期信息 → 快速更新这就实现了和 Self-Attention 一样的能力1内容感知 动态聚焦但计算量完全不一样1AttentionO(n²)MambaO(n)[工程体现] 为什么 Mamba 推理这么快[工程体现 - 结论] 推理阶段纯递归不需要 KV CacheTransformer 推理为什么慢、吃显存因为它必须存KV Cache每来一个新 token要把之前所有 token 的 K、V 都存着序列越长 → 缓存越大 → O (L) 增长每一步都要和所有历史算一次注意力Q:[1,d] * K:[max_length, d]T 和 V:[[max_length, d] 加权求和 → O (L) 计算Mamba 推理完全不一样Mamba只存一个固定大小的记忆 h不管序列多长h的大小永远不变很小一般几十到一百多维度所以序列越长Mamba越有优势新来一个token→ 只用h递推一步计算量 O (1)和长度无关, 不用 KV Cache[工程体现 - 结论] 训练阶段并行扫描Parallel Scan这里解决一个经典矛盾RNN/SSM天然是串行的一步一步来但训练需要并行不然太慢而传统的Transformer虽然是 N2 复杂度但是可以通过并行优化速度h1 a*h0 b*x1 h2 a*h1 b*x2 h3 a*h2 b*x3 h4 a*h3 b*x4 ...必须算完 h1 才能算 h2算完 h2 才能算 h3[如何优化] Mamba 的解决办法用数学结合律把递归变成并行扫描通过Parallel Scan并行前缀扫描一步并行算出所有 h1, h2, h3, h4具体细节是把序列切成很多块[x1 x2] [x3 x4] [x5 x6] [x7 x8]块内并行算块之间只传一个小状态h目的是修正后面块的结果如[x3, x4]合并[x1, x2, x3, x4]最终整段可以并行跑完[不理解优化细节]后一个状态h3依赖h2那么永远都不能实现并行呀相反Transformer 架构中Q与K矩阵每一行向量之间都是独立的所以可以实现并行加速[优化细节] Parallel Scan 并行前缀扫描# 串行待优化 已知h1 h2 f(h1) h3 f(h2) h4 f(h3) h5 f(h4) h6 f(h5) h7 f(h6) h8 f(h7)对于 8 个 token分块并行计算[h1,h2] [h3,h4] [h5,h6] [h7,h8]第一次并行4组一起算但[3,4][5,6][7,8]算出来都是 “错的”因为还没拿到前面的状态。第二次串行逐层合并用[1,2]的最终状态去修正[3,4]得到[1,2,3,4]。第三次串行用[1,4]的最终状态去修正[5,6]合并得到[1,2,3,4,5,6]。第四次串行用[5,6]的最终状态去修正[7,8]合并得到[1,2,...,7,8]。类似于多叉树先通过并行求得所有叶子节点的值然后通过状态传递得到总值过程图如下所示层1: [1-2] [3-4] [5-6] [7-8] → 并行 层2: [1-4] [5-8] → 并行 层3: [1-8] → 仅1步串行结论这种方式时间消耗需要参考硬件和序列长度而Transformer那种架构可以一次暴力计算完上下文建模[时间] 并行前缀和的时间复杂度串行步数从O(n) → O(log n)n8 → 3 步n16 → 4 步n32 → 5 步n1024 → 10 步n100000 → 17 步长度越大加速越恐怖。时间复杂度为O(log n)n为序列长度[问题] 所以 Mamba 所谓的 “并行能力” 是随着分块策略大小和硬件能力来实现的吗Mamba 的并行能力来自两点块间并行扫描log 步数序列越长越有优势硬件 CUDA 核融合吃满 GPU并行优化下假设N1024dim2048传统Transformer架构计算过程如下1Q和KT计算1024 x 1024 个 2048 的计算量由于并行优化时间上和算出单次 2048 计算量等价即 O(d)2然后你还要算Softmax带来的影响它会计算每行的求和值去突出当前维度的值因此时间复杂度为 O(n)3然后和V矩阵加权求和那里也有一次 O(d) 的运算经过并行优化。所以工程真实的时间复杂度为 2 * O(d) O(n)相比mamba架构的计算过程11024 个 2048 的计算量但是后者一定依赖于前者所以进行并行前缀和模拟 1024 的并行。所以工程上真实耗时为 O(logn*d)序列长度n越大mamba效果越明显[结论 - Softmax 耗时]所以说Transformer架构在长序列下并行速度依旧比不过Mamba是因为Softmax的作用比如 n1024有 1024 行每行 1024 个值shape 为[max_length, max_length]这 1024 个行的sum完全可以同时并行计算。N 行sum→ 并行一步完成深度不随 N 变大。重点在 sum 这在每一行内部sum a1 a2 a3 ... a1024这一行内部的 1024 个数相加不可能一步完成。所以实际工程上的时间复杂度为O (N)