密度场状态空间模型：Mamba-2 中的 1 位 蒸馏、高效推理和知识组织

我们提出了密度场状态空间模型 (DF-SSM)，这是一种将 SSM 压缩为具有 int8 低秩校正的 1 位支架的框架。应用于 Mamba-2 1.3B，我们实现了 278 MB 模型（比 2.7 GB FP16 老师小 9.7 倍），在 GPU 上运行推理速度提高了 21.4 倍（批次 = 1，相对于 mamba-ssm 参考实现），同时将下游任务性能保持在 BitMamba-2 的 2-4 个百分点内，BitMamba-2 是一个在 150B 词元上从头开始训练的 1.58 位模型。 蒸馏 本身仅需要 32M 词元和单个 A100 GPU 上的 6 小时，尽管它预设了预训练的 FP16 教师。我们开发了一个优化的推理管道，结合了用于脚手架 matmul 的 cuBLAS INT8 张量核心、用于状态 SSM 和卷积操作的自定义 CUDA 内核，以及用于在 GPU 和 CPU 上高效部署的 AVX-512 CPU 后端。除了压缩之外，我们还研究了结果模型的内部知识组织，发现了三个不同的处理阶段：意图分类（第 0-3 层，在没有词汇对齐的抽象空间中操作）、知识检索（第 25-35 层，其中事实关联本地化到 5 层窗口）和输出格式化（第 36-47 层，其中类别结构消失）。通过对 19 个类别的 445 个事实提示进行系统分析，我们发现早期层分类是句法的（由模板结构驱动）而不是语义的，并且尽管事实回忆较弱，但该模型仍表现出组织良好的知识表示，这表明表示结构可能先于事实强度。