Mamba
约 1156 字大约 4 分钟
2025-01-20
作者认为，序列建模的一个基础问题是把上下文压缩成更小的状态(We argue that a fundamental problem of sequence modeling is compressing context into a smaller state)，从这个角度来看
transformer的注意力机制虽然有效果但效率不算很高，毕竟其需要显式地存储整个上下文(storing the entire context，也就是KV缓存)，直接导致训练和推理消耗算力大 好比，Transformer就像人类每写一个字之前，都把前面的所有字+输入都复习一遍，所以写的慢 RNN的推理和训练效率高，但性能容易受到对上下文压缩程度的限制 On the other hand, recurrent models are efficient because they have a finite state, implying constant-time inference and linear-time training. However, their effectiveness is limited by how well this state has compressed the context.
好比，RNN每次只参考前面固定的字数(仔细体会这句话：When generating the output, the RNN only needs to consider the previous hidden state and current input. It prevents recalculating all previous hidden states which is what a Transformer would do)，写的快是快，但容易忘掉更前面的内容 而SSM的问题在于其中的矩阵A B C不随输入不同而不同，即无法针对不同的输入针对性的推理(详见上文的2.4节) 最终，Mamba的解决办法是，相比SSM压缩所有历史记录，mamba设计了一个简单的选择机制，通过“参数化SSM的输入”，让模型对信息有选择性处理，以便关注或忽略特定的输入 这样一来，模型能够过滤掉与问题无关的信息，并且可以长期记住与问题相关的信息 好比，Mamba每次参考前面所有内容的一个概括，越往后写对前面内容概括得越狠，丢掉细节、保留大意
为方便大家对比，我再用如下表格总结下各个模型的核心特点
模型 对信息的压缩程度 训练的效率 推理的效率 transformer(注意力机制) transformer对每个历史记录都不压缩 训练消耗算力大 推理消耗算力大 RNN 随着时间的推移，RNN 往往会忘记某一部分信息 RNN没法并行训练 推理时只看一个时间步 故推理高效(相当于推理快但训练慢) CNN 训练效率高，可并行「因为能够绕过状态计算，并实现仅包含(B, L, D)的卷积核」 SSM SSM压缩每一个历史记录 矩阵不因输入不同而不同，无法针对输入做针对性推理 mamba 选择性的关注必须关注的、过滤掉可以忽略的 mamba每次参考前面所有内容的一个概括，兼备训练、推理的效率
总之，序列模型的效率与效果的权衡点在于它们对状态的压缩程度：
高效的模型必须有一个小的状态(比如RNN或S4) 而有效的模型必须有一个包含来自上下文的所有必要信息的状态(比如transformer) 而mamba为了兼顾效率和效果，选择性的关注必须关注的、过滤掉可以忽略的
为了让传统的SSM在现代GPU上也能高效计算，Mamba中也使用了Flash Attention技术
简而言之，利用内存的不同层级结构处理SSM的状态，减少高带宽但慢速的HBM内存反复读写这个瓶颈 具体而言，就是限制需要从 DRAM 到 SRAM 的次数(通过内核融合kernel fusion来实现)，避免一有个结果便从SRAM写入到DRAM，而是待SRAM中有一批结果再集中写入DRAM中，从而降低来回读写的次数(更多详见：通透理解FlashAttention与FlashAttention2：全面降低显存读写、加快计算速度)
最后，有的新闻稿会说Mamba是第一个实现匹配Transformer性能的线性时间序列模型，其实第一个是TransNormerLLM(为了更好的阐述清楚mamba本身，把原属于本文的这部分内容：“第五部分 Mamba近似工作之线性Transformer：从TransnormerLLM到RWKV”，转到了另一篇文章中《七月论文审稿GPT第1版：通过3万多篇paper和10多万的review数据微调RWKV》)
原文链接：https://blog.csdn.net/v_JULY_v/article/details/134923301