跳到主要内容
返回时间线
Amazon Science homepage来源发表:

混合架构不再从零预训练:复用Transformer权重,预训练成本降99%以上

核心概要

亚马逊云科技提出priming方法,用预训练Transformer权重初始化混合状态空间模型,把混合架构的预训练成本降低99%以上,并据此在受控条件下比较三类状态空间层。

Source-provided article image: Can we change an AI model’s architecture without retraining it from scratch?
Amazon Science

深度剖析

混合架构的预训练成本被压到源模型预训练token预算的不到0.5%,使原本因成本过高而无法在规模上做的架构对照变得可行。 此前探索这一设计空间需要每个新架构从零训练,消耗数万亿token和数十万GPU小时,受控比较极其昂贵。 整条流程消耗源模型预训练token预算的不到0.5%,并在Qwen、Llama、Mistral等Transformer家族、稠密或MoE模型类别以及80亿到320亿及以上参数规模上成立。

在相同数据、相同训练预算和相同记忆层比例下,三类状态空间层表现出可预测的强弱次序:Gated KalmaNet优于Gated DeltaNet,Gated DeltaNet优于Mamba-2。 Mamba-2对整个状态施加统一标量衰减,Gated DeltaNet在统一衰减之上加入针对特定键的定向更新,Gated KalmaNet的擦除方向还依赖所有过去键的统计量。 理论表达能力层级预测了实测层级,差距在长上下文推理上拉大、在短上下文任务上收窄;320亿参数规模下混合GKA与同数据训练的源Transformer相差不到1%,而推理时KV缓存只需一半。

基于状态空间模型的混合模型在长上下文基准上一致优于对应的滑动窗口注意力混合模型,即使后者拿到两到四倍的记忆预算。 滑动窗口注意力逐字保留最近一段token,直觉上窗口内无信息损失,堆叠多层还能扩大感受野。 所有受测的SSM混合模型在长上下文基准上都超过了对应的SWA混合模型,即使窗口的记忆是状态的两到四倍。

Gated KalmaNet提供推理时可调的计算量,同一训练好的模型可在速度与精度前沿上取不同工作点。 Mamba-2和Gated DeltaNet对每个输入token使用相同的计算预算,与内容无关。 在AIME 2025上,GKA以10次迭代代替通常的30次,精度损失约1%,解码快10%;320亿参数GKA混合模型在128000 token上下文下解码最高快两倍,达到正确答案的墙钟时间最短为源模型的三分之二。

启示与展望

报告面向关注长上下文推理与智能体负载的研究人员和工程师。已展示的范围包括:priming在Qwen、Llama、Mistral家族以及稠密或MoE模型上成立,覆盖80亿到320亿及以上参数;受测混合模型原生支持128000 token,通过分块加共享前缀、拼接KV缓存并按PICASO的加权平均合并SSM状态,可在不训练的情况下把上下文扩展到256000 token,并在RULER NIAH检索与BABILong推理任务上保持相当精度。发布物包括80亿与320亿参数的primed混合模型、训练代码、面向SSM层的点对点序列并行算法、通用序列并行模块、GKA解码与训练的Triton内核,以及支持连续批处理、分块预填充和张量并行的vLLM服务插件,均以Apache 2.0许可发布。

材料把priming描述为一种近似:只有当固定维状态的状态空间模型能产生足够接近注意力层的混合矩阵时,源Transformer的投影才可复用,而近似是否成立取决于该注意力层做的是复杂长程检索还是简单局部操作。层选择依赖一个廉价的经验代理,即用滑动窗口注意力替换该层并测量退化程度,这一代理与动力系统理论中的Hankel秩相联系。GKA以10次迭代代替30次时精度损失约1%,说明速度与精度之间存在可调但非免费的取舍。上下文扩展方法在128K到256K上保持相当精度,材料未给出更长外推的结论。