Distilling Transformers into Linear-Time Models

Pretraining a competitive linear-time language model from scratch costs trillions of tokens, and the open-weight ecosystem already contains transformers that have paid that bill. The obvious question is whether the second half of the bill can be skipped: take the trained transformer, replace its attention layers with state space layers, and spend a small fraction of the original compute repairing the damage. The answer, since 2024, is largely yes, and the way it is done says something precise about what attention layers actually store.

Reuse the projections, not just the architecture

The insight that makes conversion cheap is that linear attention and a selective state space layer are the same shape of computation. The structured state space duality writes the SSM output as \(C_i^\top B_j\) scaled by a decay product, which maps directly onto a query-key inner product. So the pretrained \(W_Q\), \(W_K\) and \(W_V\) matrices are not discarded. They initialise the new layer's \(C\), \(B\) and \(x\) projections, and only the decay parameters, \(A\) and \(\Delta\), are new. A converted layer therefore starts far closer to the teacher's function than a randomly initialised one, and the fine-tuning budget is spent on the residual rather than on relearning the representation.

Wang and colleagues built on exactly this, reusing the linear projection weights from attention to seed Mamba layers, and reported that the result is reachable with academic GPU resources rather than a pretraining cluster (Wang, Paliotta, May, Rush and Dao, 2024, The Mamba in the Llama: Distilling and Accelerating Hybrid Models, NeurIPS, arXiv:2408.15237). MOHAWK attacked the same conversion with an explicit three-stage recipe, matching the sequence-mixing matrix first, then hidden states, then the full model, and distilled a Mamba-2 variant of Phi-1.5 using 3B tokens, with a hybrid version at 5B (Bick, Li, Xing, Kolter and Gu, 2024, Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models, arXiv:2408.10189). Against the hundreds of billions of tokens the teacher consumed, that is a rounding error.

Staging, because the objective is not one objective

Distilling a whole model on output logits alone works badly here, because the student's failure is localised. Attention layers differ from state space layers in how they mix across positions, not in how they transform each position, so the productive supervision signal is at the mixing matrix, not at the output.

The staged recipes therefore proceed inward out. First align each converted layer's sequence-mixing matrix with the teacher's attention matrix, layer by layer, with the rest of the network frozen. This is a per-layer regression on an object both architectures possess, and the semiseparable-matrix view is what makes them comparable at all. Second, match hidden states block by block, which repairs the accumulated drift the per-layer fit leaves behind. Only then run end-to-end distillation on logits, followed by the ordinary post-training pipeline: supervised fine-tuning and preference optimisation, because a converted model's instruction-following degrades in ways that logit matching against a base teacher will not fix.

Keeping some attention on purpose

The strongest converted models are not pure. The Mamba-in-the-Llama hybrids retain roughly a quarter of the original attention layers and replace the rest, and that configuration matched the original transformer on chat benchmarks while beating open hybrid Mamba models that had been trained from scratch on trillions of tokens. The retained layers are there for the reason given in State Capacity and Associative Recall: exact recall is a qualitatively different memory operation, and a small number of full attention layers restores it more cheaply than any amount of extra state.

This reframes conversion. It is not "replace attention", it is "find out how much attention this model actually needed", and the empirical answer keeps coming back somewhere between an eighth and a half.

When it breaks

The teacher's context length is an upper bound, not a floor. A converted model inherits the teacher's positional behaviour through the reused projections, and the state space layers are trained to imitate attention over the lengths seen during distillation. Extrapolating a converted hybrid past that range is unreliable, and the failure is the quiet recall failure rather than a visible collapse.

Evaluation has to be recall-weighted. Conversion preserves general language modelling, chat quality and most reasoning benchmarks far better than it preserves long-context retrieval. Reporting MMLU and a chat score on a converted model, and not a needle-in-a-haystack or multi-query recall test at serving length, is the standard way to overstate what the conversion achieved.

The economics depend on the serving profile. Conversion pays off when decode dominates and batch sizes are cache-limited. For short prompts, small batches, or workloads dominated by prefill, the constant-memory advantage is small and the distillation budget plus the quality risk buys very little. The question to answer before converting anything is what fraction of the accelerator's memory is currently KV cache.