Apple distills a Pythia-1B Transformer into pure Mamba with a two-stage recipe

Converting Transformers without keeping any attention

Apple's machine learning research team has published a method for distilling an attention-based language model into a Mamba-style state-space model, removing attention entirely instead of falling back on hybrid designs. The write-up, which reached the Hacker News front page, comes from Abhinav Moudgil, Ningyuan Huang, Eeshan Gunesh Dhekne, Pau Rodríguez Lopez, Luca Zappella and Federico Danieli, with the first two authors having done the work while at Apple.

The motivation is economic as much as architectural. According to the authors, Mamba-class state-space models offer a smaller memory footprint and faster token generation than comparable attention-based models, which makes them attractive for inference-heavy deployments. Yet the ecosystem's training know-how and inventory of pretrained checkpoints are overwhelmingly Transformer-based. Distillation promises a bridge: reuse an existing Transformer's knowledge rather than training a state-space model from scratch.

The catch is that the obvious approach fails. Prior work on cross-architecture distillation, cited by the authors, shows that directly transferring a Transformer teacher into a Mamba student loses a substantial share of the teacher's performance. The usual workaround is a hybrid model that keeps some attention blocks alongside SSM blocks, but hybrids blunt the efficiency gains that motivated the switch in the first place.

The two-stage recipe

Apple's proposal routes the knowledge through an intermediate representation instead of jumping straight across architectures.

In the first stage, the original Transformer is distilled into a linearized variant of attention, constructed with an adaptation of the kernel trick. The result is an intermediate teacher whose sequence-mixing operation is much closer in form to what a state-space model can express.

In the second stage, that linearized model is distilled into an adapted Mamba that contains no attention blocks at all. The authors' central argument is about initialization: equipping the Mamba student with a mathematically motivated starting point, rather than treating the transfer as an ordinary training run, is what makes the overall recipe work.

How well it works

The headline experiment uses Pythia-1B as the teacher. According to the paper, the distilled Mamba reaches a perplexity of 14.01, close to the teacher's 13.86, a gap small enough that the student keeps most of the original model's capability.

Results on downstream tasks largely track the teacher as well, with one consistent exception: tasks that lean heavily on in-context recall still show a deficit. The authors present this as the remaining weakness of the approach rather than something the recipe solves.

The work is backed by ablations at 1B scale using 10B distillation tokens. These cover the choice of sequence-mixer architecture, a scaling analysis across model sizes and total distillation token budgets, and a sensitivity study on how tokens should be split between the two stages — the last point mattering because a practitioner following the recipe has to decide where to spend a fixed compute allowance.

Why it matters

For anyone deploying language models under memory or throughput constraints, this offers a concrete transfer path: keep pretrained Transformers as the source of capability, ship a pure Mamba at inference time, and avoid the hybrid attention blocks that dilute the efficiency benefits. That lowers the barrier to adopting SSM-based inference, since it builds on checkpoints and training practices the community already has.

Two caveats frame the result. The recipe is demonstrated at 1B parameters, so the headline numbers do not by themselves prove the approach holds at frontier scale, though the included scaling analyses address how performance moves with model size and token budget. And the in-context recall gap is a reminder that attention-free models are not yet a drop-in replacement everywhere — retrieval-heavy workloads are where the distilled student is most likely to fall short of its teacher.