Ternary Mamba:基于预训练检查点的分组QAT实现W1.58A16量化

Ternary Mamba: Grouped Quantization-Aware Training of W1.58A16 State Space Models

精选理由

把Mamba-2压缩到1比特权重不用花大钱从头训练,用预训练模型加少量微调就行,效果只差0.3%。

AI 摘要

Mamba-2 1.3B参数模型通过分组量化感知训练(QAT)从FP16教师蒸馏,仅消耗4 GPU小时(单H100)和102M tokens,将内存从2,687 MB压缩至744 MB(3.61倍)。零样本七任务平均准确率达48.1%,接近Bi-Mamba的48.4%(±0.9pp置信区间)。该方法无需从头训练150B tokens,使用预训练检查点即可。研究还发现零比例坍缩现象,即可学习量化尺度导致的不稳定性,这在从头训练中不会出现。后处理校正策略对SSM无效,因循环结构导致误差累积。

AI 翻译 · 中文

Mamba-2 1.3B参数模型通过分组量化感知训练(QAT)从FP16教师蒸馏,仅消耗4 GPU小时(单H100)和102M tokens,将内存从2,687 MB压缩至744 MB(3.61倍)。零样本七任务平均准确率达48.1%,接近Bi-Mamba的48.4%(±0.9pp置信区间)。该方法无需从头训练150B tokens,使用预训练检查点即可。研究还发现零比例坍缩现象,即可学习量化尺度导致的不稳定性,这在从头训练中不会出现。后处理校正策略对SSM无效,因循环结构导致误差累积。

arXiv cs.LGState 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 che