Apriel-H1: 通过Mamba混合架构蒸馏高效推理模型
ServiceNow AI 开发了 Apriel-H1，这是一系列 15B 推理模型，将全注意力教师模型转换为 Mamba 混合架构。这种方法在推理质量损失最小的情况下实现了最高 2.1 倍的吞吐量提升，表明可以通过有针对性的蒸馏将效率回溯到现有的强大模型中。
核心洞察：能力匹配蒸馏
蒸馏推理模型需要使用与所保留特定能力匹配的数据，而不是通用预训练数据。最初尝试使用预训练数据或预训练数据与 SFT 数据的混合进行蒸馏时，导致推理质量出现显著下降。
成功的方法利用了教师 SFT 数据集中的高质量推理轨迹，包括：
- 多步数学证明
- 具有清晰逻辑依赖关系的编码任务
- 具有详细解释链的科学分析
这是因为将注意力机制（如检索和归纳头）替换为 Mamba 的线性递归会破坏用于推理的计算路径。高质量的推理轨迹为混合模型提供了发现相同结果的新路径所需的显式结构。此外，团队在蒸馏过程中使用了 reverse KL divergence（温度 1），因为其寻求模式的行为鼓励学生模型承诺教师的高置信度、结构化预测。
Apriel-H1 性能与基准
旗舰模型 Apriel-H1-15b-Thinker-SFT 相比全注意力教师模型实现了 2.1 倍吞吐量。在推理基准上的影响如下：
| 基准 | 教师得分 | Apriel-H1-15b-Thinker-SFT 得分 |
|---|---|---|
| MATH500 | 0.90 $ | |
| ightarrow$ 0.92 | 改进 | |
| MTBench | 8.30 $ | |
| ightarrow$ 8.58 | 改进 | |
| GSM8k | 0.97 $ | |
| ightarrow$ 0.95 | 轻微回退 | |
| GPQA | 0.59 $ | |
| ightarrow$ 0.55 | 轻微回退 | |
| AIME24 | 0.70 $ | |
| ightarrow$ 0.65 | 轻微回退 |
在 Apriel-H1 系列中，检查点的 Mamba 层数范围为 25 到 40（共 50 层）。虽然 H-30 变体提供了最佳平衡，但 H-40 变体在延迟关键工作负载下最大化吞吐量（最高可达 3.4 倍）。
分阶段蒸馏过程
将全注意力模型转换为混合模型不能在一步完成。ServiceNow AI 采用了三阶段过程：
第一阶段：初始层替换
团队在 MMLU 上使用 Leave-One-Out (LOO) 分析，识别出 25 层最不重要的层。这些层被 Mamba-in-Llama (MIL) 初始化的混合器替换，并进行端到端蒸馏。
第二阶段：渐进式转换
由于 LOO 分析在超过 25 层时变得不可靠，团队实施了 MIL-Mamba-Replacement (MMR)。这种动态启发式方法初始化一个 Mamba 混合器，运行 100 训练步骤，并记录蒸馏损失。收敛到较低损失的层被优先考虑进行替换。此过程逐步递增进行（25 $ ightarrow$ 27 $ ightarrow$ 30 $ ightarrow$ 34 $ ightarrow$ 37 $ ightarrow$ 40 层）。
第三阶段：最终 SFT
达到目标 Mamba 层数后，进行最终的 SFT 训练，直至推理性能稳定。最终的 Apriel-H1-15b-Thinker-SFT 模型在 55.9B 蒸馏标记和 20.9B SFT 标记后产生。
实现与可重现性
Apriel-H1 使用 Fast-LLM 构建，这是一个开源（Apache 2.0）的训练框架，将注意力和 Mamba 视为可混合的模块接口。这使得可以通过配置中的 pattern 字段指定层顺序，从而自由交换块类型。
生产部署
Apriel-H1 已与 Hugging Face Transformers 和 vLLM 集成。vLLM 集成利用 Mamba 缓存操作进行连续批处理、前缀缓存和分块预填充。作者指出，虽然吞吐量提升显著，但由于工具尚在成熟阶段，目前部署混合模型需要自定义代码和仔细的数值验证。
技术规格
- Mamba 超参数： 状态大小 16，DT 排名 16，内部维度 4096。
- 架构： Mamba-1 和注意力层的混合。
- 训练数据： 旗舰模型的总训练标记为 76.8B。