Summary
The paper proposes a way to convert a Transformer architecture to a Mamba architecture by first initializing a Mamba layer based on a pre-trained attention layer in order to facilitate faster convergence followed by knowledge distillation from the Transformer to the Mamba model. The training procedure only involves post-training on task-specific data. The paper also uses DPO for the knowledge distillation and the authors emphasize the novelty of this method for knowledge distillation. Following their proposed method, the authors demonstrate how to quickly train a hybrid architecture with some attention and some Mamba layers based on existing LLMs, particularly Zephyr 7B. Finally, the paper proposes a way to perform speculative decoding and gain speed up using Mamba architectures.
Strengths
The paper's writing is mostly clear and can be followed. The idea of initializing a Mamba layer based on pre-trained Transformer layers is novel and interesting. The authors evaluate the models on several benchmarks.
Weaknesses
Looking at the performance in Table 1, increasing the attention percentage always improves performance in a significant way (most clear in MMLU). As such, I am not convinced that the method is actually working. There is also a lack of proper baselines. For example, no comparison is performed with random initialization. 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.
The idea of converting a pre-trained Transformer to Mamba is quite generic. In particular, there does not seem to be any specific scale needed to get the method to work. As such, it is not clear why the authors decided to mainly test the method on such a large scale. For example, a small Transformer model can be trained from scratch, used to create a Mamba model and then compared in terms of language modeling benchmarks such as perplexity. Indeed this experiment is done for one experiment in Section 7 and shows poor performance of the method in terms of perplexity. As a side note, the paragraph starting in line 275 is titled 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.
Moreover, the authors mention that they do not target pre-trained language model capabilities. Why is that the case? Furthermore, the comparison is done on a hybrid architecture which further makes it confusing to determine the success of the method. Note that large models have a lot of redundancies as suggested by recent work in pruning and this makes it harder to see whether other attention layers are stepping in to make up for the poor performance of Mamba layers or something else is happening.
Overall I find the results not convincing enough and not supporting the claims made in the paper about the success of the method.
Questions
1. The definition of # Gen in Figure 2 (right) is very vague and unclear. Can you please elaborate?
2. 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.
3. In Algorithm 2, the verify part may return the state from an earlier point, i.e. $h_j$ instead of $h_{k^\prime}$. 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, \ldots, k^\prime$ to incorporate them into the state.
Limitations
The authors mention that the performance is only measured after fine-tuning on chat corpora instead of full pre-training due to resource limitations. While at 7B scale this is acceptable, as mentioned above, there is no need to focus on that scale from the start. Further limitations, such as degradation of performance when switching from attention to Mamba layers should be more deeply discussed.