Zebra-Llama: Towards Extremely Efficient Hybrid Models
Authors: Mingyu Yang
Mehdi Rezagholizadeh
Guihong Li
Vikram Appia
Emad Barsoum
Venue: NeurIPS 2025 poster ()
Keywords: Model Efficiency
Large Language Models
Primary area: deep_learning

Abstract: With the growing demand for deploying large language models (LLMs) across diverse applications, improving their inference efficiency is crucial for sustainable and democratized access. However, retraining LLMs to meet new user-specific requirements is prohibitively expensive and environmentally unsustainable. In this work, we propose a practical and scalable alternative: composing efficient hybrid language models from existing pre-trained models.
Our approach, X-EcoMLA, introduces a family of 1B, 3B, and 8B hybrid models by combining State Space Models (SSMs) and Multi-head Latent Attention (MLA) layers, using a refined initialization and post-training pipeline to efficiently transfer knowledge from pre-trained Transformers.
X-EcoMLA achieves Transformer-level accuracy with near-SSM efficiency using only 7–11 billion training tokens (compared to the trillions required for pre-training) and an 8B teacher. Moreover, it dramatically reduces KV cache size—down to 3.9%, 2%, and 2.73% of the original for the 1B, 3B, and 8B variants, respectively—while preserving 100%, 100%, and over 97% of average zero-shot performance on LM Harness tasks.
Compared to models like MambaInLLaMA, X-EcoMLA, Minitron, and Llamba, our approach consistently delivers competitive or superior accuracy while using significantly fewer tokens, smaller teachers, and vastly reduced KV cache memory. Notably, X-EcoMLA-8B surpasses Minitron-8B in few-shot accuracy by 7%, while using 8× fewer training tokens, over 12× smaller KV cache, and a smaller teacher (8B vs. 15B). 
It also achieves 1.4x–3.3x higher throughput (tokens/s) than MambaInLlama. The source code is
released at https://github.com/AMD-AGI/AMD-Hybrid-Models.

Zebra-Llama: Towards Extremely Efficient Hybrid Models
Mingyu Yang", Mehdi Rezagholizadeh' Guihong Li', Vikram Appia, Emad Barsoum Advanced Micro Devices, Inc. (AMD)
(mingyu. yang,
mengi mehdi. rezagholi zadeh, guihong.. 1i 8amd com
Abstract
With the growing demand for deploying large language models (LLMs) across di- verse applications, improving their inference efficiency is crucial for sustainable and democratized access. However, retraining LLMs to meet new user-specific require- ments is prohibitively expensive and environmentally unsustainable. In this work, we propose a practical and scalable alternative: composing efficient hybrid language models from existing pre-trained models. Our approach, Zebra-Llame introduces family of 1B, 3B, and 8B hybrid models by combining State Space Models (SSMs) and Multi-head Latent Attention (MLA) layers, using refined initialization and post-training pipeline to efficiently transfer knowledge from pre-trained Transform- ers. Zebra-Liama achieves Transformer-level accuracy with near-SSM efficiency using only 7-1 training tokens (compared to trillions of tokens required for pre-training) and an 8B teacher. Moreover, Zebra-Llama dramatically reduces KV cache size-down to 3.9%, 2%, and 2.73% of the original for the IB, 3B, and 8B variants, respectively- wwhle preserving 100%, 100%, and 997% of average zero-shot performance on LM Harness tasks. Compared to models like MambaInL- LaMA, X-EcoMLA, Minitron, and Llamba, Zebra- consistently delivers com- petitive or superior accuracy while using significantly fewer tokens, smaller teach- ers, and vastly reduced KV cache memory. Notably, Zebra-Llama-81 surpasses Minitron-8B in few-shot accuracy by 7% while using 8x fewer training tokens, over 12x smaller KV cache, and a smaller teacher (8B vs. 15B). It also achieves 1.4x-3.3x x higher throughput (tokens/s) than MambaInL.ama. The source code is released at httptttts///. https::/tttttiiiie //tbt com/ AMD AGT AMD Hybrid Mode 1s
1 Introduction
The exponential growth of deep learning applica- tions has created an urgent demand for models that strike balance between accuracy and com- putational particeccy-prtiiullllly in scenarios con- strained by memory or limited hardware capabili- ties. Transformer- based models, despite their im- pressive performance across range of tasks, are fundamentally limited by the quadratic complexity of their self-attention mechanisms and the substan- tial memory required to store key-value (KV) caches. These bottlenecks hinder their deployment in real- world applications, especially on edge devices or in latency-sensitive settings. At the same time, the rise of large language models (LLMs) has amplified the need for customization that is, the ability adapt pre-trained models to meet diverse user needs, hard- ware configurations, and application requirements. However, developing new LLMs from scratch for each target environment is prohibitively expensive
66
64
1981
Pre buinnd Model Thininng
10011
2000
KV Cache Size (HAB)
Figure 1: Comparing SB-scale models average LM Harness score KV cache size. Zebra-I.lan (green) matches or exceeds baselines with smaller KV cache and fewer training tokens. Cirele and square sizes indicate training tokens (billions for post-training, trillions for pre-training).
Equal Contribution First Authors, with order determined randomly.
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

and environmentally unsustainable. Traditional solutions such as model compression, neural archi- tecture search (NAS), and pre-training new architectures offer potential pathways but suffer from significant limitations. Model compression often degrades quality while NAS or pre-training new models have substantial computational costs.
To overcome these challenges, a promising new paradigm has emerged: hybrid models that aim to reduce the computational cost of self-attention while maintaining generation quality. These archi- tectures typically integrate efficient state-space models [ or linear attention (2) with full attention mechanisms, leveraging the strengths of both. However, many recent hybrid approaches, including Samba (3) and Hymba [4), require extensive pre- pre- training from scratch, which is computationally expensive. Others, such as MamM [5]. experience notable performance degradation when attention is replaced too aggressively, often due to insufficient key-value cache compression or ineffective knowledge transfer from the base model.
Our goal in this work is to develop a more efficient and sustainable alternative: to compose highly efficient language models directly from existing pre-trained Transformers, avoiding the cost of full pre-training while retaining performance. Our approach, called Zebra-Llama, introduces a family of hybrid models (IB, 3B, and 8B) built on two complementary components: Multi-Latent Attention (MLA) (6). low-rank attention mechanism that compresses memory usage without sacrificing quality under moderate compression; and Mamba2, a state-space model that eliminates KV caches entirely but performs poorly when used alone (5). Specifically, we first initialize pure MLA and pure Mamba2 models from a pre-trained Transformer via a refined weight-mapp mapping procedure. We then use Intermediate Layer Distillation (IID) to align their internal representations with those of the original Transformer model, ensuring strong initialization. Finally, we strategically compose hybrid architectures from the refined MLA and Mamba2 variants using a sensitivity- aware strategy called SMART (Sensitivity Measure- Aware Replacement of Transformer layers) to select where each component is most effective. This process results in highly efficient models that retain Transformer- level quality with drastically reduced memory and compute requirements. Our Zebra-Llama family of IB, 3B, and 8B hybrid models, achieving 25x, 50 and 36x KV cache compression relative to their respective base Transformer models, while maintaining 100%. 100%, and >97% of the base model's 's average zero-shot performance perform competitively on sew-shot tasks. Notably, Zebra-Llama-83 8B improves the average few-shot accuracy over Minitron-&B by 7%, despite using 8x fewer training tokens, >12x smaller KV cache, and a smaller teacher model (8B vs. 15B). Additionally, our Zebra-Llama exhibits high inference speed. Compared to the existing work MambalnL it achieves 1.4-3.3x higher throughput with 1.8-5.5x less memory consumption. This work introduces a practical route for building efficient, customizable LLMs from existing models. Our key contributions are:
the LM Harness evaluation benchmark. Our models also
Architecture: We propose a hybrid model combining MLA and Mamba2 layers, replacing classical Transformer blocks to reduce memory usage and address the quadratic bottleneck of attention.
Training: We develop an efficient post-training pipeline including refined initialization, Intermedi- ate Layer Distillation, and SMART layer selection for hybrid model composition.
Empirical Results: Our models match or exceed Transformer-lews performance with drastically reduced KV cache and significantly improved inference throughput.
2 Related Work
Hybrid Models Prior work on hybrid models can broadly be categorized into two groups based on their training strategy: pre-training-basse based and post-training-buat
approaches. Pre-training- based methods, such as Jamba (7). Hymba (8). Samba (3] and Manba-2-Hyhrid (9), interleave heterogeneous hybrid layers during full-scale model training, allowing the development of hybrid models from scratch. While these models achieve strong performance, their training cost remains high, limiting accessibility and sustainability. In contrast, post-training approaches insert efficient modules into pre-trained Transformers, often leveraging knowledge distillation to transfer capabilities without full re-training, including MambaInLLaMa (5) MOHAWK [10). and Llamba [T]. Our work focuses the post-training setting and pushes it further by targeting extremely efficiency- both in terms of training tokens and runtime memory. We introduce a systematic method to build hybrid models with minimal training costs, making them practical for broad deployment, especially where hardware resources are critically constrained.
Efficient Post-training Several recent approaches specifically target post-training efficiency through careful initialization, distillation, and model compression. For example, MambalnLLaMA 5

---

STEP 'S Weight Iriittatiion
Skp 2- Hiffled Hriokedis
Step 2: Seseaiiiity Messure Aeas Beskcemerted Tranffomer Lapers
LD
Varner
LG
Varoet
Dave
Paxc MLA
Aiteiiion lane Madel
PL
Pare
Alc Case Model
PLY
Figure 2: Overview of our hybrid model composition pipeline. The process consists of three stages: (1) Weight Initialization - we initialize pure Mamba2 and MLA models from pre-trained Transformer via structured mapping; (2) Refined Initialization through Intermediate Layer Distillation (ILD) we refine both models by aligning their internal representations with the base model on a small dataset; and (3) SMART Layer Selection we compose the final hybrid model by selecting MLA and Mamba2 layers based on sensitivity analysis.
replaces most self-attention blocks of a pre-trained Transformer with linear RNN layers. Through the initialization of RNN layer parameters from the pre-trained Transformer's 's weights, followed by distillation- based fine-tuning, the hybrid model achieves comparable performance with the teacher model. MOHAWK [10] introduces a multi-stage cross- cross-architecture distillation strategy (aligning attention "mixing" matrices, hidden states, and outputs) to transfer a Transformer's 's knowledge into a Mamba- based model using a fraction of the original training data, Building on this, Llamba [1]] scales the distilled recurrent architecture up to 8B parameters, attaining improved accuracy, while markedly improving inference throughput and memory usage for deployment on edge devices. Orthogonally, X-BcoMLA [12] "upcycles" a pre-trained Transformer's 's attention into multi-head latent attention modules, jointly compressing key-value caches significantly with teacher- teacher-guided fine-tuning to preserve accuracy. Meanwhile, Minitron [13] compresses a 15B Nemotron Transformer by pruning its depth, width, and attention heads, then retrains on only about 2-3% of the original data via distillation; this yields 8B and 4B models that rival larger models' performance without full re- trainin..ng. Our method offers significant advantages over previous distillation techniques like Minitron, MambaInLlama, X-EcoMLA, and Llamba, which have each faced limitations in KV cache compression, training efficiency or maintaining the base models' performance.
3 Methodology
Our methodology focuses on designing and training a hybrid model architecture that achieves strong performance with significantly enhanced efficiency. The overall approach consists of two key stages: (1) constructing an extremely efficient hybrid model, and (2) applying lightweight yet effective training pipeline. An overview of the process is provided in Figure
3.1 Composing Extremely Efficient Hybrid Model
To compose our hybrid model, we combine two complementary components: Mamba2 14 and
MLA (6) blocks. Each of these blocks contributes differently to efficiency and performance: Mamba2 blocks are based on SSMs with zero KV cache usage, making them ideal for long-context
or memory-constraine settings. However, they often underperform when used exclusively (9).
MLA blocks compress standard attention process reduce KV cache requirements while maintain- ing high performance. Yet, excessive compression can lead to noticeable performance drops 12]
Our hybrid architecture interleaves Mamba2 and MLA layers to set balance between minimal memory usage and strong predictive performance. The composition process consists of two key stages: (i) we first construct pure Mamba2 model and a pure MLA model by replacing all attention blocks in the base Transformer with Mamba2 and MLA blocks, respectively, each initialized from the original pre-trained weights; (ii) we then apply intermediate layer distillation (ILD) to align the outputs of each layer in the Mamba2 and MLA models with those of the corresponding layers in the base model, thereby preserving the original model's knowledge during the transition.
Following this refined initialization, we compose the final hybrid model using our sensitivity measure- aware replacement of transformer layers (SMART) layer selection strategy, which selects the optimal placement confi guniigurations of Mamba2 and MLA layers based on laser sensitivity analysis.
3.1.1 Refined Initialization
To initialize the MLA and Mamba2 blocks in pure MLA and pure Mamba2 models derived from pre- trained base model, we introduce an enhanced initialization strategy that goes beyond conventional weight transfer.
3

---

Background: Multi-Head Attention Given an imput hidden representation H E Rixd, MHA projects it into query, key, and value matrices using learned weights:
Q Hw? = Hw = HHW,.
(1)
where WO, W WV E Rdxxndd and nA. d, denote the number of heads and per-head dimension, d is the hidden dimension, and I represents the sequence length. The attention operation computes:
A = Softmax
O AVWO
with po Raddd During inference, MHA requires caching K and for all past tokens, incurring a memory cost of 2n,
Background: Multi- Latent Attention MLA (6) introduces low-rank compression scheme to reduce memory usage. Instead of storing K and V directly, they are compressed into latent vector:
CKV HWDKV
(3)
with W DRV E R Ax and Tke < nadh, Keys and values are reconstructed via:
Kc = CKVWUK Vc = CKVWwr
(4)
where WoR and WUV are up- projection matrices. MLA also compress the query:
C² = HWPQ, QC = cQwUe
(5)
where 09 E R and ry < d. To retain compatibility with RoPE, MLA decouples positional encoding using separate projections:
Q" ( K = (oP)(HW**))
(6)
where WOR E RTaxnad, and WKR Rdxd+ The final queries and keys are then constructed as: Q = QQC; Q&1, K k= repeat( K K where repeat K denotes replication across heads. Overall, MIA reduces KV-cache size for inference from dn,, to O((rke d,)().
Initializing MLA from Pre-trained MHA To construct an MLA-based hybrid model from a pre-trained Transformer, we upcycle its attention modules by reparameterizing them into low-rank latent attention. This conversion is initialized using a structured approach based on singular value decomposition (SVD), enabling knowledge transfer while minimizing performance loss [12]. The core MLA weights (i.e. query, key, and value projections) are initialized using low-rank approximations derived from pre-trained MHA parameters. For query weights, we apply SVD on W4:
W = U
and set WDQQ Uy. The up-projection matrices W and WQR are derived from M, V. by partitioning and reshaping into the appropriate query and RoPE dimensions. For key and value projections, concatenate W and WV and apply joint SVD:
(8) We then set WDKV Ukv. and derive WUK and WUV by partitioning and reshaping E... V The shared RoPE key projection WKR is initialized from the average W across heads. We choose constant rank rg. Fike across all layers. Non-attention parameters such as the feedforward network and output projection WO are directly copied from the pre-trained model. Additional details of MLA initialization are provided in Appendix A.
[W K W = Use
Initializing Mamba2 from Attention Representations has been shown in (5 that linear attention can be reinterpreted as special case of SSMs. This connection enables the initialization of Mamba2 blocks from pre-trained attention based Transformers. In particular, the linearized form of attention without softmax resembles linear RNN update:
h4 My 1. + K, Ve, Vt = VD
Qihr,
which parallels the linear RNN structure:
hr A,hi-1 + B 91 -,h4-
(10)
To bridge this connection, LLMM (5) proposes to initialize the continuous -time Mamba2 SSM parameters (A, B, C) from the weights of the attention blocks. This includes discretizing the SSM over a learned sequence of sampling intervals A, to match the temporal dynamics of attention. Specifically, K,T Vú in attention is mapped to B, in Mamba2; Q, in attention plays the role of C₁; and the memory coefficient m1-1.t corresponds to A, in the recurrent update. Additional details of MLA initialization are provided in Appendix A.2

---

Refined Initialization through Intermediate Layer Distillation (ILD) After initializing the parameters of MLA and Mamba2 layers from the pre-trained Transformer, we further refine their weights through lightweight ILD procedure on small portion of the training data, akin to the second phase of MOHAWK [10]. This focused training aligns the internal representations between the MLA and Mamba2 layers and the pre-trained Transformer layers, ensuring a smoother start for downstream optimization and better preservation the knowledge embedded in the original model. To perform ILD, we minimize the mean squared error (MSE) between the outputs of each pre-trained Transformer attention layer lyAlin and the corresponding outputs of MLA (AMLA) and Mamba2 (ht*) layers. The ILD losses for Mamba2 and MLA are defined as:
Mamba2 LLD
Σ
h
M2 M
MLA ILD
Σ LL.
M
Amn
11
MLA
(II)
where L denotes the total number of layers and h2 refers to the output of the thh layer. Unlike MOHAWK, where the weights of MLP within each layer remain frozen during distillation, we allow training of all parameters in the Mamba2 and MLA layers. The refined initialization resulting from this ILD has proven to be crucial for enhancing the subsequent end-to-end distillation process, as evidenced in Section ??.
3.1.2 SMART: Sensitivity Measure-Aware Replacement of Transformer Layers
The final stage of our initialization process is to com- pose our hybrid model from the full MLA and full Mamba2 models. Let N denote the desired total number of MLA layers. The set of layers assigned as MLA is represented by LMLA S. 1......... L-11), To minimize the performance gap between the hybrid and fully attention-based models, we introduce a Sen- sitivity Measure-Aware Replacement of Transformer Layers (SMART) strategy, which leverages empirical sensitivity analysis guide the layer assignment. To measure the sensitivity of each layer, we begin with the full Mamba2 model after ILD. First, we compute the KL divergence between the full attention teacher model and the student model where all layers are Mamba2. Then, for each layer i, we construct a variant of the student model in which only the 7th layer is replaced with MLA while the rest remain Mamba2, and measure the KL divergence against the teacher. The sensitivity gain 8; for layer is defined as the reduction in KL divergence relative to the full Mamba2 model. The sensitivity score is formally computed as:
Layer Sensitivity
1200
Divergence Gain
800
600
0
2
2
3
10 11 Layer Index
12
13 14 15
Figure 3: Layer sensitivity scores for I Llama3.2 2-2- IB using 4096 samples from the validation dataset. Red markers indicate the MLA layer indices se- lected by our SMART strategy with N = 4.
8i = Σ KL (p(- 21 07) 1p Via, 0)] KL pp- yiur, x. 0T) || yia,, 2., 0D], (12)
where Or and 0 denote the parameters of the teacher and full Mamba2 model, and 0' corresponds to the variant with MLA inserted at layer to time step and is the input sequence and T is total number of decoding time steps. A higher score indicates that the 7th layer plays a more critical role in aligning with the teacher and should thus be prioritized for MLA replacement due to its higher representational capacity. Figure 3 Isooss shows an example sensitivity profile for the Zebra-Llama 1B model after refined initialization, where we observe that the earliest and latest layers tend to exhibit higher sensitivity, while middle layers typically less critical. While it may seem intuitive to simply select the top N layers with the highest sensitivity scores, this strategy often leads to suboptimal layer placements especiallyyy if it results in large gaps between MLA layers. To enforce a more structured distribution and preserve the most sensitive positions, adopt the following three-step procedure:
this equation, f/x1 refers to model predictions up
Terminal Preservation: We divide the total L layers into N roughly uniform partitions. We take the first and last [/N] layers in the first and last partitions. From these two partitions, we select the layer with the highest sensitivity score as the first and last MLA layers, denoted MLA and T. respectively. This ensures that the most sensitive edge layers are preserved.
Near- Uniforn Intermediate Distribution: Given the range [T.MLA + 1, I 1), we aim to place the remaining N- - 2 MLA layers such that the gaps between consecutive MLA layers are as
5

---

uniform as possible. We constrain the gap between adjacent MLA layers to lie within the range and max: LN - 1++11 We enumerate all valid configurations N-1
of min: [ [x - -
N-
{Cz}}=] that meet this spacing constraint, where each candidate Cj defines possible set of
intermediate MLA layer indices.
Maximal Sensitivity Scores: For each valid configuration, we compute the total sensitivity score and choose the one with the highest cumulative score: C arg max Eiec, S;-
C Additional details and examples of this layer selection procedure are provided in Appendix B
3.2 Efficient Training
After initialization and model composition, we follow an end-to-end knowledge distillation and DPO training stages to incrementally improve model accuracy and efficiency.
End-to-End Knowledge Distillation fine-tuning (SFT):
This stage involves an end-to-end distillation with supervised
Lo
(r))lp(-lv1a.....) 0)],
(13)
1=1
where 9 and Or are the parameters of the student model and the teacher model respectively. Such distillation step is crucial for transferring the rich, pre-trained knowledge from the teacher model.
Direct Preference Optimization (DPO) In the final training stage, we perform DPO which is binary cross entropy loss to adjust the preference of the student model. We set the distilled student model itself as the reference model as it makes the training much stabler.
4 Experiments and Results
4.1 Training Setup
All of our Zebra-Llama models are distilled from the Llama family: Llama3.2. 2-IB-Instruct, Llama3.2 2- 3B- IB-Instruct, and Llama3. 1 --8tuInttt for our experiments. For ILD and SFT, we use the same dataset as in (5) which includes multiple public datasets such as OpenHerrmes-2.511: [5)]. GenQA and Infinity-Instrue (17). with a total number of 6.8 billion tokens. The dataset is splited into 20% and 80% for ILD and SFT separately. We repeat the same training data more than one epoch to match the desired token budget. For DPO preference tuning, we adopt three datasets Llama3-ulite -alaa---rraeeeddaakk [18). orca_dpe lil_ddpppaatt _aino__arss__ 19). and altrafeedbuck_binrrree
All models were trained on single node equipped with eight AMD MI300 GPUs. Our training details are provided in Appendix [..4
4.2 Performance Evaluation
Evaluation Tasks We adopt the LM Harness Eval benchmark 21 to perform zero-shot and few-
shot evaluations on language understanding tasks, which includes ARC- ARC-Challenge (ARC) (22), ARC-
Easy (ARE) HellaSwaa (HS) [23]. MMLU (MM) (24), OpenBookQA (OB) [25), PIQA (26).
RACE (RA) [27]. and WinoGrande (WG) 28 We provide more details for model evaluations in
Appendix A.5
Zero-shot Results
The results of our zero-shot evaluations are summarized in Table I We compare our Zebra-Llama with the base Llama models and other baselines based on distillation: ManbaInLLaMA (Hybrid (a X-EcoMLA (Pure [2)) Llamba (Pure Mamba2) [I], and Minitron (Pruning) [13]. Besides the evaluation results, we list the teacher model size, number of training tokens, student model size (the number of parameters). and the KV cache size compared to the base Llama models. For the MLA layers in Zebra-Llama, we set Tke = 128 for the IB and 3B models and Tkve = 160 for the 8B models. For each model size, we tested various combinations of MLA and Mamba2 layers.
As shown in Table I compared to the base Llama models, our Zebra-Llann achieves extreme KV cache compression without noticeable performance drop. For IB and 3B models, we achieves 3.91% (25.6x compression) and 2.01% (49.78x compression) KV cache size with even higher performance than the base Llama models. For the 8B models, we reach 5. 5.47% (18.3x compression) and 2.73% (36.6x compression) KV cache size with only 1% and 2. 15% performance drop. Note that for the 8B models, we achieve such compression ratio using only same size teachers. Further more, our Zebra-Llan Llama achieves better balance between KV cache compression and performance with much

---

Model and Setting
Teacher
Tokens
Size
KV Size
ARC
ARE
HS
MM
OB
R.A
WG
Arg.
Llama3 2- 1--1--tstt
9T
1,241
100%
37.97
63.30
60.65
46.05
34,80
74.32
38,18
59.67
51.87
MumbalaL.umm-11--000
8B
7B
1.278
50%
40.78
65.57
61.4
40.17
39.2
74.32
36.
58.
$2.31
X-BooML A--18
64)
88
71
1.238
9.57%
40.02
67.17
58.4
38.53
37.8
73.83
39.43
60.93
52.01
Llamba1BB
184708
XB
1.418
0%
37.12
65.
61.31
38.11
36.8
73.78
37.61
60.62
51.34
Zobm--Lhmm-1B 18, SMLA-SM2
8B
7B
1.27B
7.81%
12:49
67.38
60.54
38.94
41.6
72.91
38.37
61.25
52.91
1obmm-Lama-111 1B, 6MLA-10M2
6B
7B
1.28B
5.86%
43.94
67.51
60.46
38.21
41.2
73.23
37.61
61.17
52.92
Zobra-Lhma-1B 18, 4MLA-12M2
8B
7B
1,288
3.91%
42.32
66.96
58.
37.91
40.6
72.96
37.7
58, 88.88
$2.03
Lluma3.2 2-30iimst
9T
3,21B
100%
16.08
67.93
70.38
60.34
36.4
75.79
40.86
67.25
58.13
MambalaLLumu-38. 38-308
708
208
3,358
5UF
54.1
80.26
74.45
$2.44
42.4
77.69
43.44
67.32
61.52
X-BooMLAA.BB (re- 816)
3B
7B
3.218
42.97%
48.
70.37
72.41
57.51
38.2
76.28
46.41
68.11
59.71
X-BooMLAA-3B (Fee 1288*
88
7B
3,2118
9.37%
52.05
75. 35
70.95
53.2
40.8
77.09
44.
66.85
60.13
Llamba-3B
3B.70B
108
3,668
0%
45.65
73.78
73.31
52.32
42.4
78.02
40.1
70.01
59.45
Zebra- Llana-Llmaa--- 3B. 14MLA-14MM
8B
4.69%
51.28
76.14
72.57
52.1
77.53
45.93
67.56
60.69
1chmm-Lhmm-33. SMLA-30M2
8B
9B
3,368
2.68%
51.96
75.97
72.38
48.
42.8
77.64
43.54
65.
54.77
L.bm--Lhmmm33 3B, 6MLA-22M2
8B
9B
3,39B
2.01%
50.77
76.09
71.16
50.06
43.4
77.26
42.19
66.16
59.75
Llama3.1 InStInsst
1ST
8.03B
100%
54.86
79.55
79.23
68.13
.4
80.9
44.69
73.88
68,53
708
208
8.3B
50%
59.73
84.81
79.69
"59.74"
'44
80.03
'2.
"14.11
66.03
Minitron-SR
1581
94B
8.38
66.67%
52.73
79.5
77.4
62.95
45.2
81.39
39.71
72.69
63.95
X-BooMLA-SBB (Pey 128)*
8B
7B
8.03B
9.37%
56.57
79.04
77.38
58.6
42.8
79.6
48.33
70.96
64.16
Llamba-8B
8R-70R
12B
8.328
0%
53.84
79:71
76.25
60.29
79.16
40.
72.77
63.3
Zehm--LJema-8. 16MLA-16M2
KB
11B
8.198
5.47%
58.68
78.
79.27
43.4
80,
49,
72.61
64,97
Zebra- Llama- 8B. SMLA-24M2
HB
11B
8,38B
2.73%
58.87
79.17
78.4
51.6
13.6
79.43
16.22
72.45
61.12
Table 1: Zero-shot evaluation on the LM Harness Eval benchmark across eight tasks:
(ARC).
ARC-Easy (ARF), HellaSwag (HS), MMLJ, OpenBookQA (OB), Pl, RACE (RA), and WinoGrande (WG). All
the teacher models are I 1lamma.1-88/70000 or 1.ama3.2-11/33 BB/3BB exeept for Minitron (which is Nemotron-1558 155B) "The X-EcoMLA results are reproduced by ourselves.
more efficient training (i.e., fewer training tokens smaller teacher size) than other distillation based methods. For example, compared with hybrid MambalnL.ama, Zebra-Llama has similar performance with 12.79x and 24.88 x smaller KV cache for 1B and 3B models. For 8B models, Zebra-Llama achieves 9.14x KV cache compression with only 1.6% performance degradation by using a much smaller teacher (8B) and fewer training tokens (1 IB). Similarly, our method has roughly the same performance as the pure-MLA X-EcoMLA models, but with up to 3.4x less KV cache size. Moreover, compared to the pure Mamba model Llamba-8E Zebra-Llama shows significantly better performance with both smaller teacher and fewer training tokens with minimal KV cache overhead. Moreover, we further compare Zebra-Llama with offte-of-the-ag art hybrid models trained from scratch in Appendix C (see Table TTI. While prior methods like SAMBA and Mamba-2-Hybrid rely on 1.5T-3.5T training tokens, our Zebra-Llama models achieve competitive or superior performance using only 7-B represetnni 214x-500x reduction in training data.
Few-shot Results In Table 2 we re-
port the results of few-shot evaluations on the Zebra-Llama-81 -B models under 25-
Model and Setting
KV
ARC
H5
MM
WG
TQ
Avg
Libmu3. 1-B Inst Minitrn
80.12
53:99
67.36
66.7%
19.30
81.61
64.39
'7.77
43.37
62.44
shot ARC-Challenge (ARC), 10-shot Hel- laSwag (HS), 5-shot Winogrande (WG),
Maanbalal.Iuaa- 100
50%
60:41
77.97
56.67
71.35
66.6
66.6
Mambalat 15m
256L
59.22
76.88
53.04
64.88
64.64
63.91
Mambalni 12.5%
12.5%
53.33
72.16
50.85
63.61
61.12
60.21
5-shot MMLU (MM), and 0-shot mc2 for Truthful-QA (TQ) tasks. Among all mod- els. our Zebra-Llama achieves the best per- formance with only 5.47% KV cache usage. The closest one to us is MambaInI. 50% but it's trained with 1.8x more tokens and has 9. 144 more KV cache usage. For pure Mamba2 models such as MambaInI.lama-88 and Llamba-8B, they don't 't use any KV cache but their performance is significantly worse than Zebra-Llama even with more training tokens and larger teachers. Complete few-shot results for Zebra-Llama-11 1B and Zebra-
Mumbalaf 10% 0%
OM
53.51
70.31
44.21
58.91
52,31
55.85
Liaabo
01S
77.13
63.44
X-FossMI A-8R (128y*
9.37%
76.9
Zahes-Llams, (16-16)
5.47%
60.49
78.29
55.84
71.08
66.7%
Zebra- (8:24)
2.73%
60.41
76.11
55.06
71.11
61.01
64.74
Table 2: Few-shot evaluation on the LM Harness Eval bench- mark across five tasks.
can be found in Appendix E
Results on GSM8k In Table we evaluate the math math-reasoning ability of Zebra-Llama via 8- shot GSM8K task, comparing its performance against baseline models. For the Zebra-Llama- IB architecture, the 8MLA-8Mamba2 -MMmba2 configuration achieves the highest accuracy, even outperforming the base Llama-3.2- 1---IButnntt model by 6.7%. Increasing the number of MLA layers to 12 results in a performance drop from 41.09% to 29.57%; however, this result remains superior to both MambalnI.lmma IB and Llamba- IB by 4. 4.63% and 6. 75%, respectively. For Zebra-Llama-33, is slightly worse than X- X-EcoMLA-33 given more than 2x smaller KV cache. However, our Zebra-Llama
7

---

6400
55004 (lr.kalae...... 1-581
oge
Throughput (Takens//s)
ScHiL buus
3200
Peak Memory (GB)
GMLL-2AN2 12ebra-Uiama-DD
L0MLA 1852 (hetra Utere EB)
1600
100
800
3RG0K (Jame)) 1-30
400
160gs 1887
EMLA-2490 (iber--iarn--a3
S0MIA- Cistro-lame-di
16k
32k
131k
4k
8k
16k
324
131k
Context Length
Context Length
Figure Inference throughput vs. context length various 8B-size models. We measure the throughput under batch size 48 and output length 1024.
Figure 5: Inference peak memory vs. context length of various 8B-size models. The out-of-memory sce- narios are marked with 'OOM'.
still 9.71 % and 15.01 better than Mambalnl.lama-3 and Llamba-3B with 14 MILA layers and 14 Mamba2 layers. Even we decrease the number of MLA layers to 8, which leads to only 2% of the original KV cache, our Zebra-Llama is still 2.1% and 7.43% better than MambaInLlama- and Llamba-3B. This consistent superiority demonstrates that our hybrid model structure and tailored training strategy effectively preserve, and many cases enhance, the math-reasoning capabilities of the underlying base model while maintaining high KV cache efficiency.
Throughput Evaluation Figures 4 and 5 detail the inference efficiency of Zebra-Llama in terms of through- put and memory consumption for the models that distilled from Ilaam-3.11- 188.. All experiments are con- ducted on single AMD MI300X GPU with 192GB memory. We fix the batch size to 48 and the output length to 1024 and measure the inference throughput across various context/input lengths. As shown, our Zebra-Llama significantly outperforms the base Llama- 3.1-8B model by 3.9x for throughput with only 17% peak memory at 16k input tokens. Similarly, Zebra- Llama outperforms the MambaInLlama model by 3.28x in throughput with only 18.2% of peak memory at 32k input tokens. Beyond 16k and 32k tokens, the Llama-3.2-88 3.B-BB model and model will both run out of memory. constrast, our Zebra-Llama model only occupied 62.68GB and 104.6GB
Model
Tokens
GSMSK
Lluma-3.2 2-1 I-Istnett
100%
9T
38.51
Mambalaa.lamaa
7B
A-HMMLA-188
9.37%
7B
3M
Llamba- 1B
0%
8B
22.82
Zehra-Llama.
Zebra-Laamas Llara- IB (4-12)
3.91%
7B
29.57
Dhina-..2 3.2-38 I-Strrnnttrcci
100%
91
70.89
Manbalal.aama
20B
X-EcoMLA-3B
9.37%
7B
65.28
Llamba-3B
0%
108
47.76
Zebra-Llama 3B (14-14)
4.69%
9B
62.77
Zehra-LLlma-388 Lama-3B (6-22)
Table 3: 8-shot GSK8K results for Zebra-
Llama and baselines.
even with 131k input tokens, showcasing
its brilliant ability for memory efficiency.
4.3 Ablation Studies
In this section, we present a series of ablation studies aimed at justifying key design decisions in approach. Specifically,
examine the impact of initialization strategies, the effectiveness of our SMART layer selection mechanism, the trade-offs between the number and size of MLA and Mamba2 layers, and the role of teacher model scaling.
4.3.1 Effect Initialization Strategies
Selection Strategy
MLA Indices
Total s2n.
Avg. Score
Figure 6 presents our assessment of various initialization methodologies through a comparison three scenar- ios: Random weight initialization without ILD, Struc- tured weight initialization without ILD, and the proposed Structured weight initialization with ILD.
Unform ti
[0.4,812]]
1787.9
48.84
Uniform 42
[3.7.11.55]
1055.8
49.72
Max suore
[0.1.2.14]
2672.5
48.95
Pussible midile ø
[0.3.9.14]
21.55.7
"68.76
Possible mudcle #2
[0.5.9.14]
2297.5
SMART (ours)
[0.5.10.44]
2500.2
Table 4: SFT results with different layer selec- tion strategies on Llama-3.2-11 IB model. Except the selected MLA indices, all models are trained with same configuration.
The results highlight that both structured weight initial- ization and ILD significantly boost SFT performance, especially when used together. For Mamba layers, ILD is crucial for aligning their outputs with the original models due to architectural differences from at- tention layers, providing the primary performance uplift. As for MLA layers, structured initialization can boost the accuracy significantly, which offers strong foundation further refined by ILD.
4.3.2 Impact of SMART Layer Selection
In Table/4 we demonstrate the benefits of our SMART layer selection strategy by comparing with other layer selections for our Labbr--Llama11B -B with 4 MLA layers, by using the configs shown in

---

Random Int. LD
Stuctared Inc, ILD
268
128
Strctured Int. WDDLD
52
50
47.87
Score
228
Avg. Score
45
40
41 00
50
134
46
30
12. 8x KV cache Compression
25. cache Compression
25
46
12
16
20
16 Mamba2
16 MLA
8 MLA 8 MambaZ
Number MLA layers
Figure 6: Performance of various initialization strategies after SFT for different model architectures. Generally, our proposed two-stage method achieves the highest average scores.
Figure 7: Tradeoff between KV rank Tke and the number of MLA layers when fixing the total KVV cache size of given model.
Figure 3 as an example. First, by comparing the first two layer selections sets and our optimal solutions, can conclude that Terminal Preservation strategy (i.e., always selecting layers from the first and last few layers) are important for preserving model accuracy. Second, we also testify the naive greed selection greedy method (the third row); the results shows that uniformly distributing the MLA layers has a major contribution to final performance as well, which verified the effectiveness of Wear-Unifin Intermediate Distribution strategy. Third, by analyzing the last three selections, together with Terminal Preservation and Near-Uniform Intermediate Distribution, Maximal Sensitiv- ity Scores is indeed good indicator for final accuracy score. In short, these outcomes underscore how each of the three pillars-when appropriately fombined-contriben selection strategy's effectiveness. While we believe that SMART offers a systematic and principled alternative to heuristic layer selection, we look forward to further exploration and advancement of this direction future work.
critically to the SMART layer
4.3.3 Tradeoff between number of MLA layers and Fkv
The KV-cache size Zebra-Illama is determined by two factors: the number MLA layers and the KV rank Tke of each MIA layer. Figure presents preseenrreeents our findings from varying these two factors for our 1B model while maintaining a constant total KV-cache size. We observed that for significant KV-cache compressions, such 12.8x x and 25.6x, optimal performance typically occurs with moderate number of MLA layers, around six, coupled with an intermediate Tke- Deviating from this balance will hurt the model's performance. On one hand, increasing the number of MLA layers excessively, like to 16 layers at a 25.6x compression, forces the Tku per layer to become very small (e.g., Tke Tke =8), which significantly degrades performance. On the other hand, reducing the number of MLA layers too much also leads to a decline in performance, since the model will then consist almost entirely of Mamba2 layers, which typically have higher capacity gap with MHA layers than MLA layers do.
4.3.4 Scaling the Teacher
Effective knowledge transfer from a well-chosen, pre-trained teacher model is crucial for our method's success. As shown in Table [ appropriately scaled teachers significantly enhance student model performance though with diminishing return when gradually increasing size of teacher models. This due the "capacity gap" phenomenon in knowledge distillation, where a smaller student model struggles to fully absorb and generalize the complex teacher's 's representations when the teacher model far exceeds student mimicry capabilities [29 30). Therefore, it's vital to select teacher model that sufficiently powerful to provide rich knowledge but not so extremely larger than the student model, thereby balancing distillation efficacy with student-teacher compatibility. Investigating adaptive teacher scaling or multi-stage distillation offers future solutions to these capacity limitations.
4.4 Extension Qwen models
We further demonstrate that the proposed initialization and distillation strategy is not limited to Llama models but also applicable other popular model families such as Qwen (3T). which we refer as Zebra-Qwen models. We choose Qwen-2.5-4 Instructt and Qwen-2.5 2.5- I-S---Insttttt as the base models. For the 0.5B model, we use the 1.5 mode as the teacher and replace layers to MLA and 20 layers to Mamba2, resulting in 6.25% KV cache size. For the 1.5B model, we use itself as the teacher model and replace 14 layers to MIA and 14 layers to Mamba2, with KV size of 12.5% of the original target model. The performance is illustrated in Table 6 The conclusion remains the same

---

Model and Setting
Teacher Size
ARC
ARE
HS
MM
OB
P1
RA
WG
Avg.
1. 1---nss
37.97
63.30
60.65
46.05
34,80
74.32
18. 18
59.67
51.87
Zebm-Llama- IB, 4MLA-12M2
1B
38.91
61.7
55.03
35.83
37.2
71.93
35.41
58.88
49.1
Zobra-Llama-IB 1MLA-12M2
39.51
62.79
57.61
37.91
38.20
72.52
36.91
7ehn-LLhum--1 Lama-18, HML.H,44MAAA11 4ML.A-12M2
88
42.32
66.96
58.95 93
37.91
40.6
72.96
37.7
58.88
52,03
Zebm-Llama-IB 4MLA--122
70B
43.17
69.57
57.77
39.45
38.80
72.80
38.09
59.83
52.44
Llama3.2-38.
46.08
67.93
70.38
60.34
75.79
40.96
67.25
58.13
28hm-1Jlama-33 8MLA-30M2
3B
45.48
69.28
69702
47.69
40.80
74.81
42.01
63.38
56.36
Sbraa-llmma-BB 8MLA-20M2
SB
51.54
75.55
71.52
47.12
43.6
77.2
42.68
65.9
59.39
Zebm-LLlum-33 SMLA-20M2
708
51.96
77.23
69 46
48.32
43.4
76.01
43.35
65.19
59.37
79.55
79.23
68.
44
65.53
Zeban-Llnma-e Liana-BB, SMLA-24M2
8B
56.48
78.79
76.8 84
53.72
44.4
79.43
44.31
70.64
63.08
7ehm-l1 Jama-8R, 8MI 2-24M2
708
58.53
80.72
76.64
53.82
45.4
80.03
43.06
69.61
63.48
Table Impact scaling up teacher size on model performance trained on 7B tokens. Except for teacher size, we use the same training configurations for the same size of student models.
Model and Setting
Teacher
Tokens
KV Size
ARC
ARE
HS
MM
OB
P1
RA
WG
Avg.
Owen2.5-0.5B-0e 5-0.5B-nstt
100%
33.11
59.05
52.26
45,86
34.2
70.62
32.06
56.27
47.93
Zobra Qwen-O.5B, (MLA-20M2
1.5B
7B
6.25%
38.74
66.92
50.83
38.43
37.2
69.91
32.31
55.09
18.68
Owenn.55-1.5B. -n-sttt
100%
47.01
75.84
68.24
60.13
41
76.01
37.8
62.67
58,59
Zobra Qwen-l SB, 1-MLA-1-MM
1.58
7B
6.25%
48.63
75.17
67.64
53.87
41.6
75.73
38.66
64.01
58.16
Table 6: Zero-shot evaluation on the LM Harness Eval benchmark for Zebra Zebra-Qwen models.
as Llama models The 0.5B Zebra-Qwen model achieves 3.65% higher accuracy than the base model and 1.5B Lebra-Qwen model only experiences 0.7% performance degradation with 8x KV cache compression.
4.5 Extended Long-Context Evaluation
Model
KV%
4K
8K
16K
For the long-context evaluation of our models, few factors should be considered. First, as- sessing long-contexi performance requires models that have been trained on datasets with extended sequence lengths. However, due to efficiency and resouree constraints, our current models were trained with a maximum sequence length of 2048 tokens. Such a training configuration generally not optimal for long-context benchmarks like RULER (3). Despite this limitation, we conducted RULER benchmark evaluations at 4K, 8K, and 16K context lengths and the results are summarized in Table 7 We observe that, at the 3B scale, Zebra-Llama surpasses Mamba across all context lengths, while at the IB scale performs on par at 4K and exceeds MambaInL at 8K and 16K. Importantly, these gains are achieved with a substantially smaller memory footprint- CooprrnttZZeba--.e cache, whereas MambaInL.ana requires nearly ten times more KV memory.
ManbalnLlama- IB
50%
38.75
21.55
3.88
Zebra-Llangr IB
3,91%
35.75
24.80
13.37
Marihaaf Ialamiaa
50%
ar.1
22.62
0.88
Zebra-Llama-31
4.69%
58.69
38.24
9.97
Table 7: RULER benchmark results at 4K, 8K, and 16K context lengths.
uses only ~5% of the KV
As future work, we aim to extend our models to support longer context lengths and reasoning.
5 Conclusion
In this work, we addressed the growing need for efficient LLMs by proposing a practical and scalable framework for composing hybrid models from existing pre- rre-trained Transformers. Motivated by the cost and environmental impact of retraining large models for downstream use, we introduced Zebra- Llama, a family of 1B, 3B, and 8B hybrid models built using SSMs and MLA layers. We developed an effective initialization scheme and a post-training knowledge transfer pipeline that enabled these models to inherit capabilities from larger teacher models with minimal additional training. Our approach significantly reduced memory while preserving the accuracy of strong baselines.
Limitations and Future Work Our work establishes several promising directions for future research. A primary next step is to move beyond single family of base models and explore hybridiza- tion strategies across diverse LLM architectures, particularly incorporating modular frameworks like Mixture-of of Experts (MoE). Scaling our training pipeline to larger models and extending the approach to reasoning-intent
architectures remain critical frontiers. Furthermore, addressing the reliance on strong teacher models is essential. In constrained scenarios, investigating alternative knowledge transfer strategies-Sus as teacher-free free or efficient self-distillation-- will be key to understanding how hybrid architectures can effectively learn with reduced supervision.
10

---

References
[1] Albert Gu and Tri Dao. "Mamba: Linear-time sequence modeling with selective state spaces". In: arXiv preprint (tti::23/2.007552 (2023).
[2] Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. "Transformers mns:
Fast autoregressive transformers with linear attention". In: International conference on machine learning. PMLR. 2020, pp. 5156-5165.
[3] Liliang Ren, Yang Liu, Yadong Lu, Yelong Shen, Chen Liang, and Weizhu Chen. "Samba: Sim- ple hybrid state space models for efficient unlimited context language modeling". In: arXiv preprint
arXiv:2406.0752. (2024).
[4] Xin Dong, Yonggan Fu, Shizhe Diao, Wonmin Byeon, Zijia Chen, Ameya Sunil Mahabaleshwarkar, Shih Yang Liu, Matthijs Van Keirsbilek, Min-Hung Chen, Yoshi Suhara, et al. "Hymba: hybrid-head architecture for small language models". In: arXiv preprini arXiv:2411. 13676 (2024).
[5] Junxiong Wang, Daniele Paliotta, Avner May, Alexander M. Rush, and Tri Dao. The Mamba in the Llama: Distilling and Accelerating Hybrid Models. 2025, arXiv: 2408 15237 [cs. LG]
URL: https
7/arxiv. org/ abs 2408 15237
[6] Aixin Liu, Bei Feng, Bin Wang. Bingxuan Wang. Bo Liu, Chenggang Zhao, Chenggi Dengr, Chong Ruan, Damai Dai, Daya Guo, et al. "Deepseek-v2: A strong, economical, and efficient mixture-of-expert language model".. In: arXiv preprint arXiv:2405.044 (2024).
[7] Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwarz, Omri Abend, Raz Alon, Tomer Asida, Amir Bergman, Roman Glozman, Michael Gokhman, Avashalom Manevich, Nir Ratner, Noam Rozen, Erez Shwartz, Mor Zusman, and Yoav Shoham. Jamba: Hybrid Transformer-Mamba Language Model, 2024, arXiv:
2403. 19887 [CS.CL] CL]
[8] Xin Dong, Yonggan Fu, Shizhe Diao, Wonmin Byeon, 7ijia Chen, Ameya Sunil Mahahaleshwarkkrr Shih Yang Liu, Matthijs Van Keirsbilek, Min
Min- Hung Chen, Yoshi Suhara. Yingyan Lin, Jan Kautz, and Pavio Molchanov. Hymba: A Hybrid-head Architecture for Small Language Models. 2024, arXiv:
2411. 13676 [cs.CL) CL]] URL: http::ttp////t /ttp/aarvvvvv///m//// org/ abs 2411 13676
[9] Roger Walefse, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadch, Sudhakar Singh, Deepak Narayaman, et al. "An empirical study of namba-based language models". In: arXiv preprint arXiv:2406. 240.0877 (2024).
[10] Aviv Bick, Kevin Y Li, Eric P Xing. Zico Kolter, and Albert Gu. "Transformers to ssms: Distilling quadratic knowledge to subquadratic models". In: arXiv preprint arXiv:2408. /0189 (2024).
[11] Aviv Bick, Tobias Katsch, Nimit Sohoni, Arjun Desai, and Albert Gu. "Llamba: Scaling distilled recurrent models for efficient language processing". In: arXiv preprint arXiv:2502. 14458 (2025).
[12] Guihong Li, Mehdi Rezagholizadeh, Mingyu Yang. Vikram Appia, and Emad Barsoum. "X-EcoMLA: Upeycling Pre-Trained Attention into MLA for Efficient and Extreme KV Compression". In: arXiv preprint arXiv:2503. //32 (2025).
[13] Saurav Muralidharan, Sharath Turuvekere Sreenivas, Raviraj Joshi, Marcin Chochowski, Mostofa Patwary, Mohammad Shocybi, Bryan Catanzaro, Jan Kautz, and Pavlo Molchanov. "Compact language models via pruning and knowledge distillation". In: Advances in Neural Information Processing Systems 37 (2024),
pp. 41076 4012..
[14] Tri Dao and Albert Gu. "Transformers are SSMs: Generalized models and efficient algorithms through structured state space duality". In: arXiv preprint arXiv:2405.2 (2024).
[15] Teknium. OpenHermes 2.5: An Open Dataset of Synthetic Data for Generalist LLM Assistants. 2023.
URL: Intps:: httpss://hhgggggggg
co/ daz asets www.sssttteettee wttps//eeiiininn um Omen.ermes-2.2.
[16] Jiuhai Chen, Rifaa Qadri, Yuxin Wen, Neel Jain, John Kirchenbauer, Tianyi Zhou, and Tom Gold- stein. "Genqa: Generating millions of instructions from handful of prompts". In: arXiv preprint 1.Xiv:2406.10333 (2024).
[17] Beijing Academy of Artificial Intelligence (BAAI). Infinity Instruct,
Cattasets/BAAl// Instrruut 2024.
[18] Junxiong Wang. Lomaa wwwggeeeeannnnaaannaaanaa dataset.
Junxxloorkaar//ieennnn--eeeeeees armorm 2024.

---

[19]
Wing Lian. Bleys Goodson, Eugene Pentland, Austn Cook, Chanvichet Vong, and "Teknium", OpenOrea: An Open Dataset of GPT Augmented FLAN Reasoning Traces.
to/Open Grea/ OpenOrca 2023.
[20] Ganqu Cui, Lifan Yuan, Ning Ding, Guamming Yao. Wei Zhu, Yuan Ni, Guotong Xie, Zhiyuan Liu, and Maosong Sun. UltraFeedback: Boosting Language Models with High-quality Feedback. 2023, arXiv:
2310. 01377 [cs.CL]
[21] Leo Gao, Jomathan Tow, Stella Biderman, Sid Black, Anthony DiPofi, Charles Foster, Laurence Golding, Jeffrey Hsu, Kyle McDonell, Niklas Muennighoff, Jason Phang, Laria Reynolds, Erie Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. "A framework for few-shot language model evaluation". In: (2023).
[22] Peter Clark, Isaac Cowhey, Oren Etzioni, Tushar Khot, Ashish Sabharwal, Carissa Schoenick, and Oyvind Tafjord. "Think you have solved question answering? try arc, the ai2 reasoning challenge". In: arXiv preprint arXiv:1803.05533 (2018).
[23] Rowan Zellers, Ari Holtzman, Yonatan Bisk, Ali Farhadi, and Yejin Choi. "Hellaswag: Can machine really finish your sentence?" In: arXiv preprint arXiv: 1905007830000000 (2019).
[24] Dan Hendrycks, Collin Burns, Steven Basart, Andy Zou, Mantas Mazeika, Dawn Song, Jacob Steinhardr. "Measuring massive multitask language understanding". In: arXiv preprint arXiv:2009.33000 (2020).
[25] Todor Mihaylov, Peter Clark, Tushar Khot, Ashish Sabharwal. "Can a suit of armor conduct electric- ity? new dataset for open book question answering". In: arXiv preprint arXiv: 18XXvv:800.000 02789 (2018).
[26] Yonatan Bisk, Rowan Zellers, Jianfeng Gao, Yejin Choi, et al. "Piqa: Reasoning about physical common- sense natural language". In: Proceedings of the AAAI conference on artificial intelligence. Vol. 34. 05. 2020, pp. 7432-7439.
[27] Guokun Lai, Qizhe Xie, Hanxiao Liu, Yiming Yang, and Eduard Hovy. "Race: Large-scale reading comprehension dataset from examinations". In: arXiv preprint 1XXiv:100.04400 (2017).
[28] Keisuke Sakaguchi, Ronan Le Bras, Chandra Bhagavatula, and Yejin Choi "Winogrande: An adversarial winograd schema challenge at scale". In: Communications of the ACM 64.9 (2021), pp. 99-106.
[29] Seyed Iman Mirzadeh, Mehrdad Farajtabar, Ang Li, Nir Levine, Akihiro Matsukawa, and Hassan Ghasemzadeh. "Improved knowledge distillation via teacher assistant". In: Proceedings of the AAAI conference on artificial intelligence. Vol. 34. 04. 2020, pp. 5191-5198.
[30] Aref Jafari, Mchdi Rezagholizadch, Pranav Sharma, and Ali Ghodsi. "Annealing knowledge distillation". In: arXiv preprint arXiv:2104.07166 (2021).
[31] Binyuan Hui, Jian Yang, Zeyu Cui, Jiaxi Yang, Dayiheng Liu, Lei Zhang, Tianyu Liu, Jiajun Zhang, Bowen Yu, Keming Lu, et al. "Qwen2. 5-coder technical report". In: arXiv preprint arXiv:2409. /2 186 (2024).
[32]
Cheng-Ping Hsieh, Simeng Sun, Samuel Kriman, Shantanu Acharya, Dima Rekesh, Fei Jia, Yang Zhang. and Boris Ginsburg. "RULER: What's the Real Context Size of Your Long-Conterr Language Models?" In: arXiv preprint arXiv:2404.06654 (2024).
[33] Stephanie Lin, Jacob Hilton, and Owain Evans. "Truthfulqa: Measuring how models mimic human falschoods". . In: arXiv preprint arXiv:2 (2021).
[34] David Rein. Betty Li Hou, Asa Cooper Stickland, Jackson Petty, Richard Yuanzhe Pang. Julien Dirani, Julian Michael, and Samuel Bowman. "Gpqa: A graduate-level google-prood q&a benchmark". In: First Conference on Language Modeling. 2024.
[35] Guilherme Penedo, Hynek Kydliček, Anton Lozhkov, Margaret Mitchell, Colin A Raffel, Leandro Von Werra, Thomas Wolf, et al. "The fineweb datasets: Decanting the web the finest text data scale". In: Advances in Neural Information Processing Systems 37 (2024). pp. 3081 30849.
12

---

A More Experimental Details for Zebra-Llama
A.1 Structured MLA Initialization
Algorithm Python-like pseudocode of the proposed SVD initialization for MLA.
MHA weights: R_Q, R_K, WVV
MIA weights: M_DQ, W_U WLOR, WLDRV, WLUR, WLRR, RLUV
Initialization of W DDDQ, U_ and 0_OR U U_q, sigma_q, q, V_q svd (W
W_DQ = U_q
__UQR_bas _arr = (sigma_ @ V ._g). view n_h, d_h)
N_QQ w W_UOR_BO barl:, ;, :d_qk] www.wwiweeee.. n gk)d_gk)
W_RR = LuroR_bbr[:: ;, -d view (r_9, r_h+d_r)
# Initialization of R_DKV, RUK,, WRRR, WUV U_kv, sigma_kv, V_kv svd (torch. cat (W_K, W_V), , -1))
W_DKV = U_kv
_avg view (d, nv,,, d mean (1)
KR = a_a_avl: -d::]
W __UKV = sigma_Ki kv @ V_kv
UK bar W_OKV[;, :d :d_h-n_kv (r_ n_k,, d_h)
W UKK = Lu_barll [ar i view (v_kv, n_k
UV U_UK[[; d h:]
Our SVD-based MLA layer initialization follows the methodology outlined in X-EcoMLA 12 for Multi-Head Attention (MHA). However, for Generalized Question Answering (GQA) models like the Llama 3 series, our approach diverges slightly by keeping the number of key/value heads from the base model for MLA while X-EcoMLA forces the number of key/value heads to be the same as the number of query head. With such modification, we notice trivial performance difference but with slightly fewer number of parameters.
Essentially, the matrices from the base MHA/GQA module could be expressed W dddands and W WV E Rdxmn dn, where denotes the hidden dimension, n, denotes the number of attention heads, nek denotes the number of key/value heads. For MHA, we have Tig. 71ke while for GQA we have n4 > Tlke-
All the matrices we need to initialize in MLA could be expressed as WDKV E BR WUK E Re xxxudae WUV E I.
XPEud, "Rrko,, DD E ddx., wie E R WKR E Rdxd, and WQR E gr.xmad, where Tke represents the latent dimension for key/value, dgp denotes the head dimension for query/key, d, denotes the head dimension for value, rg denotes the latent dimension for query, and d, represents the dimension for RoPE embeddings. For all the experiments, we keep d, = d,, and dyp + d, = dn.
Initialization of W DDQ WUQ, and WQR Given that query compression in MLA can be viewed as a low-rank approximation of attention layers, we initially perform SVD on the pre-trained weight matrix WQ
=º=
(14)
where Ug E Rdd Eg E RTaXTs, and Vg E Rann Xre, For constructing the up-projection matrices, we reshape the product E to form UUQ E R xd, and subsequently split this tensor at the last dimension into WūQ, containing the first dyp elements, and WQR, containing the remaining d, elements. The down- projeprooettion matrix W DQ is directly assigned with Ug. This method of initial assignment is formulated as:
DD Ug WUQ_ reshape ( UQR 1. duul), WQR reshape T UQR 1,;, d, :)), (15)
where the function reshape(.) is used to integrate the last two dimensions of the specified tensor.
13

---

Method
and MLA Indices
Fg
dyt
d,
Model Size
7ebra-Ilama-IB. 8MLA8M2
10,
128
1344
1.27B
Zebra-Llama-11 60LA
11,
128
1344
32
32
1,28B
Jehra-Llama-11, 4MLA
[0,5,10,14]
128
1344
32
32
1.28B
Zebra-Llama-3B. 14MLA14M2
220.446...0.2244111111822222222.
128
1536
64
64
3.27B
Zebra- Zebra-Llama-3B. SMLA20M2
(0.4..112.1.22022..26))
128
1536
3,36B
Zebra- Sabra-Llamaa3B.. 6MLA22M2
10.5.10,16.21126]
128
1536
64
64
3.39B
Zebra-Llamaa 16MLAT6M2
160
2048
64
64
8.19B
-laaa-Laama-8B 8MLA24M2
(0.4.8.12.16.20.25.300
160
2048
8.38B
Table 8: Configurations Zebra-Ilama models' architecture.
Initialization of WDKV, W WUV, and WKR The initialization of the MLA weights associ- ated with keys and values is more complicated because of the decoupled RoPE mechanism. First, a joint SVD performed on the concatenated W and WV
[W 1 Uko Exe V..V
(16)
where Un E RdxrEs, DED E RrkuXrke and V& E X For the down-projection matrix WDKV we directly set it to Uk.. To derive the up- oppprojection matrix W and W we first set ky. Since we have d, dg. simply extract the last dg columns of WUKV as WUK, For WUK, we first extract the first d columns of WUKV and reshape them into WUK K E R Xdd, Then, we select the first dys elements along the last dimension of W K reshape it back to obtain WUK In general, this process can be expressed as:
WURV Exv VT
WDKV = Uke. WUV = BDB=WWKKK
WUKV Rkodd, WUK = reshape W [:::: dgxl).
(17)
In the final step, the initialization of the RoPE key embedding matrix WKR requires distinet approach, given that all attention heads in MLA utilize the identical RoPE key embedding. First, the average key projection matrix W ang Raxdd, is calculated for all attention heads. Then, the final d, columns are extracted for initializing WKR formulated follows:
W Kn
W W dr :].
(18)
The detailed initialization algorithm can be found in Algorithm
A.2 Structured Mamba2 Initialization
The structured initialization process of Mamba2 layers follows precisely to the solution MambaIn- Llama (5). As outlined in Section 3.1.I excluding the softmax operation in attention shows direct one-to-one mapping between Bt, 24, and C, linear RNN and Kg, Vi, and Q₂ in attention layers. In the Mamba2 framework, B,, 24, and C₁ for the continuous-time SSM are derived from input or via passing through MI.P followed by I convolution layer. The MI.P is replaced directly with pre-trained transformer layer weights follows:
Ze = W. or E Ruxnsohs B, = W E Rixmtoha C, = W E R.xnaha
(19)
where b signifies the batch size. Subsequent to this, x, B, and C are processed through the Id convolutional layer for temporal fusion before undergoing discretization in Mamba SSM. It is important to highlight that for GQA and MQA scenarios where nyu < ng. It and B, are replicated after the convolution ensure that F4, B., and C, share identical dimensions. Other parameters, such as A and A,, adhere to the original initialization procedure within the Mamba2 layers.
A.3 Model Structure Details
Detailed architectures of our IB, 3B, and 8B models, as presented Table 1 are comprehensively outlined in Table g This includes MIAA layer selections, parameters specific to the MIA layers, and the overall model size. Layers not categorized as MILA are Mamba2 layers, which follows the configuration used in MambalnLlama [
14

---

Stage
Model
Teacher
Tokens
BS
LR
Time (H)
ILD
Zebra-Llama-IB-ML.
Llama3.2-18
1.36B
96
8e-5
1.8
ILD
ebra-Llaama-B--Mambaa
Llama3.2-1B
1.36B
96
8c-5
1.7
ILD
Zebra-Llama-3BB-ML.
Llama3.2-3B
1.36B
96
8e-5
3.9
ILD
Zeb Zobra-Llamaa-B--aambal
Llama3.2-3B
1.36B
96
8c-5
4.4
ILD
Zebra-Llama-BB-MLLA
Llama3,1-80
1.36B
48
4e-5
9.2
ILD
7abra-Llama-8B--Manha2
Llama3. 1-8B
1.36B
4e-5
10.3
SFT
Zebra-Llama-IB
Llama3. 1-BB
5,44B
192
8c-5
~ 13.7
SFT
Zebra-I.lmma-3B
Llama3. 1-8B
7.44B
96
4e-5
& 31.2
SFT
Zebra-Llama-SB
Llama3.1
9.44B
64
4e-5
78.1
DPO
7ebra-Llama--1B
N/A
0.2B
32
5e-7
0.5
DPO
Zebra-Llama-3B
N/A
0.2B
32
5e-7
1.2
DPO
Zebra-L.amma-8B
N/A
0.2B
32
5e-7
R 2.3
Table 9: Training hyper- of Zebra-Llama models.
Layer Index
Sensitivity Score
Layer Index
Sensitivity Score
1185.06
8
238.1
1
382.73
9
120.56
480.68
10
323.23
3
350.95
11
228.9
196.03
12
168.69
5
367.82
13
233.87
6
250.45
14
624.03
114.44
15
361.47
Table 10: Concrete sensitivity score Figure
A.4 Training Details
In Table we present the training configurations for our Zebra-Llama series models, including the number of tokens, batch size, learning rate, and total training time. All experiments are conducted a single node equipped with eight AMD MI300 GPUs, each featuring 192GB of memory. We apply learning rate warmup over the first 1% of training data, followed by cosine annealing. The models are optimized using Adam W, with hyperparameters set to 3 = (0.9, 0.8). Additionally, all models process input sequences of length 2048 through sample packing.
A.5 Evaluation Details
We evaluate all models using the Im-evaluation harness library (commit from the big-refactor branch) following the aask-speifiic few-shot configurations defined by the Open L.L. Leaderboard. For zero- shot evaluation, we report performance across broad suite of language understanding tasks: MMLU, HellaSwag. PIQA, ARC-Fasy,
Winogrande, OpenBookQA, and RACE. Evaluations performed using the command-ling interface with ROCm-enabled devices and batch size of 16. For few-shot runs targeting leaderboard comparisons, we use the officially recommended number of shots per task (e.g., 25-shot for ARC-Challenge, 10-shot for HellaSwag, 5-shot for Winogrande and MMLU, 0-shot mc2 for Truthful-QA(TQ) 33 the same protocol prior work.
We report average scores across tasks, following
B SMART Layer Selection Algorithms
We provide the pseudo code for SMART in Algorithm 2 Besides, we provide three examples for the layer selection process of our Zebra-Llama 1B models following SMART. The examples are based on the sensitivity analysis shown in Figure whose actually values are listed in Table I0
15

---

Algorithm 2 Pseudo code: SMART-Strrcture MLA Layer Selection via Sensitivity Scores
Require: Sensitivity scores (2,,82,... ST.). number of MLA layers N
Ensure: Selected MIAA layer indices {TAA M.LAA)
1: // Terminal Preservation
2: Divide layers into N equal partitions
3: LALA highest-sensitivii layer in first partition of layers {i, E [ 1 []}
4: MLA < nnighest-senstiviie layer in last partition layers {i, L + +1,
5: // Near-Uniform Intermediate Distribution
6: // Define Valid Intermediate Layer Range
7: Let start f- L +
8: Let Fend f- I. T.MLA
9: Let R < list of candidate intermediate layers in [Tstart Tend
// Compute Allowable Gap Bounds
11: Let +- LN
L LML
N +1
12: Let Jwnn N-1
13: Let Imax
T
N-1
14: // Enumerate All Valid Configurations
15: Initialize empty list of configurations C +-
16: for all combinations N layers from R do
17: Sort the selected layers in ascending order form Cj
18: Let G, < list of gaps between consecutive layers in ( U UC,
(L
19: if all gaps in Gj satisfy Jmin I gap S Imax then
20:
Append Cj to C
21: end
22: end for
23: // Maximal Sensitivity Scores
C* arg maxc, Eieg, 8
25: return {LMLA JUC U (LMIA)
Example 1: Zebra-Llama-IB with N = 4 Terminal Preservation: LMLA O, LALA 14
Intermediate Distribution:
Define intermediate layer range: Istart 1, Tend = 13, R {1, 13).
Compute allowable gap bounds: T 11, Imin = 3, Imaz = 4.
Enumerate valid Configurations: C [[4,9) (5, 9), (5, 103).
Maximal Sensitivity Scores:
Calculate total sensitivity score:
+ 316.59,
+ 488.38,
* + $10 691.05
-Layers with maximal score: C* = {5, 10].
Return {0, 5, 10, 14).
16

---

Model and Seiting
Tokens
Size
ARC
ARE
HS
MM
OB
P1
RA
WG
Mamba-2-Hyhrit
3.5T
8.668
.47.7
77.23/-
-77.68
51.46
-22.
79.65/--.
39.71
71.27
SAMBA
3.51
1.78
48.211-
79.25/
49.74/----
48.01
37.20V
77. 77.10--
72.93
Hymba
42.321-
74.54:
53.55----
76.66/
66.61
Zeben-Lloma, 8MLA-8ME
78
1.278
89.6842.83
72.35/67 3
45 59.66660.59
38.99
30.2441 6
72.91773.29
36.56
61.33
Zebrn-Lluma, 14MLA-1442
9B
3,278
50/52.6
80.47/76.5
53.52/7245 43
51.97
31 31.8/44444
76.61/76.99
46.99
67.8
Llama,
Table 11: Comparing our Zebra-Llama with state-of-the-art hybrid models that are trained from scratch. We report two accuracy scores for each model, i.e, accuracy and normalized accuracy (ace / acc_ roc_norm. The missing results ('-") in the table means the results not reported in original paper.
Example 2: Zebra-Llama-IB with N 6 Terminal Preservation: T.ALLA = 0, IL 14
Near-Unifomy Intermediate Distribution:
Define intermediate layer range: Ystart 1, rend 13, R f1,2...... 13). Compute allowable gap bounds: T' 9, Jmin Imax = 2.
Enumerate Valid Configurations: C = 12,5, 8, 11}, 13,5,8111} (3, 6,,,,8,7 11), {3, 6, 11), (3, 6, 12}].
Maximal Sensitivity Scores:
-Calculate total sensitivity score:
811 1315.5,
* 53 + 15
3 ss 811 = 1185.8,
* 83 86 ig
S11 $11 = 1068.4,
83 83 86
811 950.86,
83 + 86 8g 812 890.65
Layers with maximal score: C* = 12,
- Return (0,,,5,8,1,, 14]
Example 3: Zebra-Llama-IB with N = 8
Terminal Preservation: L.MLA LA O, L LMLA 14
Near-Uniform Intermediate Distribution:
Define intermediate layer range: Tstart 1, Fend 13, R {1, 13).
Compute allowable gap bounds: T = Jmin Jmar 1.
Enumerate Valid Configurations: ([2, 6, 8, 10, 12)]-
Only valid configuration. Return (0,2,4,6,8, 10, 12, 14).
C Comparison with Pre-training Methods
we include more comparisons of our Zebra-Llama with other hybrid models based on pre-training instead of distillation. It can be observed that our method could achieve competitive performance as the pre-trained baselines 214 x -000 fewer training tokens, demonstrating our advantage in the training efficiency. For instance, Zebra-Llama (16MLA-16MM) matches or exceeds accuracy Mamba-2-Hyybid (8.66B, 3.5T tokens) and SAMBA (1.7B, 3.5T tokens), while using only 1 IB tokens and smaller model sizes. This highlights the efficiency and practicality of our post-training hybrid composition strategy, as achieves high performance with fraction of the training budget required by craach--trinnel models.
In Table II]
17

---

12000
5000
Throughput (tokens/s)
8975
0583
2253
7073
5083
3089
3077
2200
220
0OMCON
OON OOM
800
Batch Size
Llama3, 1-88
HambainLaama 88-50%
I EsoMMLL8888
Zebra- Llamr--Laam-- 88-16MLA16G2
Zebra- Rebra--Uama- 0B- BB-8MLA
Figure 8: Inference throughput vs. batch size of various 8B-size models. We the throughput with output sequence length 8192.
200
351
145
340
130
Poak Memony (GB)
100
30
54
35 18
OOM
DON OOM
OON DOM
32
64
120
256
$12
1024
Batthh Size
E:
4ML424M2
Figure 9: Peak memory usage during inference vs. batch size of various 8B 8B-size models. We measure the memory footprint with output sequence length 8192.
D More Inference Evaluations
In Figure 8 and 9 include more inference evaluations for 8B models iunder setting where the prompt length is fixed at 1, the output length at 8192, and the batch size gradually increases from 32 to 1024. All measurements with Hipgraph. Our Zebra-I.lama models demonstrate significant throughput gains compared to Llama, X-FcoMLA (pure MI.A), and MambaInI.lama (hybrid GQA-Mamba). Additionally, the KV cache compression in our Zebra- models translates to notable memory savings, as illustrated Figure 9 Models incorporating GQA layers experience a sharp increase in memory usage as the batch size grows. In contrast, MLA-based approaches, such as X-EcoMLA and our Zebra-Llama models, demonstrate superior memory efficiency.
recorded single AMD MI300 GPU with 192GB of memory
E Extended Few-Shot Evaluation Results
In the main text, we focused on reporting the few-shot performance of the 8B model due to space constraints. However, for completeness, we include the full few-shot evaluation results for the 1B and 3B model variants as well. As shown below, our Zebra-Llama models maintain competitive few-shot performance across multiple benchmarks even at smaller scales, demonstrating that our hybrid MLA-Mamba design preserves strong generalization while significantly reducing KV cache requirements. both and 3B scales, the proposed Sebra-Llama configurations maintain or even exceed the baseline performance on several tasks, while operating with drastically reduced KV cache memory (as low as 2.9%). This demonstrates that our hybrid MLA-Mamba2 architecture effectively balances efficiency and few-shot reasoning ability smaller model sizes.
18

---

Model and Setting
KV %
Avg.
ARC(25)
HS(10)
MMLU(S)
WG(S)
TQ(0)
Llama3.2 2- 2-1B-Instruc
100%
49.98
41.38
59.80
45.48
59.35
43.88
MambainL 10-a-18-50%%
50%
48.60
42.32
60.46
35,55
59.35
45.31
X-EcoMLA-IB
9.37%
47.97
41.04
56.13
35.32
60.77
46.59
Llamba-1B
0%
47.57
41.72
60.34
31.88
60.69
43,20
Zebra-Llama-1B 8MLA-8M2 (Ours)
7.80%
49.68
45.56
59.44
37.81
60.77
44.80
Zebra-Llama- IB, 6MIA--002 (Ours)
5.86 %
49.06
44.03
59.22
36.06
60.54
45,46
Table 12: Few-Shot Performance for 1B Models. Each model is evaluated on five standard few-shot
henchmarks: ARC (25-shot), HellaSwag (10-shot), MMLU (5-shot), Winogrande (5-shot), and TruthfulQA (0-shot). KV% denotes the relative K KV-cache size compared to the dense Llama baseline.
Model and Setting
KV%
Avg.
ARC(25)
HS(10)
MMLU(S)
WG(5)
TQ(0)
Llama3 Llama3.2-3B-Bn 3-3B-Instruct
100%
60.54
52.39
73.51
59.71
67.32
49.75
MambalnLLlama-18-550
50%
61.18
51.88
74.58
52.31
67.64
59.51
X- X-BcoMLA-3B*
9.37%
57.50
49.49
69.20
52.26
66.69
49.86
Ilamba-3B
0%
58.01
50.09
74.21
49.87
70.09
45.79
Zebra-Llama- 14MLA-14M2 (Ours)
4.69 %
60.12
53.67
71.30
51.05
67.64
56.94
Zebra-Llama-23 8MLA-20M2 (Ours)
2.86%
58.49
54.52
70.44
46.43
65.98
55.07
Table 13: Few-Shot Performance for 3B Models. Same benchmarks and settings as Table 12
E.1 Extending Evaluation to More Challenging Benchmarks
In this section, evaluate models more challenging benchmarks for assessing deeper reasoning capabilities, particularly given the architectural innovations in our hybrid design. Here, we report our results on GSM8K for mathematical reasoning, GPQA (34) for graduate-level scientific question answering and RACE (27]. a well-established benchmark for reading comprehension and multi-sentence reasoning. These benchmarks provide a broader view of our models' reasoning and generalization abilities. The results are summarized in Table 14 and we can draw the following insights:
1B scale: The Zebra-LJama-11 1B (8MLA-8M22 achieves the best average score, outperform- ing the target 1-t-a---.------tttttt model with only 7.81% of the KV cache.
3B scale:
The Zebra-L.lmma-3B models show only 1.46% and 0.55% drop compared to the target model and X-FcoMLA, respectively, while achieving 21.3x and 2x KV cache compression. Even at 2.01% % KV cache, performance remains higher than both Mamba-In- Llama and Llamba.
8B scale: The Zebra-Llama-8l variants exhibit modest 7.1 7.1% and 4.2% performance gap relative to Llama-3. -88-Innrrrtt and Bamba-In-LLamme respectively, despite 18.3x and 9.1x KV cache compression. The primary contributor this gap appears to be the more challenging GSM8K task. note that Mamba-In-Laaen used 70B teacher and 20B tokens, whereas setup used 8B teacher and 1IB tokens, suggesting room for further improvement with additional training data.
Across all model scales, Rebra-Llanne consistently outperforms Llamba, confirming the effectiveness and scalability of our hybrid 1LMama architecture.
F MMLU Performance Discussion
When comparing the different size Zebra-Llama models their corresponding base models, observed that MMLU performance reaches approximately 84% of the base model at the IB scale, 86% at the 3B scale, and 83% at the 8B scale, which is pretty consistent performance across different sizes. We explain the potential reasons for the observed MMLU performance gap below:
MMLU Task Formatting Difficulty: The observed gap in MMLU performance between hybrid models and pure Transformer baselines may stem from formatting sensitivity in MMLU's multiple- choice structure. The study by (9) shows that state-space models (SSMs) like Mamba struggle with the standard MMLU format, which involves selecting single letter (A/B/C/D) corresponding the
19

---

Model & Setting
Size
Teacher
Tokens
GSMSK (8-shot)
CPQA (Main)
GPQA (Diamond)
RACE
IB Models
Llama: 3.2-iI-lnstt
100%
32.44
38.51
26.79
26.26
38.18
Maamba-ln-Liamme X-EooMLAA-18
50%
SB
7B
30.04
24.94
26.12
30.80
38.28
SB
26
39.43
1ammba-1R
1B-70R
Zebru-Llama-Il IB, SMLLA-M2 (Ours)
7.81%
SB
7B
32.64
41.09
25.67
38.56
Zebra Llama- IB, 4MLA 12M2 (Ours)
3.91%
SB
7B
29.36
29.57
23.21
26.77
37.89
3B Models
Mamba--nn-Lane Llammmn-
30%
70B
20B
38.62
53.06
27.68
30.30
43:44
X-EcoMLA-338
9.37%
8B
7B
41.93
65.28
27 A/
30.30
44,69
Llamba-3B
3B-70B
1OB
34.35
47.76
22.77
26.77
40.10
(Ours)
Sebra-Lham-33B 6MLA-22M2 (Ours)
2.01%
3B
9B
39.12
55.19
28.35
30.81
42.11
38 Models
Llama 3.1-8B-Instrue
100%
47.12
78.17
35,49
31.31
44.69
Mamba-le-LLaamas
70B
X-FooM1A M.LA-B
8B
7B
Llamba-88
SB-708
12B
38.28
57.62
28. 28.35
26.77
40.38
Dehmn-Llamz-Be 8B, 16MLA-16M2 (Ours)
5.47%
SB
11B
44.05
68.16
29 29.02
30.30
48.71
Zebra Lluma- SMLA 24L2 (Ours)
2.73%
3B
11B
40.90
63.53
27.46
28.28
44.31
Table 14: Additional Evaluation on Challenging Benchmarks, Results are reported for GSMSK (8-shot),
GPQA (Main and Diamond), and RACE. KV Size indicates the fraction of KV cache relative to the dense Llama baseline.
correct answer. While SSMs contain the same underlying knowledge as Transformers, they require more training to learn the task format, particularly the routing of information from multiple choices into single output suggests that Lebra-Llamat hybrid use of Mamba/MLL layers may inherit some of this difficulty (inefficient decoding under MMLU's standard prompting format), especially in tasks like MMLU that rely heavily on understanding structured input-output formatting, rather than open-ended generation reasoning.
something Transformers handle more naturally via self-attention. This
Importance of Training Data Selection We believe that careful curation and selection of training datasets play crucial role in enhancing model performance, particularly for knowledge-intensive benchmarks such as MMLU. Our Zebra-Llama models were trained using the same datasets as Cramba-In-Llane (our main initial baseline) namely, OpenHermes-2.55 (15). GenQA (T6). and
Infinity-Instrut [17). However, as highlighted in the Llamba paper [11], training on a carefully filtered version of Feb-Web-Edu [35] leads to substantial improvements in MMLU scores following
distillation.
To further examine this effect, we re-trained Llamba model using our dataset configuration (without FineWeb-Edih As demonstrated in Table [5] the MMIJU score dropped from 38.11 (as reported in the original paper) to 27.35. In contrast, dataset and using the same training pipeline- achieved a significantly higher MMLU score of 36.91. This result suggests that Zebra-Llama would likely surpass Llamba even further if trained with similarly curated data. We view this as an exciting opportunity for future work, and encourage the broader community to explore the interplay between hybrid architectures and dataset quality in reasoning- intensive domains.
Zebra-Llama model-trainee on the exact same
Model
Training Dataset
MMLU Score
Llamba (Original paper)
FineWeb-Edu (fittere)) + OpenHermes-2.5 OpeaHermes-22.5 + GenQA InfinityInntruct
38.11
Llamba (Our re-training)
27.35
Zebra-Llamma (Ours)
Opentermees 5 GenQA Infinity- Instruct
36.91
Table 15: MMLU Comparison Across Training Datasets. FineWeb-Ech filtering notably improves reasoming performance, and our Zebra-Llama model demonstrates strong robustness even without access curated data.
G Evaluating Llamba and Zebra-Llama Under Consistent Training Conditions
To provide fair and controlled comparison between our Zebra-Llama and Llamba models, we re-implemented both and Llamba under identical training pipelines and datasets.
20

---

We conducted this study at the 1B model scale and Both models were trained using the three-stage pipeline introduced in the Llamba paper. Since the filtered FineWeb- Edu dataset is not publicly available, we used the same dataset configuration tokens. We maintained the same token distribution across stages as in Llamba: 225M tokens for Stage 2.025B for Stage 2, and 3.75B for Stage 3. All other hyperparameters were kept identical to ensure consistency.
in Zebra-I totaling approximately 6B
Target model: Llama-3.2-1B-- 3.2.-...---IIsstttcct
Teacher model: Llama-3.2-1B- I-Istutntttuct
Initial learning rate: 8e-5
Batch size: 96
The results in Table [6 show that Zebra-Llama surpasses Llamba (pure SSM) across all eight tasks while using only 3.91% of the KV cache, highlighting the benefits of integrating efficient attention modules (MLA) with Mamba.
Model & Setting
KV Size
Avg.
ARC
ARE
HS
MMLJU
OBQA
PIQA
RA
WG
Llama-3.2 3.2-IBBInnstree
100%
51.87
37.97
63.30
60.65
46.05
34.80
74.32
38.18
59.67
Llamba-11
0%
47.23
34.98
62.12
54.41
27.35
34.60
71.82
34.83
$7.70
Zebra-Llama-IB. 4MLA-12M2 (Ours)
3.91%
49.32
37.12
62.67
55.54
37.10
35.80
72.20
35.50
58.64
Table 16: Controlled Comparison of Llamba and Zebra-Ilama (1B scale). Both models are trained with identical datasets and pipelines. Zebra- lama achieves higher average performance with over
25x KV cache compression.
21

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: We summarized contributions at the end of the introduction section in three main items: architecture, training and our results.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
The limitations of work is discussed in the conclusion section.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
22

---

Our paper is not a theory paper and we do not have theoretical assumption and proofs.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proof's provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: We have all details of our experiments in the paper and appendix. We also will release the codes for reproducing our results when the paper accpeted.
We provide detailed descriptions of our methodology, model architecture, training pipeline, and experimental setup in both the main paper and the appendix. This includes model configurations, training schedules, and evaluation procedures. While our code not released at submission time, commit making them publicly available upon paper acceptance to further support reproducibility.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or a way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may tricky in some cases, in which case authors are welcome describe the particular way they provide for reproducibility.
23

---

In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: We have described all the training and evaluation data in the paper and they are publicly available. Our codes for training ready to be released when the paper is accepted.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.cr cc/
CodeSulomi public/ / ssisiiiin ssionpolicy for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run reproduce the results. See the NeurIPS code and data submission guidelines (https
/ 7/n Inips. cc/ http:/uullllccgcc.
des Codeeu ss i onPol 1 icy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: See the appendix for hyperparamters of our experiments. We have all the details of our experiments the paper.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [No]
Justification: Due to the large scale of our models (ranging from 1B to 8B parameters), running multiple independent training runs to compute standard deviations or confidence intervals is computationally prohibitive. As result, we do not report enror bars. This is common constraint in large model research, where the focus is typically placed on extensive ablations, benchmark diversity, and comparative baselines rather than repeated trials.
24

---

Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, at least for the experiments that support the main claims of the paper.
The factors of variability that the error bars capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall
run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
The paper reports detailed information about the computational resources used for training, including GPU type, training time, and hardware configuration. This information is presented both in the main results section and in the appendix. For instance, Table includes training time and model structure details.
Justification:
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics htttp:://tttrrrr wttp:ccppppccccccccccccccc c/ csGui dee inss?
Answer: [Yes]
Justification: We have carefully reviewed the NeurlPS Code of Ethics and confirm that our research complies with its guidelines. Our work does not involve human subjects, personally identifiable information, or any sensitive data.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
25

---

10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [Yes]
Justification: Our work focuses on improving the efficiency of large language models through post-training hybridization, making them more accessible for deployment in resource- constrained environments. The positive societal impacts include reducing the environmental and financial cost of training and serving LLMs, and enabling broader access advanced language models, especially for smaller institutions or regions with limited computational infrastructure.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No. they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our paper does not release any new datasets or pretrained models that pose high risk of misuse. The research focuses on the methodology for constructing efficient hybrid language models using components from existing publicly available models. As such, it does introduce new assets that require additional safeguards for responsible release.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use fual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.
12. Licenses for existing assets
26

---

Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes)
Justification: All external assets used in our work including pretrained models, datasets, and libraries- are properly cited in the paper. We used publicly available models and datasets under their respective open-source licenses. For each asset, we include appropriate references and URLs when applicable, and ensure that all licenses and terms of use have been respected throughout our experiments.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of license (e.g., CC-BY 4.0) should be included for each asset. For scraped data from a particular source (e.g.. website), the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms of the package should be provided. For popular datasets, paperssile wetssswithooooo hcooe com/ dat Bataseta s has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We introduce hybrid model variants and will release the codebase used for training and evaluation upon paper acceptance. The release will include documentation detailing the model architecture, training configurations, and usage instructions support reproducibility and adoption. All assets will be shared under permissive open-source license and accompanied by clear guidelines.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/moodll part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember anonymize your assets (if applicable). You either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our research does not involve any crowdsourcing experiments with human subjects. All evaluations were conducted using automated benchmarks and publicly available datasets, without human annotation feedback.
Guidelines:
27

---

The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: Our paper does not involve crowdsourcing research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent)
may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors adhere the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that the LLM is used only for writing, editing, formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: While our work builds hybrid language models using components from existing pretrained LLMs, the usage of these models follows standard practices in the field (c.g., distillation, benchmarking, initialization). No LLMs were used novel or non-standard way as part of the methodology. Therefore, no declaration is required under the NeurIPS 2025 LLM policy.
Guidelines:
The answer NA means that the method development in this research does not involve LLMs as any important, original, or non-standart components.
wwt.rriisscccc..ooee cccccc/aeeeeeeeeee/2222222 L.L)) for what should should not be described.
Please refer to our LLM policy
28


Reviews:
[{"review_id": "PvfoU9GSWQ", "reviewer": "Reviewer_yD7D", "summary": "The authors introduce Zebra-Llama, a hybrid LLM designed to significantly reduce KV cache size while preserving the performance of the original pre-trained model on one-shot and few-shot benchmarks.\n\nThe proposed workflow consists of four main steps:\n\n- First, separate Multi-head Latent Attention (MLA) and Mamba2 models are initialized using weights from the pre-trained LLM.\n\n- Next, this initialization is refined by aligning the intermediate representations of the new models with those of the original.\n\n- The best layer-mixing scheme is then identified to create an optimal hybrid architecture.\n\n- Finally, the mixed model is further fine-tuned using Supervised Fine-Tuning (SFT) and Direct Preference Optimization (DPO).\n\nThis hybrid model achieves a significant reduction in test-time computation and outperforms other competitive hybrid models in few-shot settings. Further ablation studies validate the layer selection strategy, analyze the trade-offs between model configurations, and examine scaling behaviors as the size of the teacher model changes.", "questions": "Given that long Chain-of-Thought (CoT) reasoning is susceptible to compounding errors [1], how does Zebra-Llama's performance on such tasks compare to its base model? Does its hybrid architecture of attention and state-space layers potentially exacerbate error propagation in multi-step reasoning?\n\n[1] Luo, Renjie, et al. \"Through the Valley: Path to Effective Long CoT Training for Small Language Models.\" arXiv preprint arXiv:2506.07712 (2025).", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 3, "strengths_and_weaknesses": "Strengths: \n\n- The paper is well written and quite involving. I appreciate the effort the authors made to deliver the implementation details. \n- The authors report a comprehensive series of experimental results and ablation studies. The use of SMART allocation scheme is also justified by the ablation study. \n\nWeakness: \n- It seems that the authors are using 1B/3B/8B models for the proposed scheme, how does the workflow perform when the model size is extremely large/small? \n- A typo in line 180. \"MLA initialization\" -> \"Mamba2 initialization\"", "was_revised": true, "final_justification": "This hybrid model achieves a significant reduction in test-time computation and outperforms other competitive hybrid models in few-shot settings. The authors' additional experiments demonstrate the potential for improvement in long CoT setting.", "rebuttal": "We sincerely thank the reviewer for the thoughtful feedback and kind recognition of our contributions. In response to your comments, we have done our best to address the raised concerns and will incorporate the corresponding improvements in the final version of the paper. Specifically, we have provided further insights into **[Answer 1] the scalability of the proposed workflow across different model sizes, and [Answer 2] its performance on long chain-of-thought (CoT) reasoning tasks**.\n\n---\n## **[Answer 1] Scalability of the Proposed Workflow across Smaller/Larger Model Sizes**\nThank you for the suggestion. While the core components of our workflow—such as SMART layer selection, distillation, and hybrid architectural design—are scalable and model-agnostic, we focused on 1B, 3B, and 8B models in this work for two main reasons:\n\nComparison with existing baselines: To the best of our knowledge, the largest post-trained hybrid models available in the literature are at the 8B scale. To ensure meaningful comparisons, we aligned our model sizes with those baselines.  \nPractical training constraints: The 8B model is the largest that could be reasonably trained on a single node (with 8 decent GPUs). Due to time and resource limitations, we were not able to explore training at larger scales.  \nHowever, based on your comment and within the limited time available for the rebuttal, we trained a smaller language model—Qwen-0.5B—to verify the scalability of our workflow to both smaller model sizes and a different model family. \n\n|Model & Setting|Teacher|KV Size|Avg.|ARC|ARE|HS|MMLU|OBQA|PIQA|RA|WG|\n|-----------------------------|--------|--------|-----|-----|-----|-----|------|------|------|----|----|\n|Qwen2.5-0.5B-Instruct|-|100%|47.93|33.11|59.05|52.26|45.86|34.20|70.62|32.06|56.27|\n|**Zebra-Qwen-0.5B,4MLA-20M2 (Ours)**|1.5B|6.25%|48.68|38.74|66.92|50.83|38.43|37.20|69.91|32.34|55.09|\n\nOur results show that the proposed method (referred to as Zebra-Qwen) performs well on models from the Qwen family. For training, we applied the same process as used for the LLaMA series. Notably, the 0.5B Zebra-Qwen model achieved better performance than the original target model while using only 6.25% KV cache.\n\n---\n## **[Answer 2] Performance on Long CoT Reasoning**\nThank you for the insightful question and for referencing the paper [1]. The paper shows that small language models (≤3B) can suffer performance degradation when trained on limited long Chain-of-Thought (CoT) data, raising concerns about the effectiveness of such training for small models. To investigate whether our hybrid architecture exacerbates this issue, we conducted a controlled evaluation comparing Zebra-Llama-1B, 8MLA-8M2 on GSM8K, training the model with 8K, 16K, and 32K samples from OpenR1-Math-220k, following a setup similar to the referenced paper.\nDue to time constraints, our initial evaluation was conducted on a subset comprising the first 150 questions of GSM8K, with training on three distinct dataset sizes: 8k, 16k, and 32k samples. We follow the same evaluation protocol as [1], sampling 4 answers per question (temperature = 0.6, top p = 0.95) and report the average accuracy.\nOur experiments show that:  as Zebra-Llama is further trained with long CoT data, the average response length dramatically increases from approximately 390 to over 4000 tokens. However, this increased \"thinking\" does not translate to improved accuracy. The GSM8k accuracy decreases from 39.5% to 25.83% when further trained with 8k CoT samples and further declines to 25% with 16k CoT samples.  This observation strongly aligns with the hypothesis in [1] that when small Language Models (LLMs) are trained on a limited amount of long CoT data, they tend to learn superficial reasoning patterns. This can lead to an accumulation of errors during the extended \"thinking\" process, ultimately degrading accuracy. When we increased the CoT training dataset to 32k samples, the reasoning ability of the model starts to recover and the token efficiency increases, which also aligns with the findings in [1]. \n\nIn summary, we **have observed a similar trend that training small LLM on limited long CoT data hurts the reasoning accuracy, despite a marked increase in the response length**. We hypothesize this phenomenon is general for small LLMs across various model structures, including hybrid models like Zebra-Llama. We are going to add this insightful analysis to the final version of the paper. \n\n|Column 1|original|8k samples|16k samples|32k samples|\n|--------|--------|-----------|------------|------------|\n|Accuracy|39.5|25.83|25|31.3|\n|Avg Response Length|390.4|5322.59|5215.15|4533.46|\n\n\n\n[1] Luo, Renjie, et al. \"Through the Valley:  Path to Effective Long CoT Training for Small Language Models.\" arXiv preprint arXiv:2506.07712 (2025)."}, {"review_id": "oiwxGw6ymk", "reviewer": "Reviewer_TctE", "summary": "This paper studies the composition of efficient hybrid language models from existing pre-trained models. It proposes a pipeline (including weight initialization, intermediate layer distillation, the SMART layer selection strategy, etc.) to mix SSM and MLA layers, in order to greatly improve inference efficiency while maintaining foundational capabilities without significant impact.", "questions": "See Weaknesses", "limitations": "None", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "#### **Strengths**\n\n- The proposed method is very effective, capable of greatly reducing inference overhead while keeping performance almost unchanged. For instance, the 8B model achieves an impressive 36.6x KV cache compression with only a 2.8% drop in average zero-shot performance.\n    \n- The experimental evaluation is relatively comprehensive, and the ablation analysis, which validates key designs like the novel SMART strategy, is solid\n    \n\n#### **Weaknesses**\n\n- The paper requires more discussion regarding the specific performance differences between the accelerated hybrid model and the original model. For example, while the average score is high, a closer look at Table 1 reveals the 8B model has a notable performance drop on the MMLU benchmark, and a discussion of such specific trade-offs would be valuable.\n    \n- In terms of capability measurement, the benchmarks used in Table 1 are not comprehensive enough. These more challenging benchmarks are crucial for probing deeper reasoning capabilities that may be affected by the architectural changes. I recommend further testing on benchmarks such as DROP (for reading comprehension), AIME (for math reasoning), GPQA (for graduate-level scientific QA), etc.", "was_revised": true, "final_justification": "I think this paper should be accepted to Neurips 2025. My concerns are mostly regarding the performance comparison and capability measurement in the Experiment section. The responses from the authors have largely addressed them.", "rebuttal": "We sincerely thank the reviewer for acknowledging the effectiveness of our method in reducing inference overhead with minimal performance loss, as well as recognizing the strength of our experimental evaluation and the significance of the SMART layer selection strategy.\nIn the rebuttal, we have carefully addressed your comments regarding **[Answer 1] task-wise comparison between hybrid and base models, and [Answer 2] the extension of our evaluation to more challenging benchmarks** in the following.\n\n---\n## **[Answer 1] Task-Wise Comparison of Hybrid and Base Models** \nThank you for the helpful observation regarding performance differences between our hybrid models and their corresponding base models. We carefully reviewed the results across all tasks in Table 1 for the 1B, 3B, and 8B models. Across 7 out of 8 benchmarks, our hybrid models consistently achieve either superior or very close performance compared to their base models. The only notable deviation occurs on the MMLU benchmark, which we discuss in more detail below.\nWhen comparing the different size Zebra-Llama models to their corresponding base models, we observed that MMLU performance reaches approximately 84% of the base model at the 1B scale, 86% at the 3B scale, and 83% at the 8B scale, which is pretty consistent performance across different sizes. We explain the potential reasons for the observed MMLU performance gap below:\n\n**1) MMLU Task Formatting Difficulty:** \nThe observed gap in MMLU performance between hybrid models and pure Transformer baselines may stem from formatting sensitivity in MMLU’s multiple-choice structure. The study by [2] shows that state-space models (SSMs) like Mamba struggle with the standard MMLU format, which involves selecting a single letter (A/B/C/D) corresponding to the correct answer. While SSMs contain the same underlying knowledge as Transformers, they require more training to learn the task format, particularly the routing of information from multiple choices into a single output token—something Transformers handle more naturally via self-attention.\nThis suggests that Zebra-Llama's hybrid use of Mamba/MLA layers may inherit some of this difficulty (inefficient decoding under MMLU’s standard prompting format), especially in tasks like MMLU that rely heavily on understanding structured input-output formatting, rather than open-ended generation or reasoning.  \n\n[2] Waleffe, Roger, et al. \"An empirical study of mamba-based language models.\" arXiv preprint arXiv:2406.07887 (2024).\n\n**2) Importance of training data selection:**\nWe believe that more carefully curated dataset selection can significantly enhance performance, particularly on knowledge-intensive tasks like MMLU. Our Zebra-Llama models were trained using the same datasets as Mamba-In-Llama (which was our main initial baseline)—namely, OpenHermes-2.5, GenQA, and Infinity-Instruct. However, as highlighted in the Llamba paper [3], training on a carefully filtered version of FineWeb-Edu leads to substantial improvements in MMLU scores following distillation.\nTo investigate this further, we trained a Llamba model using our dataset (without FineWeb-Edu). The MMLU score dropped from 38.11 (as reported in the original paper) to 27.35. In contrast, our Zebra-Llama model, trained on the same dataset and using the same training pipeline, significantly outperformed Llamba on MMLU (with the score of 37.1). This result suggests that Zebra-Llama would likely achieve even higher scores than Llamba if trained with similarly curated data, and we consider this an exciting direction for future work.\n**Table: MMLU Comparison with Different Training Datasets**\n\n|Model|Training Dataset|MMLU Score|\n|-----|----------------|-----------|\n|Llamba (Original paper)|FineWeb-Edu (filtered) + OpenHermes-2.5|38.11|\n|Llamba (Our re-training)|OpenHermes-2.5 + GenQA + Infinity-Instruct|27.35|\n|Zebra-Llama (Ours)|OpenHermes-2.5 + GenQA + Infinity-Instruct|36.91|\n\n\n[3] Bick, Aviv, et al. \"Llamba: Scaling distilled recurrent models for efficient language processing.\" arXiv preprint arXiv:2502.14458 (2025).\n\n---\n## **[Answer 2] Extending Evaluation to More Challenging Benchmarks**\n\nThank you for the thoughtful suggestion. We fully agree that evaluating on more challenging benchmarks is important for assessing deeper reasoning capabilities, especially in light of the architectural changes introduced by our hybrid design.\nAs recommended, we have added results on GSM8K for math reasoning and GPQA for graduate-level scientific question answering. These benchmarks help provide a more complete view of the model’s capabilities beyond the core benchmarks in the following table.\nWhile we were not able to complete testing on DROP due to time and infrastructure constraints, we do include results on RACE, which is also a well-established benchmark for reading comprehension and reasoning over longer passages.\nWe summarize the additional results below (and we will add them to the final paper):\n\n|Model & Setting|KV Size|Teacher|Tokens|Avg|GSM8K (8-shot)|GPQA (Main)|GPQA (Diamond)|RACE|\n|---------------|--------|--------|------|----|---------------|-------------|----------------|-----|\n|Llama-3.2-1B-Instruct|100%|-|-|32.44|38.51|26.79|26.26|38.18|\n|Mamba-In-Llama-1B|50%|8B|7B|30.04|24.94|26.12|30.8|38.28|\n|X-EcoMLA-1B|9.37%|8B|7B|31.3|38.06|21.43|26.26|39.43|\n|Llamba-1B|0%|1B+70B|8B|27.67|22.82|25|25.25|37.61|\n|**Zebra-Llama-1B,8MLA-8M2(Ours)**|7.81%|8B|7B|32.64|41.09|25.67|25.25|38.56|\n|**Zebra-Llama-1B,4MLA-12M2(Ours)**|3.91%|8B|7B|29.36|29.57|23.21|26.77|37.89|\n\n\n|Model & Setting|KV Size|Teacher|Tokens|Avg|GSM8K (8-shot)|GPQA (Main)|GPQA (Diamond)|RACE|\n|---------------|--------|--------|------|----|---------------|-------------|----------------|-----|\n|Llama-3.2-3B-Instruct|100%|-|-|42.32|70.89|29.24|28.28|40.86|\n|Mamba-In-Llama-3B|50%|70B|20B|38.62|53.06|27.68|30.3|43.44|\n|X-EcoMLA-3B|9.37%|8B|7B|41.93|65.28|27.46|30.3|44.69|\n|Llamba-3B|0%|3B+70B|10B|34.35|47.76|22.77|26.77|40.1|\n|**Zebra-Llama-3B,14MLA-14M2 (Ours)**|4.69%|8B|9B|41.7|62.77|27.23|29.8|46.99|\n|**Zebra-Llama-3B,6MLA-22M2 (Ours)**|2.01%|8B|9B|39.12|55.19|28.35|30.81|42.11|\n\n|Model & Setting|KV Size|Teacher|Tokens|Avg|GSM8K (8-shot)|GPQA (Main)|GPQA (Diamond)|RACE|\n|---------------|--------|--------|------|----|---------------|-------------|----------------|-----|\n|Llama-3.1-8B-Instruct|100%|-|-|47.42|78.17|35.49|31.31|44.69|\n|Mamba-In-Llama-8B|50%|70B|20B|45.98|75.05|29.46|33.3|46.12|\n|X-EcoMLA-8B|9.37%|8B|7B|44.78|70.81|29.69|30.3|48.33|\n|Llamba-8B|0%|8B+70B|12B|38.28|57.62|28.35|26.77|40.38|\n|**Zebra-Llama-8B,16MLA-16M2 (Ours)**|5.47%|8B|11B|44.05|68.16|29.02|30.3|48.71|\n|**Zebra-Llama-8B,8MLA-24M2 (Ours)**|2.73%|8B|11B|40.9|63.53|27.46|28.28|44.31|\n\nKey observations from the results:\n* Our 1B model excels, achieving the best average score and outperforming the target Llama-3.2-1B-Instruct model with only 7.81% KV cache.\n* Our 3B model shows a 1.46% and 0.55% performance drop compared to the target model and X-EcoMLA, respectively, with 21.32x and 2x KV cache compression. Even at 2.01% KV cache, our overall performance surpasses Mamba-In-Llama and Llamba.\n* Our 8B model exhibits a 7.11% and 4.2% performance drop compared to the target model and MambaInLlama, with 18.28x and 9.14x KV cache compression, respectively. For the performance gap, the reasoning task GSM8K is identified as a primary contributor. We hypothesize this is due to our efficient training scheme (8B teacher + 11B training tokens). Although Mamba-In-Llama exhibits higher GSM8K score, they adopted **a 70B teacher and 20B training tokens**. Since math reasoning tasks are generally harder, we expect Zebra-Llama to achieve further improvements with more training tokens. \n* **Our Zebra-Llama consistantly outperforms Llamba across all model sizes**"}, {"review_id": "ebU40UTCQ4", "reviewer": "Reviewer_GxoZ", "summary": "This paper introduces Zebra-Llama, a hybrid language model designed for high inference efficiency. The hybrid architecture is composed from existing pre-trained Transformers. Through a post-training pipeline a new architecture is created that interleaves Mamba2 layers with Multi-head Latent Attention (MLA) layers. This process involves three key stages: refined weight initialization from a teacher model, intermediate layer distillation to align representations, and a sensitivity-aware layer selection strategy called SMART that places the MLA layers where they are most necessary. The authors demonstrate that this method dramatically reduces KV cache size and improves throughput, while maintaining the performance of the original teacher model.", "questions": "The results in Table 1 indicate that the performance gap between the distilled hybrid models and the teacher baselines is larger for bigger models. Could you elaborate on why this might be the case?", "limitations": "The authors sufficiently discuss the limitations of their work.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths:**\n- The paper is generally well-written and easy to follow.\n- The approach to construct the hybrid model is intuitive and the KV cache reduction is significant.\n- The authors strengthen their claims by reporting direct inference throughput metrics.\n\n**Weaknesses:**\n- **Limited Novelty**: While the hybrid architecture seems to yield meaningful efficiency gains, the novelty of the approach is limited. The use of distillation to transfer knowledge from a pre-trained transformers to hybrid architectures is common practice and upcycling pre-trained attention into MLA to enable extreme KV cache compression was already proposed in X-EcoMLA, reducing the novelty of this paper's architectural contributions.\n\n- **Sensitivity Measure-Aware Replacement Method:** The Sensitivity Measure-Aware Replacement of Transformer Layers (SMART) strategy is central to the paper's contribution but is presented as a heuristic and is not well-motivated. This raises a major concern that the strategy may be overfitted to the distilled Llama architecture which is the only architecture family considered in this work. To substantiate the claim that SMART is a generally applicable technique, its effectiveness should be demonstrated on other architectural families, such as Qwen.\n\n- **Incomplete Results**: The few-shot performance is only shown for the 8B model. Why not reporting the results for all models?\n\nMinor points:\n*   Line 92: should be \"targeting extreme efficiency\".\n*   Table 1 Caption: There is a typo in the word \"except\".\n*   Line 307: The word \"accuracy\" is repeated.", "was_revised": true, "final_justification": "After the rebuttal, the authors addressed my concerns regarding applying SMART Layer Selection to a new architecture family and provided the complete few-shot results.", "rebuttal": "We sincerely thank the reviewer for the constructive feedback. In this revised version, we have carefully addressed the main concerns raised in your review. Specifically, we clarified **our key contributions [Answer 1], extended the evaluation of our SMART layer selection method to additional architectural families [Answer 2], completed the previously missing few-shot results [Answer 3], and provided further analysis to clarify the performance gap between the larger distilled and target models [Answer 4]**. We hope these substantial updates address your concerns, and we would be grateful if you could revisit your evaluation based on the additional evidence presented.\n\n---\n## **[Answer 1] Our Key Contributions**\nThank you for your comment. Given the importance of hybrid model architectures for balancing performance and efficiency, our paper introduces several key contributions spanning architecture design, training methodology, and empirical results:\n\nArchitecture: We propose the first hybrid model that combines MLA with Mamba2 layers to replace classical Transformer blocks. This design not only reduces memory usage but also addresses the quadratic attention bottleneck, enabling high-throughput inference with minimal performance degradation.  \nTraining Methodology: We placed particular emphasis on improving initialization and layer selection, as both have a significant impact on model performance. Notably, layer selection has received limited attention in prior work. To this end, we developed careful initialization strategies to enable stable integration of Mamba2 and MLA layers, introduced auxiliary intermediate layer distillation, and proposed our SMART layer selection algorithm. \n Empirical Results: Using our training pipeline, we successfully trained a family of hybrid models that achieve competitive performance with baseline Transformers across a range of benchmarks. Notably, our models reduce KV cache usage to as low as 2–5% and improve inference throughput by up to 2.1×, while also requiring less training data and shorter training time, surpassing prior state-of-the-art hybrid methods.\n \n--- \n## **[Answer 2] SMART Layer Selection on other Architectural Families** \nThanks for the comment. We agree that the proposed SMART layer selection is only tested in Llama families and it is necessary to test it with another model series such as Qwen. To show the generalizability of our SMART layer selection algorithm, we apply it to another scenario where we start from the Qwen-2.5-0.5B model (24 layers in total) and we would like to pick up 4 layers as MLA and 20 layers as Mamba2. \nSince we cannot include figures in the rebuttal. we include the layer sensitivity scores for each layer in the table below. The bold layers are the ones selected by SMART {Layer Indices: 0, 8, 15, 22}.\n\n|Layer Index|0|1|2|3|4|5|6|7|8|9|10|11|12|13|14|15|16|17|18|19|20|21|22|23|\n|-----------|--|--|--|--|--|--|--|--|--|--|---|---|---|---|---|---|---|---|--|---|----|----|----|---|\n|Sensitivity score|98.02|11.33|96.01|-1539.14|32.71|26.37|31.51|-68.93|26.05|-287.4|53.66|-201.49|-2.02|107.16|80.8|385.9|59.76|67.29|4|-1.07|164.41|221.56|243.83|129.34|\n\n\n\nWe also evaluated several alternative layer selection schemes, and the results remained consistent with those reported in the paper (Table X), further demonstrating the generalizability of the proposed SMART layer selection algorithm.\n\n|Model & Setting|Avg.|ARC|ARE|HS|MMLU|OBQA|PIQA|RA|WG|\n|---------------|----|---|---|--|-----|-----|-----|---|---|\n|Uniform [0,6,12,18]|46.86|34.73|64.56|48.39|38.81|33.2|69.04|31.58|54.54|\n|Uniform [5,11,17,23]|47.27|35.24|65.91|49.08|37.46|33.4|69.42|31.1|56.51|\n|Max score [15,20,21,22]|47.38|34.22|65.66|48.38|38.01|34.60|69.64|32.44|56.12|\n|SMART [0,8,15,22]|47.52|34.98|66.50|48.37|38.56|34.40|69.04|32.54|55.80|\n\n\n---\n## **[Answer 3] Complete Few-shot Results**\nWe focused on reporting few-shot performance for the 8B model in the main text due to space and time constraints; however, we acknowledge the importance of completeness and now provide the full few-shot performance results for both 1B and 3B models (and will include them in the final version of our paper). As shown, our Zebra-Llama models consistently demonstrate competitive few-shot performance across multiple benchmarks at smaller scales as well.\n\n|Model and Setting|KV%|Avg.|ARC(25)|HS(10)|MMLU(5)|WG(5)|TQ(0)|\n|------------------|----|-----|--------|------|--------|-----|-----|\n|Llama3.2-1B-Inst|100%|49.98|41.38|59.8|45.48|59.35|43.88|\n|MambaInLlama-1B-50%*|50%|48.60|42.32|60.46|35.55|59.35|45.31|\n|X-EcoMLA-1B|9.37%|47.97|41.04|56.13|35.32|60.77|46.59|\n|Llamba-1B|0%|47.57|41.72|60.34|31.88|60.69|43.2|\n|**Zebra-Llama-1B,8MLA-8M2(Ours)**|7.80%|49.68|45.56|59.44|37.81|60.77|44.8|\n|**Zebra-Llama-1B,6MLA-10M2(Ours)**|5.86%|49.06|44.03|59.22|36.06|60.54|45.46|\n\n\n|Model and Setting|KV%|Avg.|ARC(25)|hs(10)|mmlu(5)|WG(5)|TQ(0)|\n|------------------|----|-----|--------|------|--------|-----|-----|\n|Llama3.2-3B-Inst|100%|60.536|52.39|73.51|59.71|67.32|49.75|\n|MambaInLlama-3B-50%|50%|61.184|51.88|74.58|52.31|67.64|59.51|\n|X-EcoMLA-3B*|9.37%|57.500|49.49|69.20|52.26|66.69|49.86|\n|Llamba-3B|0%|58.010|50.09|74.21|49.87|70.09|45.79|\n|**Zebra-Llama-3B,14MLA-14M2(Ours)**|4.69%|60.120|53.67|71.30|51.05|67.64|56.94|\n|**Zebra-Llama-3B,8MLA-20M2(Ours)**|2.86%|58.488|54.52|70.44|46.43|65.98|55.07|\n\n---\n## **[Answer 4] Performance Gap between Larger Distilled and Target models**\nThank you for the comment. The larger performance gap observed for bigger hybrid models (especially the 8B models) is likely due to two main factors: \n* the use of the same size 8B teacher for training our Zebra-Llama 8B models—unlike the 1B and 3B models, which had larger teachers—and \n* the lack of scaling in training data proportional to model capacity. Due to limitations in time and computational resources, we kept the DPO dataset size fixed and were only able to use up to 11B tokens for SFT when training the 8B models. We believe that with a larger teacher and more extensive training (both in terms of steps and data volume), the performance of the 8B hybrid models can be further improved. We consider this an important direction for future work."}, {"review_id": "0LSMaZbjEv", "reviewer": "Reviewer_DZxh", "summary": "The paper introduces Zebra-Llama, a novel approach for creating efficient hybrid language models by combining existing pre-trained Transformer models with State Space Models and Multi-head Latent Attention layers. The method significantly reduces computational overhead and memory requirements, KV cache size, achieving performance competitive with traditional Transformers while drastically cutting down the resources required for model deployment and training.", "questions": "Any insights about why split $W_{UKV}$ in such way?", "limitations": "1. **Most importantly**, Llamba-8B performs better than Zebra-Llama on MMLU. Could the authors provide insights into this result? It would be helpful to conduct a controlled study with Llamba distillation pipeline using the same training dataset, since Llamba is a pure RNN model and it only has fixed size which has more saving compared with the proposed approach.\n\n2. Scaling to larger models and different model families: The paper could be improved by exploring different models, such as Qwen3, or by scaling to larger model sizes.\n\n3. Consider evaluating on more challenging tasks, such as GSM8K and MATH, to better demonstrate the effectiveness of the proposed approach.\n\nI can increase my score if some of these concerns are addressed.", "rating": 4, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 4, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths:\n\nS1. Zebra-Llama demonstrates remarkable efficiency in reducing KV cache size (up to ~50x compression), significantly improving inference throughput (2.6x-3.8x faster than comparable models).\n\nS2. The intermediate layer distillation and layer selection strategy strategies are methodologically robust, clearly contributing to model effectiveness. Ablation studies clearly demonstrates the effective of each idea (layer selection and intermediate layer distillation).\n\nWeaknesses:\n\nW1. Llamba-8B performs better than Zebra-Llama on MMLU. Could the authors provide insights into this result? It would be helpful to conduct a controlled study with Llamba distillation pipeline using the same training dataset, since Llamba is a pure RNN model and it only has fixed size which has more saving compared with the proposed approach.\n\nW2. Scaling to larger models and different model families: The paper could be improved by exploring different models, such as Qwen3, or by scaling to larger model sizes.", "was_revised": true, "final_justification": "Given their significant amount of work during the rebuttal, including the Qwen experiment, more challenge benchmarks and long-context evaluation, I have raised my score.", "rebuttal": "We sincerely thank the reviewer for the comprehensive review and valuable feedback, particularly your recognition of our method's strengths in KV cache saving, inference speedup, and the effectiveness of our intermediate layer distillation and layer selection strategy. We have carefully addressed the insightful comments regarding **comparing MMLU scores, comparisons with Llamba, model scalability, and performance on other datasets and supported our response with additional experimental evidence**. We appreciate your consideration of these updates and would be grateful if you could kindly revisit your evaluation of our paper in light of the improvements. \n\n---\n## **[Answer 1] Comparing MMLU score with Llamba for the 8B model**\nWe acknowledge the observed MMLU score gap between our 8B Zebra-Llama and Llamba-8B. As shown in [1], the MMLU task format can be particularly challenging for State Space Models (SSMs) to adapt to, often requiring more training efforts to improve. After a thorough review of the Llamba paper, we believe this gap likely comes from two key factors: \n1. **Fewer training tokens and smaller teacher model**: Our 8B Zebra-Llama model used the 8B base model as its teacher and was trained on 11 billion tokens. Llamba's 8B model, conversely, was distilled from a more powerful 70B teacher model in its final stage and trained on an additional 1 billion tokens (12 billion total). We believe this combination of a stronger teacher and a larger training budget contributed to Llamba's better MMLU performance. \n2. **More curated dataset selection** Our Zebra-Llama model utilizes the same training datasets as Mamba-In-Llama (`OpenHermes-2.5`, `GenQA`, and `Infinity-Instruct`). The Llamba paper, however, emphasizes the critical role of dataset selection/pre-processing and demonstrates that leveraging a carefully filtered version of `Fineweb-Edu` could significantly boost the MMLU scores after distillation. We hypothesize that applying a similar rigorous data curation strategy would significantly enhance Zebra-Llama's MMLU performance and lead to higher scores than Llamba. As presented in [Answer 2], our MMLU score significantly surpasses that of Llamba when using the same dataset and training pipeline. \n\n[1] Waleffe, Roger, et al. \"An empirical study of mamba-based language models.\" arXiv preprint arXiv:2406.07887 (2024).\n\n---\n\n## **[Answer 2] Evaluating Llamba and Zebra-Llama Under Consistent Training Conditions**\nThanks for the thoughtful and constructive comment. Following the reviewer’s suggestion, we conducted a controlled comparison between Zebra-Llama and Llamba using the same dataset and training pipeline. Due to time constraints, we were able to complete this evaluation at the 1B model scale. \n\nFor both Zebra-Llama and Llamba, we adopt the training pipeline introduced in Llamba paper which contains three stages. For the dataset, as the filtered `FineWeb-Edu` is not publicly available, we adopt the same training dataset in Zebra-Llama with ~6B tokens in total. We maintained the same token distribution across stages as Llamba: 225M tokens for Stage 1, 2.025B for Stage 2, and 3.75B for Stage 3. All other training parameters were kept consistent.\n* **Target Model**: Llama-3.2-1B-Instruct\n* **Teacher Model**: Llama-3.2-1B-Instruct\n* **Initial Learning Rate**: 8e-5\n* **Batch Size**: 96\n\nThe results, presented in the table below, demonstrate that **our Zebra-Llama surpasses Llamba (pure SSM) across all 8 tasks with only 3.91% KV cache**, highlighting the advantages of integrating efficient attention modules (i.e., MLA) with Mamba.\n\n|Model & Setting|KV Size|Avg.|ARC|ARE|HS|MMLU|OBQA|PIQA|RA|WG|\n|-|-|-|-|-|-|-|-|-|-|-|\n|**Llama3.2-1B-Instruct**|100%|51.87|37.97|63.3|60.65|46.05|34.8| 74.32|38.18|59.67|\n|Llamba-1B|0%|47.23|34.98|62.12|54.41|27.35|34.6|71.82|34.83|57.7|\n|Zebra-Llama-1B, 4MLA-12M2 (Ours)|3.91%|49.32|37.12|62.67|55.54|37.1|35.8|72.2|35.5|58.64|\n\n---\n## **[Answer 3] Scaling to larger models or other model families**\nWe thank the reviewer for the comments. Theortically it is possible to scale our method to larger models but computational resource and time constraints prevented us from conducting such experiments. Scaling to larger models remains in our future plan. \n\nHowever, we did perform additional experiments on `Qwen2.5-0.5B-Instruct` and `Qwen2.5-1.5B-Instruct` to demonstrate that our proposed method (now `Zebra-Qwen`) performs well across other model families. Using the same training process as for the Llama series, our 0.5B Zebra-Qwen model achieved better performance than the target model with only 6.25% KV cache. For the 1.5B Zebra-Qwen, even when using the same model as the teacher, we achieved similar performance to the target model with 12.5% KV cache size. \n\n|Model & Setting|Teacher|KV Size|Avg.|ARC|ARE|HS|MMLU|OBQA|PIQA|RA|WG|\n|-|-|-|-|-|-|-|-|-|-|-|-|\n|**Qwen2.5-0.5B-Instruct**|-|100%|47.93|33.11|59.05|52.26|45.86|34.2|70.62|32.06|56.27|\n|Zebra-Qwen-0.5B, 4MLA-20M2 (Ours)|1.5B|6.25%|48.68|38.74|66.92|50.83|38.43|37.2|69.91|32.34|55.09|\n|**Qwen2.5-1.5B-Instruct**|-|100%|58.59|47.01|75.84|68.24|60.13|41|76.01|37.8|62.67|\n|Zebra-Qwen-1.5B, 14MLA-14M2 (Ours)|1.5B|12.5%|58.16|48.63|75.17|67.64|53.87|41.6|75.73|38.66|64.01|\n\n---\n## **[Answer 4] Extending Evaluation to More Challenging Tasks**\n\nThanks for the comments and we agree on the importance of including more challenging tasks. As the reviewer suggested, we include the results on GSM8K and GPQA and summarize the results in the following.\n\n|Model & Setting|KV Size|Teacher|Tokens|Avg|GSM8K (8-shot)|GPQA (Main)|GPQA (Diamond)|\n|-|-|-|-|-|-|-|-|\n|Llama-3.1-8B-Instruct|100%|-|-|48.32|78.17|35.49|31.31|\n|MambaInLlama-8B|50%|70B|20B|45.94|75.05|29.46|33.3|\n|X-EcoMLA-8B|9.37%|8B|7B|43.6|70.81|29.69|30.3|\n|Llamba-8B|0%|8B+70B|12B|37.58|57.62|28.35|26.77|\n|Zebra-Llama-8B, 16MLA-16M2 (Ours)|5.47%|8B|11B|42.49|68.16|29.02|30.3|\n|Zebra-Llama-8B, 8MLA-24M2 (Ours)|2.73%|8B|11B|39.76|63.53|27.46|28.28|\n\n|Model & Setting|KV Size|Teacher|Tokens|Avg|GSM8K (8-shot)|GPQA (Main)|GPQA (Diamond)|\n|-|-|-|-|-|-|-|-|\n|Llama-3.2-3B-Instruct|100%|-|-|42.8|70.89|29.24|28.28|\n|MambaInLlama-3B|50%|70B|20B|37.01|53.06|27.68|30.3|\n|X-EcoMLA-3B|9.37%|8B|7B|41.01|65.28|27.46|30.3|\n|Llamba-3B|0%|3B+70B|10B|32.43|47.76|22.77|26.77|\n|Zebra-Llama-3B, 14MLA-14M2 (Ours)|4.69%|8B|9B|39.93|62.77|27.23|29.8|\n|Zebra-Llama-3B, 6MLA-22M2 (Ours)|2.01%|8B|9B|38.12|55.19|28.35|30.81|\n\n|Model & Setting|KV Size|Teacher|Tokens|Avg|GSM8K (8-shot)|GPQA (Main)|GPQA (Diamond)|\n|-|-|-|-|-|-|-|-|\n|Llama-3.1-8B-Instruct|100%|-|-|48.32|78.17|35.49|31.31|\n|MambaInLlama-8B|50%|70B|20B|45.94|75.05|29.46|33.3|\n|X-EcoMLA-8B|9.37%|8B|7B|43.6|70.81|29.69|30.3|\n|Llamba-8B|0%|8B+70B|12B|37.58|57.62|28.35|26.77|\n|Zebra-Llama-8B, 16MLA-16M2 (Ours)|5.47%|8B|11B|42.49|68.16|29.02|30.3|\n|Zebra-Llama-8B, 8MLA-24M2 (Ours)|2.73%|8B|11B|39.76|63.53|27.46|28.28|\n\nKey observations from the results:\n* Our 1B model excels, achieving the best average score and outperforming the target Llama-3.2-1B-Instruct model with only 7.81% KV cache.\n* Our 3B model shows a 6.7% and 2.63% performance drop compared to the target model and X-EcoMLA, respectively, with 21.32x and 2x KV cache compression. Even at 2.01% KV cache, our overall performance surpasses Mamba-In-Llama and Llamba.\n* Our 8B model exhibits a 12% and 7.5% performance drop compared to the target model and MambaInLlama, with 18.28x and 9.14x KV cache compression, respectively. For the performance gap, the reasoning task GSM8K is identified as a primary contributor. We hypothesize this is due to our efficient training scheme (8B teacher + 11B training tokens). Although Mamba-In-Llama exhibits higher GSM8K score, they adopted **a 70B teacher and 20B training tokens**. Since math reasoning tasks are generally harder, we expect Zebra-Llama to achieve further improvements with more training tokens. \n* **Our Zebra-Llama consistantly outperforms Llamba across all model sizes**\n\nTo test the potential of math reasoning of Zebra-Llama, we conducted additional experiments on our 1B model where we insert **more samples from OpenMathInstruct dataset during the SFT stage**. As shown in the table, GSM8K accuracy consistently gets improved as we add more math tokens. We believe the math reasoning performance of Zebra-Llama will be much improved with more data engineering efforts.\n\n| |Our Original SFT| Our Original SFT + 1M Math samples| Our Original SFT + 2M Math samples| Our Original SFT+5M Math samples|\n|-|-|-|-|-|\n|GSM8K (8-shot)|43.44|45|47.92|49.81|\n|Avg of 8 tasks in Table 1|51.35|51.08|51.2|50.94| \n\n---\n## **[Answer 5] Splitting of $W_{ukv}$**\n\nThe main goal for the split of $W^{UKV}$ is to align with operation when creating the $W^{UKV}$. The process begins by concatenating the original trained key ($W_K\\in\\mathbb{R}^{d\\times n_{kv}d_{h}}$) and value ($W^{V}\\in\\mathbb{R}^{d\\times n_{kv}d_{h}}$) weight matrices [$W^K$ ; $W^V$] and creating a single low-rank approximation of them called $W^{UKV}\\in \\mathbb{R}^{r_{kv}\\times 2d_{kv}n_{kv}}$ via SVD. \n\n[$W^{K}$ ; $W^{V}$] = $U_{kv}\\Sigma_{kv}V_{kv}^T$ = $W^{DKV}$ $W^{UKV}$\n\nConsidering how we concatenate $W^{K}$ and $W^{V}$, we know that $W^{UKV}$ can be divided into two parts where the upper part maps the latent to keys and the lower part maps the latent to values. \n\n$W^{UKV}$ = [ $\\tilde{W}^{UK}$ ; $W^{UV}$ ]   \n\nIt is straightforward to get $W^{UV}$. However, since MLA only decompress the nope dimensions for each key head from the latent, we need to select $d_{qk}$ dimensions out of $d_h$ from the upper part $\\tilde{W}^{UK}$. In this work, we assume that we always take the first $d_{qk}$ dimensions from $d_h$ as the nope dimensions.\n\nSuch procedure can be expressed as:\n\n$\\tilde{W}^{UK}$ --> $\\bar{W}^{UK}$ (reshape to $r_{kv}\\times n_{kv} \\times d_h$) --> $W^{UK}$ (take the first $d_{qk}$ from the last dimention and reshape to $r_{kv}\\times d_{qv}n_{kv}$)"}]

Author rebuttal:

Meta-review: Summary:
The paper introduces Zebra-Llama, a family of 1B, 3B, and 8B hybrid models that combine State Space Models (SSMs) and Multi-head Latent Attention (MLA) layers. The approach leverages a refined initialization and post-training pipeline to efficiently transfer knowledge from pre-trained Transformers. Zebra-Llama achieves Transformer-level accuracy with near-SSM efficiency using only 7–11B training tokens, a substantial reduction compared to the trillions of tokens typically required, with distillation from an 8B teacher.

Strengths highlighted by reviewers:
1. The hybrid model achieves significant reductions in test-time computation while maintaining competitive performance, outperforming other hybrid models in few-shot settings.
2. The method is effective in reducing inference overhead without sacrificing accuracy.
3. Additional experiments (e.g., long CoT settings) highlight the potential for further improvements.
4. The SMART layer selection method generalizes beyond Llama to Qwen, showing transferability.

Weaknesses (partially unresolved after rebuttal and discussions):
1. More challenging benchmarks are needed (e.g., DROP, MMLU), and the LLaMBA results shared in rebuttal should be fully integrated into the paper.
2. Long-context evaluation requires clearer presentation and interpretation in the final draft.
3. Unclear scalability to very large model sizes—evaluation is limited to 1B/3B/8B models.
4. The differences between the accelerated hybrid model and the original Transformer on harder benchmarks need more discussion.
5. Proposed  architectural contributions have incremental novelty, as knowledge distillation into hybrid architectures is well-studied, and prior work (e.g., X-EcoMLA) already explored MLA-based KV cache compression.
6. Larger performance gaps between larger (8B) distilled and target models remain unexplored, with the rebuttal attributing them to limited teacher size and training data. These claims need validation with larger teachers and datasets for 8B models. 

Discussion and Decision:
The paper received two accept and two borderline accept ratings. Reviewers acknowledged that the rebuttal addressed many concerns by including MMLU comparisons, LLaMBA baselines, scalability analysis, RULER benchmark results, and evaluations on additional datasets. The Area Chair agrees that the paper has clear merits and that the additional results strengthen the contribution. However, some weaknesses remain, particularly regarding scalability, novelty, and evaluations on harder reasoning tasks.

The Area Chair recommends acceptance, provided that the authors:
1. Incorporate all new results and clarifications from the rebuttal into the final version (main text or appendix).
2. Provide deeper analysis of long-context evaluations and challenging benchmarks.
3. Address novelty concerns more explicitly in the final draft.
4. Release code and checkpoints as promised in the submission.
Decision: Accept (poster)