将Mamba-2压缩至1比特,实现超快推理与知识结构可解释性。
Density Field State Space Models: 1-Bit Distillation, Efficient Inference, and Knowledge Organization in Mamba-2

- 用1比特骨架+8位低秩修正压缩模型,大幅降低存储与计算开销。
- 模型仅278MB,推理速度提升21.4倍,性能接近1.58比特原生训练模型。
- 揭示模型内部三阶段知识处理机制,展现结构化表征能力。
我们提出密度场状态空间模型(DF-SSM),将SSM压缩至1比特骨架并辅以int8低秩修正。应用于Mamba-2 1.3B模型时,得到278MB的小模型(比2.7GB的FP16教师模型小9.7倍),在GPU上单批推理速度达参考实现的21.4倍,同时下游任务性能仅比从头训练的1.58比特BitMamba-2低2-4个百分点。蒸馏过程仅需3200万词元和单个A100 GPU上6小时,但依赖预训练的FP16教师模型。我们开发了优化推理流水线,结合cuBLAS INT8 Tensor Core进行骨架矩阵乘法、自定义CUDA核处理状态式SSM与卷积操作,并采用AVX-512 CPU后端实现跨平台高效部署。除压缩外,我们系统分析了模型内部知识组织,发现三个明确处理阶段:意图分类(层0-3,在无词汇对齐的抽象空间中运行)、知识检索(层25-35,事实关联集中于5层窗口内)、输出格式化(层36-47,类别结构消失)。通过对19类共445个事实提示的分析,发现早期层分类由句法模板驱动而非语义,且模型虽事实记忆弱,但具备良好组织的知识表征——表明表征结构可能先于事实强度出现。
原文摘要 · Abstract (English)
We present Density Field State Space Models (DF-SSM), a framework for compressing SSMs to a 1-bit scaffold with int8 low-rank correction. Applied to Mamba-2 1.3B, we achieve a 278 MB model (9.7x smaller than the 2.7 GB FP16 teacher) that runs at 21.4x faster inference on GPU (batch=1, relative to the mamba-ssm reference implementation) while maintaining downstream task performance within 2-4 percentage points of BitMamba-2, a 1.58-bit model trained from scratch on 150B tokens. The distillation itself requires only 32M tokens and 6 hours on a single A100 GPU, though it presupposes a pretrained FP16 teacher. We develop an optimized inference pipeline combining cuBLAS INT8 tensor cores for the scaffold matmul, custom CUDA kernels for stateful SSM and convolution operations, and an AVX-512 CPU backend for efficient deployment on both GPU and CPU. Beyond compression, we investigate the internal knowledge organization of the resulting model, discovering three distinct processing phases: intent classification (layers 0-3, operating in an abstract space with no vocabulary alignment), knowledge retrieval (layers 25-35, where factual associations localize to a 5-layer window), and output formatting (layers 36-47, where category structure dissolves). Through systematic analysis of 445 factual prompts across 19 categories, we find that early-layer classification is syntactic (driven by template structure) rather than semantic, and that the model exhibits well-organized knowledge representations despite weak factual recall--suggesting that representational structure may precede factual strength.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。