1. 为什么Transformer的“上下文墙”不只是算力问题,而是记忆结构问题
我一开始接触长文本任务的时候,也以为只要把显卡堆上去、把序列长度拉到128K、甚至1M,Transformer就能从容应对。直到有一次跑一个金融公告的风险因子抽取任务,输入是几十份年报拼接成的超长文档,模型在序列长度超过32K之后,困惑度不降反升,而且关键信息召回率肉眼可见地往下掉。当时我第一反应是训练不充分、学习率没调好,但排查了几天之后才意识到,这不是工程问题,是Transformer本身的记忆结构决定了它不可能在任意长的上下文里把每一条信息都精准保留。
1.1 注意力机制的数据流里,藏着一个“每步都全量回顾”的隐形成本
先看一个最基础的事实:自注意力的计算复杂度是输入长度L的平方。也就是说,当L从8K涨到16K,计算量不是翻倍,是翻四倍。很多人把这个当成长上下文的主要障碍,于是去做稀疏注意力、哈希注意力、闪存注意力,试图把复杂度压下去。这些方案确实能降低算力开销,但把复杂度压下去之后,你还会撞到另一堵墙——KV Cache。
KV Cache是什么?简单来说,你在生成第t个token时,需要让这个token和之前的t-1个token都做一次相似度计算,所以你必须把前面所有token的Key和Value都缓存下来。这个缓存的大小和序列长度严格线性增长:对于7B模型,一个token的KV缓存大约需要几十KB到上百KB,取决于层数和注意力头数。序列到128K的时候,单个序列的KV缓存就能吃掉几个GB显存,而且这还是做一次前向推理。
这意味着什么?意味着你为了让模型“看完”整个上下文,要让每一层、每一个注意力头都把整段历史的Key和Value保留在内存里。这个线性内存消耗本身也就罢了,真正致命的是下一步:当序列非常长的时候,注意力分数的分布会趋于平均化。原因也不难理解——softmax操作的对象是L个点,当L非常庞大、且大量token之间的语义关联并不紧密时,有效的注意力权重会被摊薄。模型理论上“看过”每个token,但真正“记住”的,只有那些和当前token高度相关的局部片段。
1.2 窗口越长,注意力越容易“东看西看然后忘掉”
学术界有一系列工作都指向同一个现象:把位置编码换成相对位置、窗口注意力、滑窗稀疏结构之后,模型在短序列上几乎不掉点,但一旦超过训练长度,性能急剧下降。我之前做过一个小实验:把一份50页的技术文档切成几个段落,随机打乱顺序,分别喂给一个支持64K上下文的大型模型,让它回答文档里某个具体数值出现在哪个章节。结果相当难看——有时候答对了,有时候会自信地编造一个完全不存在的数值。这就是长上下文下的“幻觉型遗忘”:模型不是没看见,而是看见了之后在注意力分配中把关键信息边缘化了。
所以如果你问我,Transformer的“上下文墙”到底在哪?我的答案是:它不在FLOPs,不在GPU内存条数,而在Attention的数据结构本身要求每一步都对全部历史做等权重的回顾。等权重意味着没有真正的重点,没有重点就意味着长序列中的信号会被噪声稀释。想要突破这堵墙,排除掉工程优化,必须在“模型如何记忆”这件事上动刀。这也是我为什么开始认真关注SSM(状态空间模型)和SSM-Attention混合架构的原因。
2. SSM的“无限记忆”不是玄学:固定状态下的压缩与代价
先澄清一个常见误区:很多人听到“无限上下文”和“状态空间模型”,第一反应是“模型能无限记住所有输入”。哪有这么好的事。SSM真正做的事情,是把“整段历史”压缩成一个固定维度的隐状态向量,然后随着输入不断滚动更新这个状态。你可以把它理解成一个带记忆的循环过程,只不过这个“循环过程”背后是一个经过精心设计的线性时变系统,而不是像RNN那样用非线性门控硬扛。
2.1 状态空间模型的本质:用固定大小的“笔记本”记录所有过去
状态空间模型的连续形式可以写成两个方程:
- 状态更新:h'(t) = A h(t) + B x(t)
- 输出投影:y(t) = C h(t) + D x(t)
其中h(t)就是模型内部维护的记忆状态,它是一个固定维度的向量,比如说2048维或者4096维。A矩阵控制着“过去的记忆如何随时间衰减或演化”,B控制着“当前输入如何写入记忆”,C控制着“如何从记忆里读取信息用于输出”,D是一个直连项。
数字化之后,尤其在S4和Mamba这类工作里,状态更新被设计成一种可以并行计算的离散扫描形式。对输入序列x_1到x_L,模型逐步计算h_1到h_L,每一步的h只依赖前一步的h和当前输入x。这意味着模型处理任意长度序列时,复杂度是线性的,而且是O(L)级别的线性,完全不像Attention那样要O(L²)。更重要的是,模型不需要记住每个token的细节,它只需要记住“到目前为止发生了什么”的压缩表示。
我用一个生活化的类比解释:Attention的做法是“每一个会议参与者都带着一个完整的录音笔,每次讨论都要回放所有录音去寻找相关信息”;而SSM的做法是“会议设了一个专职记录员,这个人只记要点,每次讨论前翻一下自己的笔记本就能接上话”。前者的回溯质量高但成本随会议时长爆炸,后者的成本恒定但取决于记录员的笔记质量——也就是状态向量h的容量和信息压缩方式。
2.2 选择性扫描的引入:Mamba如何在“记住”和“遗忘”之间找平衡
原始的S4模型用的是线性时不变系统,也就是说A、B、C矩阵不随输入变化。这带来的问题是:不管输入内容重不重要,写入记忆的方式都一样。面对一篇新闻和一个数学证明,模型的“记笔记策略”完全相同。Mamba做了一个关键改动:让B矩阵和C矩阵的内容依赖当前输入token,这就是“选择性扫描”。模型在每一步都会决定:当前这个token值不值得被强调写入,以及过去状态里哪些部分应该被强化或弱化。
这个改动的实际效果非常可观。在语言建模、DNA序列、音频波形等任务上,Mamba系列的线性注意力变体取得了与同规模Transformer接近甚至更优的结果,而推理时的显存占用和计算开销都低了一个量级。这也是混合架构火起来的直接原因——既然SSM在“线性成本下记住了大致语义”,那能不能让Attention只负责在关键时刻做“精确检索”?
不过也要说清楚SSM的软肋在哪里。因为状态向量h是固定维度的,它本质上是一个有损压缩通道。一篇文档的信息量如果远大于h能容纳的维度,那么后进来的信息就会覆盖或冲淡先前的信息。我管这个叫“压垮笔记本”现象。SSM适合处理“总体脉络重要、具体细节可丢失”的长文本;但如果任务是“在100万token里精确找出一行十六进制指纹”,纯SSM几乎一定会失败——它记住的是模糊的摘要,而不是逐字逐句的原文。
3. 混合架构的三条路线:SSM负责扫描,Attention负责定点检索
当我意识到纯Attention和纯SSM各有各的擅长区间之后,脑子里的问题就从“谁取代谁”变成了“怎么让二者做同事,而不是让二者打擂台”。理想的混合架构应该是:SSM以线性成本完成对整段历史的粗粒度理解与压缩,Attention在这个压缩后的表示之上或之中,对局部关键区域做高精度的关系建模。这样既绕开了全局Attention的平方复杂度,又弥补了SSM的“模糊记忆”问题。
现实中的混合方案大致能分成三类,我分别说一下它们的结构意图和边界条件。
3.1 分段式混合:全局线性扫描叠加局部窗口Attention
这种方案可以看作最直观的“分而治之”:把输入序列切成若干固定长度的块,比如每块2048个token。在块内使用标准或窗口化的Attention,让模型在局部范围内精确建模词与词的关系;在块与块之间,用一个SSM层(比如全局状态扫描)把前一个块的输出状态传给后一个块,从而实现跨块的信息传递。
这种结构的逻辑是:绝大多数强语义关联都发生在短距离内,比如同一个句子里主谓宾的搭配、相邻句子之间的指代关系。你不需要让每一对远程token都直接做注意力交互,只需要让“区块的摘要信息”通过SSM流动下去,让后续块在判断时能参考前序块的语义趋势。
我实测下来,这类结构训练稳定性很高,因为Attention模块和SSM模块的职责边界清晰,梯度传播相对平缓。它的问题在于:分块大小是个敏感超参。块太小,SSM的摘要信息不足以支撑远程精确推理;块太大,局部Attention的计算量又上去了。我的经验值是文本类任务块长在1K到4K之间最合适,而音频或传感器这类连续信号可以放宽到8K以上。
3.2 交错式混合:Token级别的路由分配,让模型自己决定走哪条路
分段式还是“人为主观划定”在哪里用Attention、在哪里用SSM。更聪明的做法是让模型自己去判断:当前token或当前某个区域,使用Attention路径还是SSM路径。具体实现上,可以在每一个block里并行放一个Attention子层和一个SSM子层,再用一个门控或路由函数对两者的输出做加权融合,权重由当前输入决定。
这类设计的先例可以参考混合专家模型(MoE)的思路:专家网络对应“不同能力”,路由器对应“选择能力”。在混合架构里,Attention可以理解为“擅长精确检索的专家”,SSM可以理解为“擅长持续跟踪状态的专家”。路由器的目标就是根据当前语境激活合适的专家权重。
好处是灵活性高,模型在代码任务里可能更偏Attention(代码对精确符号匹配要求高),在长篇叙事里可能更偏SSM(故事主线比字词细节更重要);坏处是训练开销比分段式大不少,因为两条路径都需要完整计算,只是输出时做了加权。而且路由器的收敛比较娇气——我开始做的时候,路由器很容易在训练初期崩溃到“永远选Attention”或“永远选SSM”的极端。后来我把门控信号的来源从单层线性投影换成带LayerNorm的双层MLP,并且把路由温度的初值设低一些,情况才稳定下来。
3.3 分层式混合:底层做全局压缩,高层做局部精修
第三种思路是分层异构:靠前的若干层使用纯SSM或大窗口稀疏Attention,负责把整个长序列逐步压缩成一系列低分辨率但覆盖全程的语义表示;靠后的若干层再对这些语义表示做高分辨率Attention,类似于“先读一遍全文画提纲,再针对重点段落精读”。这种结构在存储和计算上最省,因为底层不需要保留完整的逐token KV缓存,顶层只需要处理压缩后的状态序列,序列长度已经大幅缩短了。
我在一些多轮对话系统里试过这种设计——底层用Mamba风格的状态空间层读完整个对话历史,保持角色状态和话题演变的连续追踪;顶层用Attention对最近几轮对话做精细回复生成。效果相当符合直觉:长对话的连贯性比纯Transformer强,而短对话的局部流畅性又比纯SSM强。
4. “无限上下文”的真实瓶颈在哪里:不是架构,是信息守恒
聊到这里,你可能会有个雄心壮志:那是不是把SSM和Attention一混合,就能真的做到无限上下文、百万token甚至千万token不是梦?答案没那么简单。要理解边界,我们得从信息论的角度重新看待“上下文长度”这个词。
4.1 有限状态容量的“覆盖写入”:任何固定维度都会饱和
先回答一个基础问题:一个固定维度的状态向量,到底能装多少信息?理论上,d维的实值向量,如果每个分量有足够的数值精度,能表达的“状态种类”自然是天文数字。但注意,模型能存储信息,不代表它能有效检索信息。你的电脑硬盘也能存下一万本书,但你如果只有一个文件目录而没有任何索引机制,找一本书依然是大海捞针。SSM的状态更新本质上是把新信息和旧状态做一个加权融合,新信息到来,旧状态会被改变——这个“改变”可能是有利的更新,可能就是灾难性的覆盖。
混合架构里Attention的角色其实是在做“反覆盖补偿”:它允许模型跳过状态压缩,直接从原始token序列中检索精确的Key-Value对。但Attention的代价是只能覆盖有限长度(无论是计算上还是训练分布上),所以你实际上是在两条曲线之间找平衡点:SSM的有效覆盖长度理论上无限但精度递减,Attention的精度高但长度有限。“混合”没有消除这个矛盾,它只是把这个矛盾的爆发点向后推移了。
有一个很有趣的思考实验:假设你有一个完美的混合架构,SSM部分可以在恒定成本下记住全文的“语义骨架”,Attention部分可以在有限窗口内做“逐步滑动的精确读取”。那这个系统的能力上限是什么?你依然受制于Attention滑窗的大小——如果关键信息A出现在第10个token,关键信息B出现在第900万个token,两者之间隔着几百万token的噪声,那滑窗式的Attention根本不可能同时看到A和B。SSM也许能告诉你“B所在段落的大致主题和A一致”,但它无法恢复A的逐字形式。所以,真正的“无限上下文推理”,需要的不是更好的注意力机制,而是一种无损或近无损的层次化记忆检索体系,这已经超出了单纯混合SSM和Attention的范畴。
4.2 上下文长度不等于记忆质量:越长的上下文,越需要“遗忘策略”
我这里想强调一个常被忽略的点:真实世界的长文本任务,绝大多数不需要“记住所有内容”。举个例子,读一本500页的技术手册,你不需要记住第137页第2段的每个标点,你只需要知道“第137页讲的是错误码处理”,并且能在用户提问时快速定位过去。这本质上是一个“索引摘要+按需检索”的过程。而现有混合架构,其实都还停留在“尽力压缩所有内容”的阶段,没有真正把“应该记住什么”和“应该丢弃什么”分开处理。
我一直在想,一个好的长上下文模型,应该内置一个显式的遗忘机制。不是靠SSM的矩阵衰减来隐式遗忘,而是有类似“信息价值评估器”的东西,对输入内容打分:高价值信息写入长期状态,中价值信息写入可丢弃的短期缓冲,低价值信息直接丢掉。然后Attention只在“值不值得检索”的地方启动。这个方向还没有特别成熟的工作,但已经有一些雏形,比如用强化学习训练一个“记忆读写控制器”的尝试。
4.3 可证明的高效边界:线性复杂度里也可能藏着隐藏的平方
如果你关注理论,还要警惕一件事:很多论文声称自己的混合架构在长序列上达到了“线性复杂度”,这个说法需要仔细分清——它可能只针对“前向传播的FLOPs”是线性的,但如果你要训练,反向传播里对状态序列的梯度计算可能引入O(L²)的依赖。原因在于SSM的状态更新是一个顺序递归过程,你要计算梯度L时的状态关于初始状态的变化,需要回溯整个状态链。有些论文用“并行扫描”技巧在训练时强行把复杂度压到O(L log L)甚至O(L),但这是以额外的显存和通信为代价的。也就是说,工程上的“无限上下文”往往不是免费的,而是把计算成本换成了别的东西,比如多机通信或更复杂的数值稳定性处理。
5. 我做混合架构实验的一些体会与更远的设想
前面讲了不少理论和结构层面的分析,最后分享一些我在实际代码实验中踩过的坑,以及我现在认为值得投入精力继续做的方向。这部分不保证全对,但至少能帮想动手复现的人少走弯路。
5.1 训练混合模型最容易忽视的数值问题
第一个坑是状态矩阵的初始化。SSM里有一个关键矩阵A,它的初始化方式直接决定了记忆衰减速度。如果A的特征值模长太接近1,模型会倾向于“什么都记住”,状态很快被填满,后续输入进来时发生严重的相互干扰;如果模长太小,模型变成“金鱼记忆”,两句话之后就忘了前文。Mamba论文里建议用HiPPO初始化或S4D初始化,把A初始化成一种能对多项式函数做压缩表示的形式。我自己实践后的感受是:混合架构里的SSM层的A矩阵初始化,可以比纯SSM更倾向于“慢衰减”一些,因为Attention已经承担了精确检索职责,SSM可以放心大胆地把长期趋势信息保存得更久一点。这个“更久一点”的度,我是通过A矩阵的尺度缩放参数调出来的,建议你也做一下简单扫描——在验证集困惑度上画一条曲线,通常能找到明显的性能峰值区间。
5.2 序列并行与状态传播:长到一定程度,通信成为瓶颈
当序列长度真正到了兆级别,单卡肯定放不下,你需要把序列切到多张卡上做序列并行。Attention的序列并行比较成熟,因为每个token的计算只依赖同一块内的其他token(或者通过Ring Attention做分块通信)。SSM的序列并行则要小心:状态h的更新是逐token的,跨卡时你必须把前一块的最终状态传给后一块,这个依赖没法完全异步化。在千卡规模上,每一层都要做一次全环通信,虽然通信量本身不大(只是一个向量),但层数一多、迭代一步的总通信次数就上去了。我当时在256卡上跑8B混合模型,实测下来如果不对通信做pipeline重叠,吞吐量会掉三成左右。把状态传递嵌到通信算子之间的空隙里(overlap)之后,才基本追平了纯Attention的并行效率。
5.3 更强的方向:我目前看好的三个演进思路
一是层次化记忆压缩。先让SSM把原始token压缩成若干语义片段,再对这些片段做第二遍压缩,形成多层索引结构;查询时按层级逐步下钻,而不是从头到尾线性扫描。这其实更像人类阅读长篇材料的方式。二是动态状态维度扩展。目前的SSM用固定维度,无论处理多复杂的历史都背同样的行囊。我设想可以和MoE思想结合——根据当前内容的复杂度动态调整参与状态更新的维度子集,复杂段调用更多状态通道,平淡段省下容量。三是查询导向的记忆写入。Attention在推理时需要什么信息,这个需求可以反馈给前面的SSM层,让它在写入新内容时优先保留“未来可能会被查询”的那部分信息。如果这个闭环能做出来,那混合架构就从“被动压缩”走向了“主动记忆”,也就更接近我所理解的真正“无限上下文”。
5.4 从零开始复现,我会怎么做
如果你刚接触这个方向,我建议的路线很具体:先不追最新的100B超大混合模型,而是拿一个1B到3B规模的模型做实验,在20B到50B token的干净英文语料上做预训练。混合结构可以选分段式起步,因为它的训练稳定性最好,实现门槛最低——你只需要在一个Transformer层旁边加一个可选的Mamba层,再用残差连接把它们叠加起来就行。我首推复现Mamba的官方开源实现,结合标准FlashAttention,跑通一个“Attention块套SSM块”的交替堆叠结构。等你在小规模上验证了长文本任务的收益,再逐步引入路由机制和MoE化,这样每一步都有基线可对比,不会让自己陷入“多个变量同时变化但不知道是谁在起作用”的泥潭。
我自己的实战心得是:不要一上来就追求“百万上下文”这种指标,而是先定义清楚你想解决的测评任务。如果你的目标是“从10万token文档中准确找答案”,那混合架构不够,还需要配合检索增强;如果你的目标是“在超长多轮对话中保持一致的角色人格”,那混合架构在不做任何检索增强的情况下,就已经能比纯Transformer和纯SSM都做得好。想清楚这个区别,你才不会在跑完一堆实验之后发现自己造了一台没有用武之地的机器。