AWS Priming 怎样用层级蒸馏,把 Transformer 改造成混合架构

如果只是拿到一个模型的 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 论文。

相关推荐
AOI小白新手上路1 小时前
ICME23 知识蒸馏复现 · SLDNet阶段性沉淀
深度学习
loulanyue_1 小时前
芯片设计:从最底层讲起——读MatX CEO Reiner Pope激辩逻辑门、脉动阵列与GPU架构演进
架构
张彦峰ZYF2 小时前
一次触发、一把锁与一个号码:分布式基础组件的统一设计推演
架构·先立六条时间纪律·再拆解编号的位分配·节点编号分配与回拨四策·辨析互斥、租约与领导权三种语·调度拆成七个子问题·五层架构
码上上班2 小时前
基于 K3s 的微服务容器化与 GitOps 持续交付平台
微服务·云原生·架构
Dawson Zhu2 小时前
《Agentic Design Patterns》第 17 章导读:推理技术
人工智能·语言模型·架构·aigc·agi
各类产品分享13 小时前
制药厂 GMP 车间 AR 巡检系统落地架构与实施要点
架构·ar·ar巡检
陈嘿萌3 小时前
学术速递|2026年8月 arXiv 图像融合论文汇总(08.01–08.31)29 篇最新成果
图像处理·人工智能·深度学习·计算机视觉·图像融合·学术研究与理论基础
云和恩墨3 小时前
从通用规范到系统级联动:以Skill为核心的Mopheus扩展架构与落地实践
架构
行者全栈架构师3 小时前
给开源 Rust 工具加一个新检查项,TRAE 帮我从想法跑到构建产物:全链路实战
前端·算法·架构