The Mamba in the Llama: Distilling and Accelerating Hybrid Models

Authors: junxiong  wang
daniele  paliotta
avner  may
alexander m rush
tri  dao
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-12 23:04:57

Abstract: Linear RNN architectures, like Mamba, can be competitive with Transformer models in language modeling while having advantageous deployment characteristics. Given the focus on training large-scale Transformer models, we consider the challenge of converting these pretrained models for deployment. 
We demonstrate that it is feasible to distill large Transformers into linear RNNs by reusing the linear projection weights from attention layers with academic GPU resources. The resulting hybrid model, which incorporates a quarter of the attention layers, achieves performance comparable to the original Transformer in chat benchmarks and outperforms open-source hybrid Mamba models trained from scratch with trillions of tokens in both chat benchmarks and general benchmarks. Moreover, we introduce a hardware-aware speculative decoding algorithm that accelerates the inference speed of Mamba and hybrid models. Overall we show how, with limited computation resources, we can remove many of the original attention layers and generate from the resulting model more efficiently. 
Our top-performing model, distilled from Llama3-8B-Instruct, achieves a 29.61 length-controlled win rate on AlpacaEval 2 against GPT-4 and 7.35 on MT-Bench, surpassing the best 8B scale instruction-tuned linear RNN model.

Main paper: 
 introduction: 
 While Transformers [72] have been an essential architecture in deep learning and have driven the success of large language models such as GPT [9] , Llama [70] , and Mistral [37] , they are prohibitively slow for very long sequence generation due to their quadratic complexity with respect to sequence length and large key-value (KV) cache requirement. Recent linear RNN models (Mamba [26] , Mamba2 [18] , GLA [78] , RetNet [67] , Griffin [19] ) beat Transformers in controlled experiments at small to medium scale, although the best Transformers still significantly outperform these models on downstream tasks. We note that the training times of linear RNN models are similar to those of highly optimized Transformers [78] , and therefore scaling up either of these models requires substantial computational resources.The primary benefit of linear RNN models (Mamba [26] , Mamba2 [18] ) is that they have faster inference (5× higher throughput) than Transformers. Efficient inference is emerging as a critical need for LLM systems such as new applications currently bottlenecked by the large KV cache of Transformers, e.g. reasoning over multiple long documents [30, 64, 56] and files in large codebases [60, 42] ). Emerging workflows with agents [80, 76] also require large-batch inference to explore more trajectories and long-context to model complex environments. These properties motivate the goal of distilling a large pretrained Transformer model into a linear RNN in order to generate as efficiently as possible. The technical challenges are two-fold: how to map pretrained Transformer weights to linear RNN weights for distillation, and how to adapt best-practice Transformer inference techniques, such as speculative decoding, to the new architecture. We make the following contributions:• We show that by reusing weights from attention layers, it is possible to distill a large transformer into a large hybrid-linear RNN with minimal additional compute while preserving much of its generation quality. We propose a modified Mamba architecture that can be directly initialized from the attention block of a pretrained model. • We propose a multistage distillation approach that mirrors the standard LLM pipeline combining progressive distillation, supervised fine-tuning [39] , and directed preference optimization [58] . This approach shows better perplexity and downstream evaluation compared with vanilla distillation. • We develop a hardware-aware speculative sampling algorithm and a fast kernel for speculative decoding on Mamba and hybrid architectures. We achieve a throughput of over 300 tokens/second for a Mamba 7B model. Additionally, we show that speculative decoding can be effectively applied to our hybrid architecture.Our experiments distill different large-scale open chat LLMs, Zephyr-7B [71] , Llama-3 8B [21] to linear RNN models (hybrid Mamba and Mamba2), using only 20B tokens of training.Results show that the distilled approach matches the teacher model in standard Chat benchmarks [83, 43] . We also show that it performs on par or better with all similarly sized pretrained-from-scatch Mamba models including Mamba 7B models [52, 26] trained from scratch with 1.2T tokens or NVIDIA Hybrid Mamba2 models [73] trained from scratch with 3.5T tokens in multiple tasks (e.g., MMLU [34] , TruthfulQA [47] ) in the LM evaluation [25] .Concurrent with this work, MOHAWK [6] distills a Mamba-2 variant based on the Phi-1.5 architecture with limited computation costs and performance loss.2 From Transformer to MambaQuestion: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes]Justification: The claims made in the abstract and introduction are supported by our theoretical analysis in section 2 and experimental results in section 5.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in 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. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
relationship between attention and linear rnns: 
 We begin by reviewing multihead attention to clarify the shapes of intermediate objects. Notationally, we use explicit subscripts for the sequence position instead of matrix representation, to better highlight similarities between the two models.Attention is computed in parallel for multiple differently parameterized heads. Each head takes sequence o with hidden size D as an argument and computes,Q t = W Q o t , K t = W K o t , V t = W V o t for all t, α 1 . . . α T = softmax [Q ⊤ t K 1 . . . Q ⊤ t K T ]/ √ D y t = t s=1 m s,t α s V s where o t ∈ R D×1 , W ∈ R N ×D Q t , K t , V t ∈ R N ×1 m s,t = 1(s ≤ t)Recent work has argued that linear RNNs can be serious competitors to attention in large language models. Several different linear RNN formulations have been proposed with similar formulations. For now, we leave the shapes of the parameters A t , B t , C t abstract, and note that linear RNNs all take the following form, where h is a matrix-valued hidden state,EQUATIONLinear RNNs have several computational advantages over attention. During training, all y t values can be computed more efficiently than attention since there is no softmax non-linearity. During inference, each next y t can be computed serially without requiring a cache.Despite the superficially different form, there is a natural relationship between linear RNNs and attention. Linearizing the attention formula by removing the softmax yields:y t = t s=1 m s,t α s V s = 1 √ D Q t t s=1 (m s,t K ⊤ s V s ) = 1 √ D Q t t s=1 m s,t K ⊤ s W V o s .This implies that there exists a linear RNN form of linearized attention, specifically:h t = m t-1,t h t-1 + K t V t y t = 1 √ D Q t h t ↓ h t = A t h t-1 + B t x t , y t = C t h t A t = m t-1,t , B t = W K o t , C t = W Q o t , x t = W V o tNote though that this version uses a hidden state of size h ∈ R N ×1 . Effectively tracking only one scalar over time per hidden dimension. Naively applying this transformation leads to poor results. The issue is that linearizing attention produces a degraded representation of the original model, as the softmax nonlinearity is critical to attention.The key to improving these models is to increase the capacity of the linear hidden state to better capture long-term structure. For example, previous work has shown the use of kernel methods to improve this approximation [62, 36, 82] . These approaches expand the size of the hidden state representation to h to R N ×N ′ to better match the modeling capacity of softmax.To design a effective distilled linear RNN, we aim to stay as close as possible to the original Transformer parameterization, while also expanding the capacity of the linear RNN in an efficient manner. We will not attempt to have the new model capture the exact original attention function, but instead use the linearized form as a starting point for distillation. for each head W k , W q , W v , W o : (N, D) 7:expanding grouped KVs do 8:Head Parameter: A : (N, N ′ ) 9:for all positions t: 10:xt : (B, N ) ← W V ot 11: Bt : (B, N ) ← W K ot 12: Ct : (B, N ) ← W Q ot 13: ∆t : (B, N ′ ) ← MLP(xt) 14: A1:T , B1:T , C1:T : (B, N, N ′ ) ← Disc(A, B, C, ∆) 15: y ← LinearRNN(A, B, C, x) 16:output ← output + W O⊤ y 17: return output Specifically, we adapt the parameterization from Mamba, [26] to increase the hidden state size, while initializing from the attention representation. Mamba uses a continuous time state-space model (SSM) to parameterize a linear RNN at run time, described by the differential equation,h ′ (k) = Ah(k)+B(k)x(k) y(k) = C(k)h(k)Where A is a diagonal matrix and other values are continuous signals. To apply this formulation to a discrete-time problem like language modeling, we use a neural network to produce a sequence of sampling intervals ∆ t and samples of the signals at these time steps. Given these sampling intervals, and T samples of B, C, Mamba approximates the continuous-time equation using a linear RNN as a discretization. We use an overbar to indicate the discrete-time form, which is reconstructed dynamically.A 1...T , B 1...T , C 1...T = Discretize(A, B 1...T , C 1...T , ∆ 1...T )In this simplest case, with N ′ = 1 and an identity discretization, this approach recovers the linear attention to linear RNN conversion discussed in the previous section. The benefit of Mamba is that with N ′ > 1 the continuous-time parameterization allows the model to learn significantly richer functions, without many more parameters or decreased efficiency. Specifically the only additional learned parameters will be the sampling rate ∆ and the dynamic A. These new parameters will control the constructed linear RNN through the discretization function yielding the new matrix valued linear RNN. Specifically, we take in the sameB t , C t ∈ R N ×1 and ∆ t ∈ R N ′ , but output B t , C t ∈ R N ′ ×N ×1, effectively increasing the hidden size by a factor of N ′ over the naive linear attention.A core contribution of Mamba [26, 18] is to demonstrate a hardware-aware factorization of this algorithm. Implementing the algorithm naively would be prohibitively slow as the new expanded parameters are quite large. Their approach fuses discretization, state expansion, and applying the linear RNN into a single kernel, which circumvents fully materializing the discrete parameters. This allows for large N ′ with relatively small efficiency costs.Our full approach is shown in Algorithm 1. This algorithm feeds the standard Q, K, V heads from attention directly into the Mamba discretization, and then applies the resulting linear RNN. As noted above, this can seen as roughly initializing with linearized attention and allowing the model to learn richer interactions through the expanded hidden state.Figure 1 shows the resulting architecture. Our version directly replaces Transformer attention heads directly with fine-tune linear RNN layers. We keep the Transformer MLP layers as is and do not train them. This approach also requires processing additional components like grouped query attention that shares keys and values across heads. We note that this architecture differs from the architecture used in many Mamba systems, which combines MLP-SSM layers and uses a single head. This initialization allows us to replace any attention block with a linear RNN block. We experiment with hybrid models where we keep every n attention layers. Empirically we found that replacing layers in a stepwise manner was the most effective strategy, i.e. we first keep every 2 layers, distill, and then every 4, and continue distillation.Attention Q K V Mamba C B x Δ A Figure 1:Transferring Transformer to Mamba. Weights, in orange, are initialized from the Transformer (Linear projections for Q, K, and V are initialized using linear projection for C, B, and X respectively). We replace individual attention heads with Mamba heads, and then finetune Mamba blocks while freezing the MLP blocks. Shapes are kept mainly the same. Weights in green are added. New parameters are introduced for the learned A and ∆ parameters.Question: Does the paper fully disclose all the information needed to reproduce the main experimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?Answer: [Yes] Justification: We will release our model and code. Everything will be open-sourced.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 a 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 a 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 it 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 to 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 a 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 submissions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. , with an open-source dataset or 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 reproducibility. In the case of closed-source models, it 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 to have some path to reproducing or verifying the results. 
knowledge distillation for aligned lms: 
 Knowledge distillation (KD) [35] serves as a compression technique aimed at training a smaller network that mimics the behavior of a larger teacher network. After initializing the model from the Transformer parameters, we aim to distill it to perform on par with the original language model. We assume that most of the knowledge from the transformer is maintained in the MLP layers which were transferred from the original model, and focus on distilling the fine-tuning and alignment steps of the LLM. During this stage, the MLP layers are kept frozen and the Mamba layers are trained as in Figure 1 . Supervised Fine-Tuning We first apply knowledge distillation to redo the supervised fine-tuning (SFT) stage of language model adaptation. During this stage, an LLM is trained to maximize the likelihood of a response y given an input prompt x, i.e. p(y | x). The task looks similar to conditional generation.There are two common approaches for distillation in this setting. One method is to use wordlevel KL-Divergence. In this setting, the full probability distribution of the student model p(•; θ) is trained to match the full distribution of the teacher model p(•; θ T ) by minimizing the KL divergence over the entire set of next possible tokens at position t. The second method is sequence-level knowledge distillation (SeqKD) [39] . SeqKD suggests a simple method for distillation on this style of task, by replacing the ground truth text y 1•••t with the teacher generation output ŷ1•••t , also known as pseudo-labels.EQUATIONHere θ is trainable parameters of the student model and α and β control the weights of sequence and word loss term respectively.The second stage of instruction-tuning for LLMs is to align them to a set of user preferences. During this stage, a set of desired preference pairs is used to improve the model's output. The objective is to produce outputs y to prompts x that maximize a reward model r while maintaining close to a reference model. Typically the reference model is chosen to be the model after supervised fine-tuning. For distillation, we can conveniently utilize the original teacher, i.e. This preference model is defined by a reward function r ϕ (x, y) dependent on the method used. Previous research utilizing AI feedback has primarily focused on employing reinforcement learning methods, such as proximal policy optimization (PPO) [63] , to optimize ϕ concerning this reward. Recently, methods using direct preference optimization (DPO) [58] have been effective at optimizing this objective with direct gradient updates. Specifically, DPO shows that, if we have access to preferred y w and dispreferred y l outputs for a given prompt x, we can reformulate this optimization problem as,EQUATIONThis optimization can be performed at the sequence level by scoring the preferred and dispreferred outputs of the model with the teacher and student and then backpropagating to the student. As far as we are aware this is the first use of DPO as a distillation objective. 
speculative decoding algorithms for linear rnns: 
 The main goal of the linear RNN formulation is to improve decoding efficiency. For both attention and linear RNNs, the serial dependency of autoregressive generation inherently bottlenecks efficiency. Systems cannot utilize all available compute, as they need to wait for the generation of previous tokens to proceed [66, 41, 11, 75, 10] . Speculative decoding has emerged as a method for breaking this bottleneck by spending extra compute to speculate on future generations. In this section, we consider approaches for applying this technique to large Mamba models, which can then be applied to the distilled models. When the new draft is verified, the multi-step kernel returns both the hidden state after the yellow token and the final hidden state, since verification will fall between those positions.Speculative decoding uses two models: a draft model, θ D , and a verification model, θ V . The fast draft model produces potential future completions, y * = arg max y 1:T p(y 1 , . . . , y T ; θ D ), and the larger verification model checks that these are top ranking at each time step, i.e. checking p(y * t |y * 1:t-1 ; θ V ). The longer a chain before a verification failure the faster the output. If a partial chain matches, we can rewind to the last match.Attention-based models are particularly amenable to speculation, as they are slow at generation due to sequential nature, but fast at verification due to their ability to check multiple tokens in parallel. Linear RNN models like Mamba have significantly different performance characteristics that make them less amenable to speculative decoding. Sequential decoding using recurrent-style sampling is already significantly faster than attention. Like attention, there are parallel modes for models like Mamba which are used at training. These are efficient, but are tuned for extremely long sequences. In addition, they rely on hardware-aware optimizations, such as avoiding materializing intermediate states. These properties make it difficult to use for speculation for relatively short chains when it is unknown when a conflict will occur.An additional challenge arises from caching states in RNN models. The state of an attention model is represented by the key-value cache, K 1:t , V 1:t ; whereas the state of an RNN model is simply h t . To be competitive with attention this single RNN state needs to be very large. During speculation, we need to rewind to a previous state at time step t ′ . For attention, this is simply K 1:t ′ , V 1:t ′ ; however, for RNNs this would require caching all h 1:t which would require a large memory overhead.We propose a new algorithm for linear RNN speculative decoding using hardware-aware multi-step generation. The core to the approach generation kernel that computes,y j:k , h j , h k ← MultiStep(h i , y 1:n , i, j, k; A, B, C, ∆)Where i is the starting hidden state, i ≤ j ≤ k, and j . . . k is the range of y outputs needed. The kernel is hardware-aware because it avoids materializing key terms off of the fast GPU memory. Specifically, it avoids instantiating most h 1:n as well as the discrete-time linear RNN parameters. This kernel is aimed to target the issues presented above. Specifically, it can save a snapshot of the state h j before evaluating the draft tokens. This allows recomputing the correct state on the fly after a token is rejected. The assumption is that decoding is bottlenecked by memory and not by compute, as we can compute multiple steps of decoding with very little overhead over single-step decoding.Algorithm 2 Multi-Step Linear RNN Specula- tion function Verify(y 1:k , j, h i ) // y 1:k are draft, j is last verified, // h i is a cached state with i ≤ j y ′ j:k , h j , h k ← MultiStep(h i , y 1:k , i, j, k; θ v ) k ′ ← FirstConflict(y j:k , y ′ j:k ) return k ′ , h k if k ′ = k else h j function Speculate(K) // K tokens are drafted per step h cache ← h 0 j ← 0 while y j is not end do k ← j + K y j+1:k ← arg max p(y j+1:k | y 1:j , θ D ) j, h cache ← Verify(y 1:k , j, h cache ) return y 1:jAlgorithm 2 and Figure 2 show the full algorithm. The approach maintains only one RNN hidden state in cache for verification and advances it lazily based on the success of the multi-step kernel. Since the distilled models contain transformer layers, we also extend speculative decoding to Attention/RNN hybrid architectures. In this setting, the RNN layers perform verification according to Algorithm 2, while the transformer layers simply perform parallel verification.Note that if the draft model is a Mamba or hybrid model, the speculation part of the algorithm gets more complicated, as the draft model needs to recompute the state for the tokens accepted in the previous iteration. This is done similarly to the verifier model, by caching older entries and recomputing on the fly during the next round of speculation.To verify the effectiveness of this approach we run the speculation using Mamba 7B and Mamba 2.8B as target models. Results are shown in Table 1 . Figure 3 shows the performance characteristics of the Multi-Step kernel itself. Step Size (K) Speedup on H100 GPUs. A naive implementation of our algorithm already shows strong performance on Ampere GPUs as shown in Table 1 . However, achieving strong performance on H100 GPUs is much more challenging. This is mainly due to GEMM operations being much faster, which makes the overhead incurred from the caching and recomputation operations more visible. In practice, the naive implementation of our algorithm, with several different kernel calls, achieves a decent speedup on 3090 GPUs (1.5x for Mamba 2.8B with 60% acceptance rate) but no speedup at all on H100s.We optimized our implementation by fusing kernels, and by adapting the implementation to easily allow caching and recomputing old steps. Specifically, the verifier model performs i) recomputation of previous steps from the cache, ii) multistep decoding for the new sequence of draft tokens and iii) caching within a single kernel 2. For the draft model, recomputation, decoding and caching are also fused in a single kernel. The resulting implementations archives speedups on H100s GPUs, as shown in Table 1 . 
experimental setup: 
 Target models. We perform experiments using two LLM chat models: Zephyr-7B [71] , which is a chat fine-tuned Mistral 7B [37] , and Llama-3 Instruct 8B [21] . For the linear RNN models, we use hybrid versions of Mamba and Mamba2 with 50%, 25%, 12.5%, and 0% attention layers. We refer to 0% as a pure Mamba model. Mamba2 is a variant architecture of Mamba that is designed to be more targeted to recent GPU architectures. Zephyr-Mamba refers to a distillation from Zephyr [71] , while Llama3-Mamba / Llama3-Mamba2 indicates distillation from Llama-3 instruct 8B [70] .Training. Distillation does not require any language modeling pretraining data, but instead uses the post-training process to adapt the new model. We use a three-stage process. In the first stage, we use UltraChat [20] and UltraFeedback [17] as seed prompts and use the teacher model to generate pseudo-labels. The student model is trained in one epoch using the loss L in Eq 2 with α = 1 and β = 0.1. Models are trained using AdamW optimizer with β = (0.9, 0.98) with a batch size 64. We use a linear learning rate warm-up (for the first 500 steps) followed by cosine annealing. In the second stage, we use supervised finetuning with our model on the GenQA [12] , InfinityInstruct [3] and OpenHermes 2.5 [69] datasets using SFT in one epoch, with the same hyperparameters as Zephyr [71] . In the final stage, for models distilled from Zephyr, we do distilled alignment with our model using DPO on the UltraFeedback [17] dataset which is consistent with teacher model. While models distilled from Llama-3 instructed 8B, we use datasets from SimPO [51] and Zephyr [71] . We only freeze Gated MLP (FFN) in the first stage, while in the second and final stage all parameters are trained. The total distillation process for each hybrid model (e.g., Mamba-Llama3 (50% att)) takes less than five days in 8x80G A100.Baselines. In addition to the core Transformer architectures, the main baselines we compare against are other large-scale linear RNN models. We compare with both pure SSM architectures, such as TRI Mamba 7B [52] trained with 1.2T tokens and Falcon Mamba 7B3 trained with more than 5T tokens, hybrid SSM architectures, such as Nvidia Hybrid Mamba 2 [73] trained with 3.7T tokens, and other linear hybrid RNN models, such as Recurrent Gemma-9B Instruct [8, 19] .We evaluate our models using both single-turn, AlpacaEval [43] and multi-turn chat benchmarks, MT-Bench [83] . These benchmarks assess the model's ability to follow instructions and respond to challenging prompts across a wide variety of domains.Table 2 shows the performance of our models on chat benchmarks compared with large transformer models. The distilled hybrid Mamba model (50%) achieves a similar score in the MT-benchmark as the teacher model, and slightly better than the teacher model on the AlpacaEval benchmark in both LC win rate and overall win rate. The distilled hybrid Mamba (25% and 12.5%) performance is slightly worse than that of the teacher models in the 2Additionally, we implement the convolutional part of the Mamba block using a circular buffer which allows us to keep track of the old entries and include them in the convolution when they are needed for recomputation.3https://huggingface.co/tiiuae/falcon-mamba-7b MT benchmark but still surpasses some large transformers even with more parameters in AlpcaaEval. The distilled pure (0%) model does degrade significantly in accuracy. Notably, the distilled hybrid model performs better than Falcon Mamba, which was trained from scratch with more than 5T tokens.Zero Shot Evaluation. We utilize the open-source LM Evaluation Harness library [25] (branch big-refactor) to assess 10 tasks, with the following evaluation metrics: WinoGrande (WG) accuracy [61] , PIQA (PQ) accuracy [7] , HellaSwag (HS) normalized accuracy [81] , ARC-Easy and ARC-Challenge (AE and AC) accuracy and normalized accuracy, [15] , MMLU (MM), accuracy [33] , OpenBookQA (OB) normalized accuracy [54] , TruthFulQA (TQ) accuracy [46] , PubMedQA (PM) accuracy [38] , and RACE (RA), accuracy [40] . Each task is evaluated by analyzing the probability assigned by the model to each potential answer choice. Benchmark Evaluation. We also report few-shot evaluations on OpenLLMLeaderboard by conducting 25 shots on ARC-Challenge [14] , 10 shots on HellaSwag [81] , 5 shots on MMLU [34] , and 5 shots on Winogrande [61] . For TruthFulQA, the mc2 metric is reported in this benchmark. For GSM8K [16] , we follow the evaluation for instruct tuned model [51] , which uses ZeroEval [45] , a benchmark designed for chat models. We also include the CRUX [29] from that benchmark, which is designed for evaluating reasoning on code. All models are evaluated with greedy decoding in the ZeroEval. and CRUX, we chose the zero-shot evaluation using ZeroEval, which is designed for evaluating instruct models. We evaluated the corresponding instruct-tuned models for Falcon Mamba-7b and RecurrentGemma-9B, specifically Falcon Mamba-7b-instruct and RecurrentGemma-9B-it.Table 4 shows that the performance of our distilled hybrid models matches that of the best open-source linear RNN models on the Open LLM Leaderboard, while outperforming their corresponding open-source instruct models in GSM8K and CRUX.Setup We perform speculative decoding using the distilled hybrid models. We run experiments using both Hybrid Mamba 50% and Hybrid Mamba 25% as main models. For the draft models, we train 2 and 4-layer Transformer Draft models on the OpenHermes2.5 dataset [69] , for approximately 3 full epochs, following the "shrink and fine-tune" approach from [65] . Specifically, we initialize the draft layers using layers from the Zephyr-7B model (we take layers at indices [0, 31] for the 2-layer model and [0, 10, 20, 31] for the 4-layer model), and the embeddings and language model head also from the Zephyr-7B model [71] . We perform loss masking on the prompt, thus only considering next token prediction loss (cross-entropy) on the chat continuations from the training set. Speculative decoding experiments are run on a single NVIDIA RTX 3090 on data from OpenHermes2.5. Table 5 : Performance metrics for different draft and target model configurations for K = 4 on data from OpenHermes2.5. # Gen is the average number of generated tokens per speculative decoding step and includes an additional token from the last verifier logits.Table 5 shows results for hybrid speculative decoding with, using both the Zephyr and Llama hybrid models with different configurations. For both the 50% and 25% distilled models, we achieve speedups of over 1.8x on the Zephyr-Hybrid compared to the nonspeculative baseline. We also show that the 4-layer draft model we trained achieves a higher acceptance rate, but it adds some additional overhead due to the increased draft model size.For the Llama-hybrid models, the speedups are more modest since the draft model is larger due to the large embedding table of Llama 3. In subsequent work, we will focus on making these draft models smaller.We only train our model using chat corpora due to academic budget constraints.Training on general corpora, such as those referenced in [55] , may help close the gap between teacher models and is worth exploring further. Additionally, our model is in 7B scale. Further work still needs to be done with models that have more parameters.Broader Impacts Our models are trained using a collected chat corpus. Recent research has uncovered potential societal biases embedded within many established corpora. While it is beyond the scope of this paper to delve deeply into these biases, we acknowledge the potential risk that our distilled trained models may inherit these biases. Necessity of Linear RNN. We train a model that removes Mamba blocks from the model entirely using the same recipe to see if the model can adapt. Both models are trained using Zephyr as the teacher model, with the OpenHermes 2.5 [69] dataset as the SFT dataset and UltraFeedback [17] as the DPO dataset. Table 9 compares the performance of the model with and without Mamba blocks. The model with Mamba performs significantly better than the one without it. This confirms that adding Mamba layers is critical and that the improved performance is not solely attributable to the remaining attention mechanism.Attention-free models. Attention-free models offer improved computational and memory efficiency, making them increasingly popular for various language processing tasks, including autoregressive language modeling. Models like S4 [27] and its subsequent variants [31, 28] have shown promising results in long-range synthetic tasks [68] . Gated SSM architectures, such as GSS [50] and BiGS [74] , incorporate a gating mechanism into SSMs for (bidirectional) language modeling. The recently introduced Mamba model [26] argues that the static dynamics of these methods fail to incorporate input-specific context selection within the hidden state, which could be crucial for tasks like language modeling. Mamba has been shown to outperform Transformers across different model sizes and scales. Additionally, several other sub-quadratic model architectures [57, 78, 19, 1, 2, 22, 4, 79] and hybrid architectures [23, 44] have also been proposed. Distillation from Transformers. There have been relatively few attempts to distill on to linear RNN style models. Laughing Hyena [49] proposes to distill the long convolution into a state space representation, enabling constant time inference in Hyena [57] . Ralambomihanta et al. [59] introduces a progressive knowledge approach to distill small transformer models (70M) into Hyena models. [6] Speculative Decoding. Speculative decoding [66, 41, 11, 75, 10] has recently emerged as a promising method to accelerate the inference process of large language models, particularly Transformers. This approach utilizes a smaller draft model to speculatively generate candidate tokens, which the larger target model then verifies. Leviathan et al. [41] , Chen et al. [11] proposed a rejection sampling scheme to improve inference quality, while Spector and Re [66] organized candidate tokens into a tree structure to enable more efficient verification. Subsequent work has examined both trained draft models [5, 13, 48] and training-free draft models [32, 77, 24] .The checklist is designed to encourage best practices for responsible machine learning research, addressing issues of reproducibility, transparency, research ethics, and societal impact. Do not remove the checklist: The papers not including the checklist will be desk rejected. The checklist should follow the references and follow the (optional) supplemental material. The checklist does NOT count towards the page limit.Please read the checklist guidelines carefully for information on how to answer these questions. For each question in the checklist:• You should answer [Yes] , [No] , or [NA] .• [NA] means either that the question is Not Applicable for that particular paper or the relevant information is Not Available.• Please provide a short (1-2 sentence) justification right after your answer (even for NA).The checklist answers are an integral part of your paper submission. They are visible to the reviewers, area chairs, senior area chairs, and ethics reviewers. You will be asked to also include it (after eventual revisions) with the final version of your paper, and its final version will be published with the paper.The reviewers of your paper will be asked to use the checklist as one of the factors in their evaluation. While "[Yes] " is generally preferable to "[No] ", it is perfectly acceptable to answer "[No] " provided a proper justification is given (e.g., "error bars are not reported because it would be too computationally expensive" or "we were unable to find the license for the dataset we used"). In general, answering "[No] " or "[NA] " is not grounds for rejection. While the questions are phrased in a binary way, we acknowledge that the true answer is often more nuanced, so please just use your best judgment and write a justification to elaborate. All supporting evidence can appear either in the main paper or the supplemental material, provided in appendix. If you answer [Yes] to a question, in the justification please point to the section(s) where related material for the question can be found.IMPORTANT, please:• Delete this instruction block, but keep the section heading "NeurIPS paper checklist",• Keep the checklist subsection headings, questions/answers and guidelines below.• Do not modify the questions and only use the provided macros for your answers.Question: Does the paper provide open access to the data and code, with sufficient instructions to faithfully reproduce the main experimental results, as described in supplemental material? Answer: [Yes] Justification: We will definitely release our model and code. Everything will be open-sourced. we will discuss the detailed steps and resources required for reproduction in Section 5.Guidelines:• The answer NA means that paper does not include experiments requiring code. Justification: We report the training details, include dataset, hyper-parameters and optimizer and etc in section 5.Guidelines:• The answer NA means that the paper does not include experiments.• The experimental setting should be presented in the core of the paper to a 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, or as supplemental material. 
experiment statistical significance: 
 Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?Answer: [Yes] Justification: We present the error bars in the evaluation benchmarks in Tables 2 and 4 , following the standard for these benchmarks.Guidelines:• The answer NA means that the paper does not include experiments.• The authors should answer "Yes" if the results are accompanied by error bars, confidence 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.) • 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. 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 circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: We discuss Broader Impacts in section 5.4.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [Yes] Justification: We discuss risk of pretrained language models in section 5.4.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort. 
licenses for existing assets: 
 Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We cite the original paper that produced the code package or dataset.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset. • The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset.• We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution. • For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review. 
 
 Appendix:  


Decision: Reviewers all agreed that the initial version of the paper had inadequate experimentation, and some found the exposition inadequate.  It is clear from the revised version however that these deficiencies have been largely corrected, and although the reviewer scores remain numerically low, the revised paper will be a valuable contribution to the conference.

Reviews: review: paper_summary: paper_summary: The paper proposes a way to convert a Transformer architecture to a Mamba architecture by first initializing a Mamba layer based on a pre-trained attention layer in order to facilitate faster convergence followed by knowledge distillation from the Transformer to the Mamba model. The training procedure only involves post-training on task-specific data. The paper also uses DPO for the knowledge distillation and the authors emphasize the novelty of this method for knowledge distillation. Following their proposed method, the authors demonstrate how to quickly train a hybrid architecture with some attention and some Mamba layers based on existing LLMs, particularly Zephyr 7B. Finally, the paper proposes a way to perform speculative decoding and gain speed up using Mamba architectures.
strength_weakness: strength_weakness: Looking at the performance in Table 1, increasing the attention percentage always improves performance in a significant way (most clear in MMLU). As such, I am not convinced that the method is actually working. There is also a lack of proper baselines. For example, no comparison is performed with random initialization. Similarly, the results for 100% Mamba architecture is not included. While this might not be as good as 50% architecture, it is still very much needed to understand the trade-off. 

The idea of converting a pre-trained Transformer to Mamba is quite generic. In particular, there does not seem to be any specific scale needed to get the method to work. As such, it is not clear why the authors decided to mainly test the method on such a large scale.  For example, a small Transformer model can be trained from scratch, used to create a Mamba model and then compared in terms of language modeling benchmarks such as perplexity. Indeed this experiment is done for one experiment in Section 7 and shows poor performance of the method in terms of perplexity. As a side note, the paragraph starting in line 275 is titled Does PPL correspond to ability?. However, the discussion in that paragraph does not answer this question at all. Moreover there is actually a correspondence and we see degradation in task specific performance similar to PPL. 

Moreover, the authors mention that they do not target pre-trained language model capabilities. Why is that the case?  Furthermore, the comparison is done on a hybrid architecture which further makes it confusing to determine the success of the method. Note that large models have a lot of redundancies as suggested by recent work in pruning and this makes it harder to see whether other attention layers are stepping in to make up for the poor performance of Mamba layers or something else is happening.

Overall I find the results not convincing enough and not supporting the claims made in the paper about the success of the method.
questions: questions: 1. The definition of # Gen in Figure 2 (right) is very vague and unclear. Can you please elaborate?

2. In line 182 it is mentioned that the parallel mode of Mamba is faster than a Transformer. Is there a reference for this? If so I think it would be useful to include it.

3. In Algorithm 2, the verify part may return the state from an earlier point, i.e. $h_j$ instead of $h_{k^\prime}$. How is this ok for the rest of the algorithm? This means a set of tokens will be ignored as the algorithm will never go over $j + 1, \ldots, k^\prime$ to incorporate them into the state.
limitations: limitations: The authors mention that the performance is only measured after fine-tuning on chat corpora instead of full pre-training due to resource limitations. While at 7B scale this is acceptable, as mentioned above, there is no need to focus on that scale from the start. Further limitations, such as degradation of performance when switching from attention to Mamba layers should be more deeply discussed.
paper_summary: paper_summary: The paper presents a mechanism to distill from llama to mamba.
strength_weakness: strength_weakness: The experiments are a bit limited, only performed on chat corpora. Further data mixture is needed to understand the effectiveness of the approach. (Yet due to other strength, the reviewer thinks it should receive a positive score).
questions: questions: Please see the weakness section.
limitations: limitations: Please see the weakness section.
paper_summary: paper_summary: The paper focuses on distilling transformers into SSMs for fast inference after training. Specifically, the authors initialize states of SSMs by pretrained transformer components and utilize speculative decoding for multi-step predictions. Experiments show that the distilled hybrid model can achieve competitive performance compared to the teacher network, i.e. Zephyr-7B on standard Chat benchmarks after training with 3B tokens only. The distilled hybrid model can outperform Mamba 7B which is trained from scratch with 1.3T tokens.
strength_weakness: strength_weakness: 1. The writing is not well-organized and involves a lot of unclear parts. In Section 2.1, the derivations from softmax-based attention to linear RNNs (Mamba) are not clearly stated, e.g. unexplained symbols in equations. Also the derivations should obey some assumptions. Please refer to Mamba-2.
2. Some settings of the proposed method are not clear. For example, in Section 2.2, Why is A set as fixed over time? In Section 2.3, which components in the model are from attention and which parts are trainable. Figure 1 is also not informative without explanations of the colors and arrows. In section 4, what are the differences between speculative decoding used in transformer and that used in Mamba here? What are the details of hare-aware designs?
3. The experiments are not enough. For comparison methods, Hybrid Mamba which is trained from scratch should be included. Besides Mamba, other SSMs and linear RNNs should be included. The benchmarks are only on chat and academic tasks.
questions: questions: see the weaknesses.
limitations: limitations: The paper should be generalized to all SSM models (including Mamba-1/2) or linear RNN models, instead of only focusing on Mamba-1.

Comments: comment: We deeply appreciate all the supportive and constructive feedback. We are running a smaller scale (110M) ablation experiment to address your concerns. We will definitely discuss this and include it in the next version.
Dear Authors,

Thank you for the additional results. 

One of my concerns remain and I think it would be a great addition if the authors could show whether the model relies on the language models being large scale or is applicable on smaller scales as well. As I mentioned originally in my review, it is possible to train a smaller transformer from scratch for a reasonable number steps, then do the same approach for converting layers to Mamba layers and compare the two models in terms of perplexity (or any other metric). Currently the paper is only evaluating on 7B models but it is not clear whether there is a limitation for smaller sizes or not. Either way I think it should be discussed and would be a nice addition.

Still, given the new results provided in the rebuttal, I think there is sufficient evidence to show the method has merit and can be useful. As such, I have increased my score.
We are very grateful for all the encouraging and constructive reviews. Please find our responses below. 

$\mathcal{Q}.$ The paper’s story seems to rely on a biased assumption that Mamba will replace transformers—an appealing idea, but speculative at this stage. 

$\mathcal{A}.$ We appreciate the reviewer's concerns. We do not claim in the paper that Mamba will replace transformers and in fact, believe that people will continue mainly training transformers. We instead think that some organizations may want prioritize fast and lower-memory inference, and therefore want to distill their transformer models to Mamba for some applications.

$\mathcal{Q}.$ Additionally, the paper does not introduce any new method or theory. 

$\mathcal{A}.$ We think there are two new methods introduced by this paper. 1) Our distillation approach is novel. To the best of our knowledge, this is the first work demonstrating that transformations between large transformers and linear RNN models is possible through a mathematical transformation. And in our experiments, we distill different large transformers to support this claim.  2) We introduce a new algorithm for speculative decoding in the Mamba model which is efficient in practice. We agree that there is no new theory introduced in this work, although that is generally true for most deep learning papers.  

$\mathcal{Q}.$ We don't understand that it is just an application of proven techniques on mamba.

$\mathcal{A}.$ While distillation is a proven technique, as far as we know there is no work published on distilling from transformer to Mamba or related linear RNNs. The details of what distillation approaches work and how they can reuse parameters is a the contribution. We are a bit perplexed by this comment, and it would be great help if you can point us any thing related to this. (Previous work on this topic does not do anything like distillation.)
My main concern is the narrow scope of the methods and experiments of this paper. The paper’s story seems to rely on a biased assumption that Mamba will replace transformers—an appealing idea, but speculative at this stage. The focus on distilling and accelerating the Mamba hybrid model, while potentially applicable to other models (like linear RNN or linear attention), still requires further demonstration. And the paper does not propose any new methods or theory, it is just a application of proven techniques on mamba. Given the limited scope and novelty, I don’t believe the paper meets the NeurIPS standard in its current form.
We thank the reviewer for raising his score. However, we wonder why the reviewer believes we shouldn't focus on Mamba, since it's clear from the title that this is the goal and scope of our work. We would like to ask the reviewer to reconsider their score of rejection, unless there are other doubts or issues about our work, which we are happy to engage with.
Thanks for the effects. I appreciate the mamba2 experiments provided, but I am still doubt the generalization of the proposed method to other models like linear RNN or linear attention. The paper should not only focus on mamba. I would like to raise my score to 4.
We thank again the reviewer for the helpful feedback.

$\mathcal{Q}.$ I could not find the new PPL results. Could you please share them as well?

$\mathcal{A}.$ Since our models are alignment using DPO, we evaluated ppl of our final models on lambada_openai. Here are ppl of models in the ablation results.

| Model                         | PPL |
|-------------------------------|-----|
| Hybrid Mamba (50%) w/o Attention init |  55.01 |
| Hybrid Mamba (50%) w Attention init  | 6.20  |

| Model                         | PPL |
|-------------------------------|-----|
| 50% Attention w/o Mamba Block | 151.98 |
| 50% Attention w Mamba Block   | 6.20  |

$\mathcal{Q}.$ Just to confirm, you still keep the FFN parts of the layers, and only remove the attention, then do fine-tuning? I think if this is the case, this is an important result as it shows the efficacy of using Mamba blocks and the projection. Combined with the obtained speedup I think this makes the results quite interesting. However, the trend that using Mamba blocks lead to decreased accuracy is completely observable in all the results. As such, I think it is paramount that this limitation be clearly specified and discussed in the paper. Do authors agree on this or have I misinterpreted the results?

$\mathcal{A}.$ Yes, we still keep those three Gated MLPs (FFN components) and only remove half of the attention layers before fine-tuning. We acknowledge that changing attention to Mamba may still lead to decreased performance as more Mamba layers are introduced. At the same time, we believe that as computational resources scale up, this impact will be mitigated. Additionally, the inductive bias learned from attention can help Mamba scale up, as our model has already outperformed other hybrid Mamba models trained from scratch with significantly more tokens. We definitely will address and discuss this limitations in the paper.

Additionally, we will update the speculative decoding section to improve and correct the explanation of the algorithm.
Dear Authors,

Thank you very much for the additional experiments and your clarifications and replies. I have asked for some additional clarifications below.

>  In the new ablations that we ran, we found that initializing using attention weights not only gives lower PPL, but also significantly improves task-specific performance.

I could not find the new PPL results. Could you please share them as well? 

> About Random Init: We added a comparison with random initialization. Please refer to 3a in the general response. With the Mamba default initialization and the same training tokens, the model gets a very low score in MT-bench and AlpacaEval.
> Regarding no Mamba layers:  This is a good question, and we were curious as well. Please refer to 3b in our general response. This result confirms that adding Mamba layers is very critical and that the results are not just from the remaining attention.

Thank you for providing these results. Just to confirm, you still keep the FFN parts of the layers, and only remove the attention, then do fine-tuning? I think if this is the case, this is an important result as it shows the efficacy of using Mamba blocks and the projection. Combined with the obtained speedup I think this makes the results quite interesting. However, the trend that using Mamba blocks lead to decreased accuracy is completely observable in all the results. As such, I think it is paramount that this limitation be clearly specified and discussed in the paper. Do authors agree on this or have I misinterpreted the results?

>  We are sorry for the confusion. Yes you are correct in that in speculative coding we may throw away generated tokens. In the case where a state from an earlier point is returned by the verify part, the draft model will also resume decoding from a previous snapshot of the state, recomputing the states for the accepted tokens from its cache.
Could you please fix this in the next revision of the paper? As I think the current Algorithm is incorrectly specified and does not include the recomputation part which can be confusing.

Assuming the above are addressed, I will be increasing my score as I think the new results provide enough supporting evidence that the method is working.