AWS 提出 priming 方法，可将预训练 Transformer 改造成混合状态空间架构

AI 智能体正在快速扩展：更长的会话、更深的推理链、更强的自主性。但它们底层的 Transformer 架构仍然要支付一笔随处理 token 数量增长的内存税。随着智能体能力越来越强，这笔成本也随之增长。如今，混合架构大量涌现，它们组合使用多种方法，以在内存、速度和质量之间取得更好的平衡。这些方法包括稀疏注意力（在选择下一个 token 时只考虑过去 token 的一个子集）、滑动窗口注意力（只保留最近的 token），以及状态空间模型（将完整历史的信息压缩为固定大小的状态）。这些混合架构可以理解为序列模型设计空间中的一些点：它们以不同程度在遍历序列、经过不同状态时混合信息。它们彼此之间的区别在于如何保留和使用过去。迄今为止，探索这一设计空间需要从头训练每一个新架构，使用数万亿 token 和数十万 GPU 小时。这使得对混合架构进行受控比较变得极其昂贵。在一篇新论文《Priming: Hybrid state space models from pre-trained Transformers》中，我们 Amazon Web Services 的团队解决了这个问题。我们引入了 priming，一种利用预训练 Transformer 权重作为起点的方法，将混合架构的预训练成本降低 99% 以上。借助 priming，模型开发者可以提出以前因成本过高而无法大规模回答的问题：哪些注意力层可以被更高效的层替换？在不牺牲质量的前提下可以替换多少注意力层？哪种状态空间层效果最好？以及速度与准确率之间的权衡是什么样子？在我们的论文中，我们比较了使用三种不同类型状态空间模型（Gated KalmaNet、Gated DeltaNet 和 Mamba-2）的混合模型。我们的一些发现：并非所有状态空间模型都相同。在受控条件下，Gated KalmaNet（GKA）表现最好，其次是 Gated DeltaNet（GDN），然后是 Mamba-2。GKA 独特地提供了无需重新训练的可变长度推理计算，并让用户可以直接控制其在速度-准确率前沿上的位置。对全部内容做一个好的有损摘要，胜过对最近过去的完美记录。状态空间模型始终优于滑动窗口注意力（SWA），即使 SWA 获得两倍的内存预算。相同质量，两倍吞吐量。一个 320 亿参数的基于 GKA 的混合模型在推理基准上与源 Transformer 相当，同时在 128,000 token 上下文下解码速度最高可达两倍，并在最短仅三分之二的墙钟时间内得到正确答案。为促进对混合架构的进一步研究，我们发布了 80 亿参数和 320 亿参数规模的 primed 混合模型库，用于长上下文推理和指令遵循，并附带训练代码、优化内核以及一个 vLLM 服务插件，均采用 Apache 2.0 许可证。注意力机制。Transformer 的注意力机制——它决定先前见过的 token 与当前正在处理的 token 有多相关——依赖三种不同的向量，它们的名称来自数据库设计：query、key 和 value。对于每个 token，query 捕捉该 token“在寻找什么”，key 捕捉它“提供什么”，value 携带它将传递的实际信息。相关性通过将当前 token 的 query 与每个其他 token 的 key 进行比较来计算；匹配越接近，该 token 的 value 对注意力输出的贡献就越大。这些向量中的每一个都通过将 token 乘以一个学习到的权重投影矩阵来产生。Transformer 一次输出一个 token，而每个输出 token 都会成为新的输入，因此上下文也一次增长一个 token。为节省计算，Transformer 因此将每个新输入的 key 和 value 投影存储在“KV 缓存”中。两类内存层：eidetic 与 fading。注意力机制赋予 Transformer 我们所说的 eidetic 记忆：对每个过去 token 的逐字回忆。能够在 LLM 的整个上下文窗口内进行回忆，正是上下文学习、长文档推理和检索增强生成成为可能的原因。但它的成本随输入 token 数量增长。然而，并非所有注意力层都做同样的工作。有些执行跨越数万 token 的长程检索。另一些则更偏局部，跟踪当前从句、复制附近 token 等等。局部层不需要不断增长的缓存：一个紧凑、固定大小的过去摘要——我们称之为 fading 记忆——就足够了。混合模型将两者混合：对需要全局回忆的层使用注意力层（eidetic，缓存增长），其余层使用状态空间层（fading，固定大小）。结果是一个内存占用降低、同时在关键之处仍能进行逐字回忆的模型。Priming：知识并不锁定于架构。探索混合模型的设计空间一直不切实际，因为每个混合配置通常都需要从头预训练。Priming 通过复用预训练 Transformer 权重中存储的知识来绕开这一点。Priming 之所以可行，是因为 Transformer 本身就是状态空间实现的一个特例。注意力层和 SSM 层的输出都可以表示为 value 序列乘以一个“混合”矩阵的结果，该矩阵决定每个过去 token 对每个当前 token 的贡献程度。注意力通过将 query 和 key 矩阵相乘来计算混合矩阵；SSM 则通过一个所有过去信息都必须经过的固定维度循环状态产生类似输出。Priming 利用了这种结构相似性：如果一个具有固定维度状态的 SSM 能够产生一个足够接近注意力层的混合矩阵，那么源 Transformer 的 query、key 和 value 投影就可以复用为 SSM 自身的输入投影，只需要学习 SSM 特有的参数。这种近似是否成立，取决于该注意力层执行的是复杂的长程检索还是简单的局部操作。我们用一个廉价的经验代理来区分两者：将每一层替换为 SWA——它只保留上下文的一部分——并测量性能退化。能够容忍这种替换的层，是 SSM 转换的自然候选：它们的混合矩阵不需要捕捉长程关系，因此 SSM 的状态可以忠实表示它们所做的事情。这个代理与动力系统理论中的一个形式化量相关，即该层混合矩阵的 Hankel 秩，它衡量必须通过状态瓶颈的信息量，因此也衡量精确实现该层所需的最小 SSM 状态维度。Priming 分三个阶段进行：层选择：确定哪些层可以被忠实压缩为固定大小状态，同时将其余层保留为注意力；权重迁移与对齐：从对应的注意力权重初始化 SSM 层，然后运行一个简短的蒸馏过程，使每个 SSM 对齐到它所替换层的行为；端到端适配：在下一 token 预测上微调整个混合模型，让所有层共同适配。整个流程消耗的 token 预算不到源模型预训练的 0.5%，并且适用于不同 Transformer 家族（Qwen、Llama、Mistral）、模型类别（dense 或 MoE）以及规模（80 亿到 320 亿甚至更多）。借助 priming，我们现在可以廉价且系统地探索设计空间、层类型、混合比例和层放置，因为与预训练相比，每个变体都从相同的强初始化开始，并且只需要极少的训练。对 fading 记忆层的受控比较。例如，考虑这个问题：在相同数据、相同训练预算、相同 eidetic 与 fading 层比例下，哪种 SSM 层在混合模型中效果最好？在我们的论文中，我们评估了三类状态空间模型，每一类代表一种使用固定大小状态的不同方法：Mamba-2 在每一步对整个状态施加统一的标量衰减，这是一种以相同速率遗忘一切的运行平均。当新 token 到达时，它被写入状态，而所有先前条目被统一衰减。Mamba-2 简单高效，但不加区分。Gated DeltaNet（GDN）在统一衰减之上增加了有针对性的更新。它不是统一衰减一切，而是在写入新值之前，可以选择性地擦除先前与特定 key 关联的 value（delta 规则）。这使它比 Mamba-2 更具表达力：它可以修订特定关联，而不必统一衰减其余部分。Gated KalmaNet（GKA）更进一步。它的擦除方向不仅取决于当前 key，还取决于所有过去 key 的统计量，这是一种历史感知更新，考虑了模型所见一切的累积结构。理论表达力层级预测了经验层级：GKA > GDN > Mamba-2。这种差距在长上下文推理中扩大，因为压缩质量在那里最重要，而在短上下文任务中缩小。此外，在 320 亿参数（32B）规模下，混合 GKA 与在相同数据上训练的源 Transformer 相差不到 1%，尽管推理时只需要一半的 KV 缓存。Fading 记忆胜过逐字窗口记忆。我们研究中的一个有趣发现来自对使用 SSM 层的混合模型与使用滑动窗口注意力（SWA）的混合模型的直接比较，后者逐字保留最近 token 的固定窗口，而不是将过去压缩为状态。人们可能会认为，逐字回忆即使在有限窗口内，也会优于有损压缩。毕竟，在滑动窗口范围内没有信息丢失。而且堆叠多个 SWA 层会增加端到端模型的“感受野”。相反，我们测试的每一个基于 SSM 的混合模型在长上下文基准上都优于对应的 SWA 混合模型，即使窗口获得的状态内存是状态的两到四倍。速度与准确率的运行时旋钮：推理时的可变计算。GKA 具有 Mamba-2 和 GDN 都不具备的特性：可变测试时计算。Mamba-2 和 GDN 对每个输入 token 都有相同的计算预算，无论其内容如何。相比之下，GKA 迭代地求解一个回归问题，并且迭代次数可以在服务时调整而无需重新训练。在实践中，以 10 次迭代而不是通常的 30 次来服务 GKA，会在 AIME 2025 上牺牲大约 1% 的准确率，同时带来 10% 更快的解码。一个训练好的模型可以服务速度-准确率前沿上的多个点，因此模型提供商可以提供不同的延迟层级，而无需维护单独的模型变体。这与我们在早期工作中讨论过的一个原则相关：智能是以与难度相称的方式分配努力的能力。GKA 的可变计算机制是这一原则的架构体现：要求速度的部署可以运行更少迭代，而要求准确率的部署可以运行更多迭代。更快的推理，而不仅仅是更快的 token。原始每 token 吞吐量只能说明部分情况。一个每 token 成本更低的模型，如果其推理更弱、需要更多 token 来补偿，端到端仍可能更慢。与部署相关的问题是，以墙钟时间衡量，混合模型达到目标准确率要快多少。我们在 AIME 2025 这一具有挑战性的数学竞赛基准上测量了这一点，将我们的 GKA-Primed-HQwen3-32B-Reasoner 与在相同推理配方上训练的源 Qwen3-32B 进行比较。两个模型生成相似长度的思考轨迹（完整集上中位数约 14,000 token，最难问题上为 28,000 token）——也就是说，它们产生同样深入的回答——但混合模型的吞吐量高得多。在现实并发水平下，当 Transformer 的 KV 缓存使 GPU 内存饱和时，混合模型在完整基准上达到目标准确率阈值所需时间快 1.2 到 1.4 倍，在最难问题上最高快 1.6 倍。这是 SSM 层固定大小状态的结果，它将每序列内存占用大致减半，使同一硬件上能够进行更多并发操作。吞吐量优势随上下文长度增长：32B 模型（GKA 10 次迭代）在 16K token 时为 1.50×，到 128K 时扩展到 2.06×。对于 8B 参数模型，收益更大：从 16K 时的 1.94× 到 128K 时的 2.29×。在长上下文下，KV 缓存读取主导 Transformer 的前向传播，而在内存饱和时将缓存减半大致可使有效批大小翻倍。这不仅对推理服务有影响。在强化学习（RL）后训练流程中，吞吐量决定策略生成正确轨迹的速度。混合模型的优势直接转化为在固定 GPU 预算下 RL 训练数据速率提高 20% 到 100%。超越训练窗口：通过状态组合扩展上下文。我们的 primed 混合模型原生支持 128,000 token。但智能体会话、代码库和长法律文件经常超过这一长度。进一步扩展上下文通常需要在精选数据上进行另一轮持续预训练。混合架构本身能否在没有额外训练的情况下泛化到更长上下文？现有的无需训练扩展上下文窗口的方法分为两类：用于纯 Transformer 的方法（如 YaRN 这样的位置嵌入重缩放方案）和用于纯 SSM 的方法，例如我们去年在国际学习表征会议（ICLR）上提出的 PICASO。两者都没有解决混合设置，在混合设置中 KV 缓存和循环状态必须联合协调。在我们的新 arXiv 论文中，我们描述了一种无需训练即可将混合模型上下文扩展到其原生限制之外的方法。过程很简单：将长输入划分为短于原生上下文长度的块，并在每个块前添加一个共享前缀；并行独立处理所有块；跨块拼接 KV 缓存，同时通过我们在 PICASO 论文中引入的加权平均方案合并 SSM 状态。这种方法使我们能够将 primed 混合模型的上下文长度从 128K 扩展到 256K token，同时在模型从未训练过的长度上，在检索（RULER NIAH）和推理（BABILong）任务上保持可观准确率。在我们的实验中，我们表明混合状态组合保留了 SSM 表达力层级（GKA > GDN > Mamba-2）。混合设计空间现已开放。本博客中讨论的所有工件均以 Apache 2.0 许可证发布在我们的 GitHub Hybrid Model Factory 和 HuggingFace 模型库中：一组用于长上下文推理和指令遵循的 primed 混合模型（8B 和 32B）；priming 训练代码；用于 SSM 层的点对点序列并行算法（在多节点规模下比基于 all-to-all 的替代方案最高快 2.3 倍）；通用序列并行，用于在长序列上原型设计新层，而无需自定义内核或分布式训练专业知识；用于 GKA 解码和训练的优化 Triton 内核；一个 vLLM 服务插件，支持 primed 混合架构，具备连续批处理、分块预填充和张量并行，用于大规模推理部署。如果你服务长上下文工作负载，带有 vLLM 插件的 32B GKA 推理模型可在长上下文推理任务上为你提供更高吞吐量且无质量损失。如果你正在扩展 RL 后训练，混合模型的吞吐量优势意味着在固定 GPU 预算下每单位时间有更多正确 rollout，从而在不改变训练配方的情况下更快进行策略迭代。如果你正在研究新架构，priming 让你能够以预训练成本的一小部分，在受控条件下评估任何层类型，同时从强初始化开始。通用 SP 模块意味着你不需要自定义分布式内核即可在长上下文下训练。Priming 不绑定于特定 Transformer 家族或模型类别。将其扩展到 Llama、Mistral 或 MoE 模型无需算法更改。我们希望这些工具能让社区探索我们测试之外的架构，并开发新的算法和训练配方，为 AI 智能体构建可扩展的下一代架构。完整细节请参见我们的论文《Priming: Hybrid state space models from pre-trained Transformers》。代码、模型和服务插件可在 GitHub 上获取，模型也在 Hugging Face 上。