我们能否在不从头重新训练的情况下改变 AI 模型的架构？

人工智能代理正迅速规模化：会话时长更长、推理链条更深、自主性更强。然而，其底层的Transformer架构仍需承受与处理的标记数成比例增长的内存开销。

随着代理能力的提升，这一成本也随之攀升。如今，涌现出大量混合架构，它们通过多种方法的组合，在内存、速度与质量之间实现更优的平衡。

这些方法包括：稀疏注意力机制——在生成下一个标记时仅选取部分历史标记进行计算；滑动窗口注意力机制——仅保留最近的若干标记；以及状态空间模型——将完整历史信息压缩为固定大小的状态表示。

这些混合架构可被视为序列模型设计空间中的不同点：它们在遍历序列、经历不同状态的过程中，以不同程度地融合信息。区分它们的关键在于如何保存并利用历史信息。

迄今为止，探索这一设计空间通常需要从头训练每一种新架构，耗费数万亿个标记和数十万GPU小时，使得对混合架构的可控比较成本极高。

在一篇新论文《Priming：基于预训练Transformer的混合状态空间模型》中，我们亚马逊云科技团队提出了一个解决方案。我们引入了“Priming”方法，利用预训练Transformer的权重作为起点，将混合架构的预训练成本降低99%以上。

借助Priming，模型开发者得以提出此前因成本过高而难以回答的问题：哪些注意力层可以被更高效的层替代？在不牺牲质量的前提下，最多可替换多少个注意力层？

哪种状态空间层效果最佳？速度与准确率之间的权衡又是什么？在论文中，我们对比了采用三种不同状态空间模型（Gated KalmaNet、Gated DeltaNet和Mamba-2）的混合模型。

部分研究发现如下：并非所有状态空间模型都等效。在受控条件下，Gated KalmaNet（GKA）表现最佳，其次是Gated DeltaNet（GDN），最后是Mamba-2。

GKA的独特之处在于无需重新训练即可实现变长推理计算，并让用户直接调控其在速度—准确率权衡曲线上的位置。对整体信息的良好有损摘要，胜过对近期历史的完美记录。

即使滑动窗口注意力机制（SWA）的内存预算翻倍，状态空间模型的表现仍持续优于SWA。同等质量，吞吐量却提升一倍。一款基于GKA、参数规模为320亿的混合模型，在推理基准测试中与源Transformer性能相当，同时以12.8万标记的上下文长度解码速度最高可达前者的两倍，并可在仅为三分之二的墙钟时间内得出正确答案。

为推动混合架构的进一步研究，我们已在Apache 2.0许可下发布了八亿和三十二亿参数规模的混合模型库，用于长上下文推理与指令遵循任务，并附带训练代码、优化内核及vLLM服务插件。

注意力机制 Transformer的注意力机制决定了当前正在处理的标记与先前见过的各标记之间的相关性，其核心依赖于三个向量，分别对应数据库设计中的查询、键和值。

对于每个标记，查询表示该标记“寻找什么”，键表示它“提供什么”，而值则携带其将传递的实际信息。相关性通过比较当前标记的查询与所有其他标记的键来计算；匹配度越高，相应标记的值对注意力输出的贡献就越大。

这三个向量均由标记与学习得到的权重投影矩阵相乘后生成。Transformer每次输出一个标记，且每个输出标记随即成为新的输入，因此上下文也随时间逐标记增长。

为节省计算资源，Transformer会将每个新输入的键和值投影存储于“KV缓存”中。两种记忆层：全忆与渐逝 Transformer赋予模型我们所称的“全忆”：对每一个历史标记的逐字复现。

正是这种能够回溯整个LLM上下文窗口的能力，才使上下文学习、长文档推理及检索增强生成成为可能。但这也带来了与输入标记数成比例增长的内存开销。

不过，并非所有注意力层的作用相同。有些负责跨数万标记的长距离检索，另一些则更为局部，仅追踪当前子句、复制邻近标记等。后者无需不断增长的缓存：一份紧凑的、固定大小的历史摘要——即我们所说的“渐逝记忆”——便已足够。

混合模型将两者结合：对于需要全局记忆的层使用注意力层（全忆，缓存增长），而对于其余部分则采用状态空间层（渐逝，固定大小）。

由此得到的模型不仅内存占用更低，且在关键场景下仍能实现逐字复现。Priming：知识不局限于特定架构 探索混合架构的设计空间一直不具可行性，因为每种配置通常都需要从零开始预训练。

Priming通过复用预训练Transformer中已存储的知识，绕过了这一限制。Priming之所以可行，是因为Transformer本身即是一种特殊的状态空间实现。

无论是注意力层还是SSM层的输出，均可表示为值序列与“混合矩阵”相乘的结果，该矩阵决定每个历史标记对当前标记的贡献程度。注意力层通过查询矩阵与键矩阵的乘积计算混合矩阵；而SSM则通过一个固定维度的递归状态产生类似结果，所有历史信息均需经过该状态。

Priming利用了这种结构上的相似性：若某个具有固定维度状态的SSM能够生成与注意力层相近的混合矩阵，则源Transformer的查询、键和值投影便可直接用作SSM的输入投影，只需学习SSM特有的参数即可。

这种近似是否成立，取决于注意力层执行的是复杂的长距离检索，还是简单的局部操作。我们采用一种经济的实证指标加以区分：将每一层替换为仅保留部分上下文的SWA，并测量性能下降情况。

那些能容忍这种替换的层，便是SSM转化的天然候选者：它们的混合矩阵无需捕捉长距离关系，因此SSM的状态足以忠实再现其功能。这一指标与动力系统理论中的一个正式概念——混合矩阵的汉克尔秩——相关联，该秩衡量必须通过状态瓶颈的信息量，从而确定精确实现该层所需的最小SSM状态维度。

Priming分为三个阶段：层选择：识别哪些层可在保持其余层为注意力层的同时，被忠实地压缩至固定大小的状态；权重转移与对齐：以对应的注意力权重初始化SSM层，随后运行一段简短的蒸馏过程，使每个SSM的行为与其所替代的层保持一致；端到端适配：在下一标记预测任务上微调整个混合模型，让各层协同适应。

整个流程仅消耗源模型预训练标记预算的不到0.5%，且适用于不同Transformer家族（Qwen、Llama、Mistral）、不同模型类别（密集型或MoE）以及不同参数规模（80亿至320亿及以上）。

借助Priming，我们现在可以低成本、系统性地探索设计空间、层类型、混合比例及层位布局，因为每种变体均始于相同的强大初始化，且相较于从零预训练所需，训练量极小。

渐逝记忆层的受控比较 例如，我们可以提出这样一个问题：在数据、训练预算及全忆层与渐逝层的比例均相同的情况下，哪种SSM层在混合模型中表现最佳？

在论文中，我们评估了三类状态空间模型，它们代表了利用固定大小状态的不同方式：Mamba-2在每一步对整个状态施加统一的标量衰减，形成一种按相同比例遗忘一切的移动平均机制。

每当有新标记进入时，它会被写入状态，而所有旧条目则被均匀衰减。Mamba-2简单高效，但缺乏针对性。Gated DeltaNet（GDN）在统一衰减的基础上增加了定向更新。

它并非均匀衰减所有内容，而是能够在写入新值之前，选择性地抹去与特定键相关的旧值（即“delta规则”）。这使其比Mamba-2更具表达力：可在不均匀衰减其余内容的前提下，修正特定关联。

Gated KalmaNet（GKA）更进一步。它的擦除方向不仅取决于当前键，还取决于所有历史键的统计信息，是一种面向历史的更新，能够考虑模型所见一切的累积结构。

理论上的表达力层级与实证结果一致：GKA > GDN > Mamba-2。在长上下文推理任务中，这一差距进一步拉大，而在短上下文任务中则有所收窄。

此外，在320亿参数规模下，混合GKA在推理时仅需一半的KV缓存，却与在同一数据上训练的源Transformer性能相差不到1%。渐逝记忆胜过逐字窗口记忆 我们研究中的一项有趣发现来自对采用SSM层的混合模型与采用滑动窗口注意力机制（SWA）的混合模型的直接比较。

SWA通过逐字保留固定窗口内的最新标记来代替将历史压缩为状态的做法。人们可能会认为，即使是有限窗口内的逐字复现，也应优于有损压缩。

毕竟，滑动窗口内的信息并未丢失，而且叠加多层SWA还能扩大端到端模型的“感受野”。然而，我们测试的每一款基于SSM的混合模型，都在长上下文基准测试中胜过相应的SWA混合模型，即便SWA的窗口内存是SSM状态的两到四倍。

推理时的速度—准确率调节旋钮：变长推理计算 GKA具备Mamba-2和GDN所不具备的特性：推理时的变长计算。Mamba-2和GDN对每个输入标记的计算预算固定不变，无论其内容如何。

相比之下，GKA以迭代方式解决回归问题，且在服务时无需重新训练即可调整迭代次数。实践中，将GKA的服务迭代次数从常规的30次降至10次，会在AIME 2025竞赛中损失约1%的准确率，但解码速度却提升10%。

同一款训练好的模型可在速度—准确率曲线上服务于多个点，因此模型提供商无需维护不同的模型变体即可提供不同的延迟等级。这与我们在早期工作中提出的观点相呼应：智能即根据难度合理分配精力的能力。

GKA的变长计算机制正是这一原则的架构体现：对速度要求高的部署可减少迭代次数，而对准确率要求高的部署则可增加迭代次数。更快的推理，而不只是更快的标记 单个标记的原始吞吐量只能说明部分问题。

如果模型的推理能力较弱，需要更多标记才能弥补，那么即便单个标记的成本较低，其端到端速度也可能更慢。部署层面的关键在于，混合模型在墙钟时间上达到目标准确率的速度究竟有多快。

我们在AIME 2025这一高难度数学竞赛基准上进行了测量，将我们的GKA-Primed-HQwen3-32B-Reasoner与在同一推理配方上训练的源Qwen3-32B进行了对比。

两款模型生成的思维轨迹长度相近（全集的中位数约为1.4万标记，最难题则达2.8万标记），即它们的回答同样深入，但混合模型的吞吐量要高得多。

在实际并发水平下，当Transformer的KV缓存占满GPU内存时，混合模型在全基准上的达标速度是源模型的1.2至1.4倍，而在最难题上甚至可达1.6倍。

这得益于SSM层的固定大小状态，它将每序列的内存占用大致减半，从而在同一硬件上支持更多并发操作。吞吐量优势随上下文长度增长：在1.6万标记时为1.50倍，而在12.8万标记时，320亿参数的GKA（10次迭代）可达2.06倍。

对于80亿参数的模型，增益更为显著：在1.6万标记时为1.94倍，而在12.8万标记时高达2.29倍。在长上下文场景下，KV缓存读取主导了Transformer的前向传播，而将缓存减半相当于在内存饱和时将有效批大小几乎翻倍。

这一效应不仅影响推理服务，在强化学习（RL）的后训练流水线中，吞吐量也决定了策略生成正确轨迹的速度。混合模型的优势直接转化为在固定GPU预算下，RL训练数据采集速率提升20%至100%。

超越训练窗口：通过状态组合扩展上下文 我们的Priming混合模型原生支持12.8万标记。然而，代理会话、代码库以及长篇法律文件往往远超此限。

若要进一步扩展上下文，通常需要在精选数据上进行新一轮的持续预训练。混合架构本身能否在无需额外训练的情况下推广至更长的上下文？

现有的无训练扩展上下文方法可分为两类：一类适用于纯Transformer（如YaRN等位置嵌入缩放方案），另一类适用于纯SSM，例如我们去年在国际表征学习大会（ICLR）上提出的PICASO。

这两类方法均无法应对KV缓存与递归状态需协同处理的混合场景。在我们发表于arXiv的新论文中，我们描述了一种无需训练即可扩展混合模型上下文的方法。

步骤很简单：将长输入分割为短于原生上下文长度的块，并在每块前添加共享前缀；并行独立处理所有块；在合并SSM状态时，通过我们在PICASO论文中提出的加权平均方案，同时拼接各块的KV缓存。

这一方法使我们能够将Priming混合模型的上下文从12.8万标记扩展至25.6万标记，并在检索（RULER NIAH）与推理（BABILong）任务上保持较高准确率，即便是在模型从未接受训练的长度下亦然。

实验表明，混合状态组合能够维持SSM的表达力层级（GKA > GDN > Mamba-2）。混合设计空间现已开放 本博客中讨论的所有成果均在GitHub上的Hybrid Model Factory及HuggingFace上的模型库中以Apache 2.0许可发布：用于长上下文推理与指令遵循的Priming混合模型集合（8B与32B）；Priming训练代码；SSM层的点对点序列并行算法（在多节点规模下速度可达全对全方案的2.3倍）；用于在长序列上原型化新型层的通用序列并行模块，无需定制内核或分布式训练专业知识；用于GKA解码与训练的优化Triton内核；支持Priming混合架构的vLLM服务插件，具备连续批处理、分块预填充及张量并行等功能，适用于大规模推理部署。

若您处理长上下文工作负载，配备vLLM插件的320亿参数GKA推理模型可在长上下文推理任务中实现更高吞吐量且不损失质量。若您在扩展RL后训练规模，混合模型的吞吐优势意味着在固定GPU预算下单位时间内可完成更多正确的模拟，从而加快策略迭代，无需更改训练配方。

若您从事新型架构研究，Priming让您能在受控条件下以远低于预训练的成本评估任意层类型，并从强大的初始化出发。通用SP模块使您无需定制分布式内核即可在长上下文中进行训练。

Priming并不局限于特定的Transformer家族或模型类别。将其扩展至Llama、Mistral或MoE模型无需任何算法变更。我们希望这些工具能帮助社区探索我们尚未测试过的架构，并开发用于构建可扩展下一代AI代理架构的新算法与训练配方。

详情请参阅我们的论文《Priming：基于预训练Transformer的混合状态空间模型》。代码、模型和服务插件均可在GitHub上获取，模型亦可在Hugging Face上找到。