[Research Frontier] Attention to Mamba: A Principled Recipe for Converting Transformers to SSMs

TL;DR

Researchers have finally cracked the code on how to turn an existing Transformer into a Mamba-like State Space Model (SSM) without losing significant accuracy. By using a two-stage distillation process—first linearizing attention via the "Hedgehog" technique, then mapping those weights into a Mamba structure—the team successfully distilled a Pythia-1B model into a linear-time student that retains 98% of the teacher's performance using only 10 billion tokens.

The "Quadratic Wall" and the Distillation Dilemma

Transformers are the kings of NLP, but they are prisoners of their own complexity. As sequence lengths grow, the cost of the Softmax Attention mechanism grows quadratically ($O(L^2)$). SSMs like Mamba offer a way out with linear scaling ($O(L)$), but they are notoriously hard to train from scratch.

The industry has spent millions of dollars pretraining Transformers. Naturally, we want to "recycle" that knowledge. However, previous attempts at Cross-Architecture Distillation (Transformer $\to$ SSM) usually resulted in "brain-dead" models unless some attention blocks were kept in a hybrid form. HedgeMamba changes this.

Methodology: The Two-Stage Bridge

The central insight of this work is that forcing a direct leap from Transformers to Mamba is too extreme. Instead, they build a bridge.

Stage 1: The Hedgehog Alignment

Before touching Mamba, the authors replace Softmax Attention with Linear Attention. They use the Hedgehog approach, which uses an MLP-based feature map to approximate the "spikiness" of the softmax kernel. This stage aligns the behavior of the attention mechanisms.

Stage 2: Linear Attention $\to$ HedgeMamba

The authors then mathematically map the Linear Attention parameters into the Mamba recurrence variables ($\Lambda ,B,C$).

• Initialization: They initialize Mamba to act as an identity to the linearized attention learned in Stage 1.
• Augmentation: They unlock Mamba-specific features like the 1D Convolution and the SiLU Gating branch to provide the student with more expressivity than the original teacher.

Fig 1: The two-stage recipe. Green represents Softmax Attention, Blue is the Hedgehog Linearization, and Yellow is the final Mamba integration.

Experimental Battleground

The team used Pythia-1B as the teacher and OpenWebText for distillation. The results were striking:

The "Aha!" moment comes from the Ablation Studies. The authors found that the Gating Branch (SiLU) was the most critical Mamba component for closing the performance gap, consistent with modern "Gated Attention" theories.

Fig 2: Validation Perplexity curves during training. Stage 1 (the initial dot) sets a foundation that prevents the model from stalling during Mamba finetuning.

Critical Insights & Takeaways

1. Stage 1 is Mandatory: Attempting to finetune a Mamba model directly into a Transformer (the 0/100 split) results in poor perplexity. You need the "Softmax mimicry" phase to align the representations.
2. Budget Efficiency: The fact that this works with only 10B tokens—a fraction of the 300B+ tokens usually needed for pretraining—is a massive win for "green AI."
3. Limitations: The study primarily focuses on the Pythia architecture. Whether this recipe holds for more complex "Llama-style" models (using GQA or SwiGLU) remains an open research question.

Conclusion

HedgeMamba represents a shift in how we think about model evolution. Instead of discarding old architectures, we can "transplant" their knowledge into more efficient skeletons. For engineers looking to deploy long-context LLMs without the quadratic memory overhead, this two-stage distillation is now the state-of-the-art starting point.