MambaInLlama: Hybrid Transformer Conversion

• MambaInLlama is a post-training architectural conversion that hybridizes Llama Transformer decoders by replacing selected self-attention layers with Mamba-style recurrent blocks.
• The framework preserves key components like embedding layers and MLPs while reusing attention projections, achieving a balance between task performance and inference efficiency.
• A multi-stage process involving pseudo-label distillation, supervised fine-tuning, and KL-regularized preference optimization ensures the hybrid model retains high performance.

Searching arXiv for the core paper and closely related hybrid Mamba–Llama work. MambaInLlama is a post-training architectural conversion framework in which a pretrained Llama-family Transformer is reconfigured into a hybrid Transformer–Mamba LLM by replacing selected self-attention layers with Mamba-style linear recurrent or state-space blocks, while preserving most of the surrounding decoder scaffold and reusing pretrained attention projections for initialization (Wang et al., 2024). In its original formulation, the method targets instruction-tuned teachers such as Llama-3-8B-Instruct and Zephyr-7B, retains a fraction of attention layers in an interleaved pattern, and recovers quality through pseudo-label distillation, supervised fine-tuning, and preference optimization (Wang et al., 2024). In later literature, “MambaInLlama” refers both to that specific method and to a broader design pattern for Llama-to-Mamba hybridization, especially in efficiency-oriented reasoning and deployment settings (Paliotta et al., 27 Feb 2025, Wang et al., 14 Apr 2025, Yang et al., 22 May 2025).

1. Historical setting and problem formulation

MambaInLlama emerged from the practical mismatch between the quality of pretrained Transformer checkpoints and the deployment advantages of linear RNN or state-space sequence models. The motivating observation is that strong open models are overwhelmingly available as Transformers, whereas Mamba-style models avoid Transformer KV-cache growth and are attractive for long-context generation, large-batch serving, and agentic rollouts (Wang et al., 2024). The method therefore addresses a conversion problem rather than a from-scratch pretraining problem: how to begin with a strong Llama checkpoint and remove many attention layers without discarding the teacher’s function.

The original work frames this as architectural transplantation. A teacher Transformer is not merely distilled into a smaller student; instead, substantial parts of the teacher’s internal sequence mixer are replaced with a different class of operator, namely Mamba-style linear recurrent blocks, and the resulting hybrid is then aligned back to the teacher through post-training (Wang et al., 2024). This distinguishes MambaInLlama from pure-Mamba pretraining efforts and from hybrids trained from scratch.

Later work situates MambaInLlama within a broader spectrum of Transformer–SSM designs. “Thinking Slow, Fast” explicitly uses a hybrid model called MambaInLlama as one of its distilled reasoners, with the central question being whether subquadratic architectures can exploit higher generation throughput to outperform similarly sized Transformers under fixed inference budgets (Paliotta et al., 27 Feb 2025). “M1” adapts the same conversion idea to mathematical reasoning, starting from Llama3.2-3B-Instruct and replacing most attention layers with Mamba blocks before further SFT and RL (Wang et al., 14 Apr 2025). “Zebra-Llama” treats MambaInLlama as a baseline post-training hybrid approach that replaces most self-attention blocks with linear RNN layers and uses structured initialization plus distillation (Yang et al., 22 May 2025).

2. Architectural construction and weight reuse

Architecturally, MambaInLlama preserves the global Llama-style decoder scaffold while changing selected sequence-mixing sublayers. The embedding layer, output LM head, residual stream, and MLP or FFN sublayers are inherited from the teacher, and the paper is explicit that the Transformer MLP layers are kept “as is” and frozen in the first distillation stage (Wang et al., 2024). The decisive intervention is at the attention sublayer: a chosen subset of self-attention layers is replaced by Mamba-style linear recurrent blocks.

The conversion is motivated by writing attention and linear RNNs in compatible forms. For an attention head, the original paper writes

$$q_t = W^Q x_t,\qquad k_t = W^K x_t,\qquad v_t = W^V x_t$$

and then defines the output through causal attention weights. It then recalls the generic linear-RNN form

$$h_t = \bar A_t h_{t-1} + \bar B_t x_t,\qquad y_t = \bar C_t h_t.$$

After linearizing attention, the paper rewrites the attention computation in a recurrent form and uses that correspondence as the initialization bridge from attention parameters to recurrent parameters (Wang et al., 2024).

Algorithm 1, “Attention-Initialized Mamba,” gives the concrete mapping. For each attention head and each position $t$,

• $\mathbf{x}_t \gets W^V x_t$,
• $\mathbf{b}_t \gets W^K x_t$,
• $\mathbf{c}_t \gets W^Q x_t$,
• $\Delta_t \gets \text{MLP}(x_t)$,

followed by dynamic discretization

$$\bar A_{1:T}, \bar B_{1:T}, \bar C_{1:T} \gets Disc(A, B, C, \Delta),$$

a recurrent rollout

$y \gets LinearRNN(\bar A, \bar B, \bar C, x),$

and accumulation through the reused output projection

$\text{output} \gets \text{output} + W^{O\top} y.$

In operational terms, $$h_t = \bar A_t h_{t-1} + \bar B_t x_t,\qquad y_t = \bar C_t h_t.$$0, $$h_t = \bar A_t h_{t-1} + \bar B_t x_t,\qquad y_t = \bar C_t h_t.$$1, $$h_t = \bar A_t h_{t-1} + \bar B_t x_t,\qquad y_t = \bar C_t h_t.$$2, and $$h_t = \bar A_t h_{t-1} + \bar B_t x_t,\qquad y_t = \bar C_t h_t.$$3 are reused from the teacher attention layer, while new Mamba-specific parameters are introduced for the state dynamics, especially $$h_t = \bar A_t h_{t-1} + \bar B_t x_t,\qquad y_t = \bar C_t h_t.$$4 and the network predicting $$h_t = \bar A_t h_{t-1} + \bar B_t x_t,\qquad y_t = \bar C_t h_t.$$5 (Wang et al., 2024).

The original paper experiments with retaining 50%, 25%, 12.5%, or 0% of the attention layers, and it states that retained attention layers are kept in an interleaved pattern by “keep[ing] every $$h_t = \bar A_t h_{t-1} + \bar B_t x_t,\qquad y_t = \bar C_t h_t.$$6 attention layers” (Wang et al., 2024). This makes MambaInLlama a sparse-attention hybrid rather than a pure recurrent model in its strongest configurations.

Subsequent works refine this conversion template. “M1” states that the linear projections for $$h_t = \bar A_t h_{t-1} + \bar B_t x_t,\qquad y_t = \bar C_t h_t.$$7, $$h_t = \bar A_t h_{t-1} + \bar B_t x_t,\qquad y_t = \bar C_t h_t.$$8, $$h_t = \bar A_t h_{t-1} + \bar B_t x_t,\qquad y_t = \bar C_t h_t.$$9, and $t$0 are initialized from the corresponding projections for $t$1, $t$2, $t$3, and $t$4, respectively, and further adds two linear layers to expand from $t$5 to $t$6 because the source Transformer uses GQA while Mamba does not use a KV cache (Wang et al., 14 Apr 2025). “Thinking Slow, Fast” likewise presents hybrid students as Llama decoder stacks in which many self-attention sublayers are swapped for Mamba sublayers while the MLP stack is reused directly (Paliotta et al., 27 Feb 2025).

3. Distillation, alignment, and post-training

The original MambaInLlama pipeline is explicitly multi-stage. It is not a one-shot architectural surgery but a conversion-and-alignment procedure tailored to instruction-tuned models. The three stages are pseudo-label distillation, supervised fine-tuning, and preference optimization, intended to mimic the standard post-training stack of aligned LLMs (Wang et al., 2024).

For supervised distillation, the paper combines sequence-level pseudo-labeling and token-level KL:

$t$7

In the first stage, $t$8 and $t$9 (Wang et al., 2024). Preference alignment is then formulated through a KL-regularized reward-maximization objective and instantiated as DPO with the teacher as reference:

$\mathbf{x}_t \gets W^V x_t$0

The paper notes this use of DPO specifically as a distillation objective (Wang et al., 2024).

The data schedule is also specified. Pseudo-label distillation uses seed prompts from UltraChat and UltraFeedback; supervised finetuning uses GenQA, InfinityInstruct, and OpenHermes 2.5; and DPO uses UltraFeedback for Zephyr-derived models and datasets from SimPO and Zephyr for Llama-3-8B-Instruct-derived models (Wang et al., 2024). Optimization details include AdamW, $\mathbf{x}_t \gets W^V x_t$1 betas, batch size 64, linear warmup for the first 500 steps, and cosine annealing. Each hybrid model distillation is reported to take less than five days on $\mathbf{x}_t \gets W^V x_t$2GB A100 GPUs (Wang et al., 2024).

A notable procedural feature is progressive replacement. The paper states that stepwise replacement is best: first keep every 2 attention layers and distill, then every 4, and continue (Wang et al., 2024). This suggests that the optimization difficulty lies not only in knowledge transfer but also in stabilizing the functional transition from attention to recurrence.

Later variants extend rather than discard this logic. “Thinking Slow, Fast” modifies the earlier procedure by replacing all target layers in a single round and using token-level reverse KL on OpenMathInstruct-2, then adding post-distillation SFT (Paliotta et al., 27 Feb 2025). “M1” follows a three-stage reasoning recipe—distillation, math SFT and reasoning SFT, then GRPO RL—after initializing a hybrid from Llama3.2-3B-Instruct (Wang et al., 14 Apr 2025). “Zebra-Llama” explicitly states that its structured Mamba2 initialization follows MambaInLlama “precisely,” while adding intermediate layer distillation and SMART layer selection on top of the same attention-to-SSM conversion idea (Yang et al., 22 May 2025).

4. Benchmark behavior and ablation findings

The original paper’s strongest model is a hybrid distilled from Llama-3-8B-Instruct with 50% retained attention. It reports MT-Bench $\mathbf{x}_t \gets W^V x_t$3, AlpacaEval 2 length-controlled win rate $\mathbf{x}_t \gets W^V x_t$4, and overall AlpacaEval win rate $\mathbf{x}_t \gets W^V x_t$5. The teacher Llama-3-Instruct 8B reports MT-Bench $\mathbf{x}_t \gets W^V x_t$6 and AlpacaEval 2 length-controlled win rate $\mathbf{x}_t \gets W^V x_t$7. More aggressive removal of attention reduces quality: Mamba-Llama3 (25%) reports MT-Bench $\mathbf{x}_t \gets W^V x_t$8 and AlpacaEval 2 length-controlled win rate $\mathbf{x}_t \gets W^V x_t$9; Mamba-Llama3 (12.5%) reports $\mathbf{b}_t \gets W^K x_t$0 and $\mathbf{b}_t \gets W^K x_t$1; and pure Mamba2-Llama3 (0%) reports $\mathbf{b}_t \gets W^K x_t$2 and $\mathbf{b}_t \gets W^K x_t$3 (Wang et al., 2024).

Configuration	MT-Bench	AlpacaEval 2 LC
Llama-3-Instruct 8B	8.00	22.90
Mamba-Llama3 (50%)	7.35	29.61
Mamba-Llama3 (25%)	6.86	25.85
Mamba-Llama3 (12.5%)	6.46	20.76
Mamba2-Llama3 (0%)	5.64	14.49

The same pattern appears in general benchmarks. The best reported hybrid, Mamba2-Llama3 (50%), reaches average $\mathbf{b}_t \gets W^K x_t$4 on the 10-task zero-shot LM Eval aggregate, compared with $\mathbf{b}_t \gets W^K x_t$5 for Nvidia Hybrid Mamba-8B and $\mathbf{b}_t \gets W^K x_t$6 for TRI Mamba-7B (Wang et al., 2024). The paper therefore argues that a distilled hybrid can outperform open-source Mamba or hybrid models trained from scratch on much larger corpora.

The ablations are unusually diagnostic. Attention-derived initialization is critical: for Zephyr-Mamba (50%), attention initialization versus random initialization changes LAMBADA perplexity from $\mathbf{b}_t \gets W^K x_t$7 to $\mathbf{b}_t \gets W^K x_t$8, MT-Bench from $\mathbf{b}_t \gets W^K x_t$9 to $\mathbf{c}_t \gets W^Q x_t$0, and AlpacaEval length-controlled win rate from $\mathbf{c}_t \gets W^Q x_t$1 to $\mathbf{c}_t \gets W^Q x_t$2 (Wang et al., 2024). Freezing MLPs during early pseudo-label distillation also helps substantially, and interleaving retained attention layers outperforms non-interleaved retention. The combined post-training stack matters as well: for Zephyr-derived hybrids, distillation plus SFT plus DPO outperforms distillation plus SFT or distillation plus DPO alone on MT-Bench (Wang et al., 2024).

Subsequent reasoning work reinforces the same hybrid preference. “Thinking Slow, Fast” reports that hybrid MambaInLlama models retain 4 attention layers out of 16 total layers for 1B and 6 attention layers out of 26 total layers for 3B, and concludes that the hybrid models often define the Pareto front under fixed time budgets for MATH and GSM8K (Paliotta et al., 27 Feb 2025). “M1,” a 3B hybrid initialized from Llama3.2-3B-Instruct with 6 interleaved attention layers among 28 total layers, reports AIME25 $\mathbf{c}_t \gets W^Q x_t$3, AIME24 $\mathbf{c}_t \gets W^Q x_t$4, MATH500 $\mathbf{c}_t \gets W^Q x_t$5, AMC23 $\mathbf{c}_t \gets W^Q x_t$6, and OlympiadBench $\mathbf{c}_t \gets W^Q x_t$7, matching or exceeding DeepSeek-R1-Distill-Qwen-1.5B on several benchmarks at similar scale (Wang et al., 14 Apr 2025).

5. Inference efficiency, state management, and speculative decoding

The deployment rationale for MambaInLlama is bound to the contrast between KV-cache-heavy Transformer decoding and recurrent-state decoding. The original paper motivates Mamba-like models by citing prior evidence that they can reach roughly $\mathbf{c}_t \gets W^Q x_t$8 higher throughput than Transformers, and positions the hybrid conversion as a way to inherit part of that systems advantage without abandoning pretrained Transformer quality (Wang et al., 2024).

Its distinctive systems contribution is a hardware-aware speculative decoding algorithm specialized for Mamba and hybrid models. The core difficulty is that a Transformer can conceptually rewind by truncating its KV cache, whereas a linear RNN compresses the entire history into recurrent state. The paper therefore introduces a multistep verification kernel

$\mathbf{c}_t \gets W^Q x_t$9

which fuses recomputation from a cached earlier state, multistep recurrent verification over the drafted suffix, and caching of future states in a single kernel (Wang et al., 2024).

The reported speedups are concrete. For pure Mamba on The Pile, a Mamba 2.8B verifier with a 130M Mamba draft reaches $\Delta_t \gets \text{MLP}(x_t)$0 tok/s and $\Delta_t \gets \text{MLP}(x_t)$1 speedup on H100 at $\Delta_t \gets \text{MLP}(x_t)$2. A Mamba 7B verifier with a trained Llama3-1B draft reports $\Delta_t \gets \text{MLP}(x_t)$3–$\Delta_t \gets \text{MLP}(x_t)$4 tok/s and about $\Delta_t \gets \text{MLP}(x_t)$5 speedup on H100. For hybrid models on a single RTX 3090, the reported speculative decoding speedups range from $\Delta_t \gets \text{MLP}(x_t)$6 to $\Delta_t \gets \text{MLP}(x_t)$7 depending on verifier, draft depth, and $\Delta_t \gets \text{MLP}(x_t)$8 (Wang et al., 2024).

The abstract additionally states that the distilled model has natural length extrapolation and shows almost perfect accuracy in a needle-in-a-haystack test at $\Delta_t \gets \text{MLP}(x_t)$9 the distillation length (Wang et al., 2024). The detailed discussion summarized in the same paper, however, centers mainly on chat benchmarks, general LM benchmarks, and speculative decoding, so the article’s most fully specified long-context evidence remains the deployment-side decoding algorithm rather than a fully detailed long-context evaluation setup.

Later comparative work treats MambaInLlama as a throughput baseline. “Zebra-Llama” reports that its 8B variants achieve $$\bar A_{1:T}, \bar B_{1:T}, \bar C_{1:T} \gets Disc(A, B, C, \Delta),$$0–$$\bar A_{1:T}, \bar B_{1:T}, \bar C_{1:T} \gets Disc(A, B, C, \Delta),$$1 higher throughput than MambaInLlama up to a $$\bar A_{1:T}, \bar B_{1:T}, \bar C_{1:T} \gets Disc(A, B, C, \Delta),$$2k context length, while attributing that gap largely to replacing retained GQA-like attention with MLA rather than to a different Mamba initialization (Yang et al., 22 May 2025). “M1” reports more than a $$\bar A_{1:T}, \bar B_{1:T}, \bar C_{1:T} \gets Disc(A, B, C, \Delta),$$3 speedup compared to a same-size Transformer and frames this speed as the reason the hybrid can win under fixed reasoning-time budgets through self-consistency voting (Wang et al., 14 Apr 2025).

6. Interpretation in the broader Mamba–Llama literature

In subsequent literature, MambaInLlama functions as both a concrete method and a reference point for alternative design choices. One branch preserves the hybrid premise but changes the non-SSM component. “Zebra-Llama” keeps the same general agenda—post-training a Llama-family Transformer into an efficient hybrid with structured Mamba initialization and distillation—but replaces the surviving GQA-style attention layers with MLA, arguing that MambaInLlama leaves too much KV-cache burden in the retained attention layers (Yang et al., 22 May 2025).

A second branch pushes further toward pure recurrence while keeping the Llama macro-architecture. “Llamba” is presented as a family of Llama-shaped, Mamba-based autoregressive LLMs distilled from Llama-3.x. It retains the overall Llama structure, tokenizer and vocabulary, transferred input embeddings and output head, MLP weights, and normalization layers, but replaces self-attention entirely with a Discrete Mamba-2 layer (Bick et al., 20 Feb 2025). This is closely related in spirit to MambaInLlama, but it is attention-free in the final architecture rather than hybrid.

A third branch focuses on task-specific hybrids. “M1” adapts MambaInLlama directly for long-chain-of-thought mathematical reasoning (Wang et al., 14 Apr 2025). “Thinking Slow, Fast” uses both pure Llamba and hybrid MambaInLlama students distilled from Llama 3.2 Instruct teachers, concluding that hybrid MambaInLlama is often the best speed–quality compromise under fixed time budgets (Paliotta et al., 27 Feb 2025). These studies collectively shift the interpretation of MambaInLlama from a chat-model conversion technique to a more general inference-scaling recipe.

Several adjacent works also clarify what MambaInLlama is not. “Jamba” is a hybrid Transformer–Mamba MoE model trained from scratch, not a post-hoc conversion of a pretrained Llama checkpoint; nevertheless, its $$\bar A_{1:T}, \bar B_{1:T}, \bar C_{1:T} \gets Disc(A, B, C, \Delta),$$4 attention:Mamba design and its finding that sparse attention is important for in-context learning and format adherence strongly support the same architectural intuition that some attention layers should remain (Lieber et al., 2024). “Falcon Mamba” takes the opposite extreme, presenting a pure Mamba 7B model and arguing that full attention removal can be competitive at scale, but only with a mature training recipe, $$\bar A_{1:T}, \bar B_{1:T}, \bar C_{1:T} \gets Disc(A, B, C, \Delta),$$5 trillion tokens, and architecture-specific stabilization such as RMSNorm after $$\bar A_{1:T}, \bar B_{1:T}, \bar C_{1:T} \gets Disc(A, B, C, \Delta),$$6, $$\bar A_{1:T}, \bar B_{1:T}, \bar C_{1:T} \gets Disc(A, B, C, \Delta),$$7, and $$\bar A_{1:T}, \bar B_{1:T}, \bar C_{1:T} \gets Disc(A, B, C, \Delta),$$8 (Zuo et al., 2024). In multimodal modeling, “VL-Mamba” is relevant chiefly as a terminological contrast: it does not insert Mamba blocks inside a LLaMA stack, but instead replaces the entire LLaVA-style language backbone with a pretrained Mamba LLM and keeps only the outer vision-encoder-plus-connector recipe (Qiao et al., 2024). In domain-specific transfer, “Mamba4Net” uses Llama2-7B as a teacher for a separate hybrid Mamba student with 10 Mamba layers and 2 Transformer layers, but frames this explicitly as cross-architecture distillation rather than an in-place Llama modification (Xia et al., 20 Oct 2025).

A persistent misconception is therefore that “MambaInLlama” means any use of Mamba near a Llama-family model. In the strict sense established by the original paper, it denotes replacement of selected Llama attention layers with Mamba blocks inside an otherwise preserved Llama scaffold, initialized from the original attention projections and then recovered through staged distillation (Wang et al., 2024). Broader usage in later work extends the label to related hybrids and baselines, but the original concept remains specific: attention-to-Mamba conversion of a pretrained Llama-like decoder, with sparse retained attention as the empirical compromise between Transformer quality and recurrent efficiency.