Attention to Mamba: A Recipe for Cross-Architecture Distillation

1. Problem Statement

The paper addresses the challenge of bridging the performance gap between Transformers (which rely on quadratic-complexity Softmax Attention) and State Space Models (SSMs) like Mamba (which offer linear-complexity inference and reduced memory usage).

• The Dilemma: While Mamba offers superior efficiency, it often underperforms compared to pre-trained Transformers on downstream tasks. Conversely, training SSMs from scratch is computationally expensive and requires massive datasets.
• The Distillation Gap: Direct knowledge distillation from a Transformer teacher to a Mamba student has historically failed. Naïve attempts result in significant performance degradation (e.g., perplexity >100) because the two architectures operate on fundamentally different computational paradigms (quadratic attention vs. linear recurrence). Previous hybrid solutions often rely on complex, multi-stage pipelines or retaining some Attention blocks, which defeats the purpose of pure efficiency.

2. Methodology: The Two-Stage Distillation Recipe

The authors propose a principled, two-stage distillation approach that leverages mathematical connections between Attention and SSMs to facilitate knowledge transfer without training from scratch.

Stage 1: Softmax Attention $\to$ Linear Attention (The "Hedgehog" Bridge)

• Goal: Approximate the non-linear Softmax Attention mechanism with a learnable Linear Attention variant.
• Mechanism: The paper utilizes the Hedgehog approach (Zhang et al., 2024). Instead of dropping the Softmax (which loses critical features like "spikiness"), they learn a feature map $\phi(x)$ via a single-layer MLP to approximate the exponential kernel using Mercer's theorem:
  $$\exp(x^T x') \approx \phi(x)^T \phi(x')$$
• Process: The Transformer's Softmax Attention is replaced with this learned Linear Attention. The model is distilled from the teacher Transformer using cosine embedding matching loss on the outputs of each layer.
• Outcome: A "Linearized Transformer" that retains the teacher's performance but operates with linear complexity.

Stage 2: Linear Attention $\to$ Mamba (The "HedgeMamba" Initialization)

• Goal: Translate the Linear Attention weights into an initialization for the Mamba architecture and fine-tune for expressivity.
• Theoretical Connection: The authors map the parameters of Linear Attention ($\hat{Q}, \hat{K}, \hat{V}$) to the SSM parameters ($C, B, X$). Specifically:
  • $B(X) \leftarrow \phi_{MLP}(K(X))$
  • $C(X) \leftarrow \phi_{MLP}(Q(X))$
  • $X \leftarrow V(X)$
  • The state decay matrix $\Lambda$ is initialized to the Identity matrix.
• Architecture Adaptation (HedgeMamba):
  • The Mamba block is modified to include a Value projection ($V$) before the SSM mixer (which standard Mamba lacks).
  • Normalization: An unnormalized attention score issue is resolved by expanding the SSM state to compute a normalization term (similar to the denominator in Softmax) in a single pass.
  • Fine-tuning: The entire model (except embeddings) is fine-tuned using Cross-Entropy loss on the ground truth data. Crucially, the Gate branch and Convolution layers of Mamba are unlocked to boost expressivity.

3. Key Contributions

1. Principled Initialization: Demonstrates that equipping Mamba with a "Hedgehog" initialization (derived from Linear Attention) solves the cross-architecture alignment problem, avoiding the failure of naïve distillation.
2. Two-Stage Recipe: Proposes a theoretically grounded pipeline:
   • Stage 1: Align Attention weights via Linear Attention (Hedgehog).
   • Stage 2: Transfer to Mamba and fine-tune with Mamba-specific components (Gates, Convolutions).
3. HedgeMamba Architecture: Introduces a hybrid layer that combines the Hedgehog feature map with Mamba's selective scan, including necessary modifications (Value projection, normalization) to ensure functional equivalence at initialization.
4. Comprehensive Evaluation: Provides extensive ablation studies, scaling analyses, and sensitivity tests on token budget allocation.

4. Results

The experiments were conducted using Pythia-1B as the teacher and OpenWebText (10B tokens) for distillation.

• Performance Preservation:

  • Teacher (Pythia-1B): Perplexity (PPL) = 13.86.
  • Naïve Distillation: PPL > 100 (Failed).
  • Hedgehog Baseline (Linear Attention only): PPL = 14.89.
  • HedgeMamba (Ours): PPL = 14.11.
  • Result: The distilled Mamba model achieves performance nearly identical to the Transformer teacher (within 0.25 PPL) while maintaining linear inference complexity.

• Downstream Tasks: The model outperforms the Hedgehog baseline and approaches the teacher on tasks like ARC-Challenge, BoolQ, and WinoGrande.

• Ablation Insights:

  • Token Budget: The optimal split is 10% Stage 1 / 90% Stage 2. Allocating all tokens to Stage 2 (fine-tuning only) results in poor performance, proving the necessity of the initialization step.
  • Components: Adding the Gate branch provided the largest performance boost, followed by the SSM mixer parameters and Convolution.
  • Scaling: Performance improves with model size (160M $\to$ 1B) and token budget (1B $\to$ 10B), with no saturation observed at 10B tokens.

5. Significance

• Efficiency without Sacrifice: This work proves that it is possible to migrate the knowledge of massive, pre-trained Transformers into efficient, linear-complexity SSMs without retraining from scratch or sacrificing significant accuracy.
• Theoretical Bridge: It establishes a concrete mathematical link between Linear Attention and SSMs, showing that Mamba can be viewed as a specialized, expressive form of Linear Attention when properly initialized.
• Practical Recipe: The proposed "HedgeMamba" recipe offers a reproducible, computationally efficient path for the community to adopt SSMs, leveraging the vast ecosystem of existing Transformer models.
• Resource Efficiency: The distillation process required only ~10B tokens (approx. 2.7% of the teacher's training budget), demonstrating that high-quality distillation is feasible with limited resources compared to pre-training.

In summary, the paper provides a robust "recipe" for converting Attention-based models into Mamba-based models, solving the initialization mismatch that previously hindered cross-architecture distillation.