Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models

We appreciate the reviewer finding our method well-designed and the insights from the ablations valuable. We respond to the reviewer’s questions and concerns, which mainly focus on additional evaluations and comparisons to other distillation methods, below.

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.

Despite MOHAWK’s overhead, MOHAWK provides better performance than other cross-architectural distillation methods, as SUPRA’s minimum gap between SUPRA distilled models and their teacher models is greater than 6 percentage points. In addition, SUPRA is confined to converting attention into RNN-based models; however, MOHAWK can work with any architecture that can be represented using matrix mixers, as seen in our ablations in Table 4 when we distilled a Toeplitz and causal low-rank based Phi model. Compared to a more standard technique of weight transfer and knowledge distillation, we show in the second table in our shared response that only weight transfer and knowledge distillation leads to larger performance gaps in downstream metrics compared to the MOHAWK method.

There are several typos throughout the paper.

There is an inconsistency in the reported model size of Phi-1.5.

We thank the reviewer for pointing out the inconsistency of the Phi-1.5 model. As mentioned in the shared notes, we have fixed it and the other typos in the paper.

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.

The metrics we reported are a standard subset of the ones reported in the baseline papers in both Figure 1 and Table 1 of the paper (Mamba, xLSTM, RWKV, GLA, etc.). These seldom report MMLU, partly because it is known that it is hard to achieve above-random-guessing at the 1B model scale. Phi-1.5 has above-chance on MMLU which is hypothesized to be a byproduct of the special dataset that Phi-1.5 is trained on. We believe that disentangling the role of data is an important consideration for our distillation methods that we would like to explore in future work; however, we were unable to control this systematically at the 1B model scale due to a lack of strong and open-weight/open-data models.

We appreciate the reviewer finding our analysis thorough and our method novel and important. The reviewer also raised an important point about the entanglement of the architecture, data, and MOHAWK method when looking at the distillation results. We would like to preface by reiterating that our key contribution in this paper is the MOHAWK method which enables the effective distillation from Transformers to alternative architectures, and distilling from Phi-1.5 to a Mamba model is validation of the method’s effectiveness. We will structure the following response by disentangling the architecture first and then the data from the method.

What are the fundamental limitations of this approach, what Mamba cannot do that transformers can?

The limitations of the Mamba architecture is an orthogonal line of work [1-5]. Our distillation method is not directly tied to Mamba, since we have distilled to other architectures in Table 4, and as stated in the shared response, our key contribution is the distillation method while Phi-Mamba is our validation of the effectiveness of the method. Our paper’s goal is to introduce a distillation method that can create performant alternative architecture-based models which we show via the strong downstream metrics of Phi-Mamba and the ability of MOHAWK to fully distill Phi-from-Phi.

how MOHAWK performs with more standard transformer architectures (e.g., LLAMA)

MOHAWK should work with any architecture that has a valid matrix mixer representation. With more standard Transformer architecture like Llama, the matrix mixer layer is similar to that of Phi’s, which means that Stage 1 and Stage 2’s strong performance should translate. The reason we did not use a more standard Transformer architecture was that there are few strong and open-weight models at the 1B scale.

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

We believe that our MOHAWK method will be able to scale to larger models but will require scaling to industrial resources. We distilled to a 1.5B model to validate our new distillation method on a smaller scale. The next step would be to distill a 7-8B model, but that is beyond the scope of this work. However, preliminary evidence (in Section 5.4) on just Stage 1 indicates potential for scaling to larger models as the experiments show that the SSD matrix mixer can effectively approximate the attention matrices of a Llama-7B model.

how much data would we need to close the gap to Phi-1.5 completely, or up to negligible levels?

In the second table of the shared response, we can close the gap to Phi-1.5 completely with only 5 billion tokens if the student model is also a Transformer architecture. We believe that the ability to close the gap with the teacher model stems from the differences in expressiveness between the teacher and student matrix mixers and not from the MOHAWK distillation method itself.

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.

When comparing our Phi-Mamba to the Transformer-Mamba hybrid Samba-1.7B and Mamba-SWA-MLP 1.6B model, which are trained on a stronger dataset (the Phi-2 dataset) compared to the student model and a comparable dataset compared to the teacher model, we achieve comparable performance with less than 2% of their reported 230B total training tokens without employing any Transformer based blocks. The Samba model uses the same model dimension but uses 12 attention-based layers while the Mamba-SWA-MLP model uses 18 attention-based layers. In addition, our Phi-Mamba is comparable to a Mamba-1.8B trained with the aforementioned Phi-2 dataset [6]. When compared to the other alternative models that use The Pile or SlimPajama, we still outperform them when the datasets are reported to be somewhat similar compared to C4 (all three are within ~2 percentage points) [7].

There are certain problems with the presentation in this paper:

Minor issues/typos:

As mentioned in the shared response, we have addressed these.

[1] https://arxiv.org/abs/2404.08819

[2] https://arxiv.org/abs/2402.04248

[3] https://arxiv.org/abs/2402.03170

[4] https://arxiv.org/abs/2402.01032

[5] https://arxiv.org/abs/2402.18510

[6] https://arxiv.org/pdf/2406.07522

[7] https://arxiv.org/abs/2406.17557

We are glad the reviewer found our experimental results impressive and liked our extensive comparison to other SSM and sub-quadratic models. The reviewer’s main concern is the presentation, clarity, and terminology. Based on our shared response, we have fixed the issues with the presentation, and the terminology used in this paper is standard and based on an existing line of work [1-3]. To make the paper self-contained, we decided to summarize some of the key terms in our Preliminaries section. We appreciate the reviewer for their comments which have led us to make the paper more readable and self-contained, and we will try to clarify the questions left further below.

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.

A mixing matrix/matrix mixer is any matrix that represents a sequence transformation when applied to an input sequence X. Hence, the definition of a mixing matrix is quite broad, as many algorithms, such as self-attention, Mamba, and other structured matrices, can be viewed as applying their unique matrix mixer to the input. In general, one can use naive matrix multiplication to transform an input sequence with a matrix mixer; however, certain classes of matrix mixers, like Toeplitz or Mamba, have special structures that allow for more efficient matrix multiplication.

What does "matrix mixer layer" mean?

The matrix mixer layer is the layer that “hosts” the matrix mixer. In the case of the Transformer Phi teacher model, this would refer to the self-attention layer, which includes the input projections, the actual self-attention mechanism, and the output projection. For the Phi-Mamba student model, this would refer to the Mamba-2 layer, which encompasses the projections, convolution, SSM mechanism, gate, output projection, etc.

What are TeacherMixer and StudentMixer? Is this done layer by layer, block by block or are they all aligned at the same time?

We refer to the TeacherMixer and StudentMixer as the extracted matrix mixer from the matrix mixer layer. The Matrix Orientation step has the student model learn to approximate the teacher’s mixing matrix; in Phi-Mamba’s case, the stage would be minimizing the difference between the “unraveled” SSM matrix mixer, i.e., StudentMixer, (Equation 2 in the paper) and the self-attention matrix Softmax(QK.T), i.e., TeacherMixer, found in the respective teacher layer.

This is done layer-by-layer but can be seen as block-by-block in our Phi case, as each Phi block only has one self-attention matrix mixer layer.

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?

$**x**$ is the input to the teacher matrix mixer layer which is then set as the input to the respective student layer as well. The input is passed through both layers and then the difference between the extracted matrix mixer is minimized. The layer parameters for the student layer, $\phi$, are then updated via gradient descent. The dimensions of the model match Phi-1.5 in terms of layers, state size, attention heads, etc.

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.

The Hidden-State Alignment aims to minimize the difference between the output of the entire block, not just the matrix mixers. In our case, $\mathbf{u}$ is the same as $\mathbf{x}$, but depending on the design of the teacher block, this might not always be the case, so we decided to utilize a different variable. We apologize for any confusion this may have caused.

No code or trained model artifacts were provided nor were mentioned that will be provided.

As mentioned in the shared response, we are planning to publicly release the model code and the pretrained weights. The code for our model architecture will contain certain flags that allow for the return of elements, such as the matrix mixer of Mamba-2, which will make replicating the experiments easier.

Minor comment…

As mentioned in the shared response, we have fixed the presentation issues, e.g., typos, broken links, and unclear portions, of the paper.

[1] https://arxiv.org/abs/2405.21060

[2] https://arxiv.org/abs/2407.09941

[3] https://arxiv.org/abs/2402.18668