---
title: 'Llamba: Recurrent Llama Distillation'
url: https://www.emergentmind.com/topics/llamba
type: topic
---

# Llamba: Recurrent Llama Distillation

Llamba is a family of efficient recurrent language models distilled from Meta’s Llama‑3.x Transformers into Mamba‑2 state‑space architectures. The series comprises Llamba‑1B, Llamba‑3B, and Llamba‑8B, and its stated objective is to improve the tradeoff between speed, memory efficiency, and performance by replacing self‑attention with selective state‑space layers while retaining the Llama tokenizer, vocabulary, and overall block topology. The paper presents Llamba as a cross‑architecture distillation result built with MOHAWK, using 8B, 10B, and 12B tokens for the 1B, 3B, and 8B variants respectively, and as an optimized recurrent alternative for resource‑constrained devices such as smartphones and edge platforms [2502.14458].

## 1. Definition, lineage, and model family

Llamba is defined as a family of recurrent language models whose teachers are Llama‑3.2‑1B‑Instruct, Llama‑3.2‑3B‑Instruct, and Llama‑3.1‑8B‑Instruct. In the reported design, the teacher remains a Transformer, while the student adopts a modified, multi‑head Mamba‑2 architecture that replaces self‑attention with selective state‑space layers. The series includes Llamba‑1B, Llamba‑3B, and Llamba‑8B, and all three models share Llama‑3.1’s tokenizer and vocabulary [2502.14458].

| Model | Teacher | Core scale |
|---|---|---|
| Llamba‑1B | Llama‑3.2‑1B‑Instruct | 16 residual Mamba‑2 blocks; hidden size 2048 |
| Llamba‑3B | Llama‑3.2‑3B‑Instruct | 28 residual Mamba‑2 blocks; hidden size 3072 |
| Llamba‑8B | Llama‑3.1‑8B‑Instruct | 32 residual Mamba‑2 blocks; hidden size 4096 |

The reported motivation is not to retrain a language model family from scratch, but to preserve much of the teacher’s language capability while replacing the sequence mixer with a recurrent state‑space mechanism. This suggests that Llamba should be understood as an architectural translation of the Llama‑3.x line into a Mamba‑2 student, rather than as a separate pretraining lineage.

## 2. Architectural design

Each Llamba block is built around a Multi‑Head variant of Mamba‑2 with 32 heads and head dimensions of 64, 96, or 128; state size is 64. The architecture interleaves Llama‑style Gated MLPs between Mamba‑2 mixing layers, rather than using a pure SSM stack. The paper also reports three explicit modifications relative to vanilla Mamba‑2: removal of the pre‑output projection normalization, removal of the post‑convolution activation, and use of a “Discrete Mamba‑2” variant that directly projects the matrix $A$ from the input and eliminates the discretization parameter $\Delta$ [2502.14458].

The block pattern is described procedurally as:

1. RMSNorm $\rightarrow$ Discrete Mamba‑2 layer $\rightarrow$ residual  
2. RMSNorm $\rightarrow$ feed‑forward Gated MLP $\rightarrow$ residual

Conceptually, the selective SSM behavior is summarized in the paper as

$$
h_t = A(x_t)\,h_{t-1} + B(x_t)\,x_t,\quad y_t = C(x_t)\,h_t
$$

with input‑dependent projections $A(x_t)$, $B(x_t)$, and $C(x_t)$. The stated systems consequence is linear‑time inference and constant activation memory, in contrast to attention mechanisms that require decoding caches. The multi‑head design is also explicitly contrasted with Llama‑3.x MQA/GQA: Llamba uses 32 independent heads because recurrent layers do not need decoding caches.

Architecturally, Llamba therefore retains a recognizably Llama‑like residual and MLP organization while substituting the central temporal mixer. A plausible implication is that the design tries to minimize the representational discontinuity between teacher and student so that distillation remains well conditioned.

## 3. Distillation method and training procedure

The reported training method is MOHAWK, organized into three progressive stages: Matrix Orientation, Hidden‑State Alignment, and Weight Transfer plus Knowledge Distillation. Initialization is identity‑based: the convolution layer uses an identity kernel and the multiplicative skip passes inputs unchanged [2502.14458].

The three optimization stages are given as follows.

**Matrix Orientation** aligns the student’s materialized matrix mixer to the teacher’s self‑attention matrix:

$$
L_{\text{orient}} = \|M_S - M_T\|_F^2
$$

**Hidden‑State Alignment** minimizes an $L_2$ distance between student and teacher hidden states:

$$
L_{\text{hidden}} = \sum_l \|h_S^{(l)} - h_T^{(l)}\|_2^2
$$

**Knowledge Distillation** applies cross‑entropy on logits after transferring MLP weights, RMSNorms, embeddings, and output heads from the teacher:

$$
L_{KD} = - \sum_i [\mathrm{softmax}(z_T)]_i \log [\mathrm{softmax}(z_S)]_i
$$

The total token budgets reported in Table 2 are 8B for Llamba‑1B, 10B for Llamba‑3B, and 12B for Llamba‑8B, split across the three stages as 300M/2.7B/5B, 500M/4B/5.5B, and 500M/5B/6.5B respectively. The paper emphasizes that this is less than 0.1% of the training data typical for models of similar size. Stage 3 initially uses fineweb‑edu‑4.0 and then Open‑Hermes‑2.5 for 4 epochs at 200M tokens per epoch with sequence length 4096 [2502.14458].

The optimization setup is specified as mixed precision training with FSDP and activation checkpointing on a single node with 8× NVIDIA H100 GPUs, AdamW with $\beta_1 = 0.9$, $\beta_2 = 0.95$, and weight decay $= 0.1$, batch size 64 in stage 1 and 128 in stages 2 and 3, and a WSD scheduler with min LR $1\times10^{-8}$ and warm‑up and decay each 10% of total steps. Maximum learning rates are $1\times10^{-4}$ for stages 1 and 2 for all models, while stage‑3 learning rate is $5\times10^{-5}$ for Llamba‑1B and Llamba‑3B and $1\times10^{-5}$ for Llamba‑8B. Distillation stages 1 and 2 use packed sequences of length 2048.

## 4. Reported benchmark performance

The main zero‑shot evaluation covers ARC‑C, ARC‑E, PIQA, Winogrande, HellaSwag, Lambada, MMLU, and OpenBookQA, with averages reported over the eight tasks. Llamba‑1B reaches an average of 53.2, Llamba‑3B 63.9, and Llamba‑8B 68.8. The corresponding teachers score 55.3 for Llama‑3.2‑1B, 61.9 for Llama‑3.2‑3B, and 69.4 for Llama‑3.1‑8B [2502.14458].

For Llamba‑8B, the taskwise scores reported in Table 1 are ARC‑C 54.6, ARC‑E 82.5, PIQA 80.9, Winogrande 73.3, HellaSwag 77.6, Lambada 69.4, MMLU 61.0, and OpenBookQA 43.4. The teacher Llama‑3.1‑8B records 55.1, 81.7, 81.1, 73.9, 79.3, 73.0, 68.0, and 43.0 respectively. At 3B scale, Llamba‑3B slightly exceeds Llama‑3.2‑3B on the average score, 63.9 versus 61.9.

The paper also compares Llamba with other recurrent or hybrid baselines. Table 1 reports average scores of 67.0 for Falcon3‑Mamba‑7B, 65.7 for RecurrentGemma‑9B, 70.1 for Zamba2‑7B, and 69.8 for Qwen2.5‑7B, placing Llamba‑8B close to strong Transformer baselines while using substantially less distillation data than the teacher pretraining totals.

An ablation reported in Figure 4 states that stage‑3 knowledge distillation on fineweb‑edu improves MMLU notably relative to fineweb or C4. This suggests that, within the Llamba training recipe, dataset quality in the final distillation phase materially affects reasoning‑heavy evaluation.

## 5. Throughput, memory behavior, and on‑device deployment

The systems argument for Llamba is central to its design. On a single NVIDIA H100 80GB GPU, Llamba‑8B is reported to achieve the highest throughput across batch sizes 8–2048 with `prompt_len=1` and `gen_len=8192`, compared against Llama‑3.1‑8B at `gen_len=2048` and `gen_len=8192`. The experimental setup used `torch.compile(model, fullgraph=True)` with CUDA graph. The paper does not report exact tokens per second, but it states that Transformer baselines either slow down or run out of memory at larger batch sizes, whereas Llamba scales robustly to very large batches [2502.14458].

A second deployment result is reported on Apple Silicon M3 Pro (36GB) via MLX with 4‑bit quantization. In that setting, Llamba‑8B maintains near‑constant high throughput and low memory consumption as context grows, while Llama‑3.1‑8B’s throughput drops linearly and memory grows with context. The implementation basis is an optimized Mamba‑2 SSM and Conv1D kernel stack using Apple’s Metal framework, integrated into MLX for dynamic graphs and efficient tensor operations. The paper characterizes this as enabling real‑time on‑device decoding on iPhones and MacBooks.

The implementation release is given as:

`https://github.com/cartesia-ai/edge`

In practical terms, the contribution is not only architectural. It is also a deployment‑oriented software package for recurrent inference on unified‑memory Apple hardware and other edge settings.

## 6. Position in efficient language modeling and stated limitations

Llamba belongs to a broader shift toward state‑space or hybrid sequence models, but its strategy is distinct. Jamba, for example, interleaves attention and Mamba layers in a hybrid Transformer‑Mamba MoE decoder, whereas Llamba distills Llama‑3.x into a recurrent Mamba‑2 student that replaces self‑attention with selective state‑space layers [2403.19887]. The Llamba paper therefore represents one branch of the post‑Transformer efficiency literature: not hybridization, but teacher‑student architectural transfer into a recurrent backbone [2502.14458].

The stated limitations are also specific. The paper notes that SSM‑based models underperform Transformers on some algorithmic tasks, and the reported benchmark table shows that Llamba‑8B still trails Llama‑3.1‑8B on MMLU, 61.0 versus 68.0. The paper also identifies future work on improving long‑context handling and dataset quality. In addition, while implementations and quantized models are said to be released, the specific license is not detailed in the text [2502.14458].

The term “Llamba” should also be distinguished from unrelated names in adjacent literatures. “LLEMamba” is a low‑light image enhancement method based on Retinex optimization, deep unfolding, and Mamba priors, not a language model [2406.01028]. This suggests that “Llamba” has become a cross‑domain spelling collision, but in the language‑model literature the name is specifically associated with the distilled recurrent Mamba‑2 family introduced in 2025.

In summary, Llamba denotes a Llama‑derived recurrent language‑model family whose significance lies in demonstrating that cross‑architecture distillation can move a modern instruction‑tuned Transformer line into a Mamba‑2 student with 8B–12B distillation tokens, competitive benchmark performance, high batch‑throughput, and strong on‑device efficiency [2502.14458].

Source: https://www.emergentmind.com/topics/llamba