Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models

Authors: aviv  bick
kevin y li
eric p xing
j  zico kolter
albert  gu
phi  -mamba
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 01:00:27

Abstract: Transformer architectures have become a dominant paradigm for domains like language modeling but suffer in many inference settings due to their quadratic-time self-attention. Recently proposed subquadratic architectures, such as Mamba, have shown promise, but have been pretrained with substantially less computational resources than the strongest Transformer models. In this work, we present a method that is able to distill a pretrained Transformer architecture into alternative architectures such as state space models (SSMs). The key idea to our approach is that we can view both Transformers and SSMs as applying different forms of mixing matrices over the token sequences. We can thus progressively distill the Transformer architecture by matching different degrees of granularity in the SSM: first matching the mixing matrices themselves, then the hidden units at each block, and finally the end-to-end predictions. Our method, called MOHAWK, is able to distill a Mamba-2 variant based on the Phi-1.5 architecture (Phi-Mamba) using only 3B tokens. Despite using less than 1% of the training data typically used to train models from scratch, Phi-Mamba boasts substantially stronger performance compared to all past open-source non-Transformer models. MOHAWK allows models like SSMs to leverage computational resources invested in training Transformer-based architectures, highlighting a new avenue for building such models.

Main paper: 
 introduction: 
 Large language models based upon Transformer architectures have become a staple of natural language processing but suffer from their reliance on quadratic self-attention -the need to compute inner products between tokens at all positions up to the context length. This has motivated the development of several alternative subquadratic models, either approximations of self-attention [23] or entirely different architectures, such as state space models (SSMs) [13, 12, 30, 35] . Training strong subquadratic models such as SSMs can benefit the community through their cheaper finetuning and inference costs; however, they have not benefitted from the same amount of community effort in the form of training and compute as for Transformers. This raises a natural question: is it possible to leverage the vast amounts of resources that have been invested in training quadratic-time Transformers and use these models to produce stronger alternative models, such as state-space models?In this paper, we present an approach for training subquadratic state-space models (specifically from the class of Mamba SSMs [12] ) through the distillation of different elements of a pretrained Transformer model. The key intuition is viewing both Attention and SSMs as sequence transformations that mix different token embeddings by applying different classes of matrices across them. Sequence model architectures can then be factored into separate (i) sequence mixing and (ii) channel mixing blocks, e.g., a Transformer is composed of Attention (sequence mixer) and MLP (channel mixer) Figure 1 : Plot of trained token budget to averaged accuracy on Winogrande, Arc-E, Arc-C, PIQA, and Hellaswag on various open-source models (mainly non-Transformer-based models). Our model (Phi-Mamba) uses more than 33× less token budget to achieve 5% higher average accuracy than the next best model. blocks. Using this breakdown, we can separately distill the mixing elements of each model explicitly at different levels of granularity. Specifically, we propose a three-phase distillation process that progressively targets higher levels of supervision from the teacher model: (1) a matrix orientation phase that aligns the sequence transformation matrices themselves; (2) a hidden-state distillation that aligns the hidden-state representations of each individual layer of the network without sacrificing preexisting learned representations; and (3) an end-to-end training phase with weight transfer that finally distills the final output of the network using only a fraction of training data. We term our approach MOHAWK after these three stages (Matrix Orientation, Hidden-State Alignment, Weight-Transfer and Knowledge Distillation).We apply our approach to a modified instantiation of the Mamba-2 architecture [8] , termed Phi-Mamba, which is aimed at more directly corresponding to the different architectural blocks of the Phi-1.5 language model [15] -a very strong Transformer model at the 1.3B parameter scale. Using our approach, the Phi-Mamba model achieves performance on benchmarks stronger than any previous Mamba models of similar size. Although performance still lags behind that of the base Phi-1.5 model on these benchmarks, the model is distilled with only 3.0B tokens, less than 1% of the data used to train either the previously best-performing Mamba models and 2% for the Phi-1.5 model itself. For instance, our Phi-Mamba achieves a 71.7% accuracy on the Winogrande dataset, compared to the pretrained Mamba-2 model's 60.9% accuracy, and 44.1% accuracy on the ARC-C dataset, compared to Mamba-2's 33.3% accuracy. Our results highlight the benefit of our three-phase distillation approach: we show in ablation experiments that each phase is highly beneficial for the eventual performance of the model, and that, e.g., only attempting to directly distill the Phi-1.5 model (i.e., Phase 3 alone) substantially underperforms the full MOHAWK method. Moreover, our findings emphasize the benefits of state-space models while training on fewer than 100× tokens than the original pretrained Mamba model. 
related work: 
 Sequence Models. State-of-the-art autoregressive language models have been pretrained on massive amounts of data, resulting in models that exhibit extensive downstream capabilities, such as zeroshot translation and long-range reasoning [3, 15, 37] . Recent work has focused on addressing the quadratic complexity of Transformers by developing subquadratic alternatives based on RNN [30, 1] , SSM [12, 35] , and linear attention mechanisms [23, 43, 25, 7, 32] , highlighting the importance of efficient sequence models in the era of large-scale autoregressive language models. Distillation. Knowledge distillation can be used to transfer knowledge from a large teacher model to a smaller student model, resulting in a more efficient model that retains the performance of the teacher model [18] . Distillation has been applied to various language modeling tasks, such as text generation [5, 17] , machine translation [16, 48, 36] , and question-answering system [19, 44] . Distillation in language models has been largely focused on compression: turning a larger pretrained Transformer into a smaller one by utilizing the weights of the teacher model [41, 21, 42] . Some of the techniques proposed look similar to ours; for example, [41] match attention matrices in a step similar to our matrix orientation, and [24] align outputs of each block (i.e., the hidden states). However, these differ in subtle and important ways because of our setting; for example, the former uses a different loss function than us that relies on softmax attention, and the latter is an end-to-end objective while our hidden state alignment occurs completely independently block-per-block. Consequently, prior work has observed that combining these objectives does not actually help or even hurts distillation [21] , whereas we show that our techniques all significantly help improve the student model.A smaller body of work has focused on our objective of distilling across architectures, in particular, turning a pretrained Transformer into a different architecture (usually some recurrent model) of the same size. [22] converted a pretrained softmax attention into linear attention by directly transferring weights and continuing fine-tuning. A similar approach was taken by concurrent works for converting Attention into linear RNNs [27, 40] . Recently, [47, 46] also proposed distilling into linear attention by first matching attention matrices. Our approach differs by using a different loss function that works beyond linear attention; incorporating more stages (e.g., the hidden state alignment step); and using recent, more expressive classes of efficient student models (Mamba-2), which we show are significantly easier to distill into (Table 4 ). 
preliminaries: 
 To facilitate a clear understanding of our distillation approach, we start with the necessary background and definitions. An overview of the Mamba-2 architecture, which forms the foundation of our Phi-Mamba model, is also provided.Following [8] , we refer to an equivalent function that represents the input and output of a sequence model as a sequence transformation or a sequence mixer. Formally, Definition 1 (Sequence Transformation). We use the term to refer to a parameterized map on sequences Y = f θ (X) where X, Y ∈ R (T,P ) and θ is an arbitrary collection of parameters. T represents the sequence or time axis; subscripts index into the first dimension, e.g. X t , Y t ∈ R P .To put it differently, sequence mixers combine tokens at various time steps, facilitating the model's comprehension of temporal information and interactions. Sequence transformations form the foundation of deep sequence models, being integral components of neural network frameworks such as Transformers. A particular family of sequence transformations can be represented by Y = MX for a matrix M ∈ R (T,T ) , which we refer to as a sequence transformation matrix or matrix mixer. An example of such a matrix mixer is the vanilla self-attention, Softmax(QK ⊤ ), which is applied to the input-dependent V resulting in the familiar Softmax(QK ⊤ )V. Similarly, Linear Attention [23] has a sequence transformation matrix of the form K ⊤ . In addition, we can easily obtain their causal variants by multiplying by L, a lower triangular matrix filled with 1s, to obtain L • Softmax(QK ⊤ ) and L • QK ⊤ , respectively. Another example is a Toeplitz matrix T used to perform discrete convolution on input X, resulting in TX [31] .A naive approach to computing the output of a sequence transformation is to multiply the input sequence X by the matrix M. However, this approach has a time complexity of O(T 2 ), which is prohibitive for long sequences. Subquadratic sequence transformations, such as Mamba-2, have been developed to address such inefficiencies through structured matrix multiplication.Mamba-2 [8] , a type of structured state space models (SSMs) [13, 11] , was recently introduced. Similarly to the original Mamba model [12] , Mamba-2 uses a time-varying state-space model which can selectively focus on or ignore inputs due to its input-dependent parameterization of the system components. The time-varying SSM is defined as follows:EQUATIONHere, B t and C t are input-dependent projections of the system, as in Mamba-1; however, A t is the identity matrix I multiplied by a scalar α t . The above formulation also differs from the previous one by treating the underlying sequence as originating from a discrete signal instead of a continuous one and therefore omits the sampling component ∆t from the original Mamba model.Importantly, Mamba-2 draws a new connection between SSMs and Transformers, termed Structured State Space Duality (SSD), which shows that a special case of SSMs can be viewed as a form of causal linear attention. In particular, fixing A t = I (a further restriction of Mamba-2 to α t = 1) results in the formulation of causal linear attention [23] with the matrices B and C representing the projections of the key and the query, respectively, while the input projection X corresponds to the projection of the value.Mamba-2 as a matrix sequence transformation. Inspired by the aforementioned connection between SSMs and Transformers, [8] shows that Mamba-2's SSD mixer family is equivalent to sequentially-semi-separable matrices [4] . Formally, the SSD mixer family can be represented as:EQUATIONwhereα t:i = α t-1 • α t-2 • • • α i .An interesting observation is that the Mamba-2 architecture can be viewed as a causal linear attention with a learnable causal mask.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof? Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
methods: 
 Throughout this section, we will describe each phase of MOHAWK. Specifically, we will cover the stages of matrix orientation, hidden-state alignment, and knowledge distillation, all three of which are crucial for developing an effective student model from the pretrained Transformer model. Unlike traditional distillation techniques, the student model retains the overall architecture of the teacher model, differing only in the replacement of the attention matrix mixer with a subquadratic alternative. We will progressively unveil our architecture, Phi-Mamba, along with the specifics of its distillation process. This section concludes with an in-depth description of the Phi-Mamba architecture and its hybrid version, which surpasses the performance of other subquadratic matrix mixers. Further examinations of the effectiveness of the method and ablation studies are discussed in Section 5.For clarity, the term block refers to the repeating components that form the end-to-end model. The blocks are composed of layers, such as the self-attention layer (including projections), the SSM layer (including the mixer and convolution), and the convolutional layer. In this manner, many Transformer models, like Llama [37] , are viewed as a stack of alternating self-attention and MLP blocks, whereas the Phi and Phi-Mamba models are comprised of Phi blocks that have parallel Attention/SSM and MLP blocks.The first stage of MOHAWK aims to align the student matrix mixer with the teacher's self-attention matrix. Achieving this alignment is a two-step process: first, at every mixing layer, the student components preceding the matrix mixer are set to match the teacher's components. This ensures that each layer's input undergoes the same transformation up to the matrix mixer section. Consequently, the only variation from the input to the mixing process is the matrix calculation. We then minimize the distance between the matrix mixer, e.g., the self-attention matrix and the materialized SSM matrix (2), of each layer within the student and teacher models:EQUATIONwhere ϕ denotes the parameters within the student's sequence mixing layer, and u indicates any arbitrary input. In our experimental setup, u was chosen as the output from the teacher model's preceding layer to better mimic the input distribution to the layer. This stage ensures that the student and teacher models have roughly similar mixing layers and sets the foundation for the subsequent stages of the distillation process. In particular, this stage can be done in parallel across all the student layers, as the inputs to the student and teacher blocks are identical.For Mamba-2, we begin by setting the convolution to an identity function, effectively nullifying its initial impact. This results in the computation of the semi-separable matrix being the sole distinction between the layers. We then proceed to minimize the distance between the two matrix mixers: the semiseparable scalar identity and the attention matrix (see Figure 2 ). Figure 3 demonstrates the importance of this stage in the distillation process. Furthermore, Table 3 shows that the Mamba-2 matrix mixer is more expressive than popular alternatives and can closely approximate the selfattention matrix of various data samples across all layers of a Transformer model through gradient descent, solidifying it as a strong sequence mixer.Following the optimization of Equation ( 3), we must still address the differences between the outputs of the student and teacher blocks. To achieve this, we further align the components of the two blocks using initialization and distillation. Specifically, our goal is to match each student and teacher mixing blocks by minimizing the L2 norm of their output (e.g., the entire Mamba block with the self-attention block):EQUATIONwhere similar to Section 4.1, ϕ represents student's block parameters, and u is an input. Once again, this stage can be done in parallel across all the student layers.In the case of Mamba-2, we modify the remaining components to be identical to the Phi-1.5's Attention block, so that the overall functionality is preserved from Stage 1. Concretely, we initialize the gate (see Figure 2 ) to a constant value of 1 to "open" the gate, canceling its initial effect. In addition, we remove the normalization prior to the output projection, as it cannot be set to align with the Attention block. We then minimize the distance between the output of the Mamba-2 block and the output of the teacher's self-attention block. Our analysis indicates that the distance between the Mamba-2 block and the self-attention block is strongly correlated with the model's ability to learn the teacher's distribution, as shown in Table 3 . Furthermore, Figure 3 shows that a better independent alignment of the student and teacher blocks results in performance improvements, highlighting the importance of this stage in the distillation process.The final stage of the distillation process aims to fine-tune the student model to match the performance of the teacher model. Although each student mixing block is aligned with its corresponding teacher mixing block, discrepancies are still present between consecutive blocks throughout the network To bridge these gaps and address the remaining components of the language model, we transfer the remaining weights of the teacher model to the student's respective components. For Phi-Mamba, this involves the token embedding, the final layer normalization, the Language Model head, and the MLP and input norm at each block (see Figure 2 ). We then fine-tune the complete end-to-end student model under teacher supervision. Concretely, we use a distillation loss to encourage the student model to mimic the distribution of the teacher model's logits, also known as knowledge distillation [18] :EQUATIONwhere x is the input tokens to the models.It has been hypothesized that much of the information stored in language models resides in MLP blocks [28] . To utilize the work already done pretraining the teacher, MOHAWK adjusts the structure of the student blocks to utilize the MLP in the same way as the teacher model, effectively swapping the teacher's matrix mixer with that of the student.Interestingly, during this step, the MLP weights can be kept frozen while keeping the model performant. This showcases Mamba-2's powerful expressiveness crucial for replacing Attention, cuts the [8] that omits the non-linear activation function after the convolutional operation and the layer normalization present before the output projection, so that the parts of the model outside the matrix mixer can be transferred from the teacher model. The MOHAWK distillation process involves progressively matching fine-to-coarse parts of the model to the corresponding part of the teacher model: (1) the mixer mixer itself (2) the full Mamba vs. Attention blocks, and (3) the end-to-end model. number of trained parameters by more than half, and, in larger models, helps prevent the student model from experiencing catastrophic forgetting of the teacher model's information. We validate Mamba-2's ability to do so in Table 5 .Combining the three stages of MOHAWK, we introduce the Phi-Mamba architecture, which merges the Mamba-2 model of [8] with the Phi-1.5 Transformer model of [15] . It consists of a stack of Phi-Mamba blocks (Figure 2 ), initialized and distilled as described in previous sections.Overall, the Phi-Mamba architecture, as depicted in Figure 2 , differs from the vanilla Mamba-2 architecture by modifying the structure of the SSM matrix mixer, removing components from the SSM block and incorporating dense layers from the teacher model. In particular, each Mamba-2 block was modified by removing post-convolution activation and pre-output projection normalization, while setting the gate and convolution to be identity functions. Interestingly, although these components were found to be beneficial for performance when Mamba-2 was trained from scratch [8] , we find that they are unnecessary for our distillation process.Two key changes were made to the Mamba-2 matrix mixer. The first was converting the SSM head structure from multi-value to multi-head, much like the multi-head attention mechanism found in Transformers [39] , enabling the independent distillation of each Transformer head into a Mamba head. Moreover, we handle the sequence mixer as entirely discrete-time by making the A matrix a projection of the input and eliminating the ∆ discretization parameter. Although this formulation slightly differs from Mamba-2, the original algorithm can still be applied as a black-box method. 
empirical validation: 
 We empirically validate the MOHAWK framework is able to achieve better performance on various downstream benchmarks compared to previous subquadratic models of similar size. Our final Phi-Mamba-1.5B model is distilled on 3 billion tokens (distributed as 80M in Stage 1, 160M in Stage 2, and 2.76B tokens in Stage 3 as described in Appendix A) from the C4 dataset, with a sequence length of 2048. This constitutes less than 1% of the resources used by many top-performing subquadratic open-source models (e.g., the original Mamba-1/2 models pretrain on 315B tokens).Table 1 presents a comprehensive breakdown of downstream evaluation results for our models and multiple baselines on a standard set of commonsense reasoning and language understanding tasks: WinoGrande [33] , HellaSwag [45] , PIQA [2] , ARC-Challenge and ARC-Easy [6] , and Table 1 : Downstream evaluation results for full methods, comparing Phi-Mamba against open-source models of similar sizes pretrained on standard language modeling corpuses. Phi-Mamba attains performance close to the teacher model and better than all pretrained models, while using less than 1% of the training data. LAMBADA [29] . Figure 1 shows the performance versus the training cost of Phi-Mamba compared to many open-source baselines from the literature at similar model sizes.For the remainder of this section, we will analyze the impact of the 3 stages of MOHAWK one by one. Throughout the experiments detailed in this section, we use the AdamW optimizer with β = (0.9, 0.95), a weight decay of 0.1, and a learning rate of 1 × 10 -4 , combined with a Warmup-Stable-Decay (WSD) scheduler featuring 10% warmup and 10% decay. The training law figures and the final Phi-Mamba model use the regime detailed in Appendix A.As described in Section 4.3, this phase employs a simple end-to-end distillation of teacher-model logits. It leverages the alignment among all sequence mixers and successive blocks to jointly fine-tune all components of the network. Experiments shown in Table 2 highlight the relevance of implementing this end-to-end alignment, with all three architectures achieving their highest scores only after this phase. Predictably, the impact of end-to-end alignment varies by architecture: models with more mixing layers similar to the teacher model see a reduced importance of this phase.Stage 3 is the only stage in MOHAWK that trains the student model end-to-end and can be seen as the "main" stage. Many distillation methods employ only this stage; however, Table 2 shows that using only end-to-end knowledge distillation is less than ideal. Although it is slightly advantageous to use only Stage 3 compared to only Stage 2, there is a significant gap between using only Stage 2 versus using Stage 2 + 3. As elaborated in Section 6, this phase can freeze all network components except the Mamba-2 sequence mixer without a significant performance drop. This in particular indicates that the third stage (like the other stages of MOHAWK) can operate in computationally limited settings.Following the analysis of the model's end-to-end distillation in Stage ). These lower hidden states are also correlated with lower starting perplexities, which in turn are correlated with downstream performance, as shown in Figure 5 . Furthermore, Table 2 shows the synergy between Stage 2 and Stage 3, as applying Stage 3 on top of Stage 2 outperforms vanilla knowledge distillation, highlighting the importance of incorporating both hidden-state alignment and knowledge distillation methods for the tested architectures.Motivated by our previous finding, we then analyze how matching the matrix mixers can decrease the overall mixer block's hidden-state distance with the teacher model even further. Similarly to our previous protocol, we assess the positive impact of the current stage on the following phase's metrics and final model's performance by comparing models with varying amount of Stage 1 and Stage 2 training on both stage metrics.Figure 4 shows that even with constrained budgets, performing Stage 1 for a small period can help with subsequent stages and their performances. Thus, even a small amount of Stage 1 training can help their respective Stage 2 models reach better hidden-state distances compared to the fromscratch counterpart. This is despite the phenomenon that the teacher and student mixers diverge and then re-converge in Stage 2 after mixer similarity is no longer directly optimized. Coupled with Section 5.2, which discovers that lower hidden state initializations lead to better perplexity and downstream performance, it can be inferred that Stage 1 aids the overall distillation process. We further empirically validate this intuition in Table 2 , which indicates that this stage aligns the matrix mixers to a stronger degree than only the hidden state alignment. 
self-attention approximation with structured matrix mixers: 
 We start by testing the ability of various matrix mixer families to match the empirical self-attention matrices of a pretrained Transformer. We take 1000 samples from each layer of a Llama2-7b-Chat model [37] , materialize the attention matrices, and project them onto given classes of structured matrices. The results in In particular, to describe the class of linear attention matrices (3.1), we use the fact that Q and K are projections of the input x ∈ R din onto R dout , and therefore their rank is bounded by min {d in , d out }.For multihead linear attention, d out (also known as head dimension) is typically a small value (e.g., Phi-1.5 and Llama2-7b-Chat have head dimensions of 64 and 128, respectively). Thus, we approximate this family of sequence mixers using causal low-rank matrices L • QK ⊤ , where L is a lower-triangular causal mask of 1s, and Q, K are in R n×d with d ≪ n (indicating that the head dimension is substantially smaller than the sequence length).To describe the multi-head Mamba-2 matrix family, we utilize the state space dual (SSD) layer (3.2) in a manner similar to the previous linear attention, but now the causal matrix L possesses an n-degree rolling multiplicative structure for SSD which can be seen as a more expressive mask that generalizes the causal mask (Section 3.2).Both causal low-rank and SSD matrix families were approximated with 10,000 steps of gradient descent per sample. To approximate the general class of SSM matrix mixers, we utilize balanced truncation, a gradient-free projection algorithm. This method is mainly known in the field of time-invariant Dynamical System model reduction [14] and has been modified for use in timevarying systems [34] . Similarly, for the family of causal Toeplitz matrices, representing convolution operations, we employ a simple heuristic that minimizes the error for each attention matrix.Table 3 shows that while the SSM matrix family provides the closest approximation to the selfattention matrix mixer, the Mamba-2 mixer family (SSD) has just twice the distance from the SSM matrices. This is in contrast to Linear Attention, which has three times the distance, all while keeping a computational cost on par with Linear Attention. More details can be found in Appendix C.We further validate the ability of a Mamba-2 block to replace an Attention layer within a language model. Firstly, we create two variants of our architecture, Phi-Toeplitz and Phi-LR, and run the MOHAWK process for 1B tokens at each stage (see Table 4 ) to verify that the previous finding hold in a multilayer, end-to-end model case. Secondly, we run MOHAWK while freezing various parts of the Phi-Mamba modules (refer to Table 5 ), revealing that limiting the trainable elements to the Mamba-2 blocks (excluding the embedding, head and all MLP layers) results in only a minor performance decrease during MOHAWK distillation.Interestingly, in all of the aforementioned experiments, we have found a consistent correlation between the projection distances of the matrix (Frobenius distance) in Table 3 and the downstream performance metrics (accuracy) in Table 4 . Essentially, a better matrix approximation (lower Frobenius distance) is correlated with better model performance (higher accuracy) on various tasks. This connection highlights the relationship between the quality of the matrix approximation and the performance of the model. Such findings are echoed in [20] , which find that more expressive matrix mixers lead to more performant models, e.g., Low-rank-based BERT models outperform Toeplitz-based ones. 
discussion and conclusion: 
 Our experiments shows that the Mamba-2 model can be successfully distilled from a pretrained Transformer teacher model, utilizing its extensive knowledge learned from custom datasets and higher computational resources. Despite using less than 100× data compared to many open-source models, including Mamba, our subquadratic model outperforms other subquadratic models in various benchmark tests by a wide margin.The MOHAWK framework's multi-stage process which gradually increased the scope of distillation is essential extracting the teacher model's knowledge to the fullest extent as shown in our ablations and training laws. We continue to find the effectiveness of MOHAWK when distilling hybrid Attention-SSM models and provide ablations on the number and position of Attention layers.Additionally, we demonstrate that Mamba-2's relationship to Transformers is evident not only in theory, but also in practice, as it captures interactions similar to those of Transformers, and is able to replace Attention with little drop in performance. Coupled with past research which has posited that much of a language model's knowledge is embedded in the MLP blocks, we believe that any subquadratic model with a sufficiently expressive matrix mixer can replicate the behavior of pretrained Transformers, bringing quadratic knowledge to subquadratic models. We recommend further research to explore the role of sequence mixing layers in subquadratic models and their impact on performance. Advancements in both the distillation process and the sequence mixer architecture could lead to further improved performance in a range of tasks. We propose that "trainability" and "distillability" are distinct properties of the models, and therefore, distillation techniques should be more appropriately tailored to the model. 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [NA] Justification: N/A Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: The open source assets that we use have been credited.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.• If this information is not available online, the authors are encouraged to reach out to the asset's creators. 13. New Assets Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets? Answer: [NA] Justification: N/A Guidelines:• The answer NA means that the paper does not release new assets.• Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc. • The paper should discuss whether and how consent was obtained from people whose asset is used. • We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix: Hyperparameter Search To construct Appendix A, we performed grid searches for training in Stages 1, 2, and 3 independently from scratch to find the optimal hyperparameters. We explored learning rates lr = {1, 2, 5} × 10 {-3,-4} and batch sizes 2 {15,16,17,18} . AdamW Optimizer was used with β = (0.9, 0.95), incorporating a weight decay of 0.1, gradient clipping at 1.0, and a Warmup-Stable-Decay (WSD) scheduler with 10% warmup and 10% decay utilizing linear warmup and cooldown functions. Automatic mixed precision training to bf16 was used in all stages. For Stages 1 and 2, we initially fixed the batch size at 2 16 , then varied the learning rates. After identifying the optimal learning rate, we adjusted the batch sizes and subsequently finalized the learning rate after fixing the batch size. Multi-Stage Distillation Procedure In the development of the training law (see Figure 3 ), we executed a single "continuous" run initialized from a state that included several checkpoints. The warm-up period was determined as 10% of the tokens processed during the continuous run. For instance, if the model's goal was to process 640 million tokens, and it started from a run that had processed 40 million tokens, then the warm-up would be set at 60 million tokens. The checkpoints recorded during the warm-up phase were preserved as they were, while subsequent checkpoints underwent a cooling of 10% of the current phase. To illustrate, in the scenario mentioned earlier, a checkpoint at 320 million tokens during the 40M to 640M run would maintain the original warmup, while the cooldown would span 28 million tokens. Conversely, a checkpoint at 80 million tokens within the warm-up phase would be saved without any cooldown.Training Laws on Downstream Metrics Figure 5 extends the Stage 2 versus Stage 3 comparison in Figure 3 , except we measure average accuracy on downstream metrics instead of perplexity. We observe a strong correlation between the training laws of perplexity and downstream evaluation metrics. While the general trend indicates that models exposed to more tokens during the prior stage initialization tend to perform better on both perplexity and downstream metrics, the relationship is not perfectly aligned. Specifically, the order of model performance based on perplexity does not always match the order based on downstream metrics, highlighting some differences in how these metrics capture model effectiveness.Training the Final Phi-Mamba Model After confirming the importance of the stages in Section 5.1, Section 5.2, and Section 5.3, we proceed to distill the final Phi-Mamba model using the three elements of MOHAWK. We use 80M tokens for Stage 1, due to the strong performance of the token count in both the matrix and hidden state distances (Figure 4 ). Stage 2 was distilled for 160M tokens given the apparent saturation of both hidden state distance and perplexity compared to the other initialization states, such as 10M, 20M, 40M, etc. (Figure 3 ). We employed Stage 3 to a total of 3B tokens across all stages and observed that the previously optimal learning rate applied for training training laws led to instabilities in training, particularly spikes in evaluation perplexity. Decreasing the learning rate for Stage 3 mitigated this issue as mentioned above. We hypothesize that the instability is due to the Stage 1 + 2 initialization's Mamba component being quite similar to that of the teacher model, so a large learning rate coupled with disconnect between blocks, which are mended in Stage 3, can cause training instabilities. The performance of the final model is reported in Table 1 .As noted previously Section 4.4, our Mamba-based sequence mixer is slightly modified from the original to make it more amenable for distilling from a Transformer architecture. In particular, the Mamba-2 sequence mixer is treated entirely in discrete time by projecting the input onto the matrix A and removing the discretization parameter ∆. This section serves as a complement to Section 6 and outlines the methods employed to create Table 3 . Appendices C.1 to C.5 describe our strategies for finding a matrix within the specified families that closely approximates the original attention matrix using a selected distance metric. Formally, we consider the following optimization problem:EQUATION)where M is the subspace of a specific matrix family, M is the attention matrix, and ∥•∥ corresponds to a selected distance metric. In the following sections, we explore different methods and matrix families for this optimization problem.Considering a time-varying system denoted by{A k , B k , C k , D k } k∈[l], we can describe it using the matrix mixer T (also known as the transfer matrix) as follows:T =         D 1 0 0 0 0 • • • 0 C 2 B 1 D 2 0 0 0 • • • 0 C 3 A 2 B 1 C 3 B 2 D 3 0 0 • • • 0 C 4 A 3:2 B 1 C 4 A 3 B 2 C 4 B 3 D 4 0 • • • 0 . . . . . . . . . . . . . . . . . . . . . C l A l-1:2 B 1 C l A l-1:3 B 2 C l A l-1:4 B 3 C l A l-1:5 B 4 • • • C l B l-1 D l         With A k ∈ R n×n , B k ∈ R m×n , C k ∈ R p×n , and D k ∈ R p×m ,where n is the state dimension, m the input dimension, and p the output dimension. As T's form corresponds to a semi-separable matrix (i.e., each sub-matrix has a rank of up to n), we will label this matrix form as SSM with state size n throughout the remainder of this appendix, representing a state-space model with state size n or a semi-separable matrix of order n.Every matrix can be represented as a linear combination of rank-one matrices. Thus, the attention matrix M ∈ R l×l can be interpreted as an SSM with a state size of up to l ≫ n. Consequently, we can employ prior research on time-varying model order reduction [9, 26, 38] to reduce M to an SSM with a smaller state size n. Specifically, we utilize the following SVD-based approximation:Algorithm 1 Approximation of Attention Matrix M as an SSM with State Size n Input: Attention matrix M ∈ R L×L , state size n Output: Approximated attention matrix M ∈ R L×L Procedure:1. For k = 1, . . . , L -1:1.1 Define H k as the submatrix of M below and to the left of entry M k,k : Note that the diagonal elements of M were not subject to approximation in Algorithm 1 as they remain unchanged.H k =    M k,1 • • • M k,k-1 . . . . . . . . . M l,1 • • • M l,k-1    1.Although this approximation method for the semi-separable matrix is heuristic, it has been empirically shown to deliver good results. For further details on approximation methods for semi-separable matrices, and the theoretical background behind them, we refer the reader to [10, 26] Given a set of self-attention matrices, we tried to find how close an causal low-rank matrix could approximate M = Softmax(QK ⊤ ). To ensure the state size N , or in this case rank of N , of the approximation M , we composed M = L • AB ⊤ where A, B ∈ R D,N , L is a R D,D lower triangular mask, and D = 512.We used the results from our causal low-rank (LR) experiments to inform much of our experimental design for later gradient descent-based approximations, which include both SSD classes (with and without D matrix) and RetNet. We experimented with various different low-rank approximation solvers. We found that gradient descent performed better than alternating gradient descent. Both types of gradient descent were better than alternating least-squares which often times reached less than optimal local minima. Causal low-rank matrix approximation can also be seen as a softer version of the low-rank matrix completion problem, but a semi-definite programming (SDP) approach was not able to outperform standard gradient descent.Due to our LR approximation requiring gradient descent, we selected the number of steps in relation to the time required to calculate the semi-separable approximation of the same matrix. Given the heuristic approach for converting self-attention matrices to a semi-separable form (Appendix C.1) and its ability to be parallelized, we selected the number of steps for gradient descent based on the time it took to run an entire batch of matrices (32) using gradient decent on causal low-rank versus one matrix using the semi-separable heuristic. After testing with the state sizes N = 16, 32, 64, we found that 10,000 steps suitable as it was around a factor of 5× compared to SSM. The 10,000 steps was maintained across all gradient based approximation classes (SSD, SSD without D, and RetNet). Experiments using the finalized step count showed AdamW provided better results compared to SGD/Adam, and the use of a scheduler provided little gain.During the experiments, we also found that initialization of the matrices A, B played a significant role in the resulting approximation difference. The original A, B values were sampled from [0, 1); however, given ∀M ij ≤ 1, i, j ∈ [D] due to the SoftMax operator, A, B values was then sampled from 0, 1 √ 512N to have the last row of the self-attention matrix be uniform probability. We then proceeded to vary the factor of the range exponentially, testing 0, 1√ 512N * 2 {-2,-1,0,1,2,4,8}where we found 0, 1 √ 512N * 2 4 provided the best initialization across multiple datasets. A normal distribution with µ = 0 and σ 2 with the above tested values performed worse than the uniform distribution. Initialization experiments were conducted using the AdamW optimizer with a learning rate of 0.001 and the standard 10,000 steps. This and subsequent gradient descent classes use the same initialization for their A, B matrices. For all gradient descent experiments in Table 6 , Three learning rates 0.1, 0.01, 0.001 and AdamW were used for each combination of matrix class, state size, and dataset, with the best approximation being documented. The Frobenius matrix norm was used as the loss function. Our Toeplitz approximation technique calculates the matrix approximation by setting the value of each band of the Toeplitz matrix as the average of the values of the respective band in the attention matrix. Since each band in a Toeplitz matrix is constant along its diagonal, this method ensures that the approximation preserves the structure of the original matrix while maintaining computational efficiency.To justify this approach, we observe that taking the mean per band minimizes the L2 norm (i.e., the sum of squared differences) between the original attention matrix and the approximated Toeplitz matrix. Specifically, for each band, the optimal value that minimizes the L2 difference between the two matrices is the average of the elements in that band. This is because the mean is the value that minimizes the sum of squared deviations for a set of numbers. As such, using the mean ensures that the approximation is as close as possible to the original matrix in terms of L2 distance, thereby providing a robust and efficient approximation method.As before, we assume that the approximation is input-dependent, meaning that each attention matrix has its own unique Toeplitz approximation.The segsum operator computes the sum of elements across specified segments of a matrix, which, as applied in Appendices C.3 and C.4, corresponds to summing over the columns. This operation is crucial for various matrix manipulations, including the computation of the state-space dual (refer to Equation ( 2)). Below is the Python implementation of the 'segsum' operator using PyTorch.Listing 5 PyTorch implementation of the Segmented Summation (segsum) operator. Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 


Decision: The work provides a novel method for distillation from Transformer based models to Mamba based SSMs. The method performs progressive distillation of the mixing mechanism of Transformers into an SSM, by applying a multi-stage algorithm for converting the token mixing of Transformers, which operates in quadratic time, into a sub-quadratic model. The authors demonstrate the efficacy of their method by distilling phi-1.5, a small Transformer, into an equivalent Mamba-based model, with very small amount of training data.

The reviewers raised some concerns regrading the empirical evaluation in the paper, the presentation of the results and pointed some methodological problems. However, it seems that the authors addressed most of the concerns in the rebuttal period. In light of this, I recommend that this paper is accepted.

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: - 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: 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?
limitations: 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?
paper_summary: paper_summary: The paper distills a Transformer model into the Mamba architecture by using about 1% of the pertaining dataset.
strength_weakness: strength_weakness: 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: questions: See Weaknesses.
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: 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: questions: See weaknesses part.
limitations: limitations: The authors have discussed the potential limitations in Section 6.

Comments: title: Thank you
comment: Thank you for the response! In light of most of my concerns were related to the clarity of the presentation and the authors saying that they would improve that, I increased my score. This is an interesting paper with very good results and significant potential for impact. However, whether this is realized or not depends on how easy it would be for others to apply your methods. For this reason, I would strongly urge the authors to take the opportunity to improve the presentation for the camera-ready version.
Thank you for the thorough response. After reading the rebuttal and other reviews, I decided to change my score to 6 (Weak Accept).
The reviewer thanks the authors for the discussion. It solved most of my concerns. I decide to keep my score.