Ternary Mamba: Grouped Quantization-Aware Training of W1.58A16 State Space Models
作者: Ramprasath Ganesaraja, Sahil Dilip Panse, Swathika N
分类: cs.LG, cs.AI
发布日期: 2026-06-16
💡 一句话要点
提出Ternary Mamba以解决SSM边缘部署的内存限制问题
🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture)
关键词: 状态空间模型 量化感知训练 知识蒸馏 边缘计算 模型压缩
📋 核心要点
- 现有的状态空间模型在边缘设备上的应用受到内存占用的限制，影响了其实际部署。
- 论文提出了一种基于预训练检查点的分组量化感知训练方法，结合知识蒸馏以提高模型压缩效率。
- 实验结果表明，该方法在仅使用102M标记的情况下，实现了48.1%的零-shot准确率，显著降低了内存需求。
📝 摘要（中文）
状态空间模型（SSMs）如Mamba-2提供线性时间推理，但其内存占用限制了边缘部署。先前的三元SSM工作（Slender-Mamba）在150B标记上从头训练；我们展示了预训练检查点足以减少边际标记预算1000倍。通过使用分组量化感知训练（QAT）和从冻结的FP16教师模型进行知识蒸馏，我们将Mamba-2 1.3B压缩至3.61倍（从2687MB降至744MB），并在仅102M标记（4 GPU小时，单个H100）中实现48.1%的零-shot准确率，接近Bi-Mamba的48.4%。这种基于预训练的QAT设置揭示了零比率崩溃，这是一种由可学习量化尺度引起的新型不稳定性，且在从头训练中不会出现。我们进一步表明，针对变换器的后期修正策略在SSMs中因误差累积而失效。这些结果表明，三元SSMs不需要昂贵的从头训练：基于预训练检查点的QAT与KD是数据高效的替代方案。
🔬 方法详解
问题定义：本论文旨在解决状态空间模型（SSMs）在边缘设备上部署时的内存占用问题。现有方法如Slender-Mamba需要从头训练，消耗大量计算资源和数据，限制了其应用。
核心思路：论文提出通过分组量化感知训练（QAT）结合知识蒸馏，从预训练的检查点出发，显著降低模型的内存占用，同时保持较高的准确率。这种方法避免了从头训练的高成本，提供了一种更为高效的训练策略。
技术框架：整体架构包括预训练的教师模型（FP16），通过知识蒸馏将知识传递给学生模型，同时在训练过程中应用分组量化感知策略。主要模块包括量化层、蒸馏损失计算和优化器。
关键创新：最重要的技术创新在于提出了基于预训练检查点的QAT方法，解决了从头训练带来的高成本问题，并引入了零比率崩溃这一新现象，揭示了可学习量化尺度的潜在不稳定性。
关键设计：在训练过程中，采用了特定的损失函数来平衡蒸馏损失与量化损失，同时设置了适当的量化尺度以避免零比率崩溃的发生。模型结构上，保持了Mamba-2的核心设计，但在量化层进行了优化。
🖼️ 关键图片
📊 实验亮点
实验结果显示，Ternary Mamba在仅使用102M标记的情况下，实现了48.1%的零-shot准确率，接近Bi-Mamba的48.4%。同时，模型的内存占用从2687MB减少至744MB，压缩比达到3.61倍，展现了显著的性能提升和资源效率。
🎯 应用场景
该研究的潜在应用领域包括边缘计算、移动设备上的智能应用以及资源受限环境中的机器学习模型部署。通过降低内存需求，Ternary Mamba能够使得更复杂的模型在实际应用中得以实现，推动智能设备的普及与应用。未来，该方法可能在实时推理和大规模数据处理等场景中发挥重要作用。
📄 摘要（原文）
State Space Models (SSMs) such as Mamba-2 offer linear-time inference but their memory footprint limits edge deployment. Prior ternary SSM work (Slender-Mamba) trains from scratch on 150B tokens; we show a pretrained checkpoint suffices, reducing the marginal token budget by 1,000x. Using grouped quantization-aware training (QAT) with knowledge distillation from a frozen FP16 teacher, we compress Mamba-2 1.3B to 3.61x (2,687 to 744 MB) and achieve 48.1% zero-shot accuracy (7-task average) in just 102M tokens (4 GPU-hours, single H100) -- approaching Bi-Mamba's 48.4% (within +/-0.9pp CI). This QAT-from-pretrained setting reveals zero-ratio collapse, a novel instability caused by learnable quantization scales that does not arise in from-scratch training. We further show that post-hoc correction strategies effective for Transformers fail for SSMs due to error accumulation through the recurrence. These results demonstrate that ternary SSMs do not require expensive from-scratch training: QAT from pretrained checkpoints with KD is a data-efficient alternative.