Attention to Mamba: A Recipe for Cross-Architecture Distillation

(machinelearning.apple.com)

75 points | by zerop a day ago ago

6 comments

  • kooi a day ago

    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).

    https://arxiv.org/html/2312.00752v2

  • buildbot a day ago

    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

    • yorwba a day ago

      They seem combinable to me, i.e. use the layer-selection from that paper and initialize using the Hedgehog-linearization as in this paper, then distill to smooth over inaccuracies of the approximation.

  • sehw a day ago

    [flagged]

  • thisisauserid a day ago

    I totally get why people want to reinvent LLMs so they're better/cheaper/faster and this hybrid approach with State Space Models is likely going to bear fruit.

    But I'm WAY more interested in what's going to happen with Mamba-3, et al. when it's relatively easy to predict, ya know, state.

  • vlovich123 a day ago

    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.