State Space Models (SSMs) such as Mamba have become a popular alternative to Transformer models, due to their reduced memory consumption and higher throughput at generation compared to their Attention-based counterparts. On the other hand, the community has built up a considerable body of knowledge on how to train Transformers, and many pretrained Transformer models are readily available. To facilitate the adoption of SSMs while leveraging existing pretrained Transformers, we aim to identify an effective recipe to distill an Attention-based model into a Mamba-like architecture. In prior work on cross-architecture distillation, however, it has been shown that a naïve distillation procedure from Transformers to Mamba fails to preserve the original teacher performance, a limitation often overcome with hybrid solutions combining Attention and SSM blocks. The key argument from our work is that, by equipping Mamba with a principled initialization, we can recover an overall better recipe for cross-architectural distillation. To this end, we propose a principled two-stage approach: first, we distill knowledge from a traditional Transformer into a linearized version of Attention, using an adaptation of the kernel trick. Then, we distill the linearized version into an adapted Mamba model that does not use any Attention block. Overall, the distilled Mamba model is able to preserve the original Pythia-1B Transformer performance in downstream tasks, maintaining a perplexity of 14.11 close to the teacher's 13.86. To show the efficacy of our recipe, we conduct thorough ablations at 1B scale with 10B tokens varying sequence mixer architecture, scaling analysis on model sizes and total distillation tokens, and a sensitivity analysis on tokens allocation between stages.
Figures & tables
Figure 1: We propose a two-stage recipe to distill quadratic Softmax Attention (green) in a Transformer layer to a subquadratic Mamba-based Mixer module. Our sequence mixer (HedgeMamba) is a hybrid of a learned linear Attention (Hedgehog ( Zhang et al., 2024 ) in blue) and Mamba Gu and Dao (2023) (yellow). Note that we keep the rest of the Transformer layer as it is from the teacher model (grey); we only swap Softmax Attention with our proposed HedgeMamba mixer.
Figure 2: Schematics of the overall approach followed in our two-stage Transformer-to-Mamba distillation recipe. In stage one, we distill vanilla Attention into a linearized version, by learning a feature map ϕ which approximates the action of softmax (following the Hedgehog procedure ( Zhang et al., 2024 ) ). In stage two, we introduce additional components from the Mamba block, to boost the overall expressivity of the model. The resulting hybrid layer, named HedgeMamba, is then further finetuned to close the performance gap with the original teacher model.
Model (1B)
↓ PPL
↑ Arc-C
Arc-E
SIQA
PiQA
Lambada
BoolQ
RACE
LogiQA
WinoG
HSwag
Pythia (Teacher)
13.86
27.04
56.98
39.86
70.72
42.07
60.82
32.92
22.12
53.43
47.16
Hedgehog (Baseline)
14.89
26.45
52.74
38.38
68.01
30.60
54.80
30.43
21.65
50.91
40.79
HedgeMamba (Ours)
14.11
27.13
53.66
39.76
68.72
32.31
55.20
30.91
20.89
52.17
41.87
Table 1: Comparison with teacher and prior work.
Model (1B)
#params
↓ PPL
↑ Arc-C
Arc-E
SIQA
PiQA
Lambada
BoolQ
RACE
LogiQA
WinoG
HSwag
Hedgehog
1,014M
14.89
26.45
52.74
38.38
68.01
30.60
54.80
30.43
21.66
50.91
40.79
+SSM
1,020M
14.89
26.54
52.90
38.02
68.23
31.24
55.63
30.05
22.73
51.38
40.77
+Conv
1,020M
14.89
26.62
52.74
38.28
68.93
31.63
55.84
30.14
22.43
51.78
40.74
+Gate(HedgeMamba)
1,087M
14.58
26.19
53.11
39.56
68.77
32.16
57.61
31.00
24.42
50.99
41.81
Table 2: Mixer architecture ablations: perplexity on validation set and downstream task performance.
Tokens split
↑ Downstream evals
S1 / S2 (%)
↓ PPL
Arc-C
Arc-E
SIQA
PiQA
Lambada
BoolQ
RACE
LogiQA
WinoG
HSwag
Hedgehog (no FT)
100 / 0
25.71
25.85
48.70
36.34
66.49
12.12
61.47
27.27
20.58
50.83
26.14
90 / 10
16.15
25.00
52.06
38.69
68.93
28.08
56.15
30.24
22.43
51.14
39.69
75 / 25
15.18
26.71
52.31
38.59
69.26
30.66
60.61
30.24
20.58
49.96
41.02
50 / 50
14.58
26.19
53.11
39.56
68.77
32.16
57.61
31.00
24.42
50.99
41.81
25 / 75
14.25
26.19
53.91
39.71
68.93
31.90
55.41
30.81
21.35
51.30
41.59
Table 3: Sensitivity analysis on the token budget allocation. Total 10B distillation tokens. Distillation on full HedgeMamba, with convolution layer and gate branch. Notice that a 100/0 split indicates that no fine-tuning is performed, and all tokens are used for learning the Hedgehog map in stage 1 (S1); conversely, 0/100 indicates only fine-tuning (S2) on the HedgeMamba architecture.
Token budget
↓ PPL
↑ Arc-C
Arc-E
SIQA
PiQA
Lambada
BoolQ
RACE
LogiQA
WinoG
HSwag
1B
16.56
26.19
52.27
38.74
67.68
27.32
57.49
29.76
20.43
52.25
40.67
2B
15.61
25.94
51.05
38.79
69.04
29.30
56.45
29.57
23.04
51.85
40.29
3B
15.15
25.09
52.69
38.43
69.10
30.56
56.57
29.28
23.04
51.93
41.03
10B
14.11
27.13
53.66
39.76
68.72
32.31
55.20
30.91
20.89
52.17
41.87
Table 4: Scaling study on the distillation token budget.
Appendix figures & tables4 assets
Supplementary material from the paper’s appendix.
Appendix
Model (1B)
↑ Arc-C
Arc-E
SIQA
PiQA
Lambada
BoolQ
RACE
LogiQA
WinoG
HSwag
Pythia(Teacher)
27.05 1.30
56.99 1.02
39.87 1.11
70.73 1.06
42.07 0.69
60.82 0.85
32.92 1.45
22.12 1.62
53.43 1.40
47.16 0.50
Architecture ablations
Hedgehog baseline
26.45 1.29
52.74 1.02
38.38 1.10
68.01 1.09
30.60 0.64
54.80 0.87
30.43 1.42
21.66 1.62
50.91 1.41
40.79 0.49
+SSM
26.54 1.29
52.90 1.02
38.02 1.10
68.23 1.09
31.24 0.65
55.63 0.87
30.05 1.42
22.73 1.64
51.38 1.40
40.77 0.49
+Conv
26.62 1.29
52.74 1.02
38.28 1.10
68.93 1.08
31.63 0.65
55.84 0.87
30.14 1.42
22.43 1.64
51.78 1.40
40.74 0.49
+Gate(HedgeMamba)
26.19 1.28
53.11 1.02
39.56 1.11
68.77 1.08
32.16 0.65
57.61 0.86
31.00 1.43
24.42 1.69
50.99 1.40
41.81 0.49
Appendix
Table 5: Results from experiments in Sec. 4 , including standard error for LM evaluations.
Model
↓ PPL
↑ Arc-C
Arc-E
SIQA
PiQA
Lambada
BoolQ
RACE
LogiQA
WinoG
HSwag
160M
Pythia(Teacher)
39.38
23.63
43.64
36.75
62.30
22.38
56.88
28.71
19.05
51.22
30.28
Hedgehog (Baseline)
35.95
18.26
42.47
37.05
61.15
14.48
59.66
26.41
21.04
50.43
28.93
HedgeMamba (Ours)
26.84
23.04
43.27
37.36
60.88
16.34
57.68
26.03
19.35
51.07
29.71
410M
Pythia(Teacher)
16.50
24.32
51.89
38.95
66.70
36.60
60.58
30.72
21.97
53.27
40.62
Hedgehog (Baseline)
17.66
19.97
47.31
37.97
65.23
24.04
48.44
28.04
19.35
50.28
34.63
HedgeMamba (Ours)
16.48
23.81
49.54
38.69
64.69
25.91
51.68
28.42
21.35
52.80
36.28
Appendix
Table 6: Scaling analysis with respect to model size.
Figure 3: Visualizing evolution of validation perplexity during stage 2 training, for different stage1/stage2 token allocation splits. ⧫ indicates the PPL value at completion of stage 1. All sensitivity studies were run for 200K training steps, corresponding to roughly 10B tokens.
Figure 4: Schematics of architectures discussed in this project. Top: Pythia Transformer. Bottom, from left to right: Hedgehog, HedgeMamba, and Mamba.
Cross-architecture knowledge distillation transfers learned behavior between models with different architectures to build smaller or more efficient systems.
State-space models (SSMs) have emerged as promising alternatives to Transformers for sequence modeling. However, training competitive SSMs from scratch remains computationally intensive, and the ecosystem around them is far less mature than that of Transformers. Moreover, the architectural differences between SSMs and Transformers make it challenging to efficiently transfer knowledge from pretrained Transformers. In this work, we propose Cross-architecture distillation via Attention Bridge (CAB), a distillation framework that transfers attention-related representations from Transformer teachers to state-space student models. Unlike conventional knowledge distillation that supervises only final predictions, CAB enables token-level intermediate supervision through a lightweight bridge and flexible layer-wise alignment. By aligning Transformer attention-related representations with Mamba's token-dependent state projections, CAB facilitates efficient cross-architecture knowledge transfer without inference-time overhead. Experiments on image classification and language modeling demonstrate that CAB improves Transformer-to-SSM distillation, particularly under limited-supervision settings. Overall, CAB provides an efficient pathway for transferring Transformer representations to SSM-based models, helping bridge the gap between mature Transformer ecosystems and emerging SSM ecosystems.
Penghao Wang, Yuhao Zhou, Mengxuan Wu +3
National University of Singapore · The University of Texas at Austin
Balancing fine-grained local modeling with long-range dependency capture under computational constraints remains a central challenge in sequence modeling. While Transformers provide strong token mixing, they suffer from quadratic complexity, whereas Mamba-style selective state-space models (SSMs) scale linearly but often struggle to capture high-rank and synchronous global interactions. We present a consistency boundary analysis that characterizes when diagonal short-memory SSMs can approximate causal attention and identifies structural gaps that remain. Motivated by this analysis, we propose InfoMamba, an attention-free hybrid architecture. InfoMamba replaces token-level self-attention with a concept bottleneck linear filtering layer that serves as a minimal-bandwidth global interface and integrates it with a selective recurrent stream through information-maximizing fusion (IMF). IMF dynamically injects global context into the SSM dynamics and encourages complementary information usage through a mutual-information-inspired objective. Extensive experiments on classification, dense prediction, and non-vision tasks show that InfoMamba consistently outperforms strong Transformer and SSM baselines, achieving competitive accuracy-efficiency trade-offs while maintaining near-linear scaling.
Youjin Wang, Jiaqiao Zhao, Rong Fu +5
Renmin University of China, Beijing, China · University of Macau, Macau SAR, China · Central South University, Changsha, China +1
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.
Ramprasath Ganesaraja, Sahil Dilip Panse, Swathika N