博客学术研究重点混合架构状态空间模型训练效率AWS
Can we change an AI model’s architecture without retraining it from scratch?
AI 解读已生成 · AI 生成，数字以原文为准
AWS 提出 Priming，复用预训练 Transformer 权重，把混合架构预训练成本降 超过 99%
AI agents 更长会话使 Transformer 的 KV cache 内存税随 token 数增长，而探索混合架构通常需从头预训练。AWS 在论文《Priming: Hybrid state space models from pre-trained Transformers》中提出 Priming：将 Transformer 视为状态空间实现的特殊情形，复用预训练 query、key、value 投影来初始化 SSM，仅学习 SSM 专属参数，并以滑动窗口注意力（SWA）替换测试识别可压缩层，经权重迁移对齐和端到端下一 token 适配构建混合模型；该方法将混合架构预训练成本降低 99% 以上，整个流程消耗不足源模型预训练 token 预算的 0.5%，适用于 Qwen、Llama、Mistral 及稠密/MoE 模型。作者报告在受控条件下 GKA 最优，SSM 一致胜过内存预算更高的 SWA；32B GKA 混合模型在推理基准上匹配源 Transformer，在 128,000 token 上下文下解码吞吐最高达 两倍。项目已发布模型、训练代码、优化内核和 vLLM 服务插件（Apache 2.0）。
- 1Priming 复用 Transformer 权重，不从头训练混合架构
- 2GKA 受控质量最优，可变长度推理无需重训
- 3开源 8B/32B 模型、训练代码、内核与 vLLM 插件
为什么重要 可直接把 Llama/Qwen 等改造为混合架构，降低长上下文 agent 推理成本；已开源 8B/32B 模型与 vLLM 插件适合 做长上下文 LLM 推理优化、模型架构改造或 agent 成本控制的工程与研究人员
- 来源
- amazon.science
>99%
成本降低
官方称总预训练成本
<0.5%
token 预算
源模型预训练预算占比
最高 2x
解码加速
32B GKA 混合模型
2/3墙钟
到答案时间
达到正确答案用时
注释 · 成本降为总预训练开销口径；token 预算指 术语解释从预训练 Transformer 权重创建混合 SSM 架构的方法 管线相对源模型预训练用量。出自 正文·摘要
深度总结
1 / 8Priming 三阶段管线· 流程
先通过 SWA 代理 逐层探测可压缩性，再迁移权重并蒸馏对齐，最后端到端微调让层协同
- 01Layer selection
用滑动窗口注意力逐层替换 Transformer 注意力并测退化；容忍替换的层适合转 SSM，其余保留为 attention。
- 02Weight transfer and alignment
从对应 attention 权重初始化 SSM 层，再用短蒸馏让 SSM 对齐被替换层的行为。
- 03End-to-end adaptation
在 next-token prediction 上微调整个混合模型，让所有层协同适应。
注释 · SWA 只保留部分上下文，层对 SWA 替换的退化越小，越适合换成 SSM。出自 正文·Priming proceeds in three