三元 Mamba：W1.58A16 状态空间模型的分组量化感知训练

Mamba-2 等状态空间模型 (SSM) 提供线性时间推理，但其内存占用限制了边缘部署。之前的三元 SSM 工作（Slender-Mamba）在 150B 词元上从头开始训练；我们展示了一个预训练的检查点就足够了，将边际 词元预算 减少了 1,000 倍。使用来自冷冻 FP16 教师的 知识蒸馏 的分组量化感知训练 (QAT)，我们将 Mamba-2 1.3B 压缩到 3.61x（2,687 到 744 MB），并在仅 1.02M 词元（4 个 GPU 小时，单个 H100）内实现 48.1% 的零射击精度（7 个任务平均值）——接近 Bi-Mamba 48.4%（+/-0.9pp CI 以内）。这种基于预训练的 QAT 设置揭示了零比率崩溃，这是一种由可学习的量化尺度引起的新型不稳定性，而这种不稳定性在从头开始的训练中不会出现。我们进一步表明，对于 Transformer 有效的事后纠正策略对于 SSM 来说是失败的，因为错误会通过递归累积。这些结果表明，三元 SSM 不需要昂贵的从头开始训练：使用 KD 进行预训练检查点的 QAT 是一种数据高效的替代方案。