Apple's two-stage trick lets a Mamba model nearly match its Transformer teacher
Attention to Mamba: A Recipe for Cross-Architecture Distillation
Mamba and other state space models promise lower memory use and faster generation than Transformers, but naively distilling a pretrained Transformer into Mamba loses too much of the teacher's quality. Apple researchers propose a two-stage recipe: first distill the Transformer into a linearized Attention model via a kernel-trick adaptation, then distill that into a Mamba model with no Attention blocks. The result keeps most of Pythia-1B's performance, hitting perplexity 14.01 versus the teacher's 13.86, with some gaps on in-context recall tasks.
By equipping Mamba with a principled initialization, we can recover an overall better recipe for cross-architectural distillation.
- buildbot
In my opinion, a better & more comprehensive look at a similar idea for all state space attention types is this paper from AWS: https://arxiv.org/abs/2605.08301
- kooi
Cool, I like state space models because they "feel" like they're doing proper homage to the power of control and modeling for linear dynamical systems. (Mamba is a bunch of time varying linear systems chained by non-linear activations).
I'm reading one of the OG papers now and here are some interesting snip-bits. I wonder how much of inference in data centers is done with traditional transformers vs SSM, because the efficiency gains are quite dramatic.
> We propose a new class of selective state space models, that improves on prior work on several axes to achieve the modeling power of Transformers while scaling linearly in sequence length.
> Our Mamba language model has 5x generation throughput compared to Transformers of similar size, and Mamba-3B’s quality matches that of Transformers twice its size.
Then, interestingly for training:
> Commonly, the model[SSM] uses the convolutional mode (3) for efficient parallelizable training (where the whole input sequence is seen ahead of time), and switched into recurrent mode (2) for efficient autoregressive inference (where the inputs are seen one timestep at a time).
- vlovich123
1B model only. No indication if this continues to scale to larger models / what that looks like from the Mamba side.
> This largely carries over to downstream tasks, although some gaps appear on tasks relying more on in-context recall
And yup, the famous achilles hill of SSMs - poor recall.