Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models
Authors: Aviv Bick
Kevin Li
Eric P. Xing
J Zico Kolter
Albert Gu
Venue: NeurIPS 2024 poster (2024)
Keywords: distillation
mamba
sub-quadratic
ssm
state-space-models
Primary area: deep_learning_architectures
TL;DR: We show that we can leverage pre-trained Transformers to create highly performant state space models using very little training data

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.

Transformers to SSMs: Distilling Quadratic Knowledge to Subquadratic Models
Aviv Bick Kevin Y. Li Eric Xing J. Zico Kolter', Albert Gu
CCannegie Mellon University, MBZUAI,, $Cartesia.ai fabick, kyl2)0cs cmu. cmu.ddu
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 distill pretrained Transformer architecture into alternative archi- tectures such as state space models (SSMs). The key idea to our approach is that view both Transformers and SSMs applying different forms of mixing matrices over the token sequences. We can thus progressively distill the Trans- former architecture by matching different degrees 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 distill a Mamba-2 variant based on the Phi-1.5 architecture (Phi-Mambe) using only 3B tokens. Despite using less than 1 % of the training data typically used to train mod- els from scratch, Phi- mhi-a boasts substantially stronger performance compared all past open-source non- non-Transformer models. MOHAWK allows models like SSMs to leverage computational resources invested in training Transformer-ban architectures, highlighting new avenue for building such models.
Introduction
Large language models based upon Transformer architectures have become staple of natural language processing but suffer from their reliance on quadratic self-attention - the need to compute inner products between tokens development of several alternative subquadratic models, either approximations of self-attention [23]
positions up to the context length. This has motivated the
or entirely different architectures, such state space models (SSMs) [ 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 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 state-space models?
In this paper, we present approach for training subquadratic state-space models (specifically from the class of Mamba SSMs (12)) through the distillation of different elements of a pretrained Trans- former model. The key intuition viewing both Attention and SSMs as sequence transformations that mix different token embeddings by applying different classes of matrices across them. Sequence model archiectures 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)
Authors contributed equally to this work.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Token Budget vs Common Sense and Language Understanding
65
Average Accuracy (%)
60
55
50
1
10
100
1000
Token Budget In Billions (log scale)
Phi-Mamba
RIVKV6-1.60
RWKV4-1.58
HGRN2-18 -B
Mambal: 1.4B
RWKVS 1.5B
XLSTM-1.46
GLA-1.38
Mamba2-1.38
Pythla-1.44
HGRNI-16
RetNet-11.38
Figure 1: Plot of trained token budget to averaged accuracy on Winogrande, Are-E, Arc-C, PIQA, and Hellaswag various open-source models (mainly non- monn-raanfoomr-bbeeee 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 three-phase distillation process that progressively targets higher levels of supervision from the teacher model: (1) matrix orientation phase that aligns the sequence transformation matrices themselves; (2) hidden-state distillation that aligns the hidden-state representations of each individual layer of the network without sacrificing pre- existing learned representations; and (3) an end-to-end training phase with weight transfer that finally distills the final output of the network using only fraction of training data. We term our approach MOHAWK after these three stages (Matrix Orientation, Fidden-Statt Alignment, Weight-Transfer and Knowledge Distillation).
Transfer
We apply our approach to a modified instantiation of the Mamba-2 architecture 8 termed Phi- Mamba, which is aimed more directly corresponding the different architectural blocks of the Phi-1 1.5 language model [33 - 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 71.7% accuracy on the Winogrande dataset, compared to the pretrained Mamba-2 model's 's 60.9% accuracy, and 44.1 % accuracy on the ARC-C dataset, compared 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 100x tokens than the original pretrained Mamba model.
Related Work
Sequence Models. State-of-the -he-thh-rrt autoregressive language models have been pretrained on massive amounts of data, resulting in models that exhibit extensive downstream capabilities, such as zero-
shot translation and long-range reasoning 15 Recent work has focused on addressing the
quadratic complexity of Transformers by developing subquadratic alternatives based on RNN [30 I.
SSM 12 35). and linear attention mechanisms [23 25| 7 R2,, highlighting the importance of
efficient sequence models in the of large-scale autoregressive language models.

---

Distillation. Knowledge distillation can be used to transfer knowledge from a large teacher model to smaller student model, resulting in more efficient model that retains the performance of the teacher model [8]. Distillation has been applied to various language modeling tasks, such as text
generation 177. machine translation [6
48
36], and
system [9]
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 [4] 21 42 Some of the techniques proposed look similar to ours; for example, [4 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 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-bile per--bee-........ 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.
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. 12 converted pretrained softmax attention into linear attention by directly transferring weights and continuing fine-tuning. A similar approach was taken by concurrent works for converting
also proposed distilling into linear attention by first matching attention matrices. Our approach differs by using 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
Attention into linear RNNs [27 40 Recently, 47 46
3 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.
3.1 Matrix Mixers
Following [. we refer to an equivalent function that represents the input and output of a sequence model sequence transformation or a sequence mixer. Formally, Definition (Sequence Transformation).
use the term to refer to a parameterized map on sequences Y = fx( where X, Y E (TT.P) and 0 is an arbitrary collection of parameters. T represents the sequence or time axis; subscripts index into the first dimension, e.g. X1,Yi E I...
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 foun- dation of deep sequence models, being integral components of neural network frameworks such as Transformers. particular family of sequence transformations can represented by Y MX for matrix M (TT),, 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 KT. In addition, we can easily obtain their causal variants by multiplying by L, lower triangular matrix filled with 1s, to obtain L Softmax QKT) and L QK respectively. Another example is Toeplitz matrix T used to perform discrete convolution on input X, resulting in TX [31].
A naive approach to computing the output of sequence transformation is to multiply the input sequence X by the matrix M. However, this approach has a time complexity of O which is prohibitive for long sequences. Subquadratic sequence transformations, such as Mamba-2, have been developed address such inefficiencies through structured matrix multiplication.
3.2 Mamba-2
was recently introduced. Similarly to the original Mamba model [1], Mamba-2 uses time- varying state-space model which
Mamba-2 [ a type of structured state space models (SSMs) K3
3

---

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:
h++1 = A,hn + B
//+ Cyhl
Here, B, and C₁ are input-dependent projections of the system, in Mamba-l, I; however, A, the identity matrix I multiplied by a scalar dt, The above formulation also differs from the previous one by treating the underlying sequence as originating from a discrete signal instead of continuous one and therefore omits the sampling component At from the original Mamba model.
Importantly, Mamba-2 draws State Space Duality (SSD), which shows that a special case of SSMs can viewed as a form of causal linear attention. In particular, fixing At I (a further restriction of Mamba-2 o 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 the projection of the value.
connection between SSMs and Transformers, termed Structured
Mamba-2 matrix sequence transformation. Inspired by the aforementioned connection between SSMs and Transformers, [ shows that Mamba-2's SSD mixer family is equivalent to sequentialy-sem---sppaaabee matrices Formally, the SSD mixer family can be represented as:
Q11
0
0/2:1
02
0
0
h411 04 -he + Bxg y: C. h4
033:1
0/3:2
BC.B
0n:2
an:3
where ats 0--- 04-2 ag.. An interesting observation is that the Mamba-2 architecture can be viewed as a causal linear attention with learnable causal mask.
Methods
Throughout this section, will describe each phase of MOHAWK. Specifically, we will cover the stages matrix orientation, hidden-state alignment, and knowledge distillation, all three of which crucial for developing 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 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
For clarity, the term block refers to the repeating components that form the end-to-end model. The blocks are composed of layers, such 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), viewed 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 ML.P blocks.
4.1 Stage 1: Matrix Orientation
The first stage of MOHAWK aims 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 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:
min V'eacherMMixer(uu StudentMMixerg(uJlll
(3)

---

where Φ 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. Shows that the Mamba-2 matrix mixer is more expressive than popular alternatives and can closely approximate the self- attention matrix of various data samples across all layers of a Transformer model through gradient descent, solidifying as a strong sequence mixer.
4.2 Stage 2: Hidden-State Alignment
Following the optimization of Equation (3). we must still address the differences between the outputs of the student and teacher blocks. To achieve this, 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);
min ||tttBlock(( (u) Student (u)ck,(u)ll
where 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, 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, remove the normalization prior the output projection, as it cannot be set 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 showw that better independent alignment of the student and teacher blocks results in performance improvements, highlighting the importance of this stage in the distillation process.
4.3 Stage 3: Weight-Transe Transfer and Knowledge Distillation
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, transfer the remaining weights of the teacher model to the student's '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, use a distillation loss encourage the student model to mimic the distribution of the teacher model's d logits, also known as knowledge distillation (18):
min LCx (Teacher Modeeeeoooooooo.ssstttttttadado
where x is the input tokens the models.
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 MIP weights can be kept frozen while keeping the model perfor- mant. This showcases Mamba-2's powerful expressiveness crucial for replacing Attention, cuts the
5

---

XN
projection
Soquence
SSM
Nonilnearg
Stage Marri:
Conv
Mamba
PhiMLP
Attention
PhiMLPP
Or entation
Alignmenc
Stage Knowladige Disiilation
Mamba
Phi-Mamba
Phi
Figure 2:
The Phi-Mamba architecture consists of a stack of blocks, each of which contains a Mamba block and MLP block. The Mamba block simplified version of the Mamba-2 block B that omits the non-linear activation function after the convolutional operation and the layer normalization present before the output projection, be transferred from the teacher model. The MOHAWK distillation process involves progressively matching fine-to-coarse parts of the model 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.
that the parts of the model outside the matrix mixer can
number of trained parameters by more than half, and, in larger models, helps prevent the student model from experiencing catastrophic forgetting the teacher model's information. We validate Mamba-2's ability to do so in Table
4.4 Phi-Mamba architecture
Combining the three stages of MOHAWK, introduce the Phi-Mamba architecture, which merges the Mamha-2 model of 8 with the Phi-1.5 Transformer model of [3). 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. particular, each Mamba-2 block 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 found to be beneficial for performance when Mamba-2 was trained from scratch [8), 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 Mamba head. Moreover, we handle the sequence mixer as entirely discrete-time by making the matrix a projection of the input and eliminating the A discretization parameter. Although this formulation slightly differs from Mamba-2, 2, the original algorithm can still be applied as a black- box method.
5 Empirical Validation
We empirically validate the MOHAWK framework is able to achieve better performance various downstream benchmarks compared to previous subquadratic models of similar size. Our final Phi- Mamba-1.51 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 perfoffrnming subquadratic open-source models (e.g., the original Mamba- 1/2 models pretrain on 315B tokens).
Table presents comprehensive breakdown of downstream evaluation results for 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 16), and

---

Table Downstream evaluation results for full methods, comparing Phi-Mamba against open-source source models of similar sizes pretrained on standard language modeling corpuses. Phi-Mambe attains performance close to the teacher model and better than all pretrained models, while using less than 1% of the training data.
MODEL
TOKENS
WINOG.
ARC-E
ARC-C
PIQA
HELLAS.
LAMB.
AvG. f
Phi-1.5-1.3B
150B
73.4
75.6
48.0
76.6
62.6
53.4
64.9
Phi-Mamba-1.5B
3.0B
71.7
74.0
44.1
75.5
60.2
50.1
62.6
1-1.4-
315B
61.5
65.5
32.8
74.2
59.1
64.9
59.7
Mamha-2-1.3B
315H
60.9
64.3
33.3
73.2
59.9
65.7
59.6
Finch-1.6B
1.1T
59.4
64.2
72.6
57.3
XLSTM-114B
300B
60.6
64.3
32.6
74.6
60.9
57.8
58.5
Eagle-1.5B
1.1T
59.1
64.3
33.5
71.1
55.0
65.7
58.1
Pythia-1.4B
300B
57.3
60.6
26.0
71.1
52.1
61.6
54.8
RWKV4-1.58
330B
54,6
60.5
29.4
72.4
52.5
56.4
54.3
DeltaNet-1.3B
100B
53.6
57.2
28.3
71.2
50.2
48.9
51.6
GLA-1.3B
100B
53.9
57.2
71.8
49.8
51.0
Table 2: MOHAWK distillation from Phi-1.5 teacher model to Phi-Mamba-1.5B "Stages Applied" details which of the three MOHAWK stages was performed, highlighting the importance of each stage. All experiments executed using fixed amount of 5B tokens for the entire distillation process.
STAGES TYPE
WINOG. APPLIED
ARC-E
ARC-C
PIQA
HELLAS.
LAMB. ACC f
AVG.
ACC f
ACC
ACC f
Acc
Acc
2
55.9
75.4
38.0
75.2
56.6
18.9
53.3
3
62.8
64.3
27.8
75.6
52.6
43.8
54.5
2-3
72.3
75.0
40.8
75.2
59.7
50.6
62.3
1-3
74.8
72.7
43.5
75.6
59.6
49.2
62.7
LAMBADA 29 Figure T shows
Ishows 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 3 = (0.9, 0.95), weight decay of 0. and a learning rate of 1 x 10-1 combined with a Warmup- Stable- Recy (WSD) scheduler featuring 10% warmup and 10% decay. The training law figures and the final Phi-Mamba model use the regime detailed in Appendix
5.1 Stage 3 (Weight-Trannl Transfe and Knowledge Distillation)
As described Section 4.3 this phase employs a simple end-to-end distillation teacher-model logits. It leverages the alignment among all sequence mixers and successive blocks jointly fine-tune all components of the network. Experiments shown in Table highhlght 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 reduced importance of this phase.
Stage 3 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 2able2]]hon shows that using only end-to-end knowledge distillation is less than ideal. Although is slightly advantageous to use only Stage 3 compared only Stage 2, there is significant gap between using only Stage versus using Stage 2 3. As elaborated in Section 6,, this phase can freeze all network components except the Mamba-2 sequence mixer without significant performance drop. This in particular indicates that the third stage (like the other stages of MOHAWK) can operate in computationally limited settings.
5.2 Stage 2 (Hidden-State Alignment)
Following the analysis of the model's end-to-end distillation Stage 3, we evaluate the impact of aligning the hidden-state outputs of mixer blocks (Stage 2) on both the subsequent Stage 3 process
7

---

Hidden State Distance (Log)
Training Laws Hidden State
Training Laws Perplexity
100
20
Perplexity
15-
10
10M
40M
160M
640M
2.568
10M
40M
160M
640M
2.568
Tokens Observed
Tokens Observed
Stage Initialization
Stage Finetuned
Stage Pretrained
Figure 3: Training laws comparing token budget between Stages and 3, as measured by the Stage 2 metric (hidden state distance) and Stage 3 metric (perplexity). Stage initializations are used as the starting checkpoint for their respective Stage 3 finetuning models. Stage pretrained is trained from scratch only with weight transfer and knowledge distillation. Despite training for less tokens Stage 3 than the Stage 3 from scratch, almost all Stage 2 initialized models eventually outperform the fixed budget. In general, better aligned Stage 2 initializations improve
baseline in perplexity post-Stage performance.
and overall downstream model performance. We accomplish this by training Phi-Mamba instances from scratch using Stage various token counts. From these checkpoints, we proceed to Stage 3 training, ending with different total budgets to allow us to analyze how the degree of Stage "pretraining" impacts Stage performance at various token budgets.
Figure 3 demonstrates that given an adequate training budget, models heginning with weights with lower hidden state distances (after Stage 2) outperform those that depend exclusively knowledge distillation (Stage 3). These lower hidden states are also correlated with lower starting perplexities, which in turn are correlated with downstream performance, as shown 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 aligmment and knowledge distillation methods for the tested architectures.
5.3 Stage (Matrix Mixer Orientation)
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 previous protocol, we assess the positive impact the current stage on the following phase's 's metrics and final model's performance by comparing models with varying amount of Stage and Stage training on both stage metrics.
Figure T 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 training can help their respective Stage 2 models reach better hidden-state distances compared to the from- scratch 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 [.2] which discovers that lower hidden state initializations lead to better perplexity and downstream performance, it can be inferred that Stage I aids the overall distillation process. We further empirically validate this intuition Table 2] which indicates that this stage aligns the matrix mixers stronger degree than only the hidden state alignment.
Self-Attention Approximation with Structured Matrix Mixers
start by testing the ability various matrix mixer families to match empirical self-attention
matrices of pretrained Transformer. We take 1000 samples from each layer of Llama2-7b- 2--hhat model [37], materialize the attention matrices, and project them onto given classes of structured matrices. The results in Table 3 are averaged across all layers.

---

Training Laws Matrix Mixer
Hidden State Distance (Log)
Training Laws on Hidden State
200
11
Matrix Mixer Distance
100
50
20
10M
20M
40M 80M 160M
320M 640M
10M
320M 640M
Tokens Observed
Tokens Observed
Stage Initialization
Stage Finetuned
Stage Pretrained
Figure Training laws comparing the token budget between Stages and 2, as measured by the Stage I metric (matrix mixer distance) and Stage 2 metric (hidden state distance). Even a small amount of Stage training can improve the model's hidden-state distances in subsequent stages. Notably, this improvement occurs despite increase in matrix mixer distance during Stage 2. This suggests that early Stage training provides a foundational benefit that enhances the model's performance in later stages, demonstrating the importance of initial training phases in model optimization.
Table 3: Attention matrix approximation by structured matrix mixers (Frobenius distance; lower better). Structures are Toeplitz, low-rank (LR), state space dual (SSD) model (3.2) and general semi-separable matrices (SSM), all causal. We used 1,000 samples, each 512 tokens. Samples were passed through Llama2-7B-Chan and one attention head from each layer was randomly chosen for approximation. We evaluated (LR) and SSD families with 10,000 gradient descent steps per sample.
STRUCTURE
TOBP.
LR
SSD
SSM
LR
SSD
SSM
LR
SSD
SSM
(State size N)
(16)
(16)
(16)
(32)
(32)
(32)
(61)
(64)
(64)
WT-103
12.0
0.619
0.477
0.266
0.322
0.237
0.127
0.132
0.097
0.046
OWT
12.2
0.606
0.466
0.259
0.314
0.231
0.123
0.129
0.095
0.045
C4
12.3
0.595
0,453
0.236
0.310
0.226
0.112
0.128
0.093
0.041
IMJB
12.3
0.598
0.455
0.238
0.312
0.226
0.113
0.129
0.094
0.043
In particular, to describe the class of linear attention matrices (3.1p. we use the fact that Q and K are projections of the input Rdd.. onto Rouut and therefore their rank is bounded by min {din, dows } For multihead linear attention, dout (also known as head dimension) (e.g., Phi-1.5 and Llama2-7b-Chn approximate this family of sequence mixers using causal low-rank matrices L oo QK where L is a lower-triangular causal mask of 1s, and Q K are Raxd with d < (indicating that the head dimension is substantially smaller than the sequence length).
typically a small value have head dimensions of 64 and 128, respectively). Thus,
To describe the multi-head Mamba-2 matrix family, we utilize the state space dual (SSD) layer (32) in 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 more expressive mask that generalizes the causal mask (Section 3.2)
Both causal low-rank and SSD matrix families approximated with 10,000 steps of gradient descent per sample. To approximate the general class of SSM matrix mixers, we utilize balanced Iruncation, 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 time- varying 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 shows that while the SSM matrix family provides the closest approximation to the self- attention matrix mixer, the Mamba-2 mixer family (SSD) has just twice the distance from the SSM matrices. This is in contrast 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
We further validate the ability of a Mamba-2 2 block to replace Attention layer within a language model. Firstly, we create two variants of our architecture, Phi-Toeplitz and Phi-LR, and run MOHAWK process for IB tokens at each stage (see Table 4 to verify that the previous finding

---

Table Ablations of matrix structure using the same training recipe (Stages 2 and 3). While many efficient sequence models (e.g. global convolutions, linear attention, and state space models) can be represented as structured matrix mixers (e.g. Toeplitz, low-rank, and remi-separable matrices respectively), more expressive structured matrix families can match the attention matrix more closely.
MATRIX
BLOCK OUTPUT
WINOG. ACC1 1
ARC-E 1cc
ARC-C
PIQA
HELLAS. Acc
AVG.
STRUCTURE
L2 DIST. 4
Acc 1
Acc f
Acc
Causal Toeplitz
9.6
49.3
21.2
26.2
52.3
25.9
35.0
Causal low-rank
7.6
50.2
27.9
25.0
53.3
25,6
36.4
SSD
5.5
67.2
71.0
38.6
74.2
45.0
59.2
Table 5: MOHAWK distillation for Shi-Mamba-1135B on the entire model vs just the Mamba-2 component, i.e., frozen MLP, embedding, etc. MOHAWK can be used on just the sequence mixer blocks while keeping all other components frozen without compromising performance (Section 5.1
TRAINABLE
WINOG.
ARC-E
ARC-C
PIQA
HELLAS.
LAMB.
AVG.
COMPONENTS
ACC1
Acc
Acc
ACC 1
Acc 1
ACC
ACC
All
74.8
72.7
43.5
75.6
59.6
49.2
62.7
Mamba-2
69.1
73.5
43.8
74.7
59.3
48.2
61.4
hold in multilayer, end-to-end model case. Secondly, we run MOHAWK while freezing various parts of the Phi-Mamba modules (refer to Ta)le(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 and the downstream performance metrics (accuracy) in Table Essentially, a better matrix approximation (lower Frobenius distance) 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- rank-basedd BERT models outperform Toeplitz-basee ones.
7 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 's multi-stage process which gradually increased the scope of distillation essential extracting the teacher model's knowledge to the fullest extent as shown our ablations and training laws. We continue to find the effectiveness of MOHAWK when distilling hybrid Attention-Se models and provide ablations on 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 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 language model's knowledge embedded in the MLP blocks, we believe that any subquadratic model with sufficiently expressive matrix mixer can replicate the behavior of pretrained Transformers, bringing quadratic knowledge subquadratic models. We recommend further research to explore the role of sequence mixing layers in subquadratic models and their impact on performance. Advancements both the distillation process and the sequence mixer architecture could lead to further improved performance in 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.
10

---

References
[1] M. Beck, K. Pöppel. M. Spanring, A. Auer, O. Prudnikova, M. Kopp, G. Klambauer, J. Brand- stetter, and S. Hochreiter, xlstm: Extended long short-term memory, 2024.
[2] Y. Bisk, R. Zellers, J. Gao, Y. Choi, et al. PIQA: Reasoning about physical commonsense in natural language. In Proceedings of the AAAI conference on Artificial Intelligence, volume 34, pages 7432-7439, 2020.
[3] B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert- Voss, G. Krueger. T. Henighan, R. Child, A. Ramesh, M. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, Sigler, Litwin, Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei. Language models are few-shot learners, 2020.
[4] S. Chandrasckaran, Dewilde, M. Gu, T. Pals, and A.JJ van der Veen. Fast stable solver for sequentially semi-separable linear systems of equations. In International Conference High- Perf-Peeformnnce Computing, pages 545-554. Springer, 2002.
[5] Y-C. Chen, Z. Gan, Y. Cheng, J. Liu, and J. Liu. Distilling knowledge learned in bert for text generation, 2020.
[6] Clark, I. Cowhey, O. Etzioni, T. Khot, Sabharwal, C. Schoenick, and O. Tafjord. Think you have solved question answering? try ARC, the A12 reasoning challenge. arXiv preprint arXiv:1803.05455 2018.
[7] T. Dao, Fu, S. Ermon, A. Rudra, and C. Ré. Flashattention: Fast and memory-efficient exact attention with io-awareness, 2022.
[8] T. Dao and A. Gu. Transformers SSMs: Generalized models and efficient algorithms through structured state space duality. In International Conference on Machine Learning (ICML), 2024.
[9] P. Dewilde and A. J. van der Veen. On the hankel-norm approximation of upper-triangular operators and matrices. Integral Equations and Operator Theory, 17(1):1-45, 1993.
[10] Dewilde and A.-J Veen. Time- Varying Systems and Computations. Springer Book Archive. Springer New York, NY, edition, 1998. Springer Science+Busines. Media Dordrecht 1998.
[11] A. Gu. Modeling Sequences with Structured State Spaces. Phd thesis, Stanford University, 2023.
[12] A. Gu and T. Dao. Mamba: Linear-time sequence modeling with selective state spaces, 2023.
[13] A. Gu, K. Goel, and C. Ré. Efficiently modeling long sequences with structured state spaces, 2022.
[14] S. Gugercin and A. C. Antoulas. A survey of model reduction by balanced truncation and some new results. International Journal of Control, 77(8)744-7666 2004.
[15) S. Gunasekar, Y. Zhang, J. Aneja, C. C. T. Mendes, A. D. Giorno, S. Gopi, M. Javaheripi, Kauffmann, de Rosa, O. Saarikivi, A. Salim, S. Shah, H. S. Behl, X. Wang, S. Bubeck, R. Eldan, A. T. Kalai, Y. T. Lee, and Y. Li. Textbooks are all you need, 2023.
[16] S. Hahn and H. Choi. Self-knowledge distillation in natural language processing, 2019.
[17] M. A. Haidar and M. Rezagholizadeh. Textkd-gan gan: Text generation using knowledgedistillation and generative adversarial networks, 2019.
[18] G. Hinton, O. Vinyals, and J. Dean. Distilling the knowledge in a neural network, 2015.
[19] M. Hu, Y. Peng, F. Wei, Z. Huang, D. Li, N. Yang, and M. Zhou. Attention-guided answer distillation for machine reading comprehension, 2018.
[20] S. Hwang, A. Lahoti, T. Dao, and A. Gu. Hydra: Bidirectional state space models through generalized matrix mixers, 2024.

---

[21] A. H. Jha, D. Groeneveld, E. Strubell, and I. Beltagy. Large language model distillation doesn't t need a teacher. arXiv preprint arXiv:2305. 14864, 2023.
[22] Kasai, H. Peng, Y. Zhang, D. Yogatama, Ilharco, N. Pappas. Y. Mao. W. Chen, and N. A. Smith. Finetuning pretrained transformers into rnns. arXiv preprint arXiv:2103 2021.
[23] A. Katharopoulos, A. Vyas, N. Pappas, and F. Fleuret. Transformers are rnns: Fast autoregressive transformers with linear attention, 2020.
[24] C. Liang, H. Jiang, Z. Li, X. Tang, B. Yin, and T. Zhao. Homodistil: Homotopic task-agnostic distillation of pre-trained transformers, arXiv preprint arXhv:2302.09632. 2023.
[25] H. Liu, M. Zaharia, and P. Abbeel. Ring attention with blockwise transformers for cear-infinite context, 2023.
[26] S. A. Melchior, P. Van Dooren, and K. A. Gallivan. Model reduction of linear time-varying systems over finite horizons. Applied Numerical Mathematics, 77:72-81 2014.
[27] J. Mercat, I. Vasiljevic, S. Keh, K. Arora, A. Dave, A. Gaidon, and T. Kollar. Linearizing large language models, 2024.
[28] J. Niu, A. Liu, Z. Zhu, and G. Penn. What does the knowledge neuron thesis have to do with knowledge?, 2024.
[29] D. Paperno, G. Kruszewski, A. Lazaridou, N.-Q. Pham, R. Bernardi, S. Pezzelle, M. Baroni, G. Boleda, and R. Fernández. The LAMBADA dataset: Word prediction requiring broad dis- course context. Proceedings of the 54th Annual Meeting of the Association for Computational Linguistics, pages 1525-1534 2016.
[30] B. Peng, E. Alcaide, Q. Anthony, A. Albalak, S. Arcadinho, S. Biderman, H. Cao, X. Cheng, M. Chung, M. Grella, K. K. GV, X. He, H. Hou, Lin, Kazienko, Kocon, Kong, B. Koptyra, H. Lau, K. S. I. Mantri, F. Mom, A. Saito, G. Song, X. Tang, B. Wang, J. S. Wind, Wozniak, R. Zhang, Z. Zhang, Q. Zhao, P. Zhou, Q. Zhou, Zhu, and R.-J. Zhu, Rwkv: Reinventing rnns for the transformer era, 2023.
[31] Z. Qin, X. Han, W. Sun, B. He, D. Li, D. Li, Y. Dai, L. Kong, and Y. Zhong. Toeplitz neural network for sequence modeling, 2023.
[32] Z. Qin, S. Yang, W. Sun, X. Shen, D. Li, W. Sun, and Y. Zhong. Hgm2: Gated linear rnns with state expansion, 2024.
[33] K. Sakaguchi, R. L. Bras, C. Bhagavatula, and Y. Choi. Winogrande: An adversarial Winograd schema challenge at scale. Communications of the ACM, 64(9):99-106, 2021.
[34] H. Sandberg and A. Rantzer. Balanced truncation of linear time-varying systems. IEEE Transactions on Automatic Control, 49(2)2217-22,, 2004.
[35] Y. Sun, L. Dong. S. Huang, S. Ma, Y. Xia, J. Xue, J. Wang, and F. Wei. Retentive network: A successor to transformer for large language models, 2023.
[36] X. Tan, Y. Ren, D. He, T. Qin, 7. Zhao, and T.-Y. Liu. Multilingual neural machine translation with knowledge distillation, 2019.
[37] H. Touvron, Lavril, G. Izacard, X. Martinet, M.-A. Lachaux, T. Lacroix, B. Rozière, N. Goyal, E. Hambro, F. Azhar, A. Rodriguez, A. Joulin, E. Grave, and G. Lample. Llama: Open and efficient foundation language models, 2023.
[38] A. van der Veen and Dewilde. On low-complyxity approximation of matrices. Linear Algebra and its Applications, 205-206: 1000::115-11000 1201, 1994.
[39] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin. Attention all you need, 2023.
[40] Wang, D. Paliotta, A. May, A. M. Rush, and T. Dao. The mamba in the llaaa Distilling and accelerating hybrid models, 2024.
12

---

[41] W. Wang, F. Wei, L. Dong, H. Bao, N. Yang, and M. Zhou. Minilm: Deep self alf-attention distillation for task-agnostic compression of pre-trained transformers. Advances in Neural Information Processing Systems, 33:5776-5788, 2020.
[42] M. Xia, T. Gao, Z. Zeng, and D. Chen. Sheared Ilama: Accelerating language model pre-training via structured pruning. arXiv preprint arXiv:2310.06694 2023.
[43] S. Yang, B. Wang, Y. Shen, R. Panda, and Y. Kim. Gated linear attention transformers with hardware-efficien training. 2024.
[44] Z. Yang, L. Shou, M. Gong, W. Lin, and D. Jiang. Model compression with two-stage multi- teacher knowledge distillation for web question answering system, 2019.
[45] R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi. HellaSwag: Can machine really finish your sentence? Proceedings of the 57th Annual Meeting of the Association for Computational Linguistics, 2019.
[46] M. Zhang, S. Arora, R. Chalamala, A. Wu, B. Spector, A. Singhal, K. Ramesh, and C. Ré. Lolcats: On low-rank linearizing of large language models, 2024.
[47] M. Zhang, K. Bhatia, H. Kumbong, and C. Ré. The hedgehog & the poreupine: Expressive linear attentions with softmax mimicry, 2024.
[48] C. Zhou, G. Neubig, and J. Gu. Understanding knowledge distillation in non-autoregressive machine translation, 2021.
13

---

A Experiments and Experimental Details
Hyperparameter Search To construct Appendix A we performed grid searches for training in Stages 1, 2. and independently from scratch find the optimal hyperparameters. We explored learning rates Ir {1, 2, 5) 10(-3,-1) and batch sizes 2 AdamW Optimizer was used with / = (0.9, 0.95). incorporating a weight decay of 0.1, gradient clipping at 1-0, and a Carmup-Stable-DDer
(WSD) scheduler with 10% warmup and 10% decay utilizing linear warmup and cooldown functions. Automatic mixed precision training to bfl6 was used in all stages. For Stages and 2, we initially fixed the batch size at 2 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. Consequently, Stage used bs 2 Ir 5 10 and Stage used bs = -2',,1r = 2 x 10-3, In Stage 3. we set the batch size to 2 ~ 0.5M and focused solely on varying the learning rate, resulting in 5 x 10-4, Stages and were trained to 200M steps each while Stage 3 extended to IB steps. For the Phi-Mamba ultimate model, the Stage 3 learning rate was reduced to x 10-1 to enhance stability.
Multi-Stage Distillation Procedure In the development of the training law (see Figure 3 we executed 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 320 million tokens during the 40M 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.
extends Training Laws on Downstream Metries Figure eextends 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 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 metries capture model effectiveness.
Training the Final Phi-Mamba Model After confirming the importance of the stages in Section 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 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 + 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 reported in Table T
B Applying Mamba-2 as Black Box
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 A. Even though this formulation is somewhat different from Mamba-2, the original algorithm remains applicable through a reduction expressed in Appendix B
14

---

Training Laws Downstream Metrics
Zoomed In
Average Accuracy
40M
160M
640M
2.568
640M
1.288
2.568
Tokens Observed
Stage Initialization
Stage Finetuned
Stage Pretrained
Figure 5: Training laws comparing the amount of token budget between Stages and 3, as measured by the average accuracy of downstream evaluation metrics.
Listing PyTorch example for using the Mamba algorithm for a Delta-free variation.
I: (batch, seglen, nheads, headdim)
4_ logo (batch, seqlen, nheads)
B: (batch, seqlen, nheads, dstate) (batch, seqlen, nheads, dstate)
C:
D: (nheads)
y Mamba
X A_log. unsqueeze(11),
dt rearrange (A_log, "b h -> b h ct),
A torch. ones (self, nheads), B,
C,
)
Du torch. einsum ("h, bihp->bbhpp D, X) y rearrange Du, "b 1 -> b 1 (h p)")
C Attention Matrix Approximation Details
This section serves as a complement to Section outlines the methods employed to create Table 3
Appendices C.T to C.5 describe our strategies for finding matrix within the specified families that closely approximates the original attention matrix using selected distance metric. Formally, we consider the following optimization problem:
min
M X||
(6)
XEM
where M the subspace of specific matrix family, M is the attention matrix, and ||-|| corresponds to selected distance metric. In the following sections, explore different methods and matrix families for this optimization problem.
C.1 Semi-Separable Matrix Approximation
Considering time- vireinarying system denoted by (Ak. Bk. Ck, Dalken we can describe it using the matrix mixer T (also known as the transfer matrix) as follows:
D1
0
C₂B
D2
0
0
0
0
CBArBi
CBBr
D3 CaB3
0
T
CiAsa
CaAzBr
Da
0
0
GAL11B GA GALIA By GAL11 Bg
GB
Dr
With Ak E RRxx, B& Rmmxn Cx E Roxx, and Dj E RRxxM where n is the state dimension, m the input dimension, and p the output dimension.
15

---

Table 6: Full attention matrix approximation by structured matrix mixers Structures are Toeplitz. causal low-rank (LR), RetNet, state space dual (SSD) model (3.2) with and without the diagonal D term and general semi-separable matrices (SSM). We have used 1,000 samples, each consisting of 512 tokens. Llama2-7B-CC was randomly chosen for approximation. We evaluated (LR), RetNet, and SSD families with 10,000
was applied on every sample, and one attention head from each layer
gradient descent steps per sample.
STRUCTURE
TOEPLITZ
CAUSAL LOW-RANK
RETNET
SSD WITHOUT
SSD
SEMI. SEP. MATRIX
(State size N)
(16)
(16)
(16)
(16)
(16)
WT-103
12.0
0.619
0.530
0.477
0.266
OWT
12.2
0,606
0.516
0.508
0.466
0.259
C4
12.3
0.595
0.503
0.496
0.453
0.236
IMdB
12.3
0.598
0.505
0.498
0.455
0.238
State size N)
(32)
(32)
(32)
(32)
(32)
WT-103
12.0
0.322
0.268
0.262
0.237
0.127
OWT
12.2
0.314
0.261
0.255
0.231
0.123
C4
12.3
0.310
0.255
0.249
0.226
0.112
IMdB
12.3
0.312
0.256
0.251
0.226
0.113
State size N)
(64)
(64)
(64)
(64)
(64)
WT-103
12.0
0.132
0.110
0.097
OWT
12.2
0.129
0.107
0.045
C4
12.3
0.128
0.106
0.102
0.093
0.041
IMdB
12.3
0.129
0.106
0.103
0.094
0.043
As T's form corresponds to semi-separable matrix (i.e., each sub-matrix has rank of up to n), we will label this matrix form as SSM with state size n throughout the remainder of this appendix, representing state-space model with state size semi-separable matrix of order
Every matrix can be represented as a linear combination of rank-one matrices. Thus, the attention matrix M E gxxx can be interpreted as an SSM with state size of up to / > n. Consequently, we can employ prior research on time-varying model order reduction 26 with smaller state size n. Specifically, we utilize the following SVD-based approximation:
reduce M an SSM
Algorithm 1 Approximation of Attention Matrix M as SSM with State Size
Input:
Attention matrix M E L. state size n Output: Approximated attention matrix M E RLXL Procedure:
1. For & = l,,.., 1:
1.1 Define Ha the submatrix of M below and the left of entry Ma
Ma
MR.K-1
H2
:
Ma
M
1.2 Perform the SVD Hμ and truncate it to rank 1.3 Integrate the truncated Hş back into new matrix M
Note that the diagonal elements of M were not subject to approximation in Algorithm they remain unchanged.
Although this approximation method for the semi-separable matrix is heuristic, it has been empirically shown to deliver good results. further details on approximation methods for remi-separable matrices, and the theoretical background behind them, we refer the reader to [10] 26
16

---

C.2 Causal Low-rank Matrix Approximation
Given a set of self-attention matrices, we tried 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, composed M Lo AB where A, B RRN,N L a 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 semi-defiiite 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 semi-separable form (Appendix C.T) and its ability to be parallelized, we selected the number of steps for gradient descent based on the time it took 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 around factor of 5x compared 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.
entire batch of matrices (32) using gradient decent causal low-rank versus
During the experiments, we also found that initialization of the matrices A, B played significant role the resulting approximation difference. The original A, B values were sampled from [0, 1); however, given YMy 1, E [D] due to the SoftMax operator, A, B values then sampled
from 0, V5122 to have the last row of the self-attention matrix be uniform probability. We then
proceeded vary the factor of the range exponentially, testing 0, /512N
2(-2.-1.0.12.4.8)
where we found (0, 1
2 provided the best initialization across multiple datasets. A normal
V512N
distribution with μ and o² 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 used as the loss function.
Listing 2 PoTch example for generating Causal Low-rank approximation.
n_states 512 state_size = 16 # or 32, 64
num _emdseeds 32
torch. rand seam_ n_states, state_size)) B torch. rand heads state_size, n_states))
L torch. tril (torch. ones
wwwonnnnnnnnnnnenneen ((__states. n_states)))
approximation B)
17

---

C.3 State Space Dual (SSD) Approximation
For the SSD approximation, we utilize the scalar SSM recurrent (ISS) representation introduced in [8]. A key component is the values of a, which we will refer to from here on out to avoid confusion with matrix A, that constitute the final matrix mixer M.
Listing Py Torch example for generating SSD (with and without component) approximation.
n ststates 512
state_size = 16 # or 32, 64 softplus = tor ch. nn. Softplus
jum_heads 32
A torch. rand
n_states, state_stt sitee_siz)) B torch. rand (num_heads) state_size, n_states))
D torch. rand (oum_heads))
1 torch. rand (h, n_states))
L torch. exp
exp (age@gauumaoftppuus
-1))
approxximation (A B)
if apply_D:
M_approximation M_approximation torch eye (nsttn D None, None]
Given the rolling multiplicatiie property of L and the size of states, initialization of was important to prevent the bottom-right values of L quickly reaching O. We explored the uniform initialization leads to less "decay" within the L matrix. We found sampling / from [-8, -7) resulted in the best performance and use this initialization in the SSD family and RetNet class. As expected, adding the component helps reduce the error between the approximation and actual attention matrix Table 6
of [0, 1) {-10,- 8, -, -4, 22, 0, 2} where smaller values
C.4 RetNet Matrix Approximation
The Retention mechanism, introduced by [35], is key component in RetNet models and can be represented mathematically as (QK V. Here, the matrix is defined element-wist by
n- m
T1 m
Lym
0,
T < m
where y is a decay factor. This lower triangular matrix L captures the temporal dependencies by decaying past values with respect the current position.
In our approximation, we replace the product QK with matrices A and B. The matrix L can be efficiently constructed in PyTorch using the following code, which generates RetNet approximation:
Listing 4 Py Torch example for generating the RetNet matrix approximation.
n_states 512
state_sizz size = 16 $ or 32, 64 softplus torch. nn. Softplus
nun_heads- 32
A torch. forch.rand ((nuunnnadss,a.a heads, n_states, state_size))
B torch. rand
state_size, n_states))
torch. rand (num_Haae (aum_heads))
L torch. eorch.exp (sp repeatt,, th >> h n', n=n_states) -1))
apapproximation (A B)
This implementation provides a practical method for simulating the Retention mechanism, crucial for reducing computational complexity in RetNet models.
18

---

C.5 Toeplitz Approximation
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 Toeplitz matrix 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 1.2 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 set of numbers. As such, using the mean ensures that the approximation as close as possible the original matrix terms of L2 distance, thereby providing robust and efficient approximation method.
As before, assume that the approximation is input-dependent, meaning that each attention matrix has its own unique Toeplitz approximation.
C.6 Segsum Operator
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 ().. Below is the Python implementation of the 'segsum' operator using Py Torch.
Listing 5 PyTorch implementation of the Segmented Summation (segsum) operator.
segsum(x))
"""Naive segment sum calculation. eap (segaum(A)) produces a 1-55 matria, which is equivalent to a scalar SSH, (-1)
xumsumm torch. cumsum(x, dimm-1)
x x_segaum x_ cumsum :, None] - x_ cumsum None, :] mask torch. tril (torch. ones(T, т, device-x. device, dtype=bool), diagonal=0) x_segsum s_seggum.. masked_ ("mask, -torch. sorch.inf)
return x_seguum
19

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our claims supported by concrete results and thorough ablations for each stage of our proposed method.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals are not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: See Section 7
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 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 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 few datasets or with 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, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or 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 used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
20

---

Justification: N/A.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Our proposed method is quite simple and all significant details are included.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
21

---

Answer: [No]
Justification: Our proposed method quite simple and all important details are reported. The datasets used reported and are all standard, open source datasets.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.c cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run reproduce the results. See the NeurlPS code and data submission guidelines (https /nipss.cc/pulli//guu guwdes/ddee///
for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Experimental details are reported.
Guidelines:
The answer NA means that the paper does include experiments.
The experimental setting should be presented the core the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, supplemental material.
7. Experiment Statistical Significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Our experiments are expensive for academic results; furthermore, it common knowledge that language modeling results are reproducible across runs, when only the model initialization varies and the dataset order is fixed.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
22

---

The assumptions made should be given (e.g.. Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
OK to report 1-sigma error bars, but one should state The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative eтror rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [No]
Justification: Experiments were run across different types of hardware with frequent check- pointing and resuming. but using commercially available standard resources (e.g. single- node training A100 or H100 GPUs).
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: httt::::///nurip..tt https///ccccc..cc/ccccceeeeeeee
Answer: [Yes]
Justification:
No additional harms beyond standard concerns exist from this line of research.
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [No]
Justification:
No additional harms beyond standard concerns exist from this line of research.
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.
23

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is 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 being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there 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 monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: N/A
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have 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 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 best faith effort.
12. Lieenses 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 use have been credited.
Guidelines:
The answer NA means that the paper does not 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 URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.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, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of 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 '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/moodl 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.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: N/A
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: N/A
Guidelines:
The answer NA means that the paper does involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS 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.
25


Reviews:
[{"review_id": "Io3G6rtEXX", "reviewer": "Reviewer_1qnq", "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.", "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": "The limitations are not clearly discussed. Issues that should be mentioned include:\n* How well the method will scale?\n* How well does it work with more standard transformer/mamba architectures (not Phi)?\n* What is the impact of the data?", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "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\n- 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. \n- To the best of my knowledge, the paper faithfully discusses the related work.", "weaknesses": "- I have some reservations about the empirical evaluation:\n    - 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?\n    - 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)?\n    - 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).\n    - 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.\n- There are certain problems with the presentation in this paper:\n    - There are numerous typos and grammar errors, see below.\n    - 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.\n\nMinor issues/typos:\n- Line 27: “raises a natural question: it is” → is it?\n- Line 33: “differen”\n- Line 48: “benchmarks strong than” → stronger than\n- 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? \n- Figure 3: “Despiting training” → Despite training\n- And many more I think. I suggest carefully proofreading the manuscript.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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. \n\n> What are the fundamental limitations of this approach, what Mamba cannot do that transformers can? \n\n**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.\n\n> how MOHAWK performs with more standard transformer architectures (e.g., LLAMA) \n\n**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.\n\n> 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)?\n\nWe 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.\n\n> how much data would we need to close the gap to Phi-1.5 completely, or up to negligible levels?\n\nIn 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.\n\n> 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.\n\nWhen 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].\n\n|Model|WinoG.|ARC-E|ARC-C|PIQA|HellaS.|Avg.↑|\n|-----|------|-----|-----|----|-------|-----|\n|Phi-1.5 1.3B|73.4|75.6|48.0|76.6|62.6|67.2|\n|Phi-Mamba 1.5B|71.7|74.0|44.1|75.5|60.2|65.1|\n|Mamba (Phi2 Dataset) 1.8B|73.4|78.0|45.2|77.3|49.8|64.7|\n|Mamba-SWA-MLP 1.6B|73.7|76.7|46.2|76.5|49.7|64.6|\n|Samba 1.7B|72.9|79.3|48.2|77.1|49.7|65.4|\n\n> There are certain problems with the presentation in this paper:\n\n> Minor issues/typos:\n\nAs mentioned in the shared response, we have addressed these.\n\n[1] https://arxiv.org/abs/2404.08819\n\n[2] https://arxiv.org/abs/2402.04248\n\n[3] https://arxiv.org/abs/2402.03170\n\n[4] https://arxiv.org/abs/2402.01032\n\n[5] https://arxiv.org/abs/2402.18510\n\n[6] https://arxiv.org/pdf/2406.07522\n\n[7] https://arxiv.org/abs/2406.17557", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "YXynupTudB", "reviewer": "Reviewer_2yXa", "summary": "The paper distills a Transformer model into the Mamba architecture by using about 1% of the pertaining dataset.", "questions": "See Weaknesses.", "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.\n\nThere is no potential negative societal impact to be addressed.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 1, "contribution": 2, "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.\n2. The authors compare against a large number of SSM models.\n3. 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. \n2. 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?\n3. 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. \n4. 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.\n5. Minor comment: On line 108, matrix multiplication has complexity larger than $O(T^2)$.\n6. Minor comment: On line 48, \"strong\" -> \"stronger\"\n\n\n\nOverall, 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n>  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.\n\nA 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. \n\n> What does \"matrix mixer layer\" mean?\n\nThe 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.\n\n> What are TeacherMixer and StudentMixer? Is this done layer by layer, block by block or are they all aligned at the same time?\n\nWe 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. \n\nThis 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. \n\n> 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?\n\n$\\textbf{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.\n\n> 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. \n\nThe 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.\n\n> No code or trained model artifacts were provided nor were mentioned that will be provided.\n\nAs 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.\n\n> Minor comment… \n\nAs mentioned in the shared response, we have fixed the presentation issues, e.g., typos, broken links, and unclear portions, of the paper.\n\n[1] https://arxiv.org/abs/2405.21060\n\n[2] https://arxiv.org/abs/2407.09941\n\n[3] https://arxiv.org/abs/2402.18668", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "5i05JNToia", "reviewer": "Reviewer_Kdfj", "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.", "questions": "See weaknesses part.", "limitations": "The authors have discussed the potential limitations in Section 6.", "rating": 5, "confidence": 4, "soundness": 3, "presentation": 2, "contribution": 3, "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.\n2. 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.\n3. 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.\n\n2. 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.\n\n3. 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.\n\n4. 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "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.\n\n> 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.\n\nDespite 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. \n\n> There are several typos throughout the paper.\n\n> There is an inconsistency in the reported model size of Phi-1.5.\n\nWe 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.  \n\n> 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.\n\nThe 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.", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:
We thank all reviewers for their time spent reading our paper and providing thoughtful comments and feedback. All reviewers agree that our method enables cost-effective cross-architectural distillation, which is validated by our strong downstream performance, and that the careful ablations support our claim that all three stages of MOHAWK are crucial for effective distillation. The main weaknesses that have been brought up mainly surround the reproducibility and presentation of the paper. We have used these comments, questions, and concerns to improve our submission. In addition, we directly answer them in both the shared and individual responses and provide new experiments and analyses we would like to share:

- Improving results and reproducibility: We describe our new training regime to **improve reproducibility** and distill a **new version of Phi-Mamba 1.5B using only 3B tokens**, achieving an **average metric score of 62.6**, more than 1.2 percentage points than in the original submission. 
- Disentangling MOHAWK from Architecture: We ran additional baselines where we **distill from a pre-trained Phi-1.5 to a new Phi-1.5** to show the effectiveness of the MOHAWK distillation method compared to other potential alternatives and **highlight the role of the distillation method vs that of the architecture**.
- Paper presentation and model release: We have addressed all presentation issues raised by reviewers and plan to **open-source all model code and final model weights**.

To recap, MOHAWK is a method that can distill a pre-trained Transformer mode into any alternative architectures that can be represented using matrix mixers. The three stages progressively distill the Transformer architecture by first matching the mixing matrices, then the hidden states at the end of each mixer layer, and finally the end-to-end logits. MOHAWK is validated by our Phi-Mamba 1.5B model that demonstrates substantially stronger performance than all past open-source non-Transformer models at similar sizes while using less than 1% of the typical pre-training token count. 
## Improving Results and Reproducibility
Reviewer 2yXa has mentioned increasing the reproducibility of our distillation process. Below is a paraphrased excerpt from our appendix in the revised paper.

> To train the final model, we use the AdamW optimizer with $\beta = (0.9, 0.95)$, a weight decay of 0.1 with gradient clipping of 1.0, and a Warmup-Stable-Decay (WSD) scheduler featuring 10\% warmup and 10\% decay with a linear warmup and linear cooldown function. Automatic mixed precision training to bf16 was used for all three stages as well. Stage 1 used $\text{batch size}=2^{15}, \text{lr}=5\times 10^{-4}$, Stage 2 used $\text{bs}=2^{15}, \text{lr}=2\times 10^{-3}$, and Stage 3 used $\text{bs}=2^{19} \approx 0.5\text{M}, \text{lr}=5\times 10^{-4}$. These values were determined via sweeping over the hyperparameters, where the exact process is explained in the new appendix. The previous model was trained with a constant learning rate of $1 \times 10^{-4}$ for all three stages with a functional batch size of $2^{16}$ for Stages 1 and 2 and $2^{18}$ for Stage 3. All other components were the same. 

Based on the revised training regime detailed above, we trained a refined version of Phi-Mamba 1.5B using only 300 million more tokens (3.0B tokens total) by distributing the 3B C4 tokens into 80M/160M/2.76B splits for Stage 1/2/3. The new model is **strictly better on all metrics than our previously reported Phi-Mamba**.
| Model|Tokens / Dataset|WinoG.|Arc-E|Arc-C|PIQA|HellaS.|Lamb.|Avg. ↑|
|-|-|-|-|-|-|-|-|-|
|Phi-1.5-1.3B|150B / unknown|73.4|75.6|48.0|76.6|62.6|53.4|64.9|
|**New Phi-Mamba-1.5B**|3.0B / C4|**71.7**|**74.0**|**44.1**|**75.5**|**60.2**|50.1|**62.6**|
|Old Phi-Mamba-1.5B|2.7B / C4|69.1|73.5|43.8|74.7|59.3|48.2|61.4|
|Mamba-2-1.3B|315B / The Pile|60.9|64.3|33.3|73.2|59.9|**65.7**|59.6|
## Disentangling MOHAWK from Architecture
MOHAWK is a high-level method for cross-architectural distillation. Reviewers have noted that the effectiveness of the method has been conflated with the impact of the student model’s architecture. To address this, we conducted a baseline experiment where the student model architecture is fixed to match Phi-1.5. This experiment demonstrates the **effectiveness of MOHAWK when distilling from Phi to Phi** using the same budget of 5B tokens and the same hyperparameters as the previously trained model, as detailed in the section above.

|Stages Performed|WinoG.|ARC-E|ARC-C|PIQA|HellaS.|Lamb.|Avg.|
|-|-|-|-|-|-|-|-|
|2|71.0|77.1|40.8|77.8|60.9|51.3|63.1|
|3|64.1|73.7|38.7|58.6|75.6|48.0|59.9|
|2-3|69.8|77.4|44.8|78.2|61.3|54.4|64.3|
|1-3|74.2|76.7|43.5|78.6|61.7|54.2|64.9|

Using solely Knowledge Distillation and Weight Transfer (Stage 3) can only recover parts of the original Phi’s performance (59.9 avg) while adding Stage 2 can improve it significantly (64.3). **Only with all stages of MOHAWK is the overall performance restored**, highlighting the effectiveness of our MOHAWK method over traditional knowledge distillation alternatives, and showing that **performance gaps in Phi-Mamba are due to architectural differences between Mamba and Transformer**. Understanding the fundamental limitations of subquadratic vs quadratic models is an active area of research that is independent of our distillation method, where we use MOHAWK to match or exceed the best trained-from-scratch sub-quadratic models.
## Paper Presentation and Model Release
All reviewers highlighted issues with the presentation, which we have since addressed for the camera-ready version, including typos, grammatical errors, and broken links. We thank reviewer Kdfj for noting that the base Phi-1.5 model is 1.3B, while the new Phi-Mamba model is 1.5B. Additionally, **we plan to open-source all model code and the final Phi-Mamba 1.5B model weights**.

Meta-review: 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.
Decision: Accept (poster)