Apple researchers propose two-stage Attention-to-Mamba distillation

Apple researchers have proposed a two-stage method for transferring a pretrained Transformer into an adapted Mamba model with no attention blocks. In the authors’ 1B-scale experiment, the resulting HedgeMamba model nearly matched its Pythia teacher on validation perplexity, but results were mixed across 10 downstream evaluations.

The paper’s revised experiment reported validation perplexity of 14.01 for HedgeMamba, compared with 13.86 for the Pythia-1B teacher and 14.89 for a Hedgehog linear-attention baseline. Those figures come from the authors’ OpenWebText setup; they do not establish how the model would perform in deployment.

Linear attention serves as the bridge between the two architectures. In stage one, the researchers replaced the teacher’s softmax attention with Hedgehog, which learns feature maps that approximate attention while allowing computation to scale linearly with sequence length. They copied and froze the teacher’s other parameters, then trained the new feature maps to match each student layer’s output to the teacher using a cosine-embedding loss.

In stage two, those feature maps initialized an adapted Mamba-2 mixer the authors call HedgeMamba. The researchers initialized new state-space, convolution and gating components so they would not disrupt the stage-one mapping. They then fine-tuned the model with cross-entropy loss while keeping its input and output embeddings frozen. The final model has no attention blocks, according to Apple’s research summary.

The default experiment used 1 billion tokens in stage one and 9 billion in stage two. Together, that was about 3% of the 300 billion-token budget used to train Pythia-1B. In the authors’ allocation study, perplexity reached 25.71 when all 10 billion tokens went to stage one and 15.74 with stage-two-only training, compared with 14.01 for the 10%-90% split. The authors said the result supported using both stages in their setup.

Downstream results were uneven. HedgeMamba scored slightly above Pythia-1B on ARC-Challenge, Social IQA and LogiQA, but below it on the other seven evaluations. The largest reported gap was on LAMBADA, where HedgeMamba scored 32.74 against the teacher’s 42.07. The authors attributed that difference to LAMBADA’s reliance on recalling tokens from context, which they identify as a weakness of linear attention.

The paper also reports a context-length tradeoff from a decoding experiment on one A100 80GB GPU. HedgeMamba’s recurrent and convolution-state cache stayed constant as prompt length increased, while Pythia’s key-value cache grew. At the best tested batch size for each model, HedgeMamba delivered 0.66 times Pythia’s throughput with 1,000-token prompts, 1.15 times at 2,000 tokens and 8.51 times at 16,000 tokens. The test generated 128 tokens greedily from random prompts and excluded prefill and warm-up, so it did not measure end-to-end production latency or cost.

The paper’s November 3 revision replaced the original results after the authors found an implementation error that had prevented proper training of state-space decay gates. They moved the experiments from Mamba-1 to Mamba-2 and reran them. The study covers Pythia models trained on OpenWebText. It does not evaluate instruction following, safety, factuality or application workloads, and its reported findings have not been independently replicated.