Attention to Mamba: A Recipe for Cross-Architecture Distillation

🔑 Core Idea

The paper proposes a new method to transfer knowledge from Transformer models (with Attention) into Mamba models (a type of State Space Model, SSM)—which are faster and more memory-efficient.

❗ Problem

  • Transformers (with softmax attention) are powerful but expensive (quadratic cost in sequence length).
  • Mamba/SSMs are efficient (linear cost) but harder to train and usually less accurate.
  • Directly distilling (copying knowledge) from Transformers → Mamba fails to preserve performance.

💡 Key Solution: Two-Stage Distillation

Instead of direct transfer, the authors introduce a bridge strategy:

Stage 1: Transformer → Linear Attention

  • Replace softmax attention with a learned linear approximation (via a method called Hedgehog).
  • Uses a learned feature map to mimic attention behavior.

Stage 2: Linear Attention → Mamba

  • Convert the linear attention into a Mamba-compatible initialization.
  • Then fine-tune the full model.

👉 This staged approach aligns architectures gradually, making knowledge transfer effective.

⚙️ Key Insight

  • Linear Attention and Mamba (SSMs) are mathematically related.
  • Using Linear Attention as an intermediate step avoids forcing incompatible architectures to match directly.

📊 Results

  • The distilled Mamba model almost matches the original Transformer:

    • Perplexity: 14.11 vs 13.86 (teacher)
  • Works across multiple downstream tasks with strong performance.
  • Outperforms simpler baselines and naive distillation.

🧪 Additional Findings

  • Adding Mamba components (like gating) improves results.
  • Most training budget should go to Stage 2 (fine-tuning).
  • More distillation data → better performance.

🧠 Contributions

  1. A two-stage distillation method for Transformer → Mamba.
  2. Demonstrates that proper initialization + architectural alignment is critical.
  3. Shows you can get efficient models without retraining from scratch.

🧾 Bottom Line

You can convert powerful but expensive Transformer models into fast, efficient Mamba models—without losing much performance—by:

bridging them through Linear Attention first, instead of distilling directly.