如果只是拿到一个模型的 API,我们通常只能收集它的输出,让学生学习回答或动作。拿得到权重和中间层时,能力迁移还有另一条路:保留大部分已有结构,替换部分计算模块,再让新模块对齐原模块。
Amazon Science 在 10 月 12 日发布的 Priming 解读,讨论的就是这种架构转换。1 对应论文《Priming: Hybrid State Space Models From Pre-trained Transformers》已有公开预印本。2 它并非"完全不用蒸馏的替代方法",官方流程明确包含短程逐层蒸馏,以及之后的全模型适配。
先明确改造目标
Transformer 的注意力机制可以直接利用历史 token 信息,但自回归推理中的 KV 缓存通常会随上下文增长。状态空间模型把历史信息压缩到固定维度状态里,拥有不同的内存与信息保留特性。
固定状态不意味着保留了任意历史细节,也不意味着替换后必然更好。Priming 采用混合架构,让一部分层继续使用注意力,另一部分改为状态空间层。问题从"能否全部换掉"变成了"哪些层适合换、换成什么、怎样让它们协同"。
因此,参数更少不是这项方法唯一或必要的目标。计算结构、历史状态大小和实现内核都可能影响效率。不能看到论文报告加速,就把它直接解释成"蒸馏出了一个更小的模型"。
三步流程,蒸馏在中间
第一步选择层。官方解读介绍了一种经验代理测试:临时将某层注意力换成滑动窗口注意力,观察性能损失。如果限制历史范围后损失较小,这一层就可能更适合转成固定状态模块。这个测试是选择候选层的方法,不是所有长距离任务都安全的保证。
第二步迁移权重并对齐。研究利用注意力与状态空间计算之间的结构联系,将原注意力层的投影权重用于初始化,再学习状态空间模块特有的参数,并通过蒸馏对齐被替换注意力层的实际输出。这里的监督位置在模型内部,区别于只模仿最终回答。
第三步进行端到端适配。混合模型仍需训练,让保留层和替换层重新协作。局部层看起来接近,并不意味着整个网络的误差不会累积。官方流程使用下一 token 预测来完成全模型适配,所以"无需从头预训练"不等于"无需训练"。
和常见蒸馏路线怎么比较
| 路线 | 主要可用信息 | 对齐对象 | 不能自动推出的结果 |
|---|---|---|---|
| 教师回答监督 | 输入与教师生成文本 | 回答序列 | 复制了教师内部结构 |
| 输出分布蒸馏 | 教师的概率或 logits | 输出分布 | 学生一定更小或同架构 |
| Priming 中的层级对齐 | 原权重及内部模块行为 | 被替换层 | 局部接近就代表所有任务无损 |
这些划分描述训练信号,并非互斥的产品标签。一个完整方案可以同时采用权重初始化、内部对齐和输出训练。工程选型应该问能访问什么、需要保留什么,再选择实现方式。
怎样读官方给出的数字
官方文章报告,整个转换流程消耗的 token 少于源模型预训练预算的 0.5%;也以相较从头预训练降低超过 99% 成本概括方法优势。这里比较的是复用一个已经训练好的模型,与重新做预训练的投入,不是说任意企业的整体 AI 账单能减少 99%。
文章还介绍了一个 32B 的 Gated KalmaNet 混合模型:在报告的推理测试中接近源 Transformer,在 128K 上下文设置下解码速度最高达到约两倍。模型规模、上下文长度、比较对象和测量环节都应一起保留。解码吞吐不是用户端所有请求的端到端延迟,短请求也不一定得到同样收益。
复现时应该建立哪组对照
先固定源模型、数据范围、硬件和推理配置,再比较不同替换层位置与比例。训练投入、推理缓存、吞吐、任务质量要分别记录。只拿一个原模型默认部署和一个经过优化内核的混合模型比较,难以区分架构与实现优化各自贡献。
任务评测也要覆盖不同记忆需求:需要精确引用早期内容的检索任务、主要依赖最近上下文的任务,以及综合推理。固定状态会压缩历史,某些平均分接近不代表每一类任务都接近。保留注意力层的意义之一,正是维持部分精确回忆能力。
最后再将训练成本摊入预计服务量,并计算模型转换、内核适配和维护成本。服务量不足时,即使单次推理更省,转换投入也未必值得。模型权重许可、训练数据使用权和现有推理栈兼容性,同样需要单独确认。
Priming 的启发是:教师能传递的并不只有最后一句答案,已经训练好的内部结构也能成为迁移起点。它让"换架构"不必自动等同于"重做全部预训练",但每一次替换仍要用对齐训练和任务证据承担验证责任。
参考资料
1 Amazon Science:Can we change an AI model's architecture without retraining it from scratch?,2026-10-12。
2 Priming 论文。