SpikingMamba: Towards Energy-Efficient Large Language Models via Knowledge Distillation from Mamba

TL;DR

This work addresses the energy bottleneck of large language models by distilling an energy-efficient spiking LLM, SpikingMamba, from pretrained Mamba2. It introduces TI-LIF, a ternary signed spiking neuron, and a Smoothed Gradient Compensation path to preserve semantic fidelity during training, enabling spike-driven inference with minimal accuracy loss. A single-stage distillation from Mamba2 plus reinforcement learning (DPO/KTO) transfers zero-shot abilities and further improves performance, achieving substantial energy savings (up to approximately 4.76×) with only a small accuracy gap, and additional boosts from RL. The results demonstrate practical deployment potential for edge devices, showing that SpikingMamba can maintain competitive reasoning ability while dramatically reducing energy consumption and avoiding full retraining.

Abstract

Large Language Models (LLMs) have achieved remarkable performance across tasks but remain energy-intensive due to dense matrix operations. Spiking neural networks (SNNs) improve energy efficiency by replacing dense matrix multiplications with sparse accumulations. Their sparse spike activity enables efficient LLMs deployment on edge devices. However, prior SNN-based LLMs often sacrifice performance for efficiency, and recovering accuracy typically requires full pretraining, which is costly and impractical. To address this, we propose SpikingMamba, an energy-efficient SNN-based LLMs distilled from Mamba that improves energy efficiency with minimal accuracy sacrifice. SpikingMamba integrates two key components: (a) TI-LIF, a ternary-integer spiking neuron that preserves semantic polarity through signed multi-level spike representations. (b) A training-exclusive Smoothed Gradient Compensation (SGC) path mitigating quantization loss while preserving spike-driven efficiency. We employ a single-stage distillation strategy to transfer the zero-shot ability of pretrained Mamba and further enhance it via reinforcement learning (RL). Experiments show that SpikingMamba-1.3B achieves a 4.76$\times$ energy benefit, with only a 4.78\% zero-shot accuracy gap compared to the original Mamba, and achieves a further 2.55\% accuracy improvement after RL.

Figures (7)

• Figure 1: (a) Overview of the training architecture. (b) Illustration of the SpikingMamba block. (c) Illustration of the Smoothed Gradient Compensation path.
• Figure 2: Energy efficiency ratio ($E_\text{A} / E_\text{S}$) of SpikingMamba under various configurations ($E_\text{A}$: energy of Mamba2, $E_\text{S}$: energy of SpikingMamba). Colors denote model size, striped bars indicate SGC path usage, and marker shapes indicate neuron types. Detailed data and fire rate are provided in Appendix \ref{['apx.energy.computation.']}.
• Figure 3: Activation distributions in Mamba2: (a) input projection $\boldsymbol{u}_t$ (Eq \ref{['eq.preliminary.input.projection']}), (b) output projection $\boldsymbol{y}_t$ (Eq \ref{['eq.preliminary.output.projection']})
• Figure 4: Activation statistics across channels and tokens in Mamba2: (a) input projection $\boldsymbol{u}_t$ (Eq \ref{['eq.preliminary.input.projection']}), (b) output projection $\boldsymbol{y}_t$ (Eq \ref{['eq.preliminary.output.projection']}).
• Figure 5: Activation distribution comparison for LIF-based models.