devgo.cn一个AI Agent正在处理一份上百页的合同，需要逐段比对条款，同时记住前面几十页的细节。Transformer的注意力机制给了它过目不忘的能力——每个token的key和value都被缓存下来，随时可以精确回忆——但这份记忆是有代价的：KV cache（缓存每个token的key和value向量的内存，随上下文长度线性增长）越攒越大，推理速度越来越慢。Agent越长、越自主，这个包袱就越重。
业界早就想给Transformer“减肥”了。主流思路包括稀疏注意力（只让当前token关注一部分历史token）、滑动窗口注意力（只保留最近的token）、以及状态空间模型SSM（把全部历史压缩成一个固定大小的状态向量，像Mamba那样）。这些方法各有取舍，于是出现了大量混合架构：一部分层保留完整注意力，另一部分换成更省内存的层。但问题在于，每设计一种混合架构，都得从头预训练一遍——动辄数万亿token、几十万GPU小时。想系统比较哪种组合最好，代价高到几乎没人做。
亚马逊AWS的一篇新论文给出了一个反直觉的答案：**不需要从头训练**。方法叫Priming，直接把现成预训练Transformer的权重“借用”过来，作为混合架构的起点。核心洞察是：Transformer的注意力层和SSM层在数学上其实是同构的——它们的输出都可以看成把历史token的值序列乘上一个“混合矩阵”，决定每个过去token对当前输出的贡献。注意力是动态计算这个矩阵（query和key相乘），SSM则是通过一个固定维度的循环状态来隐式逼近它。既然结构如此相似，那如果某个SSM的固定状态能足够接近原注意力层的混合矩阵，那原来Transformer的query/key/value投影矩阵就可以直接复用，只需要新学SSM特有的那几个参数。
关键问题是怎么判断哪些层适合替换。论文用一个便宜的代理实验：把某层换成滑动窗口注意力（SWA，只看最近一小段token），看模型掉多少分。如果能容忍丢失长程信息，说明这层干的活比较局部，就适合用SSM压缩；而如果一换就崩，那它是在做跨上万token的远程检索，必须保留完整注意力。这个代理与动态系统理论里的Hankel秩（衡量一个矩阵需要多大状态维度才能精确表示）有内在联系。
整个过程分三步：选层、权重迁移与对齐、端到端微调。选完层后，用原注意力权重初始化SSM层，跑一小段蒸馏让新层模仿旧层行为，最后在整个混合模型上做next-token预测微调。整个流程只用不到原模型预训练token预算的0.5%——相当于把成本砍掉了99%以上。而且跨模型家族（Qwen、Llama、Mistral）、跨规模（8B到32B）都管用。
用这个便宜的方法，论文第一次系统比较了三种SSM：Gated KalmaNet（GKA）、Gated DeltaNet（GDN）和Mamba-2。在受控条件下，GKA表现最好，而且它有一个独家优势：不需要重训就能调节推理时的计算量，让用户自己决定站在速度-精度的哪条曲线上。另一个有意思的发现是：对过去的“有损摘要”竟然完胜“忠实记录最近片段”——SSM层在只用一半内存预算的情况下，表现稳定超过滑动窗口注意力。最亮眼的数字来自一个32B参数的GKA混合模型：在12.8万token的上下文里，推理吞吐量比原版Transformer快2倍，拿正确答案的墙钟时间最短能缩短到原来的三分之二，而推理基准分数完全不掉。
为了让大家能接着往下挖，团队把这些Primed混合模型在8B和32B两个尺度上开源了，还带训练代码、优化过的kernel和一个vLLM-serving插件，Apache 2.0许可。
对我这种天天和长上下文打交道的人来说，这个思路的实用价值很直接：如果手里的Transformer模型因为上下文太长、内存扛不住了，与其换架构重训，不如先拿Priming试试哪些层能换成SSM——那些干“本地活”的层（比如跟踪当前从句、复制邻近token）通常可以被固定大小状态干净利落地替代，而真正做长程检索的层保留注意力。这样就能在几乎不损失质量的前提下把显存和延迟打下来。要注意的坑是：不是所有SSM都一样，单测的结果不能跨模型随便套，换层之前得用那个SWA代理实验先给每层做个体检。另外Priming给的是初始化，如果要让混合模型在特定任务上更强，还得做后续微调。
展示注意力机制与SSM在数学结构上的相似性：两个块都能视为把值序列乘一个混合矩阵，注意力通过query×key动态计算，SSM通过固定维状态循环生成。
Priming的三阶段流程示意：选层（用SWA代理判断哪些层可替换）、权重迁移与对齐（用注意力权重初始化SSM层并短蒸馏）、端到端微调。
对比实验图：32B GKA混合模型在长上下文推理中，吞吐量约为原Transformer的2倍，同时保持同等推理质量。
已发布的模型动物园表格：8B和32B尺度的Primed混合模型列表，涵盖Qwen、Llama、Mistral等家族。
说到底，Priming的价值不在于又一个省钱的训练技巧，而是把混合架构的设计从“凭感觉试错”变成了“可系统搜索”的工程问题——这恰好是工程实践最爱的推进方式。
内容与图片版权归原作者所有 · 原文: https://www.amazon.science/blog/can-we-change-an-ai-models-architecture-without-retraining-it-from-scratch