The Mamba in the Llama: Distilling and Accelerating Hybrid Models

Thanks for your comments! Please find our responses below.

$\mathcal{Q}.$ In Section 2.2, Why is A set as fixed over time?

$\mathcal{A}.$ The equation after line 92 has a fixed-over-time $A$ because it is the continuous SSM, which is consistent with the Mamba paper. However, after discretization, we obtain a matrix $\bar{A}$ which depends on the timestamp $t$ since $\Delta$ is not fixed over time.

$\mathcal{Q}.$ In Section 2.3, which components in the model are from attention and which parts are trainable. Figure 1 is also not informative without explanations of the colors and arrows.

$\mathcal{A}.$ In the caption of Figure 1, we state that we freeze the FFN layer and only train the projects and Mamba block. We also explain that weights in the same color are initialized from the same color in attention. Please check our updated one-page PDF for an updated explanation. We will clarify this in the final version.

$\mathcal{Q}.$ In section 4, what are the differences between speculative decoding used in transformer and that used in Mamba here? What are the details of hare-aware designs?

$\mathcal{A}.$ Speculative decoding requires verfying multiple tokens in parallel and rolling back when some of them are rejected. This is easily done in transformers by adjusting the size of the kv-cache. For SSMs and Mamba in particular, rolling back requires "undoing" the decoding and getting back to a previous state. A naive way of achieving this would be to cache the states and restore a previous state after tokens are rejected. In practice, this is very inefficent as it requires materializing several very large states. The solution we implement, which we summarize in Algorithm 2 in the paper, allows us to avoid caching the states, and instead recompute the states for the accepted tokens on-the-fly using the cached $x$, $B$ and $\Delta_t$ matrices, without ever materializing the intermediate states.

$\mathcal{Q}.$ For comparison methods, Hybrid Mamba which is trained from scratch should be included.

$\mathcal{A}.$ Please refer to the 2 and 3a for comparisons with the best Hybrid Mamba which is trained from scratch.

$\mathcal{Q}.$ The benchmarks are only on chat and academic tasks.

$\mathcal{A}.$ Please refer to section 2 in the general response. It includes 10 probability-based tasks in LM Bench, which are the same as Nvidia Hybrid Mamba.

$\mathcal{Q}.$ In Section 2.1, the derivations from softmax-based attention to linear RNNs (Mamba) are not clearly stated, e.g. unexplained symbols in equations. Also the derivations should obey some assumptions.

$\mathcal{A}.$ We apologize for the lack of clarity here. We were not trying to claim an equivalence between these equations, but instead trying to show the relationship. Please check our updated one-page PDF for an updated explanation. We will clarify this in the final version.

$\mathcal{Q}.$ Please refer to Mamba-2. Why did we choose Mamba-1 instead of Mamba-2?

$\mathcal{A}.$ Mamba 2 was released on May 31, which is after the NeurIPS submission deadline. According to NeurIPS guideline, papers appearing less than two months before the submission deadline are generally considered concurrent to NeurIPS submissions.

To show the generalization ability of our paper, we adopted our approach in Mamba 2 and report the numbers in section 2 in the general response. We would like to state that the goal of this project is to distill a large transformer into Mamba and generate fast. Indeed, our approach works for any linear RNN model and our story does not change using different linear RNN models. We hope to try out more linear RNN models as they are released.

My main concern is the narrow scope of the methods and experiments of this paper. The paper’s story seems to rely on a biased assumption that Mamba will replace transformers—an appealing idea, but speculative at this stage. The focus on distilling and accelerating the Mamba hybrid model, while potentially applicable to other models (like linear RNN or linear attention), still requires further demonstration. And the paper does not propose any new methods or theory, it is just a application of proven techniques on mamba. Given the limited scope and novelty, I don’t believe the paper meets the NeurIPS standard in its current form.

We are very grateful for all the encouraging and constructive reviews. Please find our responses below.

$\mathcal{Q}.$ The paper’s story seems to rely on a biased assumption that Mamba will replace transformers—an appealing idea, but speculative at this stage.

$\mathcal{A}.$ We appreciate the reviewer's concerns. We do not claim in the paper that Mamba will replace transformers and in fact, believe that people will continue mainly training transformers. We instead think that some organizations may want prioritize fast and lower-memory inference, and therefore want to distill their transformer models to Mamba for some applications.

$\mathcal{Q}.$ Additionally, the paper does not introduce any new method or theory.

$\mathcal{A}.$ We think there are two new methods introduced by this paper. 1) Our distillation approach is novel. To the best of our knowledge, this is the first work demonstrating that transformations between large transformers and linear RNN models is possible through a mathematical transformation. And in our experiments, we distill different large transformers to support this claim. 2) We introduce a new algorithm for speculative decoding in the Mamba model which is efficient in practice. We agree that there is no new theory introduced in this work, although that is generally true for most deep learning papers.

$\mathcal{Q}.$ We don't understand that it is just an application of proven techniques on mamba.

$\mathcal{A}.$ While distillation is a proven technique, as far as we know there is no work published on distilling from transformer to Mamba or related linear RNNs. The details of what distillation approaches work and how they can reuse parameters is a the contribution. We are a bit perplexed by this comment, and it would be great help if you can point us any thing related to this. (Previous work on this topic does not do anything like distillation.)

We sincerely appreciate your feedback. Please see our responses below.

$\mathcal{Q}.$ The experiments are a bit limited, only performed on chat corpora. Further data mixture is needed to understand the effectiveness of the approach

$\mathcal{A}.$ Please refer to 2 in general response where we hope to answer this concern. We have evaluated our model not only on generation tasks in standard chat benchmarks but also included probability-based tasks in LM eval same as Nvidia hybrid Mamba2. We also see that a model distilled from Llama-3 performs better than the Nvidia hybrid Mamba2 model, which was trained from scratch with 3.2T tokens. In addition we train on a larger set of tokens.

Thank you very much for your comments! We kindly ask you to find our responses below.

$\mathcal{Q}.$ No comparison is performed with random initialization for task specific performance.

$\mathcal{A}.$ We added a comparison with random initialization. Please refer to 3a in the general response. With the Mamba default initialization and the same training tokens, the model gets a very low score in MT-bench and AlpacaEval.

$\mathcal{Q}$ Does PPL correspond to ability?. However, the discussion in that paragraph does not answer this question at all. Moreover there is actually a correspondence and we see degradation in task specific performance similar to PPL

$\mathcal{A}.$ In the new ablations that we ran, we found that initializing using attention weights not only gives lower PPL, but also significantly improves task-specific performance.

$\mathcal{Q}.$ Similarly, the results for 100% Mamba architecture is not included. While this might not be as good as 50% architecture, it is still very much needed to understand the trade-off.

Furthermore, the comparison is done on a hybrid architecture which further makes it confusing to determine the success of the method.

$\mathcal{A}.$ We include the 1/8 Mamba model in the new result in the general response.

Here we also compare with 100% Mamba model (trained progressively). These results are less good than hybrid, but significantly better than random initialization.

$\mathcal{Q}.$ The authors mention that they do not target pre-trained language model capabilities. Why is that the case?

$\mathcal{A}.$ We apologize for the confusion in our writing. We do think this approach captures general pretraining ability. To clarify, given a limited amount of distillation compute, we hypothesized that the most effective usage of compute is to distill on chat data, which is high-quality and task relevant.

Even though we train on chat data, we want to emphasize that our model has strong performance on general benchmarks. For example, our model distilled from Llama-3 performs better than the NVIDIA hybrid Mamba2 model on general benchmarks. Please refer to 2 in the general response.

$\mathcal{Q}.$ Since there is redundancy in Transformer layers, are the Mamba layers actually useful?

$\mathcal{A}.$ This is a good question, and we were curious as well. Please refer to 3b in our general response. This result confirms that adding Mamba layers is very critical and that the results are not just from the remaining attention.

$\mathcal{Q}.$ The definition of # Gen in Figure 2 (right) is very vague and unclear. Can you please elaborate?

$\mathcal{A}.$ This is the number of tokens we generate at each generation step. It is calculated as the average number of draft tokens we accept at each iteration, plus an additional token which we is sampled from the target model, as per the speculative decoding algorithm [1].

$\mathcal{Q}.$ In line 182 it is mentioned that the parallel mode of Mamba is faster than a Transformer. Is there a reference for this? If so I think it would be useful to include it.

$\mathcal{A}.$ It is more efficient for long sequences because it doesn't require $O(L^2)$ computation. Please refer to the left figure in Figure 8 of the Mamba paper.

$\mathcal{Q}.$ In Algorithm 2, the verify part may return the state from an earlier point, i.e. $h_j$ instead of $h_{k'}$. How is this ok for the rest of the algorithm? This means a set of tokens will be ignored as the algorithm will never go over $j+1, ..., k$ to incorporate them into the state.

$\mathcal{A}.$ We are sorry for the confusion. Yes you are correct in that in speculative coding we may throw away generated tokens. In the case where a state from an earlier point is returned by the verify part, the draft model will also resume decoding from a previous snapshot of the state, recomputing the states for the accepted tokens from its cache.

[1] Fast Inference from Transformers via Speculative Decoding (https://arxiv.org/abs/2211.17192)

We thank again the reviewer for the helpful feedback.

$\mathcal{Q}.$ I could not find the new PPL results. Could you please share them as well?

$\mathcal{A}.$ Since our models are alignment using DPO, we evaluated ppl of our final models on lambada_openai. Here are ppl of models in the ablation results.

$\mathcal{Q}.$ Just to confirm, you still keep the FFN parts of the layers, and only remove the attention, then do fine-tuning? I think if this is the case, this is an important result as it shows the efficacy of using Mamba blocks and the projection. Combined with the obtained speedup I think this makes the results quite interesting. However, the trend that using Mamba blocks lead to decreased accuracy is completely observable in all the results. As such, I think it is paramount that this limitation be clearly specified and discussed in the paper. Do authors agree on this or have I misinterpreted the results?

$\mathcal{A}.$ Yes, we still keep those three Gated MLPs (FFN components) and only remove half of the attention layers before fine-tuning. We acknowledge that changing attention to Mamba may still lead to decreased performance as more Mamba layers are introduced. At the same time, we believe that as computational resources scale up, this impact will be mitigated. Additionally, the inductive bias learned from attention can help Mamba scale up, as our model has already outperformed other hybrid Mamba models trained from scratch with significantly more tokens. We definitely will address and discuss this limitations in the paper.

Additionally, we will update the speculative decoding section to improve and correct the explanation of the algorithm.