能否在不从零重新训练的情况下改变AI模型的架构？（14分钟阅读）

# 我们能否在不从零重新训练的情况下改变AI模型的架构？- Amazon Science Source: https://www.amazon.science/blog/can-we-change-an-ai-models-architecture-without-retraining-it-from-scratch AI智能体正在快速扩展：更长的会话、更深层的推理链、更高的自主性。但其底层的Transformer架构仍然面临着随处理token数量增长的内存开销。随着智能体能力的增强，这一成本也随之水涨船高。 如今，混合架构日益涌现，通过各种方法的组合来更好地平衡内存、速度与质量。这些方法包括稀疏注意力（在选择下一个token时只选取部分历史token进行考量）、滑动窗口注意力（仅保留最近的token）以及状态空间模型（将完整历史中的信息压缩为固定大小的状态）。 这些混合架构（https://www.amazon.science/publications/bmojo-hybrid-state-space-realizations-of-foundation-models-with-eidetic-and-fading-memory）可以被理解为序列模型设计空间中的不同点：它们在遍历序列并通过不同状态的过程中，以不同程度地混合信息。区分它们的关键在于如何保留和利用过去的信息。 在此之前，探索这一设计空间需要从零开始训练每一个新架构，耗费数万亿token和数十万GPU小时。这使得混合架构之间的受控对比成本极其高昂。 在一篇新论文*"Priming: Hybrid state space models from pre-trained Transformers*（https://arxiv.org/abs/2605.08301）"中，我们Amazon Web Services的团队解决了这一问题。我们提出了一种名为**priming**（启动）的方法，它利用预训练Transformer的权重作为起点，将混合架构的预训练成本降低99%以上。 通过priming，模型开发者可以提出此前在大规模验证中代价过高的问题：哪些注意力层可以被更高效的层替换？在不牺牲质量的前提下可以替换多少注意力层？哪种状态空间层表现最佳？速度与精度之间的权衡关系是怎样的？ 在我们的论文中，我们比较了使用三种不同类型状态空间模型（Gated KalmaNet、Gated DeltaNet和Mamba-2）的混合模型。以下是我们的一些发现： - ***并非所有状态空间模型都一样。*** 在受控条件下，Gated KalmaNet（GKA）表现最佳，其次是Gated DeltaNet（GDN），然后是Mamba-2。GKA独特地提供了无需重新训练即可实现的可变长度推理计算，并让用户直接控制其在速度-精度前沿上的位置。一个320亿参数的混合模型比同等规模的Transformer更快达到每个精度阈值，在更少的挂钟时间内收敛到相似的最终精度。数据集为AIME 2025数学竞赛基准（困难子集）。时间*t*处的精度是指在时间*t*之前产生正确最终答案的推演（rollout）比例。 - ***对全部信息的良好有损摘要优于对最近过去的完美记录。*** 状态空间模型始终优于滑动窗口注意力（SWA），即使SWA获得两倍的内存预算也是如此。 - ***同等质量，两倍吞吐量。*** 一个基于GKA的320亿参数混合模型在推理基准上与其源Transformer表现相当，但在128,000 token上下文下解码速度最高可达两倍，且在低至三分之二的挂钟时间内得出正确答案。 为了促进混合架构的进一步研究，我们已在Apache 2.0许可证下发布了priming混合模型系列，涵盖80亿和320亿参数规模的长上下文推理和指令跟随模型，同时提供训练代码、优化内核和vLLM服务插件。 ## 注意力机制 Transformer的注意力机制——决定先前看到的token与当前处理token的相关程度——依赖三个不同的向量，其名称源于数据库设计：查询（query）、键（key）和值（value）。对于每个token，查询捕捉该token"在寻找什么"，键捕捉它"提供什么"，值则携带它将传递的实际信息。相关性通过将当前token的查询与每个其他token的键进行比较来计算；匹配越接近，该token的值对注意力输出的贡献就越大。 这些向量中的每一个都是通过将token乘以学习到的权重投影矩阵来生成的。Transformer一次输出一个token，每个输出token成为新的输入，因此上下文也逐token增长。为了节省计算量，Transformer将每个新输入的键和值投影存储在"KV缓存"中。 ## 两种记忆层：持久记忆与衰减记忆 注意力机制赋予Transformer我们所称的***持久记忆（eidetic memory）***：对每个过去token的逐字回忆。在整个LLM上下文窗口内回忆的可能性，使得上下文学习、长文档推理和检索增强生成成为可能。但它的代价随输入token数量增长。 然而，并非所有注意力层执行相同的任务。有些进行跨越数万token的长程检索。另一些则更多在局部运作，跟踪当前子句、复制邻近token等。局部层不需要不断增长的缓存：一个紧凑的、固定大小的过去摘要——我们称之为***衰减记忆（fading memory）***——就足够了。 混合模型将两者结合：注意力层（持久记忆，增长的缓存）用于需要全局回忆的层，状态空间层（衰减记忆，固定大小）用于其余层。结果是一个内存占用降低的模型，在关键位置仍保留逐字回忆的能力。 ## *Priming*：知识不局限于架构 探索混合架构的设计空间此前并不现实，因为每种混合配置通常都需要从零开始预训练。***Priming***通过复用预训练Transformer权重中存储的知识来绕过这一障碍。 Primming之所以可行，是因为Transformer本身就是状态空间实现的一个特例（https://www.amazon.science/publications/bmojo-hybrid-state-space-realizations-of-foundation-models-with-eidetic-and-fading-memory）。注意力层和SSM层的输出都可以表示为将值序列乘以一个"混合"矩阵的结果，该矩阵决定了每个过去token对每个当前token的贡献程度。注意力通过将查询矩阵和键矩阵相乘来计算混合矩阵；SSM则通过一个所有历史信息都必须经过的固定维度循环状态来产生类似的输出。 Priming利用了这种结构相似性：如果一个具有固定维度状态的SSM能够产生足够接近注意力层混合矩阵的结果，那么源Transformer中的查询、键和值投影就可以作为SSM自身的输入投影来复用，只需学习SSM特有的参数。 这种近似是否成立取决于注意力层执行的是复杂的长程检索还是简单的局部操作。我们使用一种廉价的经验代理来区分两者：将每个层替换为仅保留部分上下文的SWA，并测量性能下降。 能够容忍这种替换的层是转换为SSM的天然候选：它们的混合矩阵不需要捕捉长程关系，因此SSM的状态可以忠实地表示它们的行为。这一代理与动力系统理论中的一个形式化量相关联，即层的混合矩阵的***Hankel秩***（https://www.researchgate.net/publication/243782056_ONTHE_STOCHASTIC_REALIZATION_PROBLEM），它衡量了必须通过状态瓶颈的信息量，从而确定精确实现该层所需的最小SSM状态维度。 Priming分三个阶段进行： 1. ***层选择：*** 识别哪些层可以忠实地压缩为固定大小的状态，同时将其余层保留为注意力层； 2. ***权重迁移与对齐：*** 从对应的注意力权重初始化SSM层，然后运行一个短程蒸馏过程，将每个SSM与其替换层的行为对齐； 3. ***端到端适应：*** 在下一个token预测任务上微调整个混合模型，让所有层协同适应。 Priming首先使用一个简单的代理（滑动窗口注意力）来识别主要在局部运作的注意力层。这些层的权重被蒸馏到状态空间模型（SSM）层中，然后与其替换层的行为对齐。最后，在下一个token预测任务上微调整个混合模型。 整个流程消耗不到源模型预训练token预算的0.5%，并且适用于不同的Transformer系列（Qwen、Llama、Mistral）、模型类别（dense或MoE）和规模（80亿到320亿及以上）。 通过priming，我们能够以低成本且系统化地探索设计空间、层类型、混合比例和层位置，因为每个变体都从相同的强初始化开始，与预训练相比只需要极小的训练量。 ## 对衰减记忆层的受控对比 例如，考虑这个问题：在相同数据、相同训练预算和相同持久记忆与衰减记忆层比例的条件下，哪种SSM层在混合模型中表现最佳？ 在我们的论文（https://arxiv.org/abs/2605.08301）中，我们评估了三族状态空间模型，每族代表了使用固定大小状态的不同方法： - ***Mamba-2*** 在每一步对整个状态施加统一的标量衰减，这是一种以相同速率遗忘所有内容的滑动平均。当新token到达时，它被写入状态，同时所有先前的条目被均匀衰减。Mamba-2简单高效，但不加区分。 - ***Gated DeltaNet（GDN）*** 在均匀衰减的基础上增加了定向更新。它不会均匀地衰减所有内容，而是可以在写入新值之前选择性地擦除先前与特定键关联的值（delta规则）。这使其比Mamba-2更具表达力：它可以修正特定的关联，而不必均匀衰减其余内容。 - ***Gated KalmaNet（GKA）*** 更进一步。它的擦除方向不仅取决于当前的键，还取决于***所有过去键***的统计量，这是一种考虑模型所见全部内容累积结构的历史感知更新。 使用Mamba-2、GDN和GKA状态空间模型的混合模型在四项任务上的结果。GKA通常优于GDN，GDN始终优于Mamba-2，与理论预测一致。 理论上的表达力层级预测了经验上的结果：GKA > GDN > Mamba-2。这一差距在长上下文推理中最为显著（压缩质量最为重要），在短上下文任务中则有所缩小。此外，在320亿参数（32B）规模下，混合GKA与其在同一数据上训练的源Transformer差距在1%以内，同时推理时只需一半的KV缓存。 ## 衰减记忆优于逐字窗口记忆 我们研究中一个有趣的发现来自对使用SSM层的混合模型与使用滑动窗口注意力（SWA）的混合模型的直接对比。SWA逐字保留一个固定窗口的最近token，而非将过去压缩为状态。人们可能预期逐字回忆即使在有限窗口内也会优于有损压缩——毕竟在滑动窗口范围内没有任何信息丢失，而且堆叠多层SWA还能增加端到端模型的"感受野"。 然而，我们测试的每个基于SSM的混合模型都在长上下文基准上优于对应的SWA混合模型，即使滑动窗口获得的状态两到四倍的内存预算也是如此。 使用有损上下文摘要的混合模型（***橙色***）始终优于使用滑动上下文窗口的混合模型（***青色***）。 ## 速度与精度的运行时调节：推理时的可变计算 GKA拥有Mamba-2和GDN都不具备的特性：可变的测试时计算量。 Mamba-2和GDN对每个输入token使用相同的计算预算，无论其内容如何。相比之下，GKA通过迭代求解一个回归问题，迭代次数可以在服务时调整而无需重新训练。 Gated KalmaNet（GKA）的精度-速度前沿，涵盖四种不同的上下文长度。几何图形表示回归问题的不同迭代次数。 在实践中，以10次迭代而非通常的30次运行GKA，会牺牲AIME 2025上约1%的精度，但解码速度提升10%。单个训练好的模型可以在速度-精度前沿上服务多个点，因此模型提供商可以提供不同的延迟层级，而无需维护单独的模型变体。 这与我们在早期工作中讨论过的一个原则相关：智能是按照难度分配与之相称的努力的能力（https://www.amazon.science/blog/intelligence-isnt-about-parameter-count-its-about-time）。GKA的可变计算机制是这一原则的架构体现：需要速度的部署可以运行较少的迭代，需要精度的部署则可以运行更多迭代。 ## 更快的推理，而非仅仅更快的token 单纯的每token吞吐量只讲述了故事的一部分。即使每token成本更低的模型，如果其推理能力较弱、需要更多token来补偿，端到端也可能更慢。与部署相关的问题是：从挂钟时间来看，混合模型能多快达到目标精度。 我们在AIME 2025这一高难度数学竞赛基准上进行了测量，将我们的GKA-Primed-HQwen3-32B-Reasoner（https://huggingface.co/collections/amazon/primed-hybrid-models-collection）与使用相同推理配方训练的源Qwen3-32B进行了对比。两个模型生成的思维链长度相似（完整集上中位数约14,000 token，最难问题上约28,000 token）——即它们产生了同等深度的回答——但混合模型的吞吐量远高于此。 在真实并发水平下（Transformer的KV缓存饱和GPU内存），混合模型在完整基准上达到目标精度阈值的速度快1.2至1.4倍，在最难问题上快至多1.6倍。这是SSM层固定大小状态的结果——它将每序列的内存占用大约减半，使得在同一硬件上实现更多并发操作。 GKA-Primed-HQwen3-32B相对于源Transformer的解码吞吐量加速比，作为上下文长度的函数。将一半注意力层替换为固定大小的SSM状态可将KV缓存大小大约减半，在128,000 token上下文长度下获得高达2.06倍的吞吐量。在服务时将GKA求解器迭代次数从30次降至10次（无需重新训练），以牺牲少量精度为代价换取额外吞吐量。 吞吐量优势随上下文长度增长：16K token时为1.50倍，32B模型（GKA在10次迭代下）在128K时扩展至2.06倍。对于80亿参数模型，增益更为显著：从16K时的1.94倍到128K时的2.29倍。在长上下文下，KV缓存读取主导Transformer的前向传播，将缓存在内存饱和时大致将有效批大小翻倍。 这一意义超越了推理服务。在强化学习（RL）后训练流程中，吞吐量决定了策略