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

Abstract: 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.

• The paper presents a grouped per-block ternary QAT method that replaces 85.7% of parameters, achieving near state-of-the-art performance with only a 24% relative increase in perplexity.
• The methodology prevents catastrophic quality collapse by using non-learnable group scaling, which stabilizes weight sparsity and enables rapid convergence.
• The experiments demonstrate a drastic reduction in training cost (102M tokens and 4 GPU-hours) while maintaining competitive accuracy across multiple evaluation benchmarks.

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

Introduction and Motivation

This paper addresses a significant open question in efficient LLM deployment on resource-constrained hardware: can large state space models (SSMs) such as Mamba-2 be quantized to the ternary regime ($\{-1, 0, +1\}$, i.e., W1.58A16 format) without catastrophic loss of accuracy, and specifically, can this be achieved using quantization-aware training (QAT) from a pretrained checkpoint rather than expensive from-scratch training? Prior work on SSM quantization (e.g., Slender-Mamba, Bi-Mamba, and Quamba) either incurs massive training costs (order 100B+ tokens, 5,000+ GPU-hours) or is limited by poor post-training quantization (PTQ) quality below 4 bits.

The core contribution is a grouped per-block ternary QAT method, evaluated on Mamba-2 1.3B, that approaches the zero-shot accuracy of state-of-the-art full-precision and binary models with dramatically lower marginal compute. The results directly address deployability of linear-time SSMs and the intersection of SSM architectures with low-bit quantization regimes that have proven effective for Transformers.

Grouped Ternary Quantization-Aware Training

The quantization procedure replaces 85.7\% of Mamba-2 1.3B parameters (the in_proj and out_proj components) with TernaryLinear modules operating at 1.58 bits per weight. Weights are partitioned into groups of $g=128$; for each group, the scale $s_g$ is computed as the non-learnable mean absolute value of the group's weights on each forward pass, and ternarization proceeds via round-and-clip to $\{-1, 0, +1\}$, with dequantization using the computed group scale. Gradients flow via the STE.

Empirically, this method entirely circumvents the catastrophic quality collapse observed with ternary PTQ, which produces random outputs (PPL $\sim$ 13M) on SSMs. With QAT, the compressed model rapidly converges and maintains sparsity at a stable equilibrium.

Zero-Ratio Collapse in Learnable Scales

A novel failure mode termed "zero-ratio collapse" is identified during QAT with learnable scale parameters. Optimization pressure drives scales to inflate, resulting in an overwhelming majority ($\sim$90\%) of weights collapsing to zero, causing sharp loss spikes and performance degradation. This is prevented by recomputing the scale as non-learnable absmean per group, which induces a self-regulating negative feedback mechanism that stabilizes sparsity at $\sim$26\% and preserves model capacity.

Figure 1: Learnable scale (exp002) collapses to 90\% zeros with a loss spike at 10k steps, while fixed absmean (exp006) self-regulates at $\sim$26\%.

Main Experimental Results and Ablations

On C4 domain-matched evaluation, the ternary QAT model achieves 21.67 PPL, a 24\% relative degradation (+4.2 PPL) versus the full-precision baseline (17.47 PPL), and a zero-shot 7-task average accuracy of 48.1\% (within $\pm$0.9pp of Bi-Mamba's 48.4\%). Crucially, this is accomplished with just 102M QAT tokens (4 GPU-hours, single H100), compared to 105B–150B tokens (5700+ GPU-hours) required by prior from-scratch ternary or binary SSMs.

Scaling investigations show that PPL on C4 improves monotonically up to at least 307M tokens, with no indication of a quality plateau, and OLS/log-linear regression suggests continued improvements are feasible under extended QAT. Train/eval variance is negligible (C4 PPL std = 0.01), and group-size ablations demonstrate near invariance of PPL to choices in $g$, with peak agreement at $g=128$0.

Figure 2: C4 PPL vs.\ training tokens (log scale). Monotonic decrease with no plateau observed; trajectory is consistent with continued improvement under additional training.

Figure 3: Group-size ablation. PPL is flat across $g=128$1; top-1 agreement peaks at $g=128$2.

Knowledge distillation (KD) is crucial: a balanced KD/CE objective ($g=128$3) outperforms CE-only and KD-only baselines by large margins, confirming the teacher signal is essential for student QAT convergence.

Robustness to Correction and Negative Results

Attempts to apply post-hoc error correction methods (Kalman gain amplification and James-Stein shrinkage), which are effective for quantized Transformers, fail to improve SSM performance in the ternary regime. In all experiments these corrections strictly worsened PPL, confirming that SSM recurrence fundamentally accumulates and entangles quantization errors such that output-linear corrections are counterproductive.

Figure 4: Both Kalman amplification and James-Stein shrinkage increase PPL monotonically from baseline. The optimum is zero correction.

Additionally, analyses of the total variation to $g=128$4 ratio (for assessing sigma-delta noise shaping applicability) reveal that SSM activations are more decorrelated than i.i.d. Gaussian signals, precluding the effectiveness of temporal noise-shaping, in stark contrast to 1-bit audio and related domains.

Figure 5: TV/L1 ratio exceeds the Gaussian baseline ($g=128$5) at every layer, ruling out sigma-delta noise-shaping for SSMs.

Layer-Selective Quantization and Compression Tradeoffs

Layer-selective quantization offers an effective compression-quality trade-off by preserving the first and last $g=128$6 blocks in FP16, with nearly linear PPL recovery proportional to the number of protected layers. No single layer dominates quantization error, and per-layer quantization error is extremely uniform.

Figure 6: Layer-selective quantization Pareto frontier. PPL recovery is approximately linear in the fraction of protected layers.

Storage compression of 3.61$g=128$7 is reported (744 MB vs 2,687 MB for the quantized projections). Actual inference acceleration and deployment footprint reduction require a custom packed ternary GEMM kernel for Mamba projection dimensions—current simulated QAT implementations consume more memory due to latent intermediates. Prior art in ternary/binary kernels for Transformers (BitBLAS, TernaryLLM) suggests similar speedups are possible for SSMs once specialized kernels are developed.

Theoretical and Practical Implications

This study demonstrates that PTQ below 4 bits is fundamentally inadequate for SSMs due to catastrophic error accumulation through the recurrence; QAT is not merely beneficial but necessary. The method presented greatly reduces the marginal computational cost and token budget required to deploy ternary SSMs in resource-constrained contexts.

Theoretically, the findings delineate key architectural distinctions between attention-based and state-space-based LLMs in quantization regimes. Practically, the method unlocks direct, data-efficient compressibility for SSMs, enabling progress toward deployable, high-throughput, low-power inference on edge hardware—contingent on kernel engineering.

Limitations and Future Work

The conclusions are currently limited to a single SSM instantiation (Mamba-2 1.3B, English, C4). Generalizing the approach to other SSM variants, larger scales, or multilingual contexts is unaddressed. The empirical evidence suggests the remaining accuracy gap may be addressable through longer QAT schedules and further KD tuning, but an asymptotic bound imposed by ternary constraints cannot be ruled out.

The absence of a custom inference kernel restricts the method's current utility to storage compression rather than true latency or throughput improvements. Furthermore, long-context generalization (i.e., sequence lengths well beyond 1024 tokens) remains unevaluated.

Conclusion

This work establishes grouped per-block ternary QAT as an effective and highly data-efficient method for compressing large pretrained SSMs to 1.58-bit weights, circumventing the unsustainable computational cost of from-scratch training. It exposes the architectural reasons for PTQ and post-hoc correction failure in SSMs, underscores the importance of non-learnable group scaling to prevent sparsity collapse, and highlights the need for further engineering and extended QAT to close the accuracy gap. The pathway to practical, low-bit SSMs for edge deployment is thus clarified, but ultimate parity with FP16 models and full realization of resource advantages will depend on future advances in quantization-aware training schedules, scale, and kernel implementations.