Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models

Summary

This paper proposes MOHAWK, a three-stage distillation method to transfer knowledge from pretrained transformer models to subquadratic models such as Mamba-2. The key idea is to view both transformers and SSMs as applying different forms of mixing matrices over token sequences. Experiments demonstrate that Phi-Mamba, a variant of Mamba-2 architecture tailored for distillation, achieves strong performance on downstream tasks using less than 1% of the training data compared to training from scratch.

Strengths

1. The paper introduces an effective way to distill knowledge from pretrained transformer models to subquadratic models, enabling the linear attention community to obtain better models at a lower cost. 2. Mamba-2 aligns well with self-attention by converting the forget gate of GLA into a data-dependent value, achieving linear attention-like structures by simply removing the softmax and adding a forget gate. 3. The three-stage distillation process is well-designed and ablation studies demonstrate the importance of each stage, providing valuable insights for future research in this direction.

Weaknesses

1. The paper does not provide a detailed comparison between the proposed distillation method and direct fine-tuning approaches like SUPRA (reference 21). While the three-stage distillation process achieves better performance with only 2.8B tokens, it is crucial to consider the additional computational overhead introduced by the need for both a teacher model and a student model during distillation. In contrast, direct fine-tuning methods, though requiring more tokens (e.g., +20B in SUPRA), may be simpler and more efficient in practice. The authors should discuss the trade-offs between these approaches and provide ablation studies to justify their choice of distillation over direct fine-tuning. 2. There are several typos throughout the paper. For example, in line 159, "Mamba-2 2" should have a "Figure" mentioned before the second "2". The authors should carefully proofread the manuscript to address these issues. 3. There is an inconsistency in the reported model size of Phi-1.5. The official documentation states that Phi-1.5 has 1.3B parameters, while this paper mentions a model size of 1.5B. The authors should clarify this discrepancy and ensure consistency throughout the paper. 4. The paper does not report the performance of the Phi-Mamba model on the MMLU benchmark, which is a standard evaluation for the Phi-1.5 model. Given that Phi-1.5 (1.3B) has reported MMLU performance, it would be informative to compare the MMLU scores of the Phi-Mamba model to assess its performance on a wider range of tasks.

Questions

See weaknesses part.

Rating

5

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

The authors have discussed the potential limitations in Section 6.

Summary

The paper distills a Transformer model into the Mamba architecture by using about 1% of the pertaining dataset.

Strengths

1. The paper successfully distills a Transformer model into the Mamba architecture by using about 1% of the pertaining dataset, reaching best performance on some of the tasks when compared with other SSM models and is close to the performance of the teacher model for most tasks. 2. The authors compare against a large number of SSM models. 3. The authors have performed extensive ablations to ensure that all three steps they propose are necessary to achieve good distillation.

Weaknesses

1. The paper uses the term "mixing matrices" repeatedly throughout the paper. However, this term is never formally defined, except as the matrix M in Section 3.1, although the next paragraph explains that such matrix multiplication is actually not used in practice. 2. The description of the "Matrix Orientation" step in Section 4.1 is also not at all clear. What does "matrix mixer layer" mean? What are TeacherMixer and StudentMixer? Is this done layer by layer, block by block or are they all aligned at the same time? What's the $x$ that we are minimizing over? Is it full sequences (previously denoted with capital $X$)? What are the dimensions of the models? 3. Similar issues are present in the description of the "Hidden-State Alignment" step in Section 4.2. There is also a new term $\mathbf u$ that appears but is never defined. 4. No code or trained model artifacts were provided nor were mentioned that will be provided. The authors declare in the Checklist that "Our proposed method is quite simple and all significant details are included", however, due to the above points, I do not think that replicating their experiments would be at all straightforward. 5. Minor comment: On line 108, matrix multiplication has complexity larger than $O(T^2)$. 6. Minor comment: On line 48, "strong" -> "stronger" Overall, while the experimental results seem impressive, the rest of the paper is not clear and do not allow for the reproducibility of the results. I would recommend that the authors clearly and formally describe the architectures of the teacher and student networks and explicitly define the objective functions that they optimize for each of the three stages.

Questions

See Weaknesses.

Rating

6

Confidence

4

Soundness

3

Presentation

1

Contribution

2

Limitations

The paper offers only one sentence of limitations in the Discussion section and it is mostly about the need for further experiments. Missing limitations are: difficulty to reproduce, investigation of only one pair of teacher and student, and failing to achieve the performance of the teacher model. There is no potential negative societal impact to be addressed.

Summary

The authors consider the problem of distilling transformers into SSM models (Mamba in particular), which results in the reduction of quadratic complexity at inference to subquadratic complexity. For this purpose, the authors propose MOHAWK, a method consisting of several steps, each aiming to match a different aspect of SSMs and transformers. The method is then tested on the Phi-1.5 transformer model and distilled into Phi-Mamba proposed by the authors. The paper shows that the resulting model performs very well and presents a thorough analysis that each step of their method is important.

Strengths

- The paper considers a novel and important problem of distilling knowledge from quadratic models to subquadratic models. This could reduce the inference costs and, as such, is of significant practical importance - The empirical results obtained by the proposed method are quite good. The authors also provide a detailed analysis of the impact of each step. The authors quite convincingly show that fine-grained matching of particular blocks is needed before we start end-to-end knowledge distillation. - To the best of my knowledge, the paper faithfully discusses the related work.

Weaknesses

- I have some reservations about the empirical evaluation: - The crucial consideration in the large language model domain is the scaling properties. Given the current experiments, I’m not convinced that the proposed method will scale well with data and available compute. It’s obviously nice that the approach works well even when using 1% of the data from other models, but how much data would we need to close the gap to Phi-1.5 completely, or up to negligible levels? What are the fundamental limitations of this approach, what Mamba cannot do that transformers can? - Also, will it work with models larger than 1.5B? For example, can we use this approach to scale Mamba up beyond its standard regime (e.g., 7B, 70B)? - The Phi architecture used in this paper works well, but I would be curious to see how MOHAWK performs with more standard transformer architectures (e.g., LLAMA). - There’s a confounding factor - Phi is trained on a very well-curated dataset, which makes it perform very well for its size. The Pile dataset used for training Mamba and other models is not that good. For a fair comparison, I think one should use models trained on (roughly) the same data as baselines. - There are certain problems with the presentation in this paper: - There are numerous typos and grammar errors, see below. - The way the Tables and Figures are referenced in the text is confusing. Some of the links are broken, and others are not referenced when they should be. In particular, Section 5 is difficult to understand because of that. Minor issues/typos: - Line 27: “raises a natural question: it is” → is it? - Line 33: “differen” - Line 48: “benchmarks strong than” → stronger than - Line 116: “$A_t h_t$ is the identity matrix I multiplied by a scalar αt” -> $A_t$, I don't think you need the $h_t$ term there? - Figure 3: “Despiting training” → Despite training - And many more I think. I suggest carefully proofreading the manuscript.

Questions

Please see the weaknesses section above. I'm particularly interested in the scaling properties of MOHAWK, and its fundamental limitations. How far can we go?

Rating

6

Confidence

4

Soundness

3

Presentation

2

Contribution

3

Limitations

The limitations are not clearly discussed. Issues that should be mentioned include: * How well the method will scale? * How well does it work with more standard transformer/mamba architectures (not Phi)? * What is the impact of the data?

Thank you

The reviewer thanks the authors for the discussion. It solved most of my concerns. I decide to keep my score.