Share your thoughts, 1 month free Claude Pro on usSee more
WorkDL logo mark

Density Field State Space Models: 1-Bit Distillation, Efficient Inference, and Knowledge Organization in Mamba-2

About

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.

Chirag Shinde• 2026

Related benchmarks

TaskDatasetResultRank
Commonsense ReasoningWinoGrande
Accuracy54.7
1581
GPU Inference ThroughputA100-SXM4-40GB
Throughput (tok/s)5.08e+3
8
Large Language Model Inference PerformanceIntel Ice Lake / AMD Zen4 AVX-512 VNNI
Memory (MB)567
4
Commonsense ReasoningPIQA
Accuracy67.1
4
Commonsense ReasoningHellaSwag
Accuracy41.4
4
Question AnsweringBoolQ
Accuracy60.8
4
Question AnsweringARC Easy
Accuracy50.2
3
Training EfficiencyDistillation (train)
Total Tokens Processed32
3
Showing 8 of 8 rows

Other info

Follow for update