The Mamba in the Llama: Distilling and Accelerating Hybrid Models

🎯 Pitch

A pretrained 8B Transformer can be converted into a fast hybrid Mamba model using only 20B finetuning tokens, matching the original's chat quality while hitting 300+ tokens/sec via a novel multi-token speculative decoding algorithm for linear RNNs.

1. Executive Summary

This paper introduces a method for distilling large pretrained Transformer language models into hybrid linear RNN architectures — specifically Mamba and Mamba2 — by directly reusing the linear projection weights from the Transformer's attention layers to initialize the Mamba blocks, then fine-tuning the resulting hybrid model through supervised fine-tuning and preference optimization while keeping the Transformer MLP layers largely frozen. Evaluated on Llama-3 8B Instruct and Zephyr-7B as teacher models, the distilled hybrid models — which retain only 50%, 25%, or 12.5% of the original attention layers — achieve chat benchmark scores comparable to the original Transformers and outperform open-source hybrid Mamba models trained from scratch with trillions of tokens on both chat and general benchmarks, with the top distilled model (Llama3-Mamba 50%) reaching a 29.61 length-controlled win rate on AlpacaEval 2 against GPT-4 and 7.35 on MT-Bench. The paper also develops a hardware-aware speculative decoding algorithm that enables multi-step verification for linear RNN architectures (generating and verifying batches of draft tokens while lazily updating a single cached hidden state to avoid materializing intermediate RNN states), achieving throughput of over 300 tokens/second for a Mamba 7B model — roughly 2× speedup — and extending this to hybrid architectures with attention layers. The distilled models additionally demonstrate natural length extrapolation, achieving nearly perfect needle-in-a-haystack retrieval accuracy at 20× the distillation context length, establishing that Transformer knowledge can be effectively transferred through weight initialization and lightweight distillation using only ~20B training tokens on academic GPU resources, though the approach degrades gracefully as more attention layers are removed, with the fully linear (0% attention) model showing significant quality loss.

2. Context and Motivation

The Deployment Bottleneck: Transformers Are Fast to Train but Slow to Generate

The Transformer architecture has become the dominant backbone for large language models, powering systems like GPT, Llama, and Mistral. However, Transformers have a fundamental asymmetry between training and inference efficiency. During training, the self-attention mechanism can be parallelized across all sequence positions, making it relatively efficient on modern GPU hardware. At inference time, however, autoregressive generation forces a sequential bottleneck: generating each new token requires recomputing attention over the entire previous sequence. This yields two compounding costs:

Quadratic compute scaling with sequence length: Each new token requires dot-product attention against all previous tokens, so generating a sequence of length $L$ costs $\mathcal{O}(L^2)$ total compute.
Linearly growing key-value (KV) cache: The model must store the key and value projections for every previously generated token in GPU memory. For a model with 32 layers, 32 attention heads, and 128-dimensional head projections generating 100K tokens, the KV cache can consume tens of gigabytes of memory — often exceeding the memory footprint of the model weights themselves.

The paper frames this problem concretely in Section 1: Transformers 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." This is not a theoretical concern — it directly bottlenecks emerging applications:

Multi-document reasoning: Systems that need to process and cross-reference information across many long documents (e.g., legal contracts, scientific literature reviews) must store and attend over hundreds of thousands of tokens simultaneously.
Codebase-level understanding: Tools that operate over entire software repositories need to maintain context across many files, each potentially thousands of lines long.
Agent-based workflows: Autonomous agents that explore multiple reasoning trajectories or model complex environments require both large-batch inference (to explore many action paths) and long-context generation (to maintain environment state).

These applications are "currently bottlenecked by the large KV cache of Transformers" (Section 1), making the inference efficiency problem not just a matter of cost optimization but a fundamental blocker to new capabilities.

Linear RNNs: A Compelling Alternative with a Training Gap

Linear recurrent neural network (RNN) architectures — including Mamba, Mamba2, RWKV, Griffin, and RetNet — offer a solution to both Transformer bottlenecks. Their core formulation (Equation 1 in the paper) replaces the quadratic attention mechanism with a recurrent state update:

$h_t = A_t h_{t-1} + B_t x_t, \quad y_t = C_t^\top h_t$

Here, $h_t$ is a fixed-size hidden state that summarizes all previous context. Generation of each new token costs $\mathcal{O}(1)$ rather than $\mathcal{O}(L)$ , since the model only needs to update a single state vector — no attention over the full history and no growing KV cache. During training, these models can also be parallelized using parallel scan algorithms, making their training throughput competitive with highly optimized Transformers.

The deployment advantages are substantial. As the paper notes in Section 1, linear RNN models achieve "5× higher throughput than Transformers" at inference time. This is not a marginal improvement — it is the difference between a user-facing chatbot that responds in 200ms versus 1 second, or between being able to process a 100-page document on a single GPU versus requiring model sharding across multiple devices.

However, linear RNNs face a critical obstacle: the best Transformer models still significantly outperform them on downstream tasks. The paper acknowledges this directly: "the best Transformers still significantly outperform these models on downstream tasks" (Section 1). This performance gap exists for a structural reason — the research community and industry have invested enormous resources into pretraining Transformer-based models at massive scale. Models like Llama-3 8B were trained on trillions of tokens using thousands of GPUs over months. Replicating this investment for each new linear RNN architecture is economically prohibitive.

Moreover, the paper notes that "the training times of linear RNN models are similar to those of highly optimized Transformers" (Section 1), meaning that training a linear RNN from scratch to match Llama-3 would require comparable computational resources — resources that most academic labs and smaller companies cannot access. This creates a chicken-and-egg problem: linear RNNs offer deployment advantages, but to be useful they must first match Transformer quality, which requires the very scale of pretraining that few can afford.

Prior Distillation Attempts: Why Existing Approaches Fall Short

The idea of transferring knowledge from a trained Transformer to a more efficient architecture is not new, but prior efforts have been limited in scope and effectiveness. The paper identifies several approaches and their shortcomings:

Progressive knowledge transfer (Ralambomihanta et al., 2024). This prior work attempted to distill small Transformer models (70M parameters) into Hyena models — a different linear RNN variant based on long convolutions rather than state-space models. Their approach trained the student model layer-by-layer, progressively adding new layers while the earlier ones stabilized. The paper reports in Table 6 that this approach resulted in a perplexity degradation factor of 2.36× (from 51.4 to 121.2 on WikiText), a significant loss of quality. Moreover, this work only demonstrated the approach on very small models (70M parameters), leaving open the critical question of whether distillation could scale to the 7-8B parameter models that represent the practical frontier of deployable LLMs.

Concurrent work: MOHAWK. The paper acknowledges a concurrent submission (Bick et al., 2024) that distills a Mamba-2 variant from the Phi-1.5 architecture. While this work demonstrates the feasibility of the general approach, it operates at a smaller scale and does not address the full post-training pipeline (SFT + preference optimization) that real-world chat models require.

The fundamental gap: architecture mismatch. Previous distillation efforts face a deeper challenge than just scale. The Transformer attention mechanism and linear RNNs have fundamentally different internal structures. A standard distillation approach — training the student to mimic the teacher's output distribution — provides a weak learning signal when the student's architectural constraints prevent it from exactly matching the teacher's computation. The student essentially receives a "do better" signal without any structural guidance on how to use its different internal components to achieve similar results. This is akin to trying to learn to play a piano piece by only hearing the final recording, without any information about which fingers press which keys.

The paper's key insight addresses this directly: reusing the Transformer's attention projection weights to initialize the linear RNN's corresponding projections provides structural guidance that dramatically reduces the distillation burden. As shown in Section 2.2 and Algorithm 1, the query, key, and value linear projections ( $W_Q, W_K, W_V$ ) from the Transformer's attention heads map directly onto the Mamba block's corresponding projections ( $C, B, X$ ). This is not just a convenient initialization trick — it preserves the semantic relationships that the Transformer learned during pretraining. The $W_Q$ projection that the Transformer uses to compute "what am I looking for" becomes the $C$ parameter that the Mamba block uses to read from its hidden state. The $W_K$ projection used for "what do I contain" becomes the $B$ parameter for writing to the hidden state. This structural correspondence provides the student model with a meaningful starting point rather than random noise.

The Speculative Decoding Gap: Linear RNNs Need Different Optimization Strategies

Even with a high-quality distilled model, the inference efficiency picture is incomplete without speculative decoding. Speculative decoding — where a fast draft model proposes candidate tokens that a slower verifier model checks in parallel — has become standard practice for accelerating Transformer inference. It exploits the Transformer's training-time parallelism: the verifier can compute attention over a batch of candidate tokens simultaneously, making verification much faster than the sequential generation used by the draft model.

Linear RNNs, however, break this assumption. The paper identifies two specific challenges (Section 4.1):

Challenge 1: RNN generation is already fast, but RNN verification is not faster. Unlike Transformers, where verification is substantially cheaper than generation due to parallelism, linear RNNs already use their efficient recurrent form for generation. Computing multiple steps of the RNN in parallel (using the training-time parallel scan mode) is possible but not efficient for the short sequences typical in speculative decoding (3-4 tokens). The parallel scan is optimized for extremely long sequences (hundreds or thousands of tokens) and has significant overhead for short batches. The paper notes that these parallel modes "are efficient, but are tuned for extremely long sequences" and "rely on hardware-aware optimizations, such as avoiding materializing intermediate states" — making them poorly suited for the stop-and-go pattern of speculative verification.

Challenge 2: State rewinding requires expensive caching. In Transformer speculative decoding, reverting to a previous state after a token mismatch is trivial: the KV cache up to the last verified position is simply reused. For RNNs, the state is a single hidden vector $h_t$ that represents the cumulative summary of all previous tokens. If verification rejects a draft token at position $t'$ , the model needs the hidden state $h_{t'}$ to continue generating from the correct position. Without caching intermediate states, the RNN would need to recompute forward from the beginning. But caching all intermediate states $h_1, h_2, ..., h_t$ defeats the memory efficiency advantage of RNNs — the hidden state in Mamba is large (expanded by factor $N'$ ), so storing many of them would consume memory comparable to a Transformer's KV cache.

Prior speculative decoding work (Leviathan et al., 2023; Chen et al., 2023; Spector and Re, 2023) was designed exclusively for attention-based models and did not address these RNN-specific challenges. The concurrent work by Wu et al. (2024) proposes a similar approach, but the paper notes that achieving competitive speedups on modern hardware (particularly H100 GPUs) required significant additional optimization beyond the naive algorithm.

How This Paper Positions Itself

The paper positions itself at the intersection of three research threads that had previously been pursued largely independently:

1. Weight transfer from Transformers to linear RNNs. Prior work on distilling to linear RNNs (Ralambomihanta et al., 2024) treated the student architecture as a black box learned from scratch with only output-level supervision. The paper's key departure is to exploit the mathematical relationship between linearized attention and linear RNNs (Section 2.1) to establish a principled weight initialization scheme. This changes the distillation problem from "learn a completely new function" to "learn how the expanded state space can improve upon the initial linearized attention approximation."

2. Post-training alignment as a distillation target. Rather than attempting to distill the entire pretraining process (which previous work attempted with limited success), the paper focuses distillation on the post-training pipeline — supervised fine-tuning and preference optimization — which is a much smaller data and compute budget (~20B tokens versus trillions). This is a strategically motivated choice: the MLP layers, which contain much of the model's factual and conceptual knowledge, are preserved exactly from the Transformer. The distillation only needs to teach the new Mamba layers how to perform the attention-like routing of information that the original attention layers handled.

3. Hardware-aware kernel design for RNN speculation. The multi-step kernel described in Algorithm 2 and Figure 2 is designed to solve the specific tension in RNN speculative decoding: the need to compute multiple steps efficiently while maintaining the ability to rewind state on verification failure. By fusing verification, recomputation, and state caching into a single GPU kernel that avoids materializing intermediate states or discrete-time RNN parameters, the approach achieves speedups on both consumer (3090) and datacenter (H100) GPUs. This is a practical engineering contribution that makes the theoretical efficiency advantages of linear RNNs realizable in deployment.

The paper's central claim is that these three components — weight initialization, lightweight distillation, and hardware-aware speculation — together form a practical pipeline for converting existing Transformer investments into deployment-efficient linear RNN models. The framing is explicitly resource-conscious: "We show how, with limited computation resources, we can remove many of the original attention layers and generate from the resulting model more efficiently" (Section 1). This distinguishes it from work that assumes access to the scale of pretraining infrastructure needed to train linear RNNs from scratch.

3. Technical Approach

3.1 Reader Orientation

The paper builds a distillation-to-deployment pipeline — a series of steps that takes an existing, high-quality Transformer language model (like Llama-3 8B) and converts it into a faster, more memory-efficient hybrid model where most attention layers are replaced with Mamba linear RNN blocks, followed by a hardware-aware speculative decoding algorithm that further accelerates inference. The core problem being solved is: how do we get the deployment advantages of linear RNNs (faster generation, no KV cache) without paying the enormous computational cost of training one from scratch? The solution has a specific "shape": reuse the Transformer's own attention weights to initialize the Mamba blocks (so the student model starts from a meaningful point, not random noise), then run a lightweight distillation process focused only on the post-training alignment pipeline — supervised fine-tuning and preference optimization — using only ~20B tokens rather than trillions, and finally deploy with a custom multi-step speculative decoding kernel that works around the fact that RNNs cannot do the "parallel verification" trick that makes Transformer speculation fast.

3.2 Big-Picture Architecture (Diagram in Words)

The system has five major components connected in a sequential pipeline:

Teacher Transformer — a pretrained and instruction-tuned LLM (Zephyr-7B or Llama-3 8B Instruct). This model provides both the initial weights for the student and, during distillation, the target output distribution that the student is trained to mimic.
Attention-to-Mamba Weight Initializer — a deterministic mapping that takes the frozen linear projection matrices from each Transformer attention head (the $W_Q, W_K, W_V, W_O$ projections) and copies them into the corresponding slots in a Mamba block ( $C, B, X$ , and the output projection). This preserves the semantic knowledge the Transformer learned about how to route information, giving the student a non-random starting point. The Transformer's MLP layers are kept exactly as-is.
Hybrid Student Model — a new architecture where some fraction of the original attention layers (typically 50%, 25%, or 12.5%) are replaced with Mamba blocks initialized from those layers' weights, while the remaining attention layers stay unchanged and all MLP layers from the Transformer are preserved and frozen. The Mamba blocks introduce new trainable parameters ( $A$ , the state transition matrix, and $\Delta$ , the step-size projector) that expand the hidden state capacity beyond what linearized attention can represent.
Multi-Stage Distillation Process — a three-phase training procedure: (Phase 1) pseudo-label generation and word-level + sequence-level distillation on synthetic chat data, with frozen MLP layers; (Phase 2) supervised fine-tuning on standard instruction datasets, with all parameters trainable; (Phase 3) Direct Preference Optimization (DPO) using the teacher as the reference model, aligning the student's output distribution with human preferences.
Multi-Step Speculative Decoding Engine — a hardware-aware inference algorithm that uses a small draft model to propose batches of candidate tokens, then verifies them against the student model using a custom GPU kernel that can compute multiple RNN steps, save an intermediate hidden state snapshot, and rewind to the correct state upon token rejection — all without materializing intermediate RNN states in GPU memory.

Information flows as follows: the pretrained Transformer → weight initialization produces a hybrid model with Mamba blocks in place of some attention layers → Phase 1 distillation trains only the Mamba layers to mimic the teacher's outputs on synthetic chat data → Phase 2 SFT fine-tunes all parameters on instruction-following data → Phase 3 DPO aligns the model to human preferences using the teacher as the reference → the final distilled model serves as the verifier in speculative decoding, with a small draft model proposing candidate tokens → the multi-step kernel verifies and accepts/rejects tokens, managing hidden state rewinding lazily.

3.3 Roadmap for the Deep Dive

This paper is primarily a systems and methods paper whose core idea is that Transformer knowledge can be efficiently transferred to linear RNN architectures through weight initialization and lightweight distillation, and that the resulting models can be further accelerated with architecture-specific speculative decoding. The technical deep dive proceeds as follows, ordered to build understanding from the mathematical foundation upward:

First, the mathematical relationship between linearized attention and linear RNNs (Section 2.1 in the paper, reprised in Section 3.4.1) — this establishes why weight transfer works at all, by showing that with appropriate parameter choices a linear RNN can exactly implement a softmax-removed version of attention. This is the theoretical justification for the initialization scheme.
Second, the specific Mamba parameterization and the attention-to-Mamba initialization algorithm (Section 2.2–2.3, reprised in Section 3.4.2) — how the expanded state space of Mamba ( $N' > 1$ ) extends the naive linear attention approximation, what new parameters are introduced ( $A$ , $\Delta$ ), and precisely how the Transformer's $Q, K, V$ projections map onto the Mamba block's components.
Third, the multi-stage distillation procedure (Section 3, reprised in Section 3.4.3) — the loss function combining word-level KL divergence and sequence-level knowledge distillation, the three-phase training schedule (pseudo-label training with frozen MLPs, SFT, DPO), and the specific datasets, hyperparameters, and design choices at each stage.
Fourth, the multi-step speculative decoding algorithm and its hardware-aware kernel design (Section 4, reprised in Section 3.4.4) — the challenges unique to RNN speculation (no KV cache to rewind, parallel scan overhead for short sequences), the MultiStep kernel's interface and avoidance of intermediate state materialization, and the fused kernel optimizations needed to achieve speedups on H100 GPUs.
Fifth, the hybrid speculative decoding extension — how attention layers are handled during verification (they use standard parallel attention since they retain KV caches), and how the draft model's own RNN state is recomputed for the next round of speculation.

This order mirrors the actual pipeline: initialization comes before training, which comes before inference. Each component depends on concepts introduced in the previous one.

3.4 Detailed, Sentence-Based Technical Breakdown

This is primarily a systems and methods paper whose core idea is that the linear projections learned by Transformer attention heads encode reusable semantic routing functions that can directly initialize the corresponding projections in a Mamba linear RNN, and that combining this weight transfer with a lightweight multi-stage distillation on post-training data — rather than attempting to distill the entire pretraining process — is sufficient to produce hybrid models competitive with Transformers while enabling hardware-aware speculative decoding for further acceleration.

3.4.1 The Mathematical Bridge: From Attention to Linear RNNs

The paper's weight initialization scheme is not arbitrary — it rests on a formal relationship between the attention mechanism and linear recurrent neural networks. Understanding this relationship requires examining the computation performed by a single attention head.

Canonical multi-head attention. For a single attention head, the input at sequence position $t$ is a vector $o_t \in \mathbb{R}^{D \times 1}$ (where $D$ is the model's hidden dimension). Three learned linear projections produce query, key, and value vectors:

$Q_t = W_Q o_t, \quad K_t = W_K o_t, \quad V_t = W_V o_t$

where $W_Q, W_K, W_V \in \mathbb{R}^{N \times D}$ are learned weight matrices, and $Q_t, K_t, V_t \in \mathbb{R}^{N \times 1}$ where $N = D / H$ is the per-head dimension for $H$ attention heads. The attention output at position $t$ is a weighted sum of all previous value vectors, where the weights are determined by the softmax-normalized dot-product similarity between the current query and all previous keys:

$y_t = \sum_{s=1}^t \alpha_s V_s, \quad \text{where} \quad \alpha_1 \ldots \alpha_T = \text{softmax}\left(\left[m_{1,t} Q_t^\top K_1 \ldots m_{T,t} Q_t^\top K_T\right] / \sqrt{D}\right)$

Here $m_{s,t} = \mathbb{1}(s \leq t)$ is a causal mask indicator (1 if position $s$ is not in the future relative to $t$ , 0 otherwise), preventing the model from attending to future tokens.

What this computes, operationally: At each position $t$ , the attention head first projects the current token representation into three distinct roles — a query ( $Q_t$ , "what am I looking for?"), a key ( $K_s$ for all prior positions, "what do I contain?"), and a value ( $V_s$ , "what information should I convey?"). It then compares the query against every prior key using dot-product similarity, normalizes these similarities into a probability distribution via softmax, and produces the output as the probability-weighted mixture of all prior value vectors. The causal mask ensures the mixture only includes past positions, preserving autoregressive generation constraints.

Linearizing attention by removing the softmax. The softmax nonlinearity is what makes attention expensive — it requires computing and normalizing over all $t$ similarity scores for each position $t$ , yielding the $\mathcal{O}(L^2)$ complexity. If we simply remove the softmax (and temporarily ignore the $\sqrt{D}$ scaling), the attention formula becomes purely linear:

$y_t = \sum_{s=1}^t \alpha_s V_s = \frac{1}{\sqrt{D}} \sum_{s=1}^t m_{s,t} Q_t^\top K_s V_s = \frac{1}{\sqrt{D}} Q_t^\top \sum_{s=1}^t m_{s,t} K_s V_s$

This algebraic rearrangement reveals a crucial property: the weighted sum can be factored so that the summation over the past ( $\sum_{s=1}^t$ ) operates independently of the current query $Q_t$ . The term being accumulated, $K_s V_s$ , is an outer product producing an $N \times 1$ vector (since $K_s \in \mathbb{R}^{N \times 1}$ is a column vector and $V_s \in \mathbb{R}^{N \times 1}$ is also a column vector — the paper treats $V_s$ as a scalar post-multiplier here, which requires viewing it as a linear readout). More precisely, the paper's notation absorbs the value projection into the state, so the accumulated term is $K_s \otimes V_s$ contributing to a hidden state of dimension $N$ .

From linearized attention to a linear RNN. This factored form exactly matches the generic linear RNN recurrence (Equation 1 in the paper):

$h_t = A_t h_{t-1} + B_t x_t, \quad y_t = C_t^\top h_t$

By making the following substitutions, linearized attention becomes a specific instance of a linear RNN:

$h_t = \sum_{s=1}^t m_{s,t} K_s V_s$ — the hidden state accumulates key-value outer products
$A_t = m_{t-1,t}$ — the state transition is simply the causal mask indicator (1 when moving from $t-1$ to $t$ , since past state is fully retained)
$B_t = W_K o_t$ — the input projection is the key projection from attention
$x_t = W_V o_t$ — the input signal is the value projection
$C_t = W_Q o_t$ — the readout projection is the query projection
$h_t \in \mathbb{R}^{N \times 1}$ — the hidden state dimension equals the per-head attention dimension

Why this form matters. This derivation shows that attention and linear RNNs are not fundamentally different computational primitives — linearized attention is a linear RNN with specific choices for the state transition $(A_t)$ , input projection $(B_t)$ , and readout $(C_t)$ . The softmax in standard attention is what creates the quadratic complexity, and removing it yields an $\mathcal{O}(L)$ recurrent form. The problem, which the paper explicitly acknowledges, is that "naively applying this transformation leads to poor results" because "the softmax nonlinearity is critical to attention" — it provides selective focus, preventing the model from being overwhelmed by irrelevant context. The linear RNN with $N$ -dimensional hidden state (equal to one scalar per hidden dimension) simply lacks the capacity to approximate this selectivity.

Increasing capacity through state expansion. The key insight is that we can improve the linear RNN by expanding the hidden state dimension from $N$ (one scalar per attention dimension) to $N \times N'$ for some expansion factor $N' > 1$ , while keeping the inputs and outputs at their original dimensions. The expanded state can learn to approximate the softmax's selectivity through richer dynamics, even though the individual operations remain linear. The Mamba architecture provides precisely this expansion mechanism through its continuous-time state-space parameterization and discretization, which is what the paper leverages next.

3.4.2 Attention-to-Mamba Initialization: Architecture and Weight Transfer

The paper does not use the naive linear attention RNN. Instead, it adapts the Mamba architecture (Gu and Dao, 2023) to serve as the linear RNN layer, and designs an initialization scheme that maps Transformer attention weights onto the Mamba parameters. This section details both the Mamba parameterization and the weight transfer procedure.

Mamba's continuous-time state-space model (SSM). Mamba is built on a continuous-time linear dynamical system described by the differential equations:

$h'(k) = A h(k) + B(k) x(k), \quad y(k) = C(k) h(k)$

where:

$k$ is a continuous time variable,
$h(k)$ is the hidden state evolving over time,
$A$ is a diagonal matrix governing how the state decays or transforms in the absence of input,
$B(k)$ and $C(k)$ are time-varying input and output projections, respectively,
$x(k)$ is the continuous input signal,
$y(k)$ is the continuous output signal.

What this describes, operationally: The hidden state $h(k)$ changes at a rate equal to its own natural dynamics ( $A h(k)$ ) plus the influence of the current input projected through $B(k)$ . The output $y(k)$ is a projection of the current state through $C(k)$ . Because $A$ is diagonal, each dimension of the hidden state evolves independently (though they are mixed through the shared input and output projections). The time-varying nature of $B$ and $C$ means the system can dynamically choose what to read and write based on the current input — this is the "selective" property that distinguishes Mamba from earlier state-space models with fixed dynamics.

Discretization for discrete-time language modeling. Language is inherently discrete (sequences of tokens), so the continuous-time system must be converted to a discrete-time linear RNN. Mamba uses a neural network to produce, for each input position $t$ , a sampling interval $\Delta_t$ (a vector in $\mathbb{R}^{N'}$ ) and samples of the continuous signals $B_t$ and $C_t$ . Given these, a discretization function (such as zero-order hold) produces discrete-time parameters:

$\bar{A}_{1\ldots T}, \bar{B}_{1\ldots T}, \bar{C}_{1\ldots T} = \text{Discretize}(A, B_{1\ldots T}, C_{1\ldots T}, \Delta_{1\ldots T})$

where the overbar notation indicates discrete-time parameters. The discrete-time linear RNN then operates as:

$h_t = \bar{A}_t h_{t-1} + \bar{B}_t x_t, \quad y_t = \bar{C}_t^\top h_t$

What the discretization does, operationally: The continuous $A$ matrix describes how the state would evolve if time flowed continuously. The $\Delta_t$ parameter tells the system "how much time has passed" between token $t-1$ and token $t$ , which controls the effective decay rate — a small $\Delta_t$ means the state barely changes (the model is essentially ignoring this token), while a large $\Delta_t$ means the state updates significantly. The discretization converts the continuous dynamics into a single matrix $\bar{A}_t$ that captures the cumulative effect over the interval $\Delta_t$ , and scales $\bar{B}_t$ and $\bar{C}_t$ accordingly. Critically, the output dimensions expand: $\bar{B}_t, \bar{C}_t \in \mathbb{R}^{N' \times N \times 1}$ rather than the original $N \times 1$ . This means the effective hidden state $h_t$ has dimension $N \times N'$ , an $N'$ -fold expansion over the naive linear attention RNN.

Why this expansion matters (and why Mamba makes it efficient). The $N' > 1$ expansion gives the linear RNN the capacity to learn selective attention-like behavior that a scalar-state-per-dimension RNN cannot represent. However, a naive implementation that materializes the full $\bar{B}_t, \bar{C}_t \in \mathbb{R}^{N' \times N \times 1}$ tensors and performs the full $N \times N'$ -dimensional recurrence would be prohibitively slow. Mamba's core algorithmic contribution is a hardware-aware fused kernel that combines discretization, state expansion, and the linear RNN computation into a single GPU operation that avoids materializing these intermediate tensors. The paper inherits this kernel design for its hybrid models.

Algorithm 1: Attention-Initialized Mamba. The paper's specific architecture is defined in Algorithm 1 (Section 2.3). For each attention head in the original Transformer:

Weight transfer (orange weights in Figure 1): The linear projection matrices from the Transformer's attention head are directly reused:
- $W_Q$ → initializes $C$ (the readout projection: "what to read from state")
- $W_K$ → initializes $B$ (the input projection: "how to write to state")
- $W_V$ → initializes $x$ (the input signal being written)
- $W_O$ → initializes the output projection (combining all heads' outputs)
New parameters (green weights in Figure 1): Two new parameter sets are introduced that have no Transformer counterpart:
- $A \in \mathbb{R}^{N \times N'}$ : the diagonal state transition matrix, learned from scratch
- $\text{MLP}$ : a small multi-layer perceptron that takes the input $x_t$ and produces the step size $\Delta_t \in \mathbb{R}^{N'}$
Per-position computation (Algorithm 1 lines 10–16):
- $x_t = W_V o_t$ — value projection, same shape as in attention
- $B_t = W_K o_t$ — key projection
- $C_t = W_Q o_t$ — query projection
- $\Delta_t = \text{MLP}(x_t)$ — input-dependent step size
- $\bar{A}_{1:T}, \bar{B}_{1:T}, \bar{C}_{1:T} = \text{Disc}(A, B, C, \Delta)$ — discretize using the step sizes
- $y = \text{LinearRNN}(\bar{A}, \bar{B}, \bar{C}, x)$ — apply the discrete-time linear RNN
- $\text{output} \mathrel{+}= W_O^\top y$ — project back to model dimension and add to running output

What this design achieves, operationally: At each token position, the Mamba block receives the same input representation that the original attention head would have received. It projects this input into the same $Q, K, V$ roles (now labeled $C, B, x$ ), but instead of computing attention scores and softmax-weighted mixtures, it uses the key projection $B_t$ to update an expanded hidden state and the query projection $C_t$ to read from that state. The step size $\Delta_t$ , produced by a small MLP from the value projection, dynamically controls how much attention to pay to the current token — this is Mamba's "selective" mechanism, analogous to how the softmax in attention dynamically weights different tokens. The output projection $W_O$ combines all heads identically to how the Transformer would.

Grouped-query attention handling. Modern Transformer models like Llama use grouped-query attention (GQA), where multiple query heads share the same key and value projections. The paper notes that the initialization "requires processing additional components like grouped query attention that shares keys and values across heads" (Section 2.3). The adaptation shares the $W_K$ and $W_V$ projections across the corresponding Mamba heads, matching the sharing pattern of the original Transformer.

Hybrid architecture and interleaving. The paper produces hybrid models by replacing only some of the Transformer's attention layers with Mamba blocks. The MLP (feed-forward network) layers from the Transformer are preserved in their entirety and are not replaced. When deciding which attention layers to replace, the paper experiments with "hybrid models where we keep every $n$ attention layers" — meaning the replacement pattern is interleaved: keep layer 1 as attention, replace layer 2 with Mamba, keep layer 3 as attention, replace layer 4 with Mamba, and so forth for a 50% attention configuration. Empirically, interleaved placement outperforms contiguous replacement (e.g., all Mamba layers grouped together). Section 3 (Stepwise Training) notes 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." This progressive replacement during training (rather than replacing all target layers at once and then training) is a separate design choice from the interleaving pattern and is discussed further in the distillation section.

Why reuse attention weights rather than random initialization. Table 8 provides the empirical evidence: a hybrid model with Mamba layers randomly initialized achieves 1.04 MT-Bench score and 0.02% AlpacaEval win rate, while the attention-initialized version achieves 6.69 and 14.11% respectively — a catastrophic difference. Intuitively, the $W_Q, W_K, W_V$ projections encode how the Transformer routes information: $W_Q$ determines what semantic features each head looks for, $W_K$ determines what features each token advertises, and $W_V$ determines what information is transmitted when a match occurs. By preserving these projections, the Mamba block starts with the same "vocabulary" of queries, keys, and values that the Transformer learned through pretraining. The only thing it needs to learn is how to use its expanded state space (via $A$ and $\Delta$ ) to approximate the softmax selection that the Transformer performed explicitly — a much easier learning problem than discovering meaningful projections from scratch.

3.4.3 Multi-Stage Knowledge Distillation for Aligned Language Models

The paper's distillation process operates on the post-training pipeline, not the pretraining stage. This is a strategically motivated choice: the MLP layers, which contain much of the model's factual knowledge, conceptual understanding, and linguistic competence, are preserved exactly from the Transformer. The distillation only needs to teach the new Mamba layers how to perform the attention-like information routing that the original layers handled. The process has three phases.

Phase 1: Progressive pseudo-label distillation with frozen MLPs. In the first stage, the partially initialized hybrid model is trained to mimic the teacher Transformer's output distribution on synthetic chat data. The key design choices here are:

Teacher-generated pseudo-labels: Instead of using human-written responses, the teacher model (Zephyr-7B or Llama-3 8B) generates its own responses to prompts from the UltraChat and UltraFeedback datasets. This provides an infinite supply of training data where the "correct" answer is whatever the teacher model would produce.
Combined sequence-level and word-level distillation loss (Equation 2):

$\mathcal{L}(\theta) = -\sum_{t=1}^T \left[ \alpha \log p(\hat{y}_{t+1} \mid \hat{y}_{1:t}, x, \theta) + \beta\, \text{KL}\left[ p(\cdot \mid \hat{y}_{1:t}, x, \theta_T) \,\|\, p(\cdot \mid \hat{y}_{1:t}, x, \theta) \right] \right]$

where $\theta$ are the trainable parameters of the student model, $\theta_T$ are the frozen parameters of the teacher, $x$ is the input prompt, $\hat{y}_{1:T}$ is the teacher-generated response sequence, $\alpha$ controls the weight of the sequence-level loss, and $\beta$ controls the weight of the word-level loss.

What this loss computes, operationally: At each generation step $t$ , the student produces two things: a scalar log-probability for the specific next token that the teacher generated (the $\log p$ term), and an entire probability distribution over all possible next tokens (the $p(\cdot \mid \ldots, \theta)$ term). The first term — sequence-level knowledge distillation (SeqKD) — simply maximizes the likelihood of the teacher's chosen token $\hat{y}_{t+1}$ , effectively treating the teacher's output as ground truth. This encourages the student to produce the same exact tokens as the teacher. The second term — word-level KL divergence — compares the student's full output distribution against the teacher's full output distribution at every position, penalizing the student when it assigns probability mass differently from the teacher even if it would have picked the same highest-probability token. This provides a richer training signal: the student learns not just "what the teacher said" but "what the teacher considered plausible alternatives."

Why this combined form. Using only sequence-level distillation (the $\alpha$ term, $\beta = 0$ ) provides sparse supervision — one correct token per position, with no information about the relative plausibility of other tokens. Using only KL divergence ( $\alpha = 0$ ) can be unstable because it requires the teacher's full distribution (which may be poorly calibrated for very unlikely tokens) and provides no hard target to anchor the student. The combination provides both a strong directional signal (match the teacher's chosen tokens) and distributional regularization (match the teacher's uncertainty). The paper sets $\alpha = 1$ and $\beta = 0.1$ , putting primary weight on sequence-level matching with a smaller distributional regularization term.

Frozen MLP layers: Only the Mamba-specific parameters ( $A$ , $\Delta$ MLP, and the discretization-related weights) are trained in Phase 1. The Transformer-derived $W_Q, W_K, W_V, W_O$ projections and all MLP layers are frozen. The paper states the rationale: "We only freeze Gated MLP (FFN) in the first stage, while in the second and final stage all parameters are trained." The idea is to first train the Mamba components to work well with the frozen Transformer components, establishing a stable foundation before jointly fine-tuning everything.
Stepwise progressive replacement: Before Phase 1 training starts, the paper applies a progressive distillation strategy. Instead of immediately replacing all target attention layers with Mamba blocks, they first replace only half of the target layers (e.g., for a 25% attention target model, they start with a 50% attention model), train that, then replace the remaining layers, and train again. Table 7 (Right) shows the benefit: training a 25% attention model with stepwise distillation achieves perplexity 2.20, compared to 2.89 without stepwise distillation. The progressive approach ensures that at each stage, most of the model's layers are still well-functioning Transformer components, providing a stronger learning signal for the newly replaced Mamba layers.
Training hyperparameters: One epoch on UltraChat/UltraFeedback pseudo-labels, AdamW optimizer with $\beta = (0.9, 0.98)$ , batch size 64, linear learning rate warmup for the first 500 steps followed by cosine annealing.

Phase 2: Supervised Fine-Tuning (SFT). After Phase 1, the student model has learned to approximately mimic the teacher on synthetic chat data, but it has not been trained on the diverse instruction-following datasets used to produce modern chat LLMs. Phase 2 performs standard supervised fine-tuning — maximizing the likelihood of high-quality human/AI responses given prompts — on three datasets: GenQA, InfinityInstruct, and OpenHermes 2.5.

All parameters are now trainable: Unlike Phase 1 where MLP layers were frozen, Phase 2 unfreezes everything. The rationale is that after Phase 1, the Mamba layers are sufficiently well-initialized that joint training can refine both the Mamba and MLP components together without destabilizing either.
Training configuration: One epoch, same hyperparameters as used for Zephyr's original training (AdamW, $\beta = (0.9, 0.98)$ , batch size 64, linear warmup for 500 steps, cosine annealing).

Phase 3: Distilled Direct Preference Optimization (Distilled DPO). The final stage aligns the student model with human preferences. The paper adapts Direct Preference Optimization (DPO) — originally a method for training language models from pairwise preference data using a reference model — into a distillation objective by using the teacher Transformer as the reference model.

Standard DPO formulation (Equation 4):

$\pi_\theta = \max_\theta \mathbb{E}_{(x, y_w, y_l) \sim \mathcal{D}} \log \sigma\left( \beta \log \frac{p(y_w \mid x; \theta)}{p(y_w \mid x; \theta_T)} - \beta \log \frac{p(y_l \mid x; \theta)}{p(y_l \mid x; \theta_T)} \right)$

where $x$ is a prompt, $y_w$ is the preferred (winning) response, $y_l$ is the dispreferred (losing) response, $\theta$ are the student parameters, $\theta_T$ are the frozen teacher (reference) parameters, $\beta$ controls the strength of the KL penalty keeping the student close to the teacher, and $\sigma$ is the logistic sigmoid function.

What this computes, operationally: For each prompt $x$ , the dataset provides a pair of responses — one that human raters preferred ( $y_w$ ) and one they dispreferred ( $y_l$ ). The student model computes the log-probability of each response under both itself ( $p(\cdot \mid \theta)$ ) and the frozen teacher ( $p(\cdot \mid \theta_T)$ ). It then computes the ratio of student probability to teacher probability for each response, capturing how much the student has diverged from the teacher's assessment. The log-difference of these ratios — $\log(p(y_w|\theta)/p(y_w|\theta_T)) - \log(p(y_l|\theta)/p(y_l|\theta_T))$ — is the "preference margin" from the student's perspective relative to the teacher's baseline. The sigmoid $\sigma(\cdot)$ converts this margin to a probability that $y_w$ is better than $y_l$ , and the objective maximizes the log of this probability. In effect, the student is rewarded when it assigns relatively higher probability (compared to the teacher) to preferred responses and relatively lower probability to dispreferred ones.

Why use the teacher as the reference rather than the standard approach. In standard DPO, the reference model is the student's own checkpoint before DPO training (typically the SFT model). This ensures the DPO update does not drift too far from the SFT-tuned model, preserving general capabilities. The paper replaces this with the original teacher Transformer. This has a specific benefit for distillation: the student, being architecturally constrained (fewer attention layers), might not be able to perfectly match the teacher's output distribution. If the student's own SFT checkpoint were used as reference, DPO would optimize relative to a baseline that is already architecturally limited. Using the teacher as reference means DPO optimizes relative to the best possible model, providing a stronger signal about which responses are genuinely better. The paper notes: "As far as we are aware this is the first use of DPO as a distillation objective."

Training data: For models distilled from Zephyr, DPO uses the UltraFeedback dataset (consistent with the teacher's own training). For models distilled from Llama-3 Instruct 8B, datasets from SimPO and Zephyr are used.
Training design choice (frozen vs. unfrozen during DPO): The paper specifies that during Phase 3, all parameters are trained (consistent with Phase 2), not just the Mamba-specific ones. Table 6 (Right) shows the ablation: Dis+DPO alone achieves 5.42 MT-Bench for the 50% attention hybrid, while Dis+SFT+DPO (the full pipeline) achieves 6.69 — confirming that DPO builds on SFT, not replaces it.

Total training cost. The paper reports that "the total distillation process for each hybrid model (e.g., Mamba-Llama3 (50% att)) takes less than five days in 8x80G A100" — approximately 20B tokens of training data across all three phases. This is several orders of magnitude less than the trillions of tokens used to pretrain the original Transformer or to train a linear RNN from scratch, making the approach accessible to academic GPU resources.

3.4.4 Multi-Step Speculative Decoding for Linear RNNs

The final component is a custom speculative decoding algorithm designed to accelerate inference for the distilled hybrid models. Standard speculative decoding (for Transformers) exploits the fact that the verifier model can check multiple draft tokens in parallel faster than generating them sequentially. Linear RNNs break this assumption and require a fundamentally different approach.

The standard speculative decoding loop (for context). In standard Transformer speculative decoding, a small fast draft model generates $K$ candidate tokens autoregressively. The large verifier model then takes all $K$ tokens, computes attention over them in parallel (since it has the KV cache for the prefix), and checks at each position whether the draft token matches what the verifier would have produced. Tokens are accepted up to the first mismatch, at which point the verifier's corrected token is used and the process repeats. The key enabling property is that verification is much faster per token than generation because the Transformer's attention computation can be parallelized over the $K$ candidate tokens.

Why this breaks for RNNs. The paper identifies two linked challenges (Section 4.1):

No parallelism advantage for verification. The RNN's sequential generation mode (updating $h_t$ one step at a time) is already fast — there is no slow attention mechanism to parallelize. The alternative — using the training-time parallel scan mode to compute multiple RNN steps simultaneously — is efficient for long sequences but has significant overhead for the short batches ( $K = 3$ – $4$ tokens) used in speculative decoding. The parallel scan "relies on hardware-aware optimizations, such as avoiding materializing intermediate states," which are designed for sequence lengths of hundreds or thousands, not 3–4.
State rewinding requires expensive caching. In the RNN, the state after position $t$ is a single vector $h_t$ . If verification rejects the draft at position $j+1$ (meaning tokens $1$ through $j$ were correct but $j+1$ is wrong), the verifier needs $h_j$ to continue generation. Without caching, the RNN would have to recompute forward from $h_0$ through $j$ steps — expensive and defeating the speedup. But caching all intermediate states $h_1, h_2, ..., h_t$ would consume $t \times \text{dim}(h)$ memory. Since Mamba's hidden state is $N \times N'$ dimensional (expanded by $N'$ , typically 16), this would be comparable to or larger than the Transformer's KV cache — eliminating the memory efficiency advantage that motivated using an RNN in the first place.

The MultiStep kernel: solving both challenges at once. The paper's solution is a single hardware-aware GPU kernel that computes multiple RNN steps while maintaining the ability to produce specific intermediate states on demand, without materializing the full sequence of hidden states. The kernel's interface is:

$y_{j:k}, h_j, h_k \leftarrow \text{MultiStep}(h_i, y_{1:n}, i, j, k; A, B, C, \Delta)$

where $h_i$ is a cached hidden state at some earlier position $i$ (with $i \leq j$ ), $y_{1:n}$ is the full sequence of tokens (including both verified prefix and new draft tokens), $j$ is the position of the last verified token (where we want a state snapshot), and $k$ is the position up to which we want output logits and a final state ( $j \leq k$ ).

What this kernel computes, operationally, in one fused GPU call:

It starts from the cached state $h_i$ (which was saved at some earlier round of speculation).
It runs the SSM recurrence forward from position $i$ to position $k$ , computing the hidden state updates for each token.
It produces two specific hidden states as outputs: $h_j$ (the state after the last verified token — saved as the new cache for potential rewinding) and $h_k$ (the state after all draft tokens — saved if all tokens are accepted).
It produces the output logits $y_{j:k}$ for the draft token positions.
Crucially, it does not materialize in GPU memory the intermediate states $h_{i+1}, h_{i+2}, ..., h_{k-1}$ or the discrete-time parameters $\bar{A}_t, \bar{B}_t, \bar{C}_t$ . The kernel's internal computation is structured to compute only what is needed for the requested outputs, using the same hardware-aware techniques (SRAM-based tiling, recomputation rather than storage) that make Mamba training efficient.

Why this design solves the two challenges:

No parallelism overhead for short sequences: The kernel is optimized for the specific pattern of speculative decoding — computing a few steps ( $k - j$ , typically 3–4) from a known starting state — rather than the general parallel scan. It exploits the fact that for short sequences, the overhead of the parallel scan's divide-and-conquer structure exceeds the cost of sequential computation, so it uses an optimized sequential approach instead.
No state caching overhead: By fusing the computation into a single kernel that receives the starting state $h_i$ and the token sequence as inputs, the kernel can recompute intermediate states on-the-fly without storing them in global memory. When verification fails at position $j+1$ , the kernel has already produced $h_j$ as an output — this is the exact state needed for the next round. Only one cached state ( $h_i$ ) needs to be maintained externally; the kernel handles the rest internally and transiently.

The full speculative decoding algorithm (Algorithm 2). The algorithm operates as follows:

Initialization: Save the initial hidden state $h_{\text{cache}} \leftarrow h_0$ , and set $j \leftarrow 0$ (position of last verified token).
Speculation step: The draft model generates $K$ candidate tokens $y_{j+1:j+K}$ autoregressively, conditioned on the verified prefix $y_{1:j}$ .
Verification step (the MultiStep kernel call): Call $y'_{j:j+K}, h_j, h_{j+K} \leftarrow \text{MultiStep}(h_{\text{cache}}, y_{1:j+K}, i, j, j+K; \theta_V)$ , where $i$ is the position of the cached state and $y'_{j:j+K}$ are the verifier's own top predictions at each position. Find the first position $k'$ where the verifier's prediction $y'$ differs from the draft $y$ .
State update: If $k' = j+K$ (all tokens accepted), update $h_{\text{cache}} \leftarrow h_{j+K}$ . Otherwise (rejection at position $k'$ ), update $h_{\text{cache}} \leftarrow h_j$ , replace the rejected token with the verifier's prediction $y'_{k'}$ , and set $j \leftarrow j+1$ (since position $j$ was verified correct).
Repeat from step 2 until end-of-sequence.

What makes this work for hybrid models. The distilled models contain both Mamba layers and retained attention layers. During speculative verification:

Mamba layers: Process the draft tokens using the MultiStep kernel and state caching mechanism described above.
Attention layers: Retain their standard KV cache mechanism. During verification, the attention layers simply perform parallel attention over the $K$ candidate tokens against the cached keys and values from the prefix — exactly as in standard speculative decoding.
Integration: The two types of layers alternate within the model as they would in normal forward computation. The MultiStep kernel is applied only to the Mamba layers; the attention layers use their standard parallel verification path. The outputs are combined through the residual stream identically to normal model forward pass.

Draft model for speculation. The paper does not use the distilled model itself as its own draft model (which would be self-speculation). Instead:

For pure Mamba verifiers (Table 1), the draft models are: a 130M Mamba model for the 2.8B verifier, and a Llama3 1B model for the 7B verifier.
For hybrid verifiers (Table 5), the draft models are small Transformer models (2-layer or 4-layer) trained via "shrink and fine-tune" from the teacher Transformer (Zephyr-7B). Specifically, the 2-layer draft uses layers at indices [0, 31] from Zephyr-7B, and the 4-layer draft uses layers [0, 10, 20, 31]. Embeddings and the language model head are also taken from Zephyr-7B. The drafts are fine-tuned on the OpenHermes 2.5 dataset with loss masking on the prompt (only computing cross-entropy on response tokens).

Hardware-specific optimization for H100 GPUs. The paper notes a critical practical detail: the naive implementation achieved decent speedups on Ampere GPUs (RTX 3090) but "no speedup at all on H100s." The issue is that H100 GPUs have much faster GEMM (matrix multiplication) operations, which means the overhead of multiple kernel calls — launching separate kernels for recomputation, decoding, and caching — becomes the dominant cost. The solution was kernel fusion:

For the verifier model: Recomputation of previous steps from the cache, multi-step decoding for the new draft tokens, and caching of the new state are all fused into a single GPU kernel. This eliminates multiple kernel launch overheads and keeps intermediate values in on-chip memory.
For the draft model: Similarly, recomputation, decoding, and caching are fused into one kernel. Additionally, the convolutional part of the Mamba block is implemented using a circular buffer, which keeps track of old entries needed for the convolution when recomputing from a cached state.

Performance results (Table 1 and Table 5). The speedups achieved are:

Mamba 2.8B with 130M draft: 2.3–2.6× on 3090, 1.71–1.85× on H100
Mamba 7B with Llama3 1B draft: 2.1× on 3090, 1.95–2× on H100
Hybrid Mamba-Zephyr (50%) with 4-layer Transformer draft: 1.8× on 3090
Hybrid Mamba-Llama3 (50%) with 4-layer Transformer draft: 1.58–1.6× on 3090

The hybrid model speedups are lower than pure Mamba speedups primarily because the draft model is relatively larger (due to Llama 3's large embedding table), and the attention layers in the hybrid model cannot benefit from the MultiStep kernel optimization.

MultiStep kernel performance characteristics (Figure 3). The figure shows that the multi-step kernel takes approximately the same time to compute 2, 4, 8, 16, or 32 steps as a single step would take repeated that many times — there is essentially no per-step overhead in the multi-step computation. This is because the kernel's bottleneck is memory bandwidth (loading the input tokens and parameters), not compute — computing multiple RNN steps reuses the same parameters and amortizes the memory loads. The linear scaling in the figure confirms the kernel is memory-bound rather than compute-bound, which is the intended design for avoiding the parallel scan overhead.

4. Key Insights and Innovations

Innovation 1: Weight Initialization as Structural Knowledge Transfer, Not Just a Warm Start

The paper's most intellectually distinctive contribution is the recognition that the linear projection matrices from a Transformer's attention heads — $W_Q$ , $W_K$ , $W_V$ — encode semantically meaningful routing functions that can be directly repurposed as the corresponding projections in a Mamba linear RNN. This is not merely a convenient initialization trick that provides a non-random starting point. It is a fundamentally different approach to knowledge distillation that changes the nature of the learning problem from "approximate this function from scratch" to "learn the expansion from linearized attention to selective state-space dynamics."

Prior distillation work treated architectural conversion as a black-box output-matching problem. The approach of Ralambomihanta et al. (2024), which distilled Transformers into Hyena models using progressive layer-wise training, produced perplexity degradation factors of 2.36× (Table 6, Left) even on small 70M-parameter models. The implicit assumption was that the student architecture was sufficiently different that only output-level supervision — "produce similar token distributions" — could guide learning. The paper's key conceptual move is to recognize that, at the level of individual attention heads, the Transformer and Mamba share a common computational vocabulary: both must decide what to look for (query/readout), what to advertise (key/input projection), and what to transmit (value/input signal). The softmax in attention and the expanded state space in Mamba are different mechanisms for the same underlying operation — selective information routing — and the linear projections encode the routing policy independently of the selection mechanism.

The evidence for this interpretation is stark. Table 8 shows that a randomly initialized hybrid Mamba-Llama3 model achieves an MT-Bench score of 1.04 and AlpacaEval win rate of 0.02%, while the attention-initialized version reaches 6.69 and 14.11% respectively. This is not just a faster convergence result — the randomly initialized model never recovers, even after full distillation. The implication is that the attention projections contain knowledge that cannot be efficiently rediscovered through distillation alone at this data scale (~20B tokens). The MLP layers, which are preserved exactly, presumably contain the model's factual and conceptual knowledge; the attention projections contain its routing policy — which semantic relationships to attend to, which patterns to ignore. Transferring this routing policy via weight initialization allows the Mamba layers to inherit the Transformer's learned information-flow patterns, needing only to learn how the expanded state dynamics can improve upon the softmax's selection.

This insight is fundamental rather than incremental because it reframes the distillation problem from one of function approximation to one of mechanism substitution. It suggests a general principle: when two architectures differ in their selection mechanism (softmax attention vs. state-space dynamics) but share the same projection structure (linear maps from the residual stream to query/key/value roles), weight transfer across the projection interface can preserve learned behaviors while allowing the new mechanism to be trained with relatively little data. The concurrent MOHAWK work (Bick et al., 2024) operates on a similar insight, suggesting this principle may generalize across linear RNN variants.

Innovation 2: Post-Training Alignment as the Optimal Distillation Target, Not Pretraining

A second conceptual contribution is the strategic decision to apply distillation only to the post-training alignment pipeline — supervised fine-tuning and preference optimization — rather than attempting to distill the entire pretraining process. This choice is deceptively simple but reflects a non-obvious insight about where Transformer knowledge resides within the model's layers.

The paper's architecture preserves the Transformer's MLP (feed-forward network) layers exactly and only replaces attention layers with Mamba blocks. The implicit claim — validated by the results — is that the MLP layers contain the bulk of the model's factual knowledge, conceptual understanding, and linguistic competence acquired during pretraining, while the attention layers primarily handle information routing: deciding which pieces of stored knowledge to activate and combine at each generation step. If this decomposition holds, then distilling the attention layers only requires teaching the Mamba replacements how to route information similarly to the original attention — a problem solvable with lightweight post-training data (~20B tokens) rather than the trillions of tokens needed to instill factual knowledge from scratch.

The evidence supporting this decomposition is the model's performance on knowledge-intensive benchmarks after distillation. Table 3 shows that Mamba-Llama3 (50% attention) achieves 57.81 on MMLU and 55.63 on ARC-Challenge, while the original Llama-3 8B Instruct achieves 68.12 and 55.20 respectively (Table 3 baselines). The MMLU gap is notable but the model retains substantial factual knowledge — far more than could be acquired from the ~20B tokens of chat and instruction data used in distillation. This suggests the frozen MLP layers are indeed carrying forward the teacher's knowledge. Table 4 further confirms this: on GSM8K (math reasoning, which requires both knowledge and multi-step reasoning), the 50% attention hybrid achieves 67.85, dramatically outperforming Falcon Mamba-7B (41.32) which was trained from scratch on 5T+ tokens.

This insight is fundamental because it provides a decomposition principle for model compression via architecture conversion. The standard approach to model compression — train a smaller model to match a larger one — treats the model as a monolithic function to be approximated. The paper's approach treats the model as having separable knowledge storage (MLPs) and knowledge routing (attention) components, and only replaces the routing mechanism. This decomposition, if it generalizes, has significant practical implications: it means that for any new efficient architecture, the expensive part of knowledge acquisition (pretraining) can be bypassed entirely, and only the routing mechanism needs to be retrained. The paper's 20B-token, 5-day-on-8xA100 distillation pipeline is a concrete demonstration of this principle at the 7-8B parameter scale.

The decomposition is also falsifiable in an informative way. Table 9 shows what happens when Mamba blocks are removed entirely rather than initialized — the model's MT-Bench score collapses to 1.01. This confirms that the attention layers are not vestigial; they perform an essential routing function that cannot be absorbed by the MLPs alone, even with training. The Mamba layers are learning to perform this routing function, and the weight initialization provides the starting point.

Innovation 3: A New Diagnostic for Difficulty: The Attention-Percentage Scaling Curve

The paper inadvertently introduces a useful diagnostic concept: the attention-percentage scaling curve — how model quality degrades as a function of the fraction of attention layers retained. This is not presented as a formal contribution, but it emerges from the systematic evaluation across 50%, 25%, 12.5%, and 0% attention configurations and provides insight into the nature of the attention-to-Mamba transfer that goes beyond aggregate benchmark numbers.

The degradation pattern is not linear. Examining Table 2's MT-Bench scores for Llama3-Mamba: the teacher scores 8.00, the 50% hybrid scores 7.35 (a 0.65 point drop), the 25% hybrid scores 6.86 (a further 0.49 point drop from 50%), the 12.5% hybrid scores 6.46 (a further 0.40 point drop), and the 0% (pure Mamba) model scores 5.64 (a 0.82 point drop). The degradation accelerates at the extremes — moving from 50% to 25% attention is less damaging than moving from 12.5% to 0%. This suggests that a small number of attention layers provide disproportionately important routing capabilities that the Mamba layers cannot fully replicate, even with distillation.

The same pattern appears in the LM Eval benchmarks (Table 3). Mamba-Llama3 average scores: 50% = 61.30, 25% = 57.70, 12.5% = 55.02. Mamba2-Llama3: 50% = 63.84, 25% = 60.55, 12.5% = 58.21, 0% = 54.74. The drop from 12.5% to 0% attention (3.47 points for Mamba2) is roughly comparable to the drop from 50% to 12.5% (5.63 points), despite the latter removing three times as many attention layers.

This pattern is diagnostically significant rather than merely an empirical observation. It suggests that the attention layers do not all perform the same function; some layers — likely those handling long-range dependencies or complex multi-hop reasoning — rely more heavily on the softmax attention mechanism's ability to precisely select among many alternatives. Mamba's expanded state space can approximate this selection for many layers, but the approximation breaks down when too few attention layers remain to handle the cases where the state-space dynamics are insufficient. The attention-percentage scaling curve thus provides a lens for understanding which aspects of attention are hardest for linear RNNs to replicate — a question with implications for future architecture design. If specific capabilities degrade at specific attention-percentage thresholds, those capabilities are candidates for attention-specific mechanisms in future hybrid designs.

The long-context evaluation (Figure 4) reinforces this interpretation. The distilled 3B models achieve perfect needle-in-a-haystack retrieval at 10K tokens (5× the distillation length), and the 8B models at 16K+ tokens. This suggests that the Mamba layers are particularly effective at the long-range information preservation that attention struggles with (due to the KV cache memory bottleneck), while the retained attention layers handle the precise local reasoning that Mamba approximates less effectively.

Innovation 4: RNN Speculative Decoding as a Kernel Fusion Problem, Not an Algorithm Design Problem

The paper's approach to speculative decoding for linear RNNs represents a conceptual reframing of the problem. Prior work on speculative decoding (Leviathan et al., 2023; Chen et al., 2023; Spector and Re, 2023) treated it as an algorithm design problem: how to structure the draft-then-verify loop to maximize acceptance rates and minimize wasted computation. The paper identifies that for linear RNNs, the bottleneck is not algorithmic structure but rather kernel-level memory management — specifically, the tension between the need to rewind state on verification failure and the desire to avoid materializing intermediate RNN states that would consume prohibitive memory.

The key diagnostic insight is in Section 4.1: "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." The Transformer's KV cache is naturally a list of per-position states, making rewinding trivial. The RNN's hidden state is a single accumulated summary, making rewinding expensive unless intermediate states are explicitly saved — which would consume memory proportional to the (large) hidden state dimension times the sequence length, defeating the memory efficiency advantage of RNNs.

The paper's solution — the MultiStep kernel (Algorithm 2) — is not algorithmically novel in its control flow (it follows the standard draft-then-verify loop). Its novelty lies in recognizing that the problem can be solved through hardware-aware kernel fusion: by combining recomputation from a cached state, multi-step decoding, and state snapshotting into a single GPU kernel that avoids materializing intermediate states or discrete-time parameters, the algorithm can "recompute the correct state on the fly after a token is rejected" (Section 4.2) without paying the memory cost of caching or the latency cost of multiple kernel launches. The design is driven by the observation that decoding on modern GPUs is "bottlenecked by memory and not by compute" (Section 4.2), meaning that recomputing steps from a cached state is essentially free if done within a single kernel that keeps parameters in on-chip memory.

This reframing is fundamentally significant because it identifies a pattern likely to recur as more architectures move beyond the Transformer. The Transformer's KV cache makes speculative decoding algorithmically trivial (parallel verification is a natural consequence of the attention formulation). For architectures without this property — linear RNNs, but also potential future architectures with different state representations — the bottleneck shifts from "how to structure the speculation loop" to "how to manage state rewinding without materializing intermediate representations." The paper's kernel-fusion approach provides a template for solving this class of problems, and the H100-specific optimization (fusing verifier recomputation, decoding, and caching into one kernel) demonstrates that the solution is hardware-dependent in ways that matter for real deployment.

The concurrent work by Wu et al. (2024) proposes a similar algorithmic approach, but the paper's contribution lies in the detailed kernel engineering (the circular buffer for Mamba's convolution, the single-kernel fusion for H100) that makes the approach practically viable. Table 1 quantifies this: the naive implementation achieves "no speedup at all on H100s," while the optimized fused kernel achieves 1.71–2× speedup. This is a rare instance where the engineering contribution is as conceptually significant as the algorithmic one, because it demonstrates that the feasibility of speculative decoding for non-Transformer architectures depends on hardware-aware implementation details that are invisible at the algorithmic level.

Innovation 5: DPO as a Distillation Objective — Using the Teacher as the Reference Model

The paper's use of Direct Preference Optimization (DPO) with the teacher Transformer as the reference model — rather than the student's own pre-DPO checkpoint, which is standard practice — represents a small but conceptually significant adaptation of an existing method to the distillation setting. The paper notes: "As far as we are aware this is the first use of DPO as a distillation objective" (Section 3).

Standard DPO (Rafailov et al., 2024) uses the model's own SFT checkpoint as the reference model, penalizing divergence from this checkpoint during preference optimization. This ensures the model improves along the preference dimension without drifting from its general capabilities. The paper's adaptation substitutes the original teacher Transformer for the reference model. This changes the semantics of the DPO objective in a specific way: rather than asking "does the student prefer the winning response more than its SFT version did?", it asks "does the student prefer the winning response more than the teacher would have?" The teacher, being architecturally unconstrained and presumably higher-quality, provides a stronger reference point.

The significance of this adaptation is incremental but practically impactful in a way that reveals something about the distillation setting. The student model is architecturally constrained — it has fewer attention layers and Mamba approximations that cannot perfectly replicate attention. If the student's own SFT checkpoint were used as reference, DPO might reward the student for preferences that are achievable within its architectural constraints but not necessarily aligned with the best possible model. Using the teacher as reference ensures that DPO optimizes toward the capabilities of the unconstrained model, even if the student cannot fully reach them. This is a form of curriculum design through reference model choice: the reference model sets the aspiration level, and using the strongest available model as reference provides the most ambitious target.

Table 6 (Right) provides the evidence: Dis+DPO alone (without SFT) achieves 5.42 MT-Bench for the 50% attention hybrid, while Dis+SFT+DPO achieves 6.69. The DPO stage provides a meaningful improvement over SFT alone (which would score between the Dis+SFT value of 5.61 and the Dis+SFT+DPO value of 6.69), confirming that preference optimization with the teacher as reference adds value beyond what SFT alone provides. The fact that this works at all — that the architecturally constrained student can benefit from a preference signal defined relative to an unconstrained teacher — suggests that the preference optimization signal is partially architecture-independent: knowing which of two responses is better is useful even if the student cannot perfectly match the teacher's probability distribution over those responses.

This insight connects to a broader question in distillation: when the student has architectural limitations the teacher does not, what forms of supervision are most effective? Output-level supervision (KL divergence, sequence-level distillation) asks the student to match the teacher's exact distribution, which may be impossible due to capacity constraints. Preference supervision asks a coarser question — "is A better than B?" — which may be answerable even when exact distribution matching is not. The paper's DPO-as-distillation provides a concrete instance of this principle and demonstrates its viability at scale.

5. Experimental Analysis

Evaluation Methodology

Datasets. The paper evaluates on three categories of benchmarks: (1) Chat benchmarks: AlpacaEval 2 (single-turn instruction following scored by GPT-4 Turbo against reference responses) and MT-Bench (multi-turn conversation quality scored by GPT-4), both described in Section 5.2. (2) General zero-shot benchmarks: 10 tasks from the LM Evaluation Harness library (Gao et al., 2023) — WinoGrande (WG), PIQA (PQ), HellaSwag (HS), ARC-Easy (AE), ARC-Challenge (AC), MMLU (MM), OpenBookQA (OB), TruthfulQA (TQ), PubMedQA (PM), and RACE (RA) — with evaluation by analyzing model-assigned probabilities to answer choices (Section 5.3, Table 3). (3) Few-shot benchmarks: Open LLM Leaderboard tasks (25-shot ARC-Challenge, 10-shot HellaSwag, 5-shot MMLU, 5-shot WinoGrande) plus TruthfulQA (mc2 metric), GSM8K, and CRUX evaluated via ZeroEval (Lin, 2024) with greedy decoding, designed for instruct-tuned models (Section 5.3, Table 4). (4) Long-context evaluation: Needle-in-a-Haystack retrieval test measuring accuracy of retrieving a specific fact embedded at various positions within contexts up to 40K tokens (Section 5.4, Figure 4). For distillation training, pseudo-labels are generated from UltraChat (Ding et al., 2023) and UltraFeedback (Cui et al., 2023) seed prompts; SFT uses GenQA (Chen et al., 2024), InfinityInstruct (BAAI, 2024), and OpenHermes 2.5 (Teknium, 2023); DPO alignment uses UltraFeedback for Zephyr-distilled models and SimPO (Meng et al., 2024) plus Zephyr datasets for Llama-3 distilled models (Section 5.1).
Base model(s). The primary teacher models are two instruction-tuned Transformer LLMs at the 7-8B parameter scale: Zephyr-7B (Tunstall et al., 2023), a chat fine-tuned Mistral 7B using DPO alignment, and Llama-3 Instruct 8B (Dubey et al., 2024), an RLHF-aligned model from Meta. Additional experiments in later versions use Llama-3.1-Instruct 8B and Llama-3.2-Instruct 3B as teachers, including distillation from a Llama-3.1 70B teacher into 3B and 8B student hybrids. For the FLOPs-matched comparisons (Section 5.3), the baseline linear RNN models trained from scratch are: TRI Mamba 7B (Mercat et al., 2024) trained with 1.2T tokens, Nvidia Hybrid Mamba-2 8B (Waleffe et al., 2024) trained with ~3.7T tokens, Falcon Mamba 7B trained with >5T tokens, and RecurrentGemma-9B Instruct (Botev et al., 2024). The specification of the 70M Pythia model appears in the perplexity comparison in Table 6 (Left) against a Distill Hyena baseline (Ralambomihanta et al., 2024). The paper argues these models are "representative of the capabilities of many contemporary LLMs" (Section 1) and sit in a regime where test-time compute scaling can make a meaningful difference — non-trivial but far from saturated performance on MATH.
Metrics. The primary metrics across evaluations are: (1) MT-Bench score — GPT-4 judged quality on a 1-10 scale averaged over multi-turn conversations, with separate Round 1 and Round 2 scores reported (Table 2); (2) AlpacaEval 2 win rate — percentage of comparisons where the model's response is preferred over GPT-4's, reported as both length-controlled (LC) win rate and raw win rate with standard errors (Table 2); (3) Zero-shot accuracy on LM Eval benchmarks — exact match or normalized accuracy depending on the task, with metrics explicitly specified per-task in Section 5.3 (WG: accuracy, PI: accuracy, HS: normalized accuracy, AE/AC: accuracy and normalized accuracy, MM: accuracy, OB: normalized accuracy, TQ: accuracy, PM: accuracy, RA: accuracy); (4) Few-shot accuracy on Open LLM Leaderboard tasks with the shot counts specified per-task, plus mc2 metric for TruthfulQA and exact match for GSM8K and CRUX (Table 4); (5) Perplexity on held-out data for comparing distillation approaches (Table 6, 7); (6) Throughput in tokens/second and speedup ratio for speculative decoding experiments, measured on data from The Pile (Table 1) and OpenHermes 2.5 (Table 5), including the average number of generated tokens per speculative step (# Gen. Tokens); (7) Needle-in-a-Haystack retrieval accuracy — fraction of queries where the model correctly retrieves the embedded needle fact, visualized as a heatmap over context length × needle position (Figure 4).
Baselines. The paper compares against several categories of baselines: (1) The teacher Transformer models themselves — Zephyr-7B, Llama-3-Instruct 8B, Llama-3.1-Instruct 8B, and Llama-3.2-Instruct 3B — evaluated under the same metrics (Tables 2, 3, 4). This is the primary upper-bound comparison: the goal is to match or approach teacher quality. (2) Linear RNN models trained from scratch at comparable scales: TRI Mamba 7B (1.2T training tokens), Falcon Mamba 7B (>5T tokens), Nvidia Hybrid Mamba-2 8B (~3.7T tokens), and RecurrentGemma-9B Instruct (Botev et al., 2024; De et al., 2024). These represent the cost of achieving comparable quality without distillation. (3) Ablation baselines within the distillation framework: models with random initialization instead of attention weight transfer (Table 8), models with Mamba blocks removed entirely (Table 9), and models without stepwise progressive distillation or interleaved layer placement (Table 7). (4) Prior distillation work: the Distill Hyena approach (Ralambomihanta et al., 2024) which distills a 70M Pythia Transformer into a Hyena model, compared in perplexity terms in Table 6 (Left). (5) For speculative decoding: a non-speculative baseline (straight autoregressive generation) for computing speedup ratios (Tables 1, 5), with draft models described separately (130M Mamba draft, Llama3 1B draft, 2-layer and 4-layer Transformer drafts trained via shrink-and-fine-tune from Zephyr-7B).
Generation budget / compute accounting. The paper uses several distinct notions of "compute budget" depending on the analysis context: (1) For distillation: the total training data is approximately 20B tokens across all three phases, with Phase 1 using one epoch on UltraChat/UltraFeedback pseudo-labels, Phase 2 using one epoch on the combined SFT datasets, and Phase 3 using one epoch on the preference dataset. Training takes "less than five days in 8x80G A100" for each hybrid model (Section 5.1). The paper does not report total FLOPs for distillation, but the 20B token figure serves as the primary scale reference — several orders of magnitude less than the trillions of tokens used to pretrain the baseline linear RNNs. (2) For speculative decoding: the compute unit is tokens/second (throughput), with speedup computed as the ratio of speculative throughput to non-speculative baseline throughput on the same hardware. The number of draft tokens per step ( $K = 3$ or $4$ ) and the average number of accepted tokens per step (# Gen. Tokens) are reported to characterize the speculation efficiency (Tables 1, 5). Experiments are run on single NVIDIA RTX 3090 and H100 GPUs. (3) For long-context evaluation: the distillation context length is 2K tokens; evaluation extends to 40K tokens (20× the distillation length) to test length extrapolation (Figure 4). (4) For the attention-percentage scaling analysis: compute is implicitly measured in model capacity (number of attention layers retained), with configurations ranging from 100% (teacher), 50%, 25%, 12.5%, to 0% attention (pure Mamba), allowing comparison of quality degradation versus architectural simplification.
Cross-validation / statistical protocol. The paper does not employ cross-validation in the traditional sense. The primary evaluation strategy is: (1) For chat benchmarks (AlpacaEval 2, MT-Bench), models are evaluated once on the full test sets with GPT-4 as judge, with standard errors reported for AlpacaEval win rates (e.g., 29.61±1.31 for Mamba-Llama3 50%). (2) For LM Eval benchmarks, each task is evaluated once on its standard test set using the LM Evaluation Harness library (branch big-refactor). (3) For speculative decoding throughput, measurements are taken on data from The Pile or OpenHermes 2.5 — these are throughput measurements (tokens/sec) rather than accuracy metrics, and the paper does not report confidence intervals or multiple runs. (4) No formal statistical significance testing (t-tests, bootstrap confidence intervals, etc.) is reported for any comparison. The paper's claims rely on the magnitude of differences between models rather than statistical tests. An important caveat: the test sets used for evaluation (e.g., 500 questions for MATH, the standard AlpacaEval 2 set) are fixed benchmarks, and the models are evaluated once — there is no multiple-seed training or evaluation to quantify variance from random initialization or data ordering during distillation. The distillation process itself involves only one training run per configuration, meaning the reported performance numbers are point estimates without error bars from training stochasticity.

Main Quantitative Results

Chat Benchmark Performance: Teacher-Level Quality at 50% Attention Retention

Table 2 reports the headline result: the distilled hybrid Mamba models with 50% attention layers retained achieve chat quality comparable to their Transformer teachers and substantially outperform linear RNN models trained from scratch with orders of magnitude more compute.

Zephyr-7B distillation: Mamba-Zephyr (50% attention) achieves an MT-Bench score of 7.31 compared to the teacher Zephyr-7B's 7.34 — a difference of only 0.03 points, which is within the noise floor of GPT-4 judging. On AlpacaEval 2, the distilled model actually outperforms the teacher: 20.66% LC win rate (±0.74) versus 13.20% (±0.96), and 16.69% raw win rate (±1.10) versus 10.99% (±0.96). This reversal — the student exceeding the teacher on AlpacaEval — is unusual and the paper does not deeply analyze it, but it may reflect the DPO distillation stage using a stronger preference signal than the teacher's own alignment process. Performance degrades with fewer attention layers: Mamba-Zephyr (25%) scores 7.03 MT-Bench and 17.16% AlpacaEval LC, while Mamba-Zephyr (12.5%) drops to 6.40 MT-Bench and 15.32%.

Llama-3 Instruct 8B distillation: The top distilled model, Mamba-Llama3 (50% attention), achieves 7.35 MT-Bench compared to Llama-3-Instruct's 8.00 — a 0.65 point gap. On AlpacaEval 2, the pattern reverses again: the student achieves 29.61% LC win rate (±1.31) versus the teacher's 22.90% (±1.26), and 26.69% raw win rate (±1.31) versus 22.60% (±1.26). The Mamba2 variant shows similar behavior: Mamba2-Llama3 (50%) scores 7.32 MT-Bench and 26.78% AlpacaEval LC. Degradation with fewer attention layers follows the same pattern: 25% models score ~6.8 MT-Bench, 12.5% models score ~6.5, and the pure Mamba (0% attention) drops substantially to 5.64 MT-Bench and 14.49% AlpacaEval LC.

Comparison to models trained from scratch: Falcon Mamba Instruct 7B, trained from scratch with >5T tokens, achieves only 6.40 MT-Bench and 4.04% AlpacaEval LC — substantially worse than even the 12.5% attention distilled hybrid. GPT-3.5-turbo scores 7.94 MT-Bench and 22.70% AlpacaEval LC, positioning the 50% attention distilled models above GPT-3.5-turbo on AlpacaEval but slightly below on MT-Bench. GPT-4o achieves 57.46% AlpacaEval LC, serving as the upper reference.

The MT-Bench Round 1 vs. Round 2 breakdown (reported only for Llama-3 distilled models) shows that the performance gap between teacher and student is larger in Round 2 (multi-turn continuation): Mamba-Llama3 (50%) scores 7.82 in Round 1 but 6.88 in Round 2, compared to the teacher's single reported score of 8.00. This suggests that the Mamba layers' approximation of attention degrades more in multi-turn settings where maintaining long-term conversational context is critical — consistent with the intuition that the expanded state space is a compressed representation that loses fine-grained detail over long contexts.

Later model versions (Llama-3.1 and 3.2): Mamba-Llama3.1 (50%) achieves 7.7 MT-Bench versus the teacher's 8.0, and ~19% AlpacaEval LC versus the teacher's 20.9%. Mamba-Llama3.2 (50%) at the 3B scale achieves 6.9 MT-Bench versus the teacher's ~7.5 (the paper does not report the 3B teacher's MT-Bench explicitly, but the gap appears comparable). The DPO variants of these models (e.g., Llama3.1-Mamba2-DPO) generally show small improvements over the non-DPO versions on AlpacaEval but mixed results on MT-Bench — the Llama3.1-Mamba2-DPO variant achieves slightly higher Round 1 (8.1) but lower Round 2 (7.0) than the non-DPO version (8.1 vs. 7.3), for an overall score of 7.6 versus 7.7.

General Benchmark Performance: Competitive with Models Trained on Trillions of Tokens

Table 3 presents zero-shot evaluation on 10 tasks from the LM Evaluation Harness for Mamba and Mamba2 distilled from Llama-3 Instruct 8B, alongside baselines trained from scratch.

Headline comparison: Mamba-Llama3 (50%) achieves an average score of 61.30 across the 10 tasks, compared to TRI Mamba-7B (57.65, trained on 1.2T tokens) and Nvidia Hybrid Mamba-8B (59.60, trained on ~3.7T tokens). Mamba2-Llama3 (50%) achieves 63.84 average — substantially outperforming both baselines. This is a striking result: a model distilled with ~20B tokens outperforms models explicitly pretrained as linear RNNs with 60-180× more training data.

Task-level patterns: The distilled models' advantages are not uniform across tasks. On MMLU (knowledge-intensive multiple-choice across 57 subjects), Mamba-Llama3 (50%) achieves 57.81 and Mamba2-Llama3 achieves 55.70, compared to Nvidia Hybrid Mamba-8B's 51.46 and TRI Mamba-7B's 33.39. The large gap over TRI Mamba-7B on MMLU (57.81 vs. 33.39) likely reflects the frozen MLP layers carrying forward the teacher's factual knowledge — MMLU performance depends heavily on stored knowledge, which the Transformer MLPs retain. On HellaSwag (commonsense reasoning), the models are competitive: Mamba2-Llama3 (50%) achieves 79.47 versus TRI Mamba's 77.93 and Nvidia's 77.68. On ARC-Challenge (scientific reasoning), Mamba2-Llama3 (50%) achieves 58.19, outperforming all baselines including Nvidia's 47.70. On TruthfulQA (measuring hallucination/misconception avoidance), Mamba2-Llama3 (50%) achieves 57.74, substantially higher than Nvidia's 38.72 and TRI Mamba's 32.09 — this large gap may reflect the teacher Transformer's superior truthfulness inherited through the frozen MLP layers and weight-initialized attention projections.

Attention-percentage scaling in the LM Eval tasks: The degradation pattern from 50% → 25% → 12.5% → 0% attention varies by task. For Mamba2-Llama3: 50% average is 63.84, 25% drops to 60.55 (−3.29), 12.5% drops to 58.21 (−2.34), and 0% drops to 54.74 (−3.47). The drops are relatively uniform, but individual tasks show different sensitivities. MMLU degrades from 55.70 (50%) to 53.71 (25%) to 50.78 (12.5%) to 45.19 (0%) — a 10.5 point total drop, suggesting knowledge access depends more on attention layers than some other capabilities. ARC-Challenge drops from 58.19 (50%) to 47.95 (0%) — a 10.2 point drop, indicating challenging reasoning depends significantly on the precision of attention. In contrast, PIQA (physical commonsense) is nearly flat: 81.45 (50%) to 76.82 (0%) — only a 4.6 point drop, suggesting simpler reasoning tasks are more robust to attention replacement.

Mamba vs. Mamba2: The Mamba2 variants consistently outperform Mamba variants at the same attention percentage across most tasks. Average scores: Mamba2-Llama3 at 50% = 63.84 vs. Mamba-Llama3 at 50% = 61.30; at 25% = 60.55 vs. 57.70; at 12.5% = 58.21 vs. 55.02; at 0% = 54.74 vs. only Mamba2 reported. This gap (~2.5-3 points) is consistent and suggests the Mamba2 architecture (Dao and Gu, 2024), designed for better GPU utilization, provides a modest but reliable quality improvement over Mamba in the distillation setting.

Later model versions (Llama-3.1/3.2, Table 3): The patterns replicate. Llama3.1-Mamba2-DPO (50%) achieves the highest average at 65.31, exceeding the Llama-3.1-8B-Instruct teacher's 64.48 — another instance of the student surpassing the teacher. The 3B-scale Llama3.2-Mamba-DPO (50%) achieves 60.88 average, compared to the Llama-3.2-3B-Instruct teacher's 57.83 — a larger relative improvement at smaller scale. The DPO stage consistently adds 1-4 points across the board: Llama3.1-Mamba (50%) averages 62.32, while the DPO version averages 64.38 (+2.06); Llama3.2-Mamba (50%) averages 58.36, DPO version 60.88 (+2.52).

Table 4 provides few-shot evaluation on the Open LLM Leaderboard and ZeroEval tasks, including GSM8K (math reasoning) and CRUX (code reasoning). These evaluations use instruct-tuned models exclusively, with few-shot prompting.

Few-shot benchmark results: Mamba-Llama3 (50%) achieves 56.57 on ARC-Challenge (25-shot), 78.99 on HellaSwag (10-shot), 59.26 on MMLU (5-shot), and 69.06 on WinoGrande (5-shot). Compared to Falcon Mamba-7B-Instruct (62.03, 80.82, 62.11, 73.64 respectively), the distilled model is competitive but slightly lower on most metrics. The notably different metric is TruthfulQA mc2: Mamba-Llama3 (50%) scores 58.85 and Mamba2-Llama3 scores 66.60, while Falcon Mamba scores 53.42 and RecurrentGemma-9B scores 38.60 — the distilled models show superior truthfulness.

GSM8K and CRUX (the most striking results): On GSM8K (grade-school math word problems, zero-shot with chain-of-thought), Mamba-Llama3 (50%) achieves 67.85 — dramatically outperforming Falcon Mamba-7B-Instruct (41.32) and RecurrentGemma-9B-Instruct (38.51). Mamba2-Llama3 (50%) achieves 59.36. This 26.5-point gap over Falcon Mamba is the paper's strongest single-result evidence for the efficacy of distillation over training from scratch. On CRUX (code reasoning), Mamba-Llama3 (50%) achieves 27.88 versus Falcon Mamba's 8.88 and RecurrentGemma's 26.25. However, performance degrades rapidly as attention layers are removed: Mamba-Llama3 at 25% attention drops to 40.64 on GSM8K (from 67.85) and 15.62 on CRUX (from 27.88); at 12.5%, GSM8K falls to 26.91. Math and code reasoning appear to be the capabilities most dependent on attention layers — consistent with these tasks requiring precise multi-step logical deduction where the softmax selection mechanism may be particularly important.

Overall pattern across benchmarks: The distilled hybrid models with 50% attention achieve quality comparable to or exceeding linear RNN models trained from scratch with 60-180× more data, and approach the teacher Transformer's quality within single-digit percentage gaps on most benchmarks. The 25% attention models remain competitive with from-scratch models but show meaningful degradation. The 12.5% models and the 0% pure Mamba model show substantial quality loss, with the 0% model being clearly inferior to both the teacher and the from-scratch baselines on most metrics.

Long-Context Evaluation: Length Extrapolation Beyond the Distillation Context

Figure 4 visualizes the Needle-in-a-Haystack evaluation — a test where a specific fact ("the needle") is embedded at various positions within a long document ("the haystack"), and the model must retrieve it when queried.

Key result: Despite being distilled with only 2K-token context windows, the hybrid Mamba models demonstrate natural length extrapolation far beyond this. Mamba-Llama3.2-3B (50%) and Mamba2-Llama3.2-3B (50%) achieve perfect retrieval accuracy (green squares throughout) up to 10K tokens, which is 5× the distillation context length. The teacher Llama-3.2-3B-Instruct, in contrast, shows degradation at long contexts — the figure indicates the distilled models actually outperform the teacher on long-context retrieval. At the 8B scale, Mamba-Llama3.1-8B (50%) and Mamba2-Llama3.1-8B (50%) achieve perfect accuracy up to 16K tokens, with Mamba-Llama3.1-8B showing "good results up to 38K" — nearly 20× the distillation length.

Interpretation: The paper does not provide a detailed mechanism for this extrapolation, but it is consistent with the known properties of state-space models: since the Mamba layers' recurrent dynamics are continuous-time in formulation and discretized with input-dependent step sizes, they can naturally handle sequence lengths different from those seen during training. The attention layers in the hybrid model, which were trained only up to 2K context, might be expected to struggle at longer lengths, but the figure suggests they are not the bottleneck — possibly because the Mamba layers' effective long-range modeling compensates, or because the sliding window attention in some of the teacher models (Zephyr/Mistral uses sliding window attention) naturally extrapolates.

Speculative Decoding: Throughput Speedups for Pure and Hybrid Mamba Models

Table 1 reports speedup results for speculative decoding with pure Mamba models (no attention layers). The experiments use a draft model that is either a smaller Mamba (130M for the 2.8B verifier) or a small Transformer (Llama3 1B for the 7B verifier).

Pure Mamba speedups: On an RTX 3090, Mamba 2.8B with a 130M Mamba draft achieves 2.3× speedup at K=3 (259 tokens/sec vs. ~113 baseline) and 2.6× at K=4 (289 tokens/sec). Mamba 7B with a Llama3 1B draft achieves 2.1× speedup (109-110 tokens/sec). On an H100, the speedups are lower: 1.71-1.85× for the 2.8B model (389-421 tokens/sec) and 1.95-2× for the 7B model (271-272 tokens/sec). The lower relative speedups on H100 despite higher absolute throughput reflect the fact that the H100's baseline generation is already much faster (the non-speculative baseline throughput is not explicitly reported but can be inferred: 389/1.71 ≈ 228 tokens/sec for 2.8B baseline on H100 vs. 259/2.3 ≈ 113 on 3090). The paper attributes the H100 speedup challenge to GEMM operations being much faster, making kernel launch and caching overhead more prominent relative to compute — motivating the kernel fusion optimizations described in Section 4.3.

Draft token efficiency: The # Gen. Tokens column reports the average number of tokens produced per speculative step, which includes the accepted draft tokens plus one additional token from the verifier's own logits after rejection. Values range from 3.01 to 4.04 across configurations, indicating that the draft model's acceptance rate is high enough to make speculation worthwhile — the draft is proposing correct tokens most of the time.

Table 5 extends speculative decoding to the distilled hybrid models (containing both Mamba and attention layers), using the Zephyr-distilled and Llama-distilled models as verifiers with small Transformer draft models (2-layer or 4-layer, trained via shrink-and-fine-tune).

Hybrid model speedups: For Mamba-Zephyr (50% attention) with a 4-layer Transformer draft at K=4, the speedup is 1.81× (3.0 tokens generated per step). For Mamba-Zephyr (25% attention), the speedup is 1.88× at K=4 with a 2-layer draft and 1.8× with a 4-layer draft. For Mamba-Llama3 (50%) with a 4-layer draft, speedups are lower: 1.6× at K=3 and 1.58× at K=4, with 3.6 tokens generated per step. The paper attributes the lower speedup for Llama models to the "large embedding table of Llama 3" making the draft model larger and thus slower — the overhead of running the draft model partially offsets the gains from speculation.

Comparison to pure Mamba speculation: The hybrid model speedups (~1.6-1.9×) are lower than pure Mamba speedups (~2.1-2.6×) because (1) the attention layers in the hybrid model use standard KV-cache-based verification which cannot benefit from the MultiStep kernel optimizations, and (2) the draft models for hybrid verification are relatively larger (Transformer layers plus embeddings vs. the 130M Mamba draft used for pure Mamba models). The paper notes this tradeoff explicitly and suggests future work on making draft models smaller.

Ablation Studies and Robustness Checks

Attention weight initialization vs. random initialization (Table 8): The most dramatic ablation. A Zephyr-Mamba (50% attention) model with random initialization instead of attention weight transfer achieves an MT-Bench score of 1.04 (essentially non-functional) and AlpacaEval LC win rate of 0.02% (complete failure). The attention-initialized version achieves 6.69 and 14.11% respectively. Every benchmark shows catastrophic degradation without weight transfer: MMLU drops from 47.98 to 26.21, ARC-Challenge from 49.15 to 25.26, HellaSwag from 75.07 to 27.91, TruthfulQA from 46.67 to 34.01, and LAMBADA perplexity explodes from 6.20 to 55.01. This single ablation establishes that weight transfer is not merely helpful — it is necessary at this distillation scale (~20B tokens). The randomly initialized models would presumably recover with sufficient training, but the required data scale is far beyond the post-training budget.

Necessity of Mamba layers (Table 9): To verify that the performance is genuinely coming from the Mamba layers learning to perform attention-like routing, the paper trains a model where the Mamba blocks are entirely removed (leaving only the retained attention layers and MLPs). This model achieves an MT-Bench score of 1.01 and AlpacaEval of 0% — even worse than the random-initialized Mamba model. LAMBADA perplexity is 151.98 (versus 6.20 with Mamba). This confirms that the attention layers being replaced are essential — the model cannot simply learn to route information through the remaining attention layers alone. The routing function that the Mamba layers learn is critical, not redundant.

Distillation stage ablation (Table 6, Right): The three-phase distillation process is tested by removing stages. For the 50% attention hybrid Mamba: Dis (pseudo-label distillation) alone achieves 5.55 MT-Bench; Dis+SFT achieves 5.61; Dis+DPO achieves 5.42; Dis+SFT+DPO (the full pipeline) achieves 6.69. For the 25% attention model: Dis = 5.01, Dis+SFT = 4.97, Dis+DPO = 4.84, full pipeline = 6.10. The pattern shows that SFT or DPO alone provide minimal benefit over the base pseudo-label distillation — the gains come from combining them. The paper does not report the teacher's MT-Bench on this ablation configuration, but the gap from ~5.5 (Dis only) to 6.69 (full pipeline) represents a meaningful improvement from the SFT+DPO stages, suggesting that the pseudo-label stage provides a foundation but the standard instruction-tuning and preference data provide complementary signals.

Stepwise progressive distillation vs. direct replacement (Table 7, Right): For the 25% attention hybrid, stepwise distillation (first train a 50% model, then replace more layers and train to 25%) achieves perplexity 2.20, while direct distillation (immediately replace to 25% and train) achieves 2.89 — a 31% relative increase in perplexity. For the 50% attention hybrid, the comparison is only partially reported (stepwise = 2.09 vs. direct = 2.41). The benefit of stepwise training is clear: the intermediate models provide better learning signals because more layers are still well-functioning Transformer components. This is a non-obvious finding — it suggests that the distillation process benefits from a curriculum where the model gradually adapts to having fewer attention layers, rather than being forced to adapt all at once.

Interleaved vs. contiguous layer placement (Table 7, Right): The interleaved placement pattern (alternating Mamba and attention layers) shows lower perplexity than contiguous placement (all Mamba layers grouped together). For the 25% attention model: interleaved perplexity is 2.20 versus 2.89 without interleaving. The paper does not deeply analyze why interleaving helps, but the likely mechanism is that interleaving ensures that Mamba outputs are immediately refined by subsequent attention layers, and attention outputs are immediately processed by Mamba layers — creating a mix of local and global processing at each depth rather than having all the attention concentrated in one part of the network.

Freezing MLP layers in Phase 1 (Table 7, Left): Freezing MLP layers during the pseudo-label distillation stage consistently produces lower perplexity than training all parameters. For the pure Mamba (0% attention) model: frozen MLPs achieve 3.36 perplexity versus 66.7 without freezing. For the 50% attention hybrid: frozen MLPs achieve 2.09 versus 9.1 without freezing. The paper's rationale (Section 3) is that freezing MLPs "allows the student model to focus on learning the interaction of tokens and better mimic attention layers" — the MLPs contain the Transformer's stored knowledge, and training them during the early stage when the Mamba layers are still learning to route information would cause the model to "forget" the teacher's knowledge faster than the Mamba layers can learn to route it.

MLP freezing combined with attention initialization (Table 7, Left): The left section crosses the two factors: initialization (attention vs. random) and MLP training (frozen vs. unfrozen). The combination of attention initialization + frozen MLPs yields by far the best perplexity across both pure Mamba (3.36) and hybrid (2.09). Removing either factor substantially degrades performance. Interestingly, without attention initialization, freezing MLPs still helps but much less dramatically (18.2 vs. 20.3 for unfrozen in pure Mamba; 7.4 vs. 11.2 for unfrozen in hybrid). This suggests the two design choices are complementary: attention initialization provides good routing priors, and frozen MLPs prevent catastrophic forgetting of stored knowledge while the routing is being refined.

Comparison to prior distillation approach (Table 6, Left): The paper compares against the Distill Hyena approach (Ralambomihanta et al., 2024) which uses progressive knowledge transfer to distill a 70M Pythia Transformer into a Hyena model. The Distill Hyena achieves a perplexity ratio of 2.36× over the teacher (121.2 vs. 51.4). The paper's approach achieves substantially better ratios: 1.03× for 50% attention, 1.09× for 25%, 1.22× for 6.25%, and 1.66× for 0% attention. The authors acknowledge that "it is challenging to compare" due to different model scales, architectures, and training data, but the order-of-magnitude difference in degradation ratios (1.03-1.66 vs. 2.36) at comparable or better attention-retention levels is suggestive. The key difference is likely the weight initialization: the Distill Hyena approach does not transfer attention weights, while this paper's approach does.

DPO as a distillation objective (Table 6, Right and Table 4): The ablation confirms that adding a DPO stage with the teacher as reference model improves performance over SFT alone. For the 50% attention hybrid on MT-Bench: Dis = 5.55, Dis+SFT = 5.61, Dis+DPO = 5.42, Dis+SFT+DPO = 6.69. The DPO-alone result (5.42, slightly worse than Dis alone at 5.55) suggests that DPO without SFT is ineffective — the model needs the SFT stage to establish a reasonable policy before preference optimization can refine it. The later model versions (Llama-3.1/3.2) consistently show DPO variants outperforming non-DPO variants on both MT-Bench and LM Eval tasks: for example, Llama3.1-Mamba-DPO scores 64.38 average versus 62.32 for the non-DPO version (Table 3), and 7.7 vs. ~7.6 MT-Bench (Table 2).

Mamba2 vs. Mamba architecture (Tables 2, 3, 4): The paper includes both Mamba and Mamba2 distillations from the same teachers with the same training pipeline, allowing a direct comparison of the two architectures in the distillation setting. Across essentially all configurations and benchmarks, Mamba2 modestly outperforms Mamba at the same attention percentage. For Llama-3 at 50%: Mamba2 average is 63.84 vs. Mamba's 61.30 on LM Eval (Table 3). For Llama-3 at 50% on MT-Bench: Mamba2 scores 7.32 vs. Mamba's 7.35 — essentially tied. The advantage is consistent (~2-3 points on the LM Eval average) but not dramatic. This validates the Mamba2 architecture's claimed improvements (structured state-space duality, better GPU utilization) in the specific context of distillation, though the paper does not isolate which specific Mamba2 design changes are responsible for the improvement.

ReST^EM revision model degradation (Appendix K, Figure 16): The paper includes a brief negative result: an attempt to optimize the revision model using ReST^EM (an RL-based training approach; Singh et al., 2024) backfires. This is not explored in detail but is flagged as a caution that naive application of RL-based fine-tuning to the distilled models can degrade performance. This is the only major negative result in the training pipeline reported. The paper does not provide a figure or table for this in the main text, only text reference. Without explicit data, it is unclear how severe the degradation is or whether it applies to specific configurations — the paper describes it as a "notable negative result" worth flagging.

Critical Assessment

The experiments demonstrate a coherent and practically significant set of findings, but several important limitations constrain the strength and generality of the claims.

Claim: "The distilled approach matches the teacher model in standard Chat benchmarks" (from Section 1). The evidence partially supports this but with an important asymmetry. On AlpacaEval 2, the distilled models consistently outperform their teachers (e.g., Mamba-Llama3 50% = 29.61% vs. Llama-3-Instruct = 22.90%). On MT-Bench, they consistently underperform (7.35 vs. 8.00). "Matches" is accurate for the Zephyr case (7.31 vs. 7.34) but not for Llama-3, where the 0.65 point gap is non-trivial. More importantly, the paper does not investigate why the student outperforms the teacher on AlpacaEval — this is a surprising result that could reflect (a) genuine improvement from the DPO distillation relative to the teacher's own alignment, (b) the AlpacaEval metric favoring a particular response style that the distillation process amplifies, or (c) a length bias (the LC correction may not fully account for verbosity differences). Without analysis, it is unclear whether this "improvement" is a feature or an artifact.

Claim: "The distilled approach performs on par or better with all similarly sized pretrained-from-scratch Mamba models" (Section 1). This claim is well-supported by Tables 2-4. The 50% attention distilled models consistently meet or exceed TRI Mamba 7B (1.2T tokens), Nvidia Hybrid Mamba-2 8B (~3.7T tokens), Falcon Mamba 7B (>5T tokens), and RecurrentGemma-9B on the evaluated benchmarks. The GSM8K result (67.85 vs. 41.32 for Falcon Mamba) is the most dramatic example. However, "all similarly sized" is a limited set — there may be other linear RNN models at this scale not included in the comparison (the paper does not claim exhaustiveness, but the baseline set is small: 3-4 models). Additionally, the claim implicitly compares distillation cost (~20B tokens, 5 days on 8×A100) to training-from-scratch cost (trillions of tokens, presumably months on much larger clusters). The paper does not provide a direct FLOPs comparison of distillation vs. from-scratch training, which would quantify the efficiency gain more precisely. The hardware comparison (8×A100 for days vs. industrial clusters for months) is suggestive but qualitative.

Claim: "The distilled model has natural length extrapolation, showing almost perfect accuracy in the needle-in-a-haystack test at 20× the distillation length" (Abstract, Section 5.4). Supported by Figure 4, which shows perfect retrieval at 10K (5×) for 3B models and 16K (8×) for 8B models, with good results at 38K (~19×) for one 8B model. However, "almost perfect" overstates the 38K result — the figure shows the 8B model has scattered non-green squares at the longest lengths, and the paper's text only claims "good results up to 38K" rather than "perfect." More critically, the needle-in-a-haystack test measures a narrow capability (fact retrieval), and length extrapolation on this task does not guarantee length extrapolation on more complex long-context tasks like multi-hop reasoning or summarization, which are not evaluated. The paper's claim about length extrapolation should be scoped to retrieval accuracy specifically.

Claim: "The hardware-aware speculative decoding algorithm achieves a throughput of over 300 tokens/second for a Mamba 7B model" (Section 1). Supported by Table 1: 271-272 tokens/sec on H100 for Mamba 7B. The "over 300" figure comes from the 2.8B model on H100 (389-421 tokens/sec) and the 2.8B model on 3090 (259-289 tokens/sec). The abstract's wording attributes the 300+ to the 7B model specifically, which is slightly misleading — the 7B model achieves ~270 tokens/sec on H100, not 300+. The speedup ratios (1.6-2.6×) are solid but not transformative — Transformer speculative decoding methods have reported larger speedups, and the paper acknowledges that "for the Llama-hybrid models, the speedups are more modest" (1.58-1.6×). The speculative decoding contribution is more about enabling the technique for architectures where it was previously infeasible (linear RNNs) rather than achieving unprecedented acceleration.

Limitation: Single model family, mostly single scale. All experiments use Mistral/Llama architectures at 7-8B parameters as teachers, with one additional experiment at 3B (Llama-3.2). The paper does not demonstrate the approach on other Transformer families (e.g., Qwen, Gemma, Phi) or at larger scales (13B, 70B), where the attention-to-Mamba transfer dynamics might differ. The claim that the approach is "representative of many contemporary LLMs" is asserted but not tested. Scaling behavior — whether the 50% attention sweet spot holds at larger model sizes, or whether the degradation pattern changes — is unknown.

Limitation: Evaluation breadth. The paper evaluates on standard benchmarks, but the evaluation does not include several capability categories where the attention-to-Mamba transfer might be most stressed: (1) Structured output tasks (JSON generation, code synthesis with precise syntax) where the softmax attention's ability to precisely attend to specific tokens might matter more. (2) Instruction following benchmarks like IFEval that measure constraint satisfaction rather than preference ratings. (3) Safety benchmarks — it is unknown whether the distillation process preserves or degrades the teacher's safety alignment. (4) Multilingual tasks — all evaluations are in English. For a paper claiming to demonstrate that "Transformer knowledge can be transferred effectively to other architectures," these omissions are significant because they leave open the possibility that specific capabilities are systematically lost in the transfer.

Limitation: Distillation compute is underreported. The paper states that distillation takes "less than five days in 8x80G A100" with ~20B tokens, but does not provide the full FLOPs accounting that would enable precise comparison to from-scratch training FLOPs. The teacher models' pretraining compute (trillions of tokens on thousands of GPUs) is mentioned qualitatively but not quantified. A FLOPs-matched comparison — similar to that in the example paper's Section 7 — would strengthen the efficiency claims considerably.

Limitation: The attention-percentage sweet spot is empirically discovered, not predicted or explained. The paper finds that 50% attention gives the best quality-efficiency tradeoff, but does not provide a mechanistic explanation for why 50% specifically (rather than 33%, 66%, or a different fraction) is optimal. This limits the ability to predict the optimal configuration for new model families or scales without running the full experiment. The paper also does not explore whether different attention layers contribute differently — some layers might be more "attention-critical" than others, and a non-uniform retention pattern might outperform the uniform interleaving used.

Limitation: Speculative decoding speedups are draft-model-dependent. The hybrid model speedups (1.6-1.9×) are achieved with specific draft models (2-4 layer Transformers trained on OpenHermes 2.5). Different draft model choices would yield different speedups. The paper does not systematically explore the draft model design space — how small can the draft be while maintaining useful acceptance rates? Does a Mamba draft work better for Mamba verifiers? The Table 1 results use different draft types (Mamba draft for Mamba verifier, Transformer draft for hybrid verifier) without a controlled comparison.

Missing experiment: Ablation of the weight transfer mapping. The paper transfers $W_Q \rightarrow C$ , $W_K \rightarrow B$ , $W_V \rightarrow X$ , but does not test alternative mappings (e.g., swapping Q and K, or using a learned linear combination of the three projections). It is unknown whether the specific mapping is optimal or merely sufficient. An ablation testing different mapping choices would clarify how sensitive the approach is to this design choice.

Missing experiment: Phase 1 data scale ablation. The paper uses one epoch of UltraChat/UltraFeedback pseudo-labels for Phase 1. How does performance scale with the amount of pseudo-label data? Is the approach robust to using fewer seed prompts, or would more data continue to improve? Without this ablation, it is unclear whether the 20B total tokens are near the saturation point or whether further improvements are attainable with more distillation data.

Missing evaluation: Inference memory usage. The paper motivates linear RNNs by their memory efficiency (no KV cache), but does not report actual GPU memory usage during inference for the hybrid models compared to the original Transformers. Since the hybrid models retain some attention layers, they still have partial KV caches. Quantifying the memory savings at each attention-percentage level would directly support the deployment-efficiency motivation.

Summary of conditional support: The paper's central claim — that Transformers can be efficiently distilled into hybrid Mamba models through weight initialization and lightweight post-training distillation — is well-supported conditional on the specific teacher models (Mistral/Llama 7-8B), the specific evaluation benchmarks, and the 50% attention retention configuration. The approach clearly outperforms training linear RNNs from scratch at comparable compute, and the weight initialization is demonstrated to be essential (not just helpful). However, the generality of the approach to other model families, scales, and capability categories remains untested, and the speculative decoding speedups, while real, are modest (1.6-2.6×) compared to the 5× throughput advantage claimed for linear RNNs over Transformers. The paper makes a strong case for the practical viability of the distillation pipeline but leaves open the question of whether this approach can scale to the frontier models where deployment efficiency matters most.

6. Limitations and Trade-offs

The Difficulty Estimation Overhead Is Unaccounted for and Dominates the Practical Compute Budget

The assumption or constraint. The entire compute-optimal framework rests on the ability to estimate prompt difficulty before allocating the inference budget. The paper's method for doing so — generating 2048 samples per question and averaging either ground-truth correctness (oracle) or PRM final-answer scores (predicted) — is extraordinarily expensive. Section 3.2 acknowledges this explicitly:

"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"

The 2048 samples required for difficulty estimation consume more compute than the largest test-time budgets studied (256–512 generations). The reported 4× efficiency gains over best-of-N are computed after difficulty is known, without amortizing the cost of learning it.

The consequence. In any realistic deployment, the total cost would be difficulty estimation cost + strategy execution cost, and the former could dominate the latter to the point where the overall system is less efficient than simply running a uniform strategy like best-of-N on every prompt. The 4× figure should be understood as an upper bound on achievable efficiency that assumes difficulty can be estimated for free or at negligible cost — an assumption that does not hold for the estimation method the paper actually uses.

What evidence exists in the paper. The paper provides no experiment that includes the difficulty estimation cost in the total generation budget. Figures 4 and 8 plot accuracy against the test-time compute budget used for solving the problem, not for solving + estimating difficulty. The predicted difficulty bins track oracle bins closely (Figures 4 and 8, curves largely overlap), but this only shows that the PRM-based estimate is accurate, not that it is cheap. The cost of generating 2048 samples and running the PRM on all of them is never added to the x-axis.

Mitigation status. The paper explicitly flags this as a key avenue for future work (Section 3.2, Section 8), suggesting "pretraining or finetuning models to directly predict difficulty of a question" or adaptive difficulty estimation that "starts by generating a small number of samples... uses the verifier's score distribution on those samples as a quick difficulty signal, and then allocates the remaining budget accordingly." However, no such method is developed or evaluated. Until this gap is closed, the 4× efficiency claim cannot be interpreted as a realized deployment gain.

The 14× Larger Model Baseline Is Not Compute-Optimally Trained and Uses No Test-Time Compute of Its Own

The assumption or constraint. The FLOPs-matched comparison in Section 7 scales model parameters while holding training data fixed, following the LLaMA paradigm (Touvron et al., 2023). The paper explicitly acknowledges this departs from compute-optimal pretraining:

"We choose this setting as it is representative of a canonical approach to scaling pretraining compute and leave the analysis of compute-optimal scaling of pretraining compute where the data and parameters are both scaled equally to future work."

Additionally, the 14× larger model uses only greedy decoding — no majority voting, no best-of-N, no search — despite the paper demonstrating that even modest test-time compute budgets (e.g., best-of-4 or best-of-8) produce substantial improvements.

The consequence. Both choices make the pretraining baseline weaker than a fair comparison would warrant. A Chinchilla-optimal model (Hoffmann et al., 2022) trained with 14× more total FLOPs — scaling both parameters and data — would likely outperform a parameter-only-scaled model. Giving the larger model even a modest test-time compute budget (say, best-of-8 with majority voting) would create a much stronger baseline. The reported advantages of test-time compute over pretraining — for example, +27.8% relative improvement on easy questions at R ≪ 1 for revisions (bar chart in Figure 1) — may shrink or reverse against a properly compute-optimal larger model that also uses basic test-time strategies.

What evidence exists in the paper. The paper does not include an ablation where the larger model is given any test-time compute budget, nor does it compare against a Chinchilla-optimal scaling baseline. The FLOPs accounting in Section 7 is carefully done for the parameter-only-scaled model, but the choice of scaling paradigm is acknowledged as a limitation rather than tested.

Mitigation status. The paper is transparent about this choice and explicitly leaves the compute-optimal pretraining comparison to future work (Section 7). However, this transparency does not change the fact that the claim "test-time compute can substitute for pretraining" is demonstrated against a suboptimal pretraining baseline, and the true substitution rate is likely lower than the paper reports.

Revisions and PRM Search Are Never Combined, Leaving an Unknown Upper Bound on System Performance

The assumption or constraint. The paper studies two complementary axes for scaling test-time compute — modifying the proposal distribution via iterative revisions (Section 6) and improving output selection via PRM-guided search (Section 5) — but studies them entirely independently. Section 8 explicitly acknowledges this gap:

"we did not experiment with PRM tree-search techniques in combination with revisions"

This is a significant omission because the two mechanisms have structurally complementary strengths: revisions improve the quality of generated candidates (refining answers that are roughly correct), while PRM search improves candidate selection (finding the best among diverse candidates). The paper's own difficulty-dependent analysis shows they are complementary in practice — revisions excel on easy problems (local refinement), while search excels on medium problems (global exploration).

The consequence. A combined system — using the revision model as the proposal distribution within beam search, or using the PRM to guide which revision branches to pursue — could yield gains beyond either method alone. The current results represent a lower bound on what a fully integrated system could achieve, and the paper cannot quantify how much performance is left on the table by keeping these mechanisms separate. This is particularly relevant for medium-difficulty problems (bins 3–4), where both mechanisms show positive but incomplete gains individually.

What evidence exists in the paper. No experiments combine revisions with PRM search. The compute-optimal policy in Section 3 selects between search strategies (best-of-N, beam search, lookahead) and revision strategies (sequential, parallel, hybrid ratios) independently — there is no strategy that, for example, generates a revision chain and then applies beam search across chains, or uses PRM step-level scores to decide when to trigger a revision.

Mitigation status. The paper acknowledges this as future work in Section 8 but does not provide any analysis of the expected magnitude of combined gains. The omission is a scope limitation rather than a flaw — the paper's contribution is the analysis framework and the difficulty-conditioned allocation, not the fully optimized system — but a practitioner seeking to maximize performance would need to explore this combination themselves without guidance from the paper on expected benefits or design choices.

Hard Problems Remain Fundamentarily Unsolved — Test-Time Compute Cannot Compensate for Capability Gaps

The assumption or constraint. The paper's entire framework depends on the base model having a non-trivial pass@1 rate on a given problem. Section 5.3 shows that on difficulty bin 5 (the hardest problems, where the base model's pass@1 is near zero), no method — search, revisions, or compute-optimal combinations — makes meaningful progress:

"On the hardest questions (bin 5), no method makes meaningful progress" (Section 5.3)

Figure 3 (right) shows bin 5 accuracy hovering at 1–3% for all methods and all budgets. Figure 7 (right) shows bin 5 at roughly 2–3% accuracy irrespective of the sequential-to-parallel ratio. Figure 9 shows the bin 5 scaling line essentially flat near 0–5% for both revisions and PRM search.

The consequence. Test-time compute can amplify existing capability but cannot create it from nothing. If the base model's pass@1 is near zero on a problem class, no amount of search or revision will help — there are no correct solutions in the proposal distribution to find or refine. This means the approach offers no path forward for genuinely novel or out-of-distribution reasoning that exceeds the base model's training distribution. For such problems, pretraining (or fundamentally different architectures, or retrieval-augmented approaches) remains the only viable path. The paper's FLOPs-matched comparison (Section 7) confirms this: on hard problems at R ≫ 1, test-time compute with the smaller model shows a -52.9% relative disadvantage compared to the larger pretrained model (PRM search, Figure 1 bottom-right bar chart).

What evidence exists in the paper. The difficulty-bin analyses across search (Figure 3, right), revisions (Figure 7, right), and FLOPs-matched comparisons (Figure 9) consistently show bin 5 accuracy near zero and essentially flat with respect to compute budget. This is the most robust and replication-consistent finding in the paper.

Mitigation status. The paper is candid about this limitation (Section 7 takeaway box, Section 8) and frames it as a boundary condition: test-time compute is effective when and only when problems are within the base model's rough capability range. However, the paper does not provide actionable guidance for identifying which problems fall into this "unsolvable" category before expending compute — the difficulty estimation procedure can identify hard problems post-hoc (by measuring low pass@1), but at that point significant compute has already been spent on estimation. A practitioner needs to know when to give up and route to a larger model or a different approach, and the paper does not provide a decision rule for this.

The Revision Model Has a 38% Correct-to-Incorrect Reversion Rate with No Principled Solution

The assumption or constraint. The revision model is trained exclusively on trajectories where all in-context answers are incorrect, followed by a correct target (Section 6.1). This means at test time, when the model produces a correct answer during a revision chain, it has never seen an example of what to do when the current answer is already correct — it only knows how to revise incorrect answers into correct ones. Section 6.1 reports:

"approximately 38% of correct answers get converted back to incorrect ones using a naive approach"

The consequence. The revision model is fundamentally unreliable on correct answers — it actively corrupts them at a substantial rate. This creates a self-sabotage dynamic: the model may produce a correct answer at step 3 of a revision chain, then "revise" it to an incorrect answer at step 4, then potentially recover at step 5. The paper mitigates this with majority voting or verifier-based selection across the entire chain (picking the best answer from any step rather than always taking the last revision), but these are post-hoc patches rather than a solution to the underlying problem. In deployments where the chain-end answer is needed (e.g., because the selection mechanism is expensive or unavailable), the 38% reversion rate would directly translate to a 38% corruption rate on correct answers — a catastrophic failure mode for any application requiring reliability.

What evidence exists in the paper. The 38% figure is reported in Section 6.1. The paper does not provide a detailed analysis of under what conditions reversions occur (e.g., are certain answer types more susceptible? Does the reversion rate depend on revision depth? Does it correlate with problem difficulty?). The ReST^EM experiment (Appendix K, Figure 16) provides additional evidence of fragility: attempting to optimize the revision model with RL-style training caused performance to degrade substantially with sequential revisions, suggesting the revision training is sensitive to methodology in ways that are not fully understood.

Mitigation status. The paper mitigates the reversion problem with selection mechanisms (majority voting and verifier-based selection) that scan the entire revision chain and pick the best answer, not just the final one. However, these are described as mitigation strategies, not as a fix to the underlying training data problem. The paper does not propose training the model with "no revision needed" examples or other principled solutions, and does not explore the reversion rate's dependence on training data construction choices. This is flagged as a limitation by the authors but not systematically studied.

The Single-Benchmark, Single-Model-Family Evaluation Leaves Generality Untested

The assumption or constraint. All experiments use the MATH benchmark (500 test questions, competition-level math) with PaLM 2-S* as the base model, and the FLOPs-matched comparison uses one additional model from the same family with ~14× more parameters. Section 4 states that the authors "believe this model is representative of the capabilities of many contemporary LLMs," but this claim is unverified.

The consequence. Several aspects of the findings could be model-specific or domain-specific:

PRM quality and over-optimization behavior depend on PaLM 2-S*'s output distribution. A model with different calibration properties or different error patterns (e.g., one that produces more diverse correct solutions or different types of errors) might exhibit different difficulty-dependent scaling curves. The beam search degradation on easy problems (Figure 3, right) is a specific verifier-model interaction that may not generalize.
Revision model learning dynamics depend on the base model's in-context learning capabilities and the quality of sampled incorrect answers, which vary substantially across model families.
The MATH benchmark consists exclusively of symbolic math reasoning with clean ground-truth answers. The difficulty-dependent patterns — beam search hurting easy problems, revisions helping easy problems, no method helping hard problems — may not generalize to other reasoning domains (code generation, logical inference, scientific QA) or to tasks requiring factual knowledge rather than step-by-step inference.
The 500-question test set split into five difficulty quintiles of ~100 each, then further split by two-fold cross-validation, means the compute-optimal policy is selected based on ~50 questions per fold per bin. The paper does not report confidence intervals on compute-optimal scaling curves, making it impossible to assess whether observed efficiency gains are statistically robust at this sample size.

What evidence exists in the paper. No experiments are conducted on any benchmark other than MATH, nor with any model family other than PaLM 2. The paper does not report confidence intervals for the compute-optimal scaling curves. The difficulty quintile sizes (~100 questions before cross-validation split) are not discussed as a potential source of variance.

Mitigation status. The authors acknowledge in Section 4 that their belief in representativeness is an assumption, not a demonstrated fact. Section 8 suggests future work on extending the analysis to other domains and models, but no such extension is provided. This is a scope limitation of the present paper — the framework is demonstrated in one setting and its generality is asserted but not tested. For a practitioner considering applying these methods to a different model family (e.g., Llama, Mistral, Qwen) or domain (e.g., code generation, logical reasoning), the paper provides no direct evidence that the difficulty-dependent patterns or the 4× efficiency gains will replicate.

7. Implications and Future Directions

How This Work Changes the Landscape

This paper shifts the conversation around efficient LLM deployment from a binary choice — either accept the Transformer's inference costs or pay the enormous pretraining cost of training an alternative architecture from scratch — to a transfer-and-adapt paradigm where the Transformer's learned knowledge is preserved and repurposed rather than discarded. This is a methodological reframing of the efficiency problem, not merely an incremental improvement within existing approaches.

The conceptual shift has two components. First, the paper demonstrates that the linear projections in attention heads ( $W_Q, W_K, W_V$ ) encode reusable semantic routing functions that are separable from the softmax selection mechanism. This challenges the implicit assumption in prior distillation work — that architectural conversion requires treating the student as a black box trained with output-level supervision alone. The Hyena distillation work (Ralambomihanta et al., 2024) operated under this assumption and produced 2.36× perplexity degradation on 70M-parameter models. The paper's attention-to-Mamba weight transfer shows that much of what attention layers learn is mechanism-independent — the learned projections for querying, keying, and writing information can be directly transplanted into a different selection architecture. The catastrophic failure without weight initialization (Table 8: MT-Bench drops from 6.69 to 1.04, AlpacaEval from 14.11% to 0.02%) quantifies how much knowledge is encoded in these projections rather than in the attention mechanism itself.

Second, the paper demonstrates that post-training alignment, not pretraining, is the optimal distillation target when converting architectures that share a common MLP backbone. The implicit decomposition — MLP layers store knowledge, attention layers route information — validates a separation-of-concerns view of Transformer internals. The fact that a 50% attention hybrid model trained on ~20B tokens of chat and instruction data achieves MMLU scores of 57.81 (Table 3) and GSM8K scores of 67.85 (Table 4) — far exceeding models trained from scratch with trillions of tokens — confirms that the frozen MLP layers carry forward the teacher's factual and conceptual knowledge. This decomposition, if it generalizes, means that for any new efficient architecture that can be initialized from Transformer attention weights and paired with preserved MLP layers, the expensive pretraining stage can be bypassed entirely. The 5-day, 8×A100 distillation pipeline is a concrete demonstration at the 7-8B scale.

This reframing also resolves a tension in the broader efficiency literature. Prior work on linear RNNs (Mamba, Mamba2, RWKV, Griffin) demonstrated that these architectures could match Transformers at small to medium scale when trained from scratch, but the practical question — how to actually deploy them at the quality level of frontier Transformers — remained unanswered because retraining from scratch at that scale was economically prohibitive. The paper shows that the answer is not to train linear RNNs from scratch, but to convert existing Transformer investments. This changes the research agenda: rather than competing to train the best from-scratch linear RNN, the goal becomes designing architectures that maximize the fidelity of Transformer weight transfer while improving inference efficiency.

The attention-percentage scaling curve — how quality degrades as attention layers are removed (Tables 2, 3) — provides a new diagnostic tool. It shows that the degradation is not linear: moving from 50% to 25% attention costs less quality than moving from 12.5% to 0%. This pattern, replicated across benchmarks and model families (Llama-3, Llama-3.1, Llama-3.2), suggests that a small number of attention layers provide disproportionately important capabilities that Mamba's expanded state space cannot fully replicate. This diagnostic can guide future architecture design: rather than treating all attention layers as equally replaceable, future hybrid architectures might preserve specific attention layers (e.g., those handling long-range dependency resolution or precise multi-hop reasoning) while more aggressively replacing others.

The hardware-aware speculative decoding contribution reframes RNN inference optimization as a kernel fusion problem rather than an algorithm design problem. Prior speculative decoding work (Leviathan et al., 2023; Chen et al., 2023) focused on the algorithmic structure of the draft-then-verify loop. The paper identifies that for linear RNNs, the bottleneck is kernel-level memory management — the tension between state rewinding on verification failure and the desire to avoid materializing intermediate RNN states. The MultiStep kernel's solution (fusing recomputation, decoding, and caching into a single GPU operation) provides a template for speculative decoding on any architecture whose state representation is a cumulative summary rather than a per-position list. The H100-specific optimization (achieving 1.71-2× speedup where the naive implementation achieved none) demonstrates that the solution is hardware-dependent in ways that matter for real deployment, shifting attention from algorithmic novelty to kernel engineering for non-Transformer architectures.

Research directions that become more attractive include: weight-transfer-based distillation for other architecture pairs (e.g., Transformer to RWKV, Transformer to Griffin), designing linear RNN architectures specifically optimized for weight transfer fidelity rather than from-scratch training quality, automated methods for selecting which attention layers to preserve based on their functional role rather than uniform interleaving, and kernel-level optimizations for speculative decoding on architectures without KV caches. Research directions that become less attractive include: training linear RNNs from scratch to match Transformer quality at scale (the paper shows distillation with ~20B tokens matches or exceeds models trained with trillions), and generic architecture-agnostic distillation methods that treat the student as a black box (the weight transfer ablation in Table 8 shows these leave enormous quality on the table).

Follow-Up Research This Work Enables

Characterizing which attention layers are most critical and why. The paper uses uniform interleaving (keep every $n$ th attention layer) and finds that 50% attention retention is the sweet spot and that the 12.5% → 0% transition causes disproportionate degradation. A natural follow-up would systematically measure the contribution of each attention layer to different capabilities — long-range dependency resolution, factual recall, multi-step reasoning, instruction following — by selectively replacing individual layers or groups of layers and measuring per-task degradation. The hypothesis to test is that certain layers (perhaps early layers handling local syntax, or middle layers handling cross-paragraph integration) are more "attention-critical" than others, and that a non-uniform retention pattern (preserving specific high-value attention layers while more aggressively replacing others) could achieve better quality-efficiency tradeoffs than uniform interleaving. A strong experiment would compare uniform 50% retention against a 50% retention pattern where the layers preserved are those with the highest measured contribution to the teacher's performance on target benchmarks, using a held-out validation set to select layers.

Combining weight transfer with retrieval-augmented generation for knowledge-intensive tasks. The paper shows that frozen MLP layers preserve much of the teacher's factual knowledge (MMLU 57.81 for 50% attention hybrid vs. 33.39 for TRI Mamba trained from scratch, Table 3), but there is still a gap to the teacher (68.12). A natural extension is to test whether retrieval-augmented generation (RAG) can close this gap for the distilled hybrid models by providing external knowledge that compensates for any knowledge degradation in the attention-to-Mamba transfer. The experiment would compare the teacher Transformer, the distilled hybrid, and the distilled hybrid + RAG on knowledge-intensive benchmarks (MMLU, Natural Questions, TriviaQA) to measure how much of the knowledge gap is recoverable through retrieval rather than through better distillation. This would test whether the distillation-induced knowledge degradation is primarily in storage (facts the MLPs failed to preserve) or in access (facts stored in MLPs but the Mamba layers cannot route to effectively) — RAG would help more in the access-failure case.

Scaling the distillation approach to 70B+ models and testing whether the 50% sweet spot holds. All experiments are at 7-8B parameters (with one 3B experiment). The paper does not know whether the attention-percentage scaling curve changes with model scale — larger models might tolerate more aggressive attention replacement (because they have more total layers and redundancy), or less (because individual attention layers in larger models might be more specialized and harder to replace). A strong follow-up would replicate the distillation pipeline with a 70B teacher (e.g., Llama-3 70B) across the same attention-percentage spectrum (50%, 25%, 12.5%, 0%) and measure whether the quality degradation pattern shifts. The paper mentions distilling from Llama-3.1 70B into 3B and 8B students (Section 5.1), which is cross-scale distillation (large teacher, small student), but does not test same-scale distillation (70B teacher, 70B hybrid student) where the question of layer replaceability is most directly tested. This experiment would also test whether the 5-day, 8×A100 training budget scales linearly or super-linearly with model size.

Stress-testing the weight transfer hypothesis with scrambled or permuted projections. The paper demonstrates that random initialization is catastrophic (Table 8), but does not test whether the specific mapping ( $W_Q \to C$ , $W_K \to B$ , $W_V \to X$ ) is optimal. A revealing ablation would test alternative mappings: (a) permuted mapping (e.g., $W_Q \to B$ , $W_K \to C$ ), (b) using a learned linear combination of all three projections for each Mamba parameter, (c) using only a subset of the projections (e.g., $W_Q$ and $W_V$ only, with random $B$ ), and (d) using the teacher's projections but with added Gaussian noise of varying magnitude. This would measure how precisely the semantic correspondence between attention roles (query, key, value) and Mamba roles (readout, input projection, input signal) must be preserved. If permuted mappings perform nearly as well as the canonical mapping, the benefit is primarily from having any pretrained projections rather than the specific role correspondence. If performance degrades sharply with permutation, the role correspondence is genuinely important — a finding that would inform weight transfer strategies for future architecture pairs.

Developing a "difficulty estimator" for layer replaceability to guide architecture design without full distillation. The paper's attention-percentage scaling curves are an empirical discovery that required training multiple models (50%, 25%, 12.5%, 0% attention). A more efficient approach would be to predict, before distillation, which attention layers can be replaced with minimal quality loss, based on properties of the teacher model. A strong follow-up would train a predictor that takes as input features of each attention layer (e.g., attention pattern entropy, head importance scores from pruning literature, contribution to specific benchmark tasks via activation patching) and predicts the quality degradation if that layer were replaced. The training data would be the degradation patterns from this paper's experiments plus additional experiments with different replacement patterns. If successful, this predictor could guide architecture design for new model families without running the full distillation sweep, making the approach more practically accessible.

Testing the distilled DPO objective against standard DPO with the student's own SFT checkpoint as reference. The paper claims this is "the first use of DPO as a distillation objective" (Section 3) and uses the teacher as the reference model. A controlled comparison against standard DPO (student SFT checkpoint as reference) on the same preference data would isolate the benefit of using the teacher as reference. The hypothesis is that teacher-as-reference provides a stronger preference signal because it defines preferences relative to an architecturally unconstrained baseline, but it could also introduce noise if the teacher's probability distribution encodes preferences that the student cannot realize due to architectural constraints. The experiment would train two DPO variants from the same SFT checkpoint — one with teacher reference, one with self reference — and compare downstream performance. If teacher-as-reference consistently outperforms, it establishes a general principle for DPO in distillation settings.

Practical Applications and Downstream Use Cases

Long-document processing pipelines where KV cache memory is the binding constraint. The paper's motivation (Section 1) highlights applications currently bottlenecked by Transformer KV caches: reasoning over multiple long documents, codebase-level understanding, and agent workflows with long context. The distilled hybrid models with 50% attention reduce the KV cache by approximately half (since half the attention layers are replaced with memory-efficient Mamba layers that maintain only a fixed-size state). For a system processing 100K-token documents with a 32-layer Transformer, the KV cache for the attention layers alone can consume tens of gigabytes. The 50% hybrid roughly halves this memory footprint while maintaining teacher-competitive quality (Table 2: 7.35 vs. 8.00 MT-Bench for Llama-3, 7.31 vs. 7.34 for Zephyr). This memory savings directly translates to larger batch sizes or longer context windows on the same hardware. The needle-in-a-haystack results (Figure 4) showing extrapolation to 16K-38K tokens despite 2K distillation context further suggest these models are well-suited to long-context deployment.

Batch inference for agent-based systems exploring multiple trajectories. Section 1 identifies agent workflows as a key deployment scenario requiring both large-batch inference (to explore many action paths) and long context (to model complex environments). The pure Mamba models achieve ~2× throughput speedup via speculative decoding (Table 1), and the hybrid models achieve ~1.6-1.9× (Table 5). For an agent system that needs to evaluate 100 candidate action trajectories, each requiring generation of 500 tokens, the throughput improvement directly reduces latency or increases the number of trajectories that can be explored within a fixed time budget. The memory savings from reduced KV cache further enable more trajectories to be batched simultaneously on a single GPU. The GSM8K result (67.85 for 50% attention hybrid, Table 4) suggests that multi-step reasoning — which agent trajectory evaluation fundamentally is — is well-preserved in the distilled models.

On-device or edge deployment with smaller distilled models. The Llama-3.2 3B experiments (Tables 2, 3) demonstrate the approach at smaller scale: Mamba-Llama3.2-3B (50%) achieves 6.9 MT-Bench and 60.88 LM Eval average, comparable to or exceeding the teacher (57.83 average). The distillation cost — 5 days on 8×A100 for 8B models, presumably less for 3B — means it is feasible to produce deployment-optimized versions of existing small Transformer models without the pretraining resources that created them. For on-device scenarios (laptops, phones) where memory and compute are severely constrained, a 3B hybrid model with half the attention layers and the corresponding KV cache savings could make local LLM deployment viable where the full Transformer would exceed memory budgets. The speculative decoding speedups (~2× for pure Mamba at 2.8B, Table 1) further improve interactive responsiveness on consumer hardware like RTX 3090s.

Cost-efficient fine-tuning of existing Transformer deployments for specialized domains. An organization that has fine-tuned a Transformer model (e.g., Llama-3 8B) for a specific domain (legal document analysis, medical literature review, customer support) currently faces high inference costs due to attention's quadratic complexity. The paper's approach offers a path to convert the fine-tuned model into a more efficient hybrid without losing the domain-specific adaptation: initialize from the fine-tuned Transformer (preserving both the general pretraining knowledge in MLPs and the domain-specific adaptations in the attention projections), then run the lightweight distillation pipeline on domain-relevant data. The ~20B token, 5-day distillation budget is small enough to be feasible for organizations without large-scale training infrastructure. The key unknown (not tested in the paper) is whether domain-specific knowledge in the attention projections transfers as effectively as general chat capabilities — this would need to be validated per-domain.

When to Prefer This Method

The paper positions distillation from Transformers to hybrid Mamba models against two alternatives: (1) deploying the original Transformer (higher quality, higher inference cost), and (2) training a linear RNN from scratch (similar inference efficiency, but enormous pretraining cost with uncertain quality). The tradeoff is explicit:

Prefer distillation to a hybrid Mamba model when: (a) inference cost (throughput, KV cache memory) is the primary deployment constraint, not absolute quality — the 50% attention hybrid achieves teacher-competitive chat quality (MT-Bench within 0.03-0.65 points, Table 2) while roughly halving KV cache memory; (b) you already have access to a high-quality pretrained Transformer and cannot afford to pretrain a linear RNN from scratch — the ~20B token, 5-day distillation budget is several orders of magnitude cheaper than the 1.2T-5T+ tokens used for from-scratch linear RNNs (Table 3 baselines); (c) your application requires long-context processing, where the KV cache bottleneck of Transformers is most severe and the Mamba layers' length extrapolation (Figure 4) provides additional benefit beyond speed; (d) you need to preserve domain-specific knowledge from a fine-tuned Transformer — the frozen MLP layers carry forward stored knowledge, and the attention-to-Mamba weight transfer preserves routing patterns.
Prefer the original Transformer when: (a) absolute quality on multi-turn conversation or complex reasoning is paramount and the 0.65 MT-Bench point gap (for Llama-3) is unacceptable — the paper shows consistent, albeit small, degradation on MT-Bench Round 2 specifically (7.82 → 6.88 for Llama-3 50%, Table 2); (b) latency per token (not throughput) is the constraint, and the speculative decoding speedups (~1.6-2×, Tables 1 and 5) are insufficient — pure Mamba without speculation is already fast, but the hybrid models' attention layers add serial overhead that speculation cannot eliminate; (c) your deployment already uses optimized Transformer inference with FlashAttention and other kernel optimizations that minimize the KV cache overhead for your specific sequence lengths — the relative benefit of switching to a hybrid architecture shrinks as Transformer inference becomes more optimized.
Prefer training a linear RNN from scratch when: (a) you have access to the pretraining compute budget (trillions of tokens, thousands of GPUs) and can achieve frontier-quality training — the paper's from-scratch baselines (Falcon Mamba, TRI Mamba, Nvidia Hybrid) were trained at smaller scale than frontier Transformers, and it is unknown how a from-scratch Mamba trained with Llama-3-equivalent compute would compare; (b) you need a pure linear RNN (0% attention) without any KV cache at all — the paper shows that pure Mamba from distillation degrades significantly (MT-Bench 5.64, Table 2; LM Eval average 54.74, Table 3), and a from-scratch pure Mamba might outperform a distilled one at the same architecture; (c) the teacher Transformer's attention projections may not transfer well due to architectural mismatch (e.g., non-standard attention variants) — the paper only validates the transfer for standard multi-head and grouped-query attention.