Prime：从预训练变压器构建混合状态空间模型
机器学习
2026-05-12 v1 人工智能
摘要
混合状态空间模型将Attention与递归状态空间模型（SSM）层相结合，平衡了来自Attention的集中记忆与来自SSM的压缩淡化记忆。这使得Key-Value缓存更小，解码更快，同时提供更丰富的架构设计空间。探索该设计空间的规模化研究迄今仍需从头训练，这一壁垒使大多数大规模混合模型研究局限于狭窄的架构范围。我们提出Prime方法，将混合架构设计从预训练问题转化为知识迁移问题。Prime从预训练变压器初始化混合模型，经过短期对齐和后训练阶段，即可在不超过源模型预训练token预算的0.5%时恢复下游质量。Prime对源变压器家族（如Qwen、Llama、Mistral）、模型类别（稠密或Mixture-of-Experts）以及模型规模均不依赖。Prime使我们能够在相同的条件下对SSM层类型进行规模化受控比较。我们评估了Gated KalmaNet（GKA）、Gated DeltaNet（GDN）和Mamba-2，表明它们的表达层次GKA>GDN>Mamba-2直接预测在长上下文推理任务中的下游性能。我们将Prime扩展至8B/32B推理模型，支持原生128K上下文。我们的混合GKA 32B在其源Qwen3-32B基础上提升了+3.8个平均推理分数，同时保持在相同数据上的Transformer后训练模型的1%以内，并实现了最高2.3倍的解码吞吐量提升。为促进混合架构研究，我们发布了一个包含Prime混合模型的模型库，支持长上下文推理和指令遵循，伴随Prime训练和推理代码（用于长上下文训练的Sequence Parallelism算法、优化GKA内核以及vLLM serving插件），全部采用Apache~2.0许可证。
引用
@article{arxiv.2605.08301,
title = {Priming: Hybrid State Space Models From Pre-trained Transformers},
author = {Aditya Chattopadhyay and Elvis Nunez and Prannay Kaul and Benjamin Bowman and Evan Becker and Luca Zancato and David Thomas and Wei Xia and Stefano Soatto},
journal= {arXiv preprint arXiv:2605.08301},
year = {2026}
}