Hymba: A Hybrid-head Architecture for Small Language Models
Authors: Xin Dong
Yonggan Fu
Shizhe Diao
Wonmin Byeon
ZIJIA CHEN
Ameya Sunil Mahabaleshwarkar
Shih-Yang Liu
Matthijs Van keirsbilck
Min-Hung Chen
Yoshi Suhara
Yingyan Celine Lin
Jan Kautz
Pavlo Molchanov
Venue: ICLR 2025 Spotlight (2025)
Keywords: hybrid model
language model
Primary area: foundation or frontier models, including LLMs
TL;DR: We propose Hymba, a family of small language models featuring a hybrid-head parallel architecture.

Abstract: We propose Hymba, a family of small language models featuring a hybrid-head parallel architecture that integrates attention mechanisms and state space models (SSMs) within the same layer, offering parallel and complementary processing of the same inputs. In this hybrid-head module, attention heads provide high-resolution recall, while SSM heads facilitate efficient context summarization. Additionally, we introduce learnable meta tokens, which are prepended to prompts to store critical meta information, guiding subsequent tokens and alleviating the “forced-to-attend” burden associated with attention mechanisms. Thanks to the global context summarized by SSMs, the attention heads in our model can be further optimized through cross-layer key-value (KV) sharing and a mix of global and local attention, resulting in a compact cache size without compromising accuracy. Notably, Hymba achieves state-of-the-art performance among small LMs: Our Hymba-1.5B-Base model surpasses all sub-2B public models and even outperforms Llama-3.2-3B, achieving 1.32\% higher average accuracy, an 11.67$\times$ reduction in cache size, and 3.49$\times$ higher throughput.

Published as a conference paper at ICLR 2025
HYMBA: A HYBRID-HEAD ARCHITECTURE FOR SMALL LANGUAGE MODELS
Xin Dong' Yonggan 2 Shizhe Diao', Wonmin Byeon', Zijia Chen', Ameya Sunil Mahabaleshwarkar Min-Hung Chen', Yoshi Suhara', Yingyan (Celine) Lin Jan Kautz', Pavlo Molchanov NVIDIA Georgia Institute of Technology Hong Kong University of Science and Technology
Yang Liu Matthijs Van Keirsbilck'
ABSTRACT
We propose Hymba, family of small language models featuring hybrid-head parallel architecture that integrates attention mechanisms and state space models (SSMs) within the same layer, offering parallel and complementary processing the same inputs. In this hybrid-head module, attention heads provide high- resolution recall, while SSM heads facilitate efficient context summarization. Ad- ditionally, we introduce learnable meta tokens, which are prepended to prompts store critical meta information, guiding subsequent tokens and alleviating the "forced-to-attend" burden associated with attention mechanisms. Thanks to the global context summarized by SSMs, the attention heads in our model can be fur- ther optimized through cross-layer key-value (KV) sharing and a mix of global and local attention, resulting in compact cache size without compromising ac- curacy. Notably, Hymba achieves state-of-the of-the-art performance among small LMs: Our Hymba- 1.5B-Base model surpasses all sub-2B public models and even out- performs Llama-3.2-30 achieving 1.32% higher average accuracy, 11.67x re- duction cache size, and 3.49x higher throughput.
Models on Hugging Face: Hymba-1.m (ommba-..5--aae 1.5e-Base
Hymba- 5B-Insttuct
INTRODUCTION
Transformers, with their attention-based architecture, have become the dominant choice for lan- guage models (LMs) due to their strong performance, parallelization capabilities, and long-term recall through key-value (KV) caches (Vaswani (2017). However, their quadratic computational cost and high memory demands pose efficiency challenges. In contrast, state space models (SSMs) like Mamba (Gu & Dao 2023) and Mamba-2 (Dao & Gu 2024) offer linear complexity and efficient hardware optimization but struggle with memory recall tasks, affecting their performance on general benchmarks (Waleffe et al. 2024 Arora et al. 2024a). While existing hybrid models that stack at- tention and SSM layers have demonstrated potential (Lieber et al. 2024 Ren et al.) 2024). they can introduce information bottlenecks when one layer type is not well-suited for specific tasks, requiring compensation from subsequent layers.
In light of this, we propose Hymba, a novel LM architecture that integrates attention heads and SSM heads within the same layer, offering parallel and complementary processing of the same inputs. This hybrid-head approach allows each layer to simultaneously harness both the high-resolution recall attention and the efficient context summarization of SSMs, increasing the model's flexibility and expressiveness in handling various types of information flows and memory access patterns.
To further enhance the achievable performance of Hymba, we introduce learnable meta tokens that prepended to the input sequences and interact with all subsequent tokens even in sliding window attention. These meta tokens appear to act as a compressed representation of world knowledge and alleviate the issue of "softmax attention not being able to attend to nothing" (Bondarenko et al. 2023 Miller Xiao et al. 2023), improving performance across both general and recall-intensive tasks. In addition, inspired by findings in (Brandon et al.) (024) that consecutive layers have high correlation in the KV cache, we propose sharing the KV cache between layers as well. Additionally,
"Equal contribution.

---

Published as a conference paper at ICLR 2025
Average Acc. (9)
Cache Size IMBI
Throughput (tok/sec)
01.1
MrDDAA/gmmb-11.
METTA Lases 12-1B
1194
ApplerOpevEEL
512
255
AupgingfianttSmma M4-179
53
840
25
54
76
06 13. 12C 186
100
Figure 1: Performance comparison of Hymba- 1.5B against sub-2B models in terms of average task accuracy, cache size (MB) relative to sequence length, and throughput (tok/sec). Specifically, the tasks include 5-shot MMLU, ARC-C, ARC-E, PIQA, Hellaswag, Winogrande, and SQuAD-C, and the throughput is measured on NVIDIA A100 with sequence length of 8k and batch size of 128 using PyTorch. For models encountering out-of-memory (OOM) issues during throughput measurement, halve the batch size until the OOM is resolved. This approach is used measure the maximal achievable throughput without OOM.
we incorporate mix of global and local (sliding window) attention, with the latter used in most of layers, to further reduce cache costs without compromising accuracy.
Comprehensive evaluations and ablation studies demonstrate that Hymba not only establishes new tate-of-thee-at (SOTA) benchmark performance across wide range of tasks but also achieves greater efficiency compared to transformers and previous hybrid models. We provide the bench- mark with other representative small LMs in Fig. [] with more comprehensive benchmarks in Fig, 6 For instance, in commonsense reasoning tasks, Hymba-1.55 1.5B can outperform Llama-3.2-3B with 1.32% higher average accuracy. while requiring 11.67x smaller cache size and being x.49 faster.
To optimize Hymba for on-device tasks, we further employ supervised finctuning and direct pref- erence optimization (DPO) (Rafailov et al., (024). Our instruction-tune model, Hymba- 1.5B- Instruct, achieves best-in-class performance on GSM8K, GPQA, and the Berkeley cenction-calling leaderboard, surpassing Llama-3.2-1 BB. Additionally, parameter-efficient finetuning shows Hymba's strong potential in this setting. For instance, a DoRA (Liu al. (024d)- finetumed version of Hymba- 1.5B outperforms Llama3.1-80D -88tttstrutt by 2.4% on RoleBench (Wang et al. 2023).
2 HYMBA: THE PROPOSED HYBRID- ARCHITECTURE
SSMs such as Mamba (Gu & Dao 2023 were introduced to address the quadratic complexity and large inference-time KV cache issues of transformers. However, due to their oowrresolution memory, SSMs struggle with memory recall and reasoning aceuracy (Waleffe et al., 2024 Jelassi al. [022]] Arora et al., (024a). To overcome these limitations, we propose roadmap for developing efficient and high-performing small LMs in Tab. and outlined as follows:
Fused hybrid-head modules. Fusing attention and SSM heads in parallel within a hybrid-head module outperforms sequential stacking (see Tab. (A)-(B) and Sec. (2T)). Both heads process the same information simultaneously, leading to improved reasoning and recall accuracy. We argue that sequential fusion lacks synergy, as both blocks operate on each set of inputs independently.
Configuration
Reasoning (%)
Commonsense
(%)
Throughput (token/sec)
Cache Size (MB)
Design Reason
Ablations on 300M model size and 100B training tokens
Transformer (Lloma)
44.08
39,98
721.1
414.7
Accurate recall while inefficient
State Space Models (Mamba)
42.98
19.23
4720.8
1.9
Efficient while inaccurate recall
Attention heads (sequential)
44.07
45.16
776.3
156.3
Enhance recall capabilitiss
B. Multi-head strecture (parallel)
45.1
49.90
876.7
148.2
Better balance of two modules
C. Local / global attention
44,56
48.79
2399.7
41.2
Boost compute/cache efficiency
cache sharing
48.04
2756.5
Cache efficiency
E. Meta tokens
45.59
$1.79
2695.8
40.0
Learned memory initialization
Scaling 1.5B model size and 1.5T training tokens
F. Size / data
60.56
64.15
664.1
78.6
Further boost task performance Improve multi-sbot recall
G.+ Extended context length (2K-+8))
661.1
78.6
Table 1: Design roadmap of our Hymba model. We evaluate the models' (1) commonsense reason- ing accuracy, averaged over 8 tasks, and (2) recall accuracy, averaged over 2 tasks. The throughput is on NVIDIA A100, sequence length 8k, batch size 128. The cache size is measured with 8k sequence length, assuming the FP16 format.
2

---

Published as a conference paper at ICLR 2025
Meta Tokens
SSM Feat.
Fading Memory (From SSM)
Gate Norm.
SSM Head
Input Tokens
Input Proj
Output Proj
Mean
Attn Feat.
Input Tokens
Attn Head
Gate Norm.
Memory
Snapshot Memory (From Attn)
(Meta Tokens)
(a)
(b)
Figure 2: (a) Visualize the hybrid-head module in Hymba; (b) Interpret from the memory aspect.
Efficiency and KV cache optimization. While attention heads improve task performance, they increase KV cache requirements and reduce throughput. To mitigate this, we optimize the hybrid- head module by combining local and global attention and employing cross-layer KV cache sharing (see Tab. T (C)- and Sec. 2.2]. This improves throughput by 3x and reduces cache by almost 4x.
Meta tokens. A set of 128 learnable embeddings is prepended input tokens, serving learned cache initialization to enhance focus on relevant information. These tokens fulfill dual purpose: (i) they mitigate attention drain by acting backstop tokens, effectively redistributing attention, and (ii) they encapsulate compressed world knowledge (see Tab. (F) and Sec. 2.3
Scaling up model size and data. Ablation studies conducted 300M-parameter model using 100B training tokens. The final models were trained with 1.5T tokens and scaled up to 1.5B- parameter models (see Tab. (F) and Sec. 2.4
2.1 A FUSED HYBRID-HEAL HEAD MODULE
SSM models are efficient but suffer from limited recall capabilities and task performance (Waleffe et al. 2024 Jelassi et al. 2024] Arora et al., 2024a] Ben- Kish et al. [002] as seen in Tab. [ Given
the high recall resolution of attention, in this step we aim (I) combine the processing efficiency
and context summarization capabilities of SSMs with the high recall resolution of attention, and (2) develop a fused building block to achieve this goal, so it can serve as a fundamental component for constructing future foundation models.
Previous hybrid models (Ren al. 2024; Glorioso et al. 2024 Lieber et al. 2024 often com- bine attention and SSMs in sequential manner. This strategy may lead information bottlenecks when a layer type that is poorly suited for a specific task cannot effectively process the informa- tion. Motivated by the multi-head attention structure in the vanilla Transformer (Vaswani 2017) where different heads undertake different roles and focus on different contexts (Lv et al., 2024 Merullo et al. 2024), we propose alternative approach: fusing atlention and SSMs in parallel into a hybrid-head module, as shown in Fig. 2 (a). The advantage of this design is that different attention and SSM heads can store, retrieve, and process the same piece of information in distinct ways, thereby inheriting the strengths of both operators.
Design formulation. We show that the hybrid-head module be represented by a unified and symmetric formulation. As shown in Fig. 2 (a). given the input sequence X, which the original input sequence X prepended with meta tokens introduced in Sec. 2.3 the input projection W Winproj = [WQ, W WV WSSMM G projects X to the query, key, and value of the attention heads using WQ W and WY, respectively, as well as the input features and gates of the SSM heads using WSSM and G respectively.
Following /Vaswani 2017), the output of attention heads Yam can formulated as:
aann softmax(Qk /V V -Mtt
where Mam softmax(QK///d) WV and Q R K = K
Similar to the attention heads, the SSM heads our model, for which we adopt Mamba (Gu & Dao 2023), can also be represented using data-controlled linear operator M following (Ali al.
3

---

Published as a conference paper at ICLR 2025
2024
Ben- Kish et al. 2024). Specifically, the SSM head output smm can be formulated as:
i
a =;
II exp(AAs) B,d, A=j+1
Y = (,,, WSSM Msm
where Msm G a WSSM, G W is an output gate, and A, B, C, A are the SSM parameters following the definition (Gu & Dao 2023). More specifically, A is a learnable matrix, B W s.m. C' Wx.ssm and A Softplus (Wa-)sm)) with Xsam WSSM X.
We observed that the output magnitudes of the SSM heads, Yssms are consistently larger than those of the attention heads, Ysu.. as visualized in Fig. 9 in Append. [ To ensure effective fusion, we normalize and re-scale them using learnable vectors to improve training stability, and then average the outputs, followed by a final output projection. The overall formulation of our fused module can be represented symmetrically:
= W Wow..proj (1 + 32norm(Mssnl
where B1 and 32 are learnable vectors that re-scale each channel of the outputs from the attention and SSM heads, respectively. We further explore the optimal ratio of SSMs and attention in hybrid heads in Append. D and amalyze the relative importance of heads in Append. E
Interpretation from the memory aspect. The components in the hybrid-heal module can be in- terpreted as analogous to human brain functions. Specifically, as shown in Fig. (b). the attention heads provide high recall resolution and thus act like snapshot memories in the human brain, stor- ing detailed recollections of a moment event. In contrast, the SSM heads summarize the context through a constant cache and thus function as fading memories, which gradually forget the details of past events while retaining their core or gist. As shown in Tab. [1 in Append. in Hymba, the summarized global context from fading memories enables allocating more snapshot memories for memorizing local information while maintaining recall capabilities. This is achieved by replacing most global attention with local attention, thus improving memory efficiency.
2.2 KV CACHE OFTIMIZATION
Our hybrid-head module improves recall and reasoning capabilities but can compromise memory and throughput efficiency due to the KV cache required by the attention heads. To address this, aim to reduce the KV cache while maintaining comparable task performance.
Combine global and local attention. Local attention, also known as sliding window attention (SWA) (Beltagy et al. 2020), offers a more efficient alternative to global full attention, though risks losing global context. However, with the presence of SSM heads in hybrid-hed module, which already summarize global context, we can more aggressively replace global full attention with local attention, achieving better balance between efficiency and performance.
Exploring the ratio of local attention and global attention. As shown in Tab. in Append. D initially replace global attention in all layers with SWA, which results in a significant degradation in recall capabilities, with accuracy dropping by over 20% on recall-intensive tasks. In response, we progressively reinstate global attention in some layers. Interestingly, as shown in Tab. (C). we find that using global attention in just three layers (i.e., the first, middle, and last layers) is sufficient to recover recall accuracy while maintaining comparable commonsense reasoning accuracy. In turn, this strategy achieves 2.7x throughput and 3.8 cache reduction.
Cross-layer KV sharing. Recent works [Luu et al. 2024a) observe that KV cache shares a high similarity between adjacent layers, suggesting that using separate KV caches for each layer leads to both cache and parameter redundancy. In light of this, we employ cross-layer KV sharing (Brandon) al. 2024), where keys and values shared between consecutive layers (e.g., every two layers share the same KV cache). This strategy reduces both KV memory usage and model parameters, allowing the saved parameters to be reallocated to other model components. As shown in Tab. (D), cross-layer KV sharing improves throughput by 1.15x while maintaining comparable recall accuracy and boosting commonsense accuracy by +0.60%
4

---

Published as a conference paper at ICLR 2025
Repeat (N-3)/2
Repeat (N-3)/2
sharing
sharing
Hymbu every layes
every Jayens
Hymbu Block
Hymbu Block
Hymbu
Hymba
LM Head
Block
Block
Full
SWA
SWA
Layernorm
Hybrid- head module
Layer
FEN
Figure 3: The overall architecture and building block of our Hymba model.
After the above optimization, Hymba's 's overall architecture is visualized in Fig. It worth noting that the proposed KV cache optimization strategies are related to the design of the hybrid-head modules. For example, if apply the same KV cache sharing strategy the Llama3 model, its average commonsense reasoning accuracy drops from 44.08% to 43.61% % and recall accuracy drops from 39.98% 28.18%. This suggests that, with the help of Mamba heads, attention heads in hybrid-heat modules are more tolerant to those lossy KV cache optimization strategies.
2.3
META TOKENS
We observed that the initial tokens, though not semantically important, often receive significant tention scores from subsequent tokens, similar to observations in prior work (Xiao et al. 2023 Han al., 2024). As shown in Fig more than 50% of the attention focused on the BOS token for I.lama.2-33 3.23B....BB To address this, we aim to guide the attention to focus more on tokens that mean- ingfully contribute to task performance. Specifically, we introduce a set of learnable meta tokens R = [r,, 72,,........ Tmm] serve the initial tokens. Given the input sequence X [x1, Zn], these meta tokens are prepended to the input sequence, forming the modified input sequence:
x [r1, Tm, I1,
Fn.
(4)
where X represents the new input sequence for our model. At inference time, since the meta tokens are fixed and appear at the beginning of any input sequences, their computation can be performed of- fline. Thus, the role of meta tokens at inference modulate the subsequent tokens, allowing subsequent tokens to focus more on those that contribute meaningfully to task performance. Similar the analogy in Sec. 2.1] the meta tokens participate in the attention and SSM calculations of all subsequent tokens, analogous to metamemory in the human brain, which helps recognize where locate needed information in other memories. We provide further analysis from the memory perspective in Append. G
also be viewed as learned cache initialization to
The role of Meta Tokens. We hypothesize that meta tokens perform the following functions. Prevent token overwriting. As shown in (Darcet et al. 2023), attention tends to overwrite and ex- cessively focus on certain tokens, functioning a garoage collector. This phenomenon was later observed in LLMS and termed "attention sinks" (Xiao et al. 2023 Han et al. 2024). Introducing learnable tokens independent of the input improves the learning of generalizable garbage collector.
Exit tokens to deal with "horced-to-attend"" Prepending tokens to the input affects the shape of the softmax function by modifying the denominator. Quiet Attention (Miller (023) alters the softmax denominator by adding one, enabling the attention to output zeros. Adding one is equivalent to prepending an all-zero token to the keys and values. Our meta tokens extend this idea by being leamable, allowing the model to optimize the softmax shape.
Initialization for KV cache and SSM state. Learning initial tokens can be seen as a form of learned prompt tuning (Lester al., 2021 Gu et al. 2021c) learned initialization. For inference, meta tokens are fixed, and the keys and values can be precomputed offline and stored. Task-specific meta tokens can be used, though in this work
one set for all tasks.
Meta tokens boost recall capabilities and commonsense reasoning accuracy. To analyze the impact of meta tokens the attention mechanism, visualize the entropy of the attention map
5

---

Published as a conference paper at ICLR 2025
Sliding Window Attention
Meta Tokens
SSM (Mamba)
Hymba (our)
+
Figure 4: Visualize Hymba's attention map the contributions of meta tokens, SWA, and Mamba.
for both attention and SSM heads (Ali et al. 2024 Bin-Kish et al., 2024) before and after introducing meta tokens. Specifically, the attention map entropy reflects the distribution of attention scores across tokens, where lower entropy indicates stronger retrieval effects (Ren et al. 2024). as the attention scores concentrated around smaller subset of tokens, and vice versa.
We provide the visualization in Fig. 13 in Append. G where we observe that, after introducing meta tokens, both the attention and SSM heads exhibit an overall reduction in entropy. Combined with the improved reasoning and recall capabilities shown in Tab. [ (E), this suggests that meta tokens may help both the attention and SSM heads focus more subset of important tokens that contribute most to task performance.
Hymba attention map. Hymba's attention pattern be viewed combination of individual components from SWA, meta tokens, and SSM, as shown in Fig. More analysis of Hymba's real attention map and a comparison with the Llama and Jamba models provided in Append.
We observe that in vanilla Transformers, attention scores are more concentrated on the beginning- of-sequence token, which is consistent with the findings (Xiao al.] 2023), and have higher proportion of local attention scores focusing on the token itself. In Hymba, meta tokens, attention heads, and SSM heads complement each other, leading to more balanced distribution of attention scores across different types of tokens.
2.4 HYMBA MODEL FAMILY
Building on the design insights explored above, we scale up the model sizes and training tokens to deliver the Hymba model family, which includes 125M model, a 350M model, and a 1.5B model.
We train Hymba- 2255/3500/ models using a mix of (CLL-BBaeelines 1.0 (Li et al. 2024) (Ben Allal et al. 2024), and proprietary high-quality dataset, with IT, 250B, and 50B tokens, respectively. We combine the Sermup-Stable-Decaae (WSD) learning rate sched- uler (Hu et al. 2024). with maximum and minimum learning rates of 3e-3 and 1e-5, and the data annealing technique (Dubey et al. 2024 Shen et al.] (024) to ensure stable pretraining. We use a se- quence length of 2k and batch size of 2M tokens throughout the training process until the last 100B tokens, where we increase the sequence length to 8k and change the ROPE base following (bloc97) 2023). The overall training pipeline is illustrated in Fig. 5 More details are provided in Append.
3 EXPERIMENTAL RESULTS
3.1 EXPERIMENT SETTINGS
Baselines. Our baselines include popular (small) LMs with quadratic attention (e.g.. Llama 3.2 (Al 2024ch, SmolLM (Allal et al., 2024b), SmolLM2 (Allal et al.] 2024a). AMD-OLMo (Liu et al. 2024b), StableLM (Bellagente et al. 2024), Olmo (Groeneveld et al. 2024). Cosmo (Huggingface)
Base Model
Instruct Model
General pretraining
LR annealing
3. SFT-1
4. SFT-2
DPO
Constant learning 2a-3
Lear 1e-5
Supervied fineturinn calling, C&A
Deneral instruction fellnsting
Peefenence Optimiation
improve
hokens taiinig
Email Peepriettry
High quality data
5000 tokeras total
samp
lingth
Figure 5: Training pipeline adapted for Hymba family. See Fig. 14 for detailed loss curve.

---

Published as a conference paper at ICLR 2025
Table 2: Benchmark Hymba with SOTA small L.M.. All models have fewer than 2B parameters, cept for Llama-3.2 3.2-3B which is marked as gray. All results are obtained through AM-EVALUATION- HARNESS (Gao et al. 2023). SQuAD-C SQuAD-Completion
indicates variant of the SQuAD question answering task proposed by Arora et al. (2024b). The throughput is measured with 8k sequence length and a 128 batch size on an NVIDIA A100 GPU. The best results are highlighted in bold, and the second-best results are highlighted in underline, where Llama-3.2-30 is not included the ranking due to its 3B model size.
Train tokens
Cache (MB)
MMLU
ARC-E ARC-C
PIQA 0-shot
Wino. O-shot
Hclla. 0-shot
SQuAD-C Avg.
Model
#arrmms.
Token/s
5-shot
O-shot
0-shot
1-shot
OpenELM-1
1.1B
1.5T
246
346
27.06
62.37
19.54
74.76
61.80
48.37
45.38
48.47
Rene-v0.1
1.3B
1.ST
800
113
32.94
67.05
31.06
76.49
62.75
51.16
48.36
52.83
Phi-1.5
1.3B
0.15T
241
1573
42.56
76.18
44.71
76.56
72.85
48.00
30.09
55.85
SmolLM
1.7B
IT
1573
27.06
76.47
43.43
75.79
60.93
49.58
45.81
54.15
Cosmo
1.8B
0.2T
244
1573
26.10
62.42
32,94
71.76
55.80
38.51
47.20
h2o-danube2
1.8B
2T
271
492
40.05
70.66
33.19
76.01
66.93
53.70
49.03
$5.65
Llama-3.2-11 IB
1.2B
9T
535
262
32.12
65.53
31.39
74.43
60.69
47.72
40.18
$0.29
Qwen2.5
1.5B
18T
469
229
60.92
75.51
41.21
75.79
63.38
50.20
49.53
59.51
AMD-OLMo
1.2B
1.31
387
1049
26.93
65.91
31.57
74.92
61.64
47.30
33.71
48.85 85
MmoILM2
1.7B
11T
238
1573
50.29
77.78
44.71
77.09
66.38
$3.55
50.50
60.04
Llama-3.2-8
3.0B
9T
191
918
56.03
74.54
42.32
76.66
69.85
55.29
43,46
59.74
Hymba
1.5B
1.ST
664
79
$1.19
76.94
45,90
77.31
66.61
53.55
55.93
61.06
2024), Phi-1.5 (Li et al. 2023), H2O-Danube (Singer al. 2024), OpenELM (Mehta et al. 2024
and MiniCPM (Hu et al. (),4)), as well as hybrid models (e.g., Rene (Al 2024a)).
Benchmark settings. We adopt two benchmarking settings: (1) In Sec. 3.2 directly bench- mark our delivered Hymba against SOTA public small LMs, and (2) Sec. 3.3 we train different architectures from scratch with the same dataset, number of layers, model size, and training recipes.
Benchmark tasks. In addition to evaluating commonsense reasoning and recall- -inteliittensive tasks on base models, we also evaluate instruction-tune tuned models downstream tasks such as math, function calling, and role- ael--plaiig in Sec. 3.4
3.2 BENCHMARK WITH SOTA SMALL LMS
We present the benchmark results of our Hymba models with parameter sizes of 125M, 350M, and 1.5B, compared to SOTA small language models within the same size range.
As highlighted in Tab 2 with only 1.5T pretraining tokens, our Hymba-1 1.5B 15B performance among all sub-2B LMs and demonstrates better throughput and cache efficiency com- pared to all transformer-basee quence length increases. For instance, compared the strongest sub-2B baseline, SmoILMM-11.7B, trained on IIT tokens, our Hymba-1.5B. trained on only 1.5T tokens, achieves a 1.02% average ac- curacy improvement, 19.91 cache size reduction, and 2.79x throughput. When comparing with small LMs trained on no more than 2T tokens, our model achieves a 1..1% 41% % average accu- racy improvement over the most competitive baselines, Phi-1.5 and h2o-danube2-1 1.8B, respectively. Additionally, our model even outperforms Llama-3.2-33 with 1.32% higher average accuracy, an 11.67x cache size reduction, and 3.49 x throughput.
model achieves the best
LMs, with this speedup becoming even more pronounced as the se-
We visualize the trade-offs between commonsense reasoning accuracy and cache size/throughput in F In addition, our delivered tiny LMs, Hymba-125M/355 125M/350M, consistently outperform all LMs of comparable model size, as summarized in Tab. and Tab. in Append. A.I We have also provided
a 1lymba-1.5B model trained exclusively on public data in Append. A.2
3.3 BENCHMARK DIFFERENT ARCHITECTURES UNDER THE SAME SETTING
General and recall-intensive tasks performance comparison. We do a comprehensive compar- ison between Hymba and other model architectures, including standard Transformer (Llama3 (Al 2024b)), pure Mamba (Gu & Dao 2023 Dao & Gul 2024). Mamba with FFN, and hybrid archi- tecture with sequential layer stacking (Samba (Ren al. (024h) several downstream tasks. All

---

Published as a conference paper at ICLR 2025
SmalLMM. 1.
Smoll 1.78
Hymba-1.50
Commonsente Reasoning (%
Hymba-1.58 [Qurs)
19.91x Cache Reduction
Phi-1.5-1388
Commonsense Reasoning (%)
Phi.115-1138
2,7 Fuster
M.1.7
?2
(mell M-1.78
Qwien2 5-1
#20 danube2 1.88
#2pxtamube2:1.88
Stablelm--11.0
Stabkmm-221.66
1c.1-
Rane-v0.1: w0.1..........
OpenFLMM M-1.1 18 AMN DLMo
La-MN-3.-1B
Qwent 1.86
@wen1.5-1.. 1-15--888
OpenELM-1 1BAMDD CLNS 412MA-32-18
Olmo 18
(0mo-1BBhh
<osmo-10 18
152
16 Cache Size (MB)
300
400
500
500
Throughput (Tokens/sec)
(b)
Figure 6: Visualize the trade-off between (a) commonsense reasoning accuracy (avr. ARC-C, ARC- E, PIQA, Hellaswag, OBQA, and Winogrande using (Gao et al. 2023)) and cache size, with through- put represented by the point size of different models, (b) commonsense reasoning accuracy and throughput, with cache size represented by the point size. The throughput is measured with 8k sequence length and a 128 batch size on an NVIDIA A100 GPU. The cache size measured with 8k sequence length, assuming the FP16 format.
Arch. Style (IB)
Mamba2 FFN
Task Type
Mamba2
Llama3
Samba
Hymba
Wiki. ppl. J LMB. ppl.
19.17
20.42
19.28
19,91
18.62
Language
12.59
14.43
13.09
12.65
10.38
SWDE f SQuAD-C
50.24
26.43
75.95
30.00
54.29
Recall Intensive
36.43
31.40
18.70
42.33
44.71
Avg.
43.34
28.92
47.33
36.17
49.50
Lambda
47.51
44.54
47.95
49.08
52.84
PIQA 1
73.94
73.07
73.45
73.23
74.97
Common- Reasoning and Question- answering
ARC-C
38.91
37.03
39.68
39.59
41.72
ARC-E
70.96
71.00
73.74
73.36
74.12
Hella. 1
57.73
55.83
57.64
58.49
60.05
Wino. 1
58.48
55.56
56.20
57.54
57.85
TruthfulQA +
30.75
29.86
31.64
28.84
31.76
SIQA
41.86
42.22
42.22
42.48
43.24
Avg.
52.32
51.14
32.82
52.83
54.57
Table 3: Apple-to-apple comparison of our Hymba, pure Mamba2 (Dao & Gu, 2024). Mamba2 with FFN, Llama3 (Dubey et al. 2024) style, and Samba-style
Ceimba-FFN-Attt-FEE (Ren al, 2024 architectures. All models have IB parameters and are trained from scratch for 100B tokens from SmolLM-Coree
(Ben Allal et al. 2024) with exactly the same training recipe. All results are obtained through L.M-EALLATION -ALLLATTIIO-HHEESSS (Gao al. 2023) using zero-shot setting. The best and second best results are highlighted in bold and underline, respectively.
models have the same number of layers and total parameters to facilitate fair comparison. Models are trained on the same data with the same hyperparameters and under the same codebase. To ensure our conclusions are generally valid, we run comparison experiments at different scales (1B and 300M) and different training datasets (SmoLL.M-crrpl) (Ben Allal et al.) 2024) and FineWeb (Penedo et al,, and Tab. 8 respectively. We evaluate the models language modeling, real-world -intesiittenivee commonsense reasoning, and pouctiion-nssweein tasks.
2024h) in Tab. and
As shown in Tab. tasks, e.g., achieving an average score 1.45% higher than the second-best model at the 300M seale and 1.74% higher at the 1B scale. The ablation study for the 300M scale is in Append. A
Hymba model consistently outperforms other 1B architectures across most
In addition, considering that Mamba models suffer from limited recall capabilities due to their constant-size cache and recurrent nature (Ben-Kish al. 2024 Arora al. 2024a JJeassi al.,
8

---

Published as a conference paper at ICLR 2025
40-
Mamba-2
03-
Llama3
60:
Hymba
10
11.0-
n-
22
20.
05
bepth Paccent
w.n-
44.
410-
05
04
ab
210-
02
Joved
Artavi
A21M5
20144
Jelse yoll
06
Tuken Limet
Token Limit
Token Limit
Figure 7: Needle-in-the -ntllll--tth- performance comparison across different architecture under apple- to-apple setting. The white vertical line represents the finetuning sequence length (4k).
2024), we test the models on two real-world recall-intensive tasks, SWDE Arora et al. 2024a Lockard al. 2019) and SQuAD (Arora et al. 2024a; Rajpurkar al., 2018), where the former is to to extract semi-structured relations from given raw HTML websites and the latter is to extract answers from given context passages. Echoing the previous findings. Mamba2 and Mamba2 with FFN architectures under-pe under-perform the Transformer model (i.e. Llama3) on these tasks (see Tab. B. Hymba model augments the Mamba heads with attention heads, which allows the model to have large effective receptive field to establish long-range dependencies and high-resolution memory to store and retrieve key information in all layers. As a result, Hymba outperforms the Transformer and Samba architectures (where the latter stacks Mamba and attention layers sequentially).
Needle-in-the-Haystee Haystack performance comparison. We further do an apple-to-apple comparison between Hymba, Mamba2, and Llama3 on the synthetic retrieval task, seelleiin-th--hayssee random and informative sentence (i.c., needle) inserted into a long document (i.c., haystack) and the model is required to retrieve the needle from the haystack to answer the questions. All models are of size IB and trained with the same setting: (i.) pretrain is done with 1 sequence length; (ii.) finetune with 4k sequence length; (iii.) test with up 16k sequence length. If model has ROPE, then we adjust the ROPE as in (Liu et al. 2023) during finetuning.
A
As shown in Fig- 7 the Hymba model significantly outperforms the Mamba2 and Llama3 models. While the Mamba2 model has good extrapolation capabilities when the needle is inserted in the end of the haystack, it struggles to retrieve the needle when the needle is in the beginning middle the haystack. In contrast, Llama3 model has limited extrapolation capabilities (Peng et al., 2023b Liu
al. 2023 Zhang al. 2024) and struggles the "lost in the middle" (Liu al. 2024c) scenario. We provide more weal-world long- context tasks evaluation in Append. where we show that Hymba has comparable or better length generalization capabilities comparing vanilla Transformers under similar training length.
3.4 BENCHMARK INSTRUCTION-TUNEE MODELS
Implementation details of post-training. We post-train Hymba-1.5B base model with two-stage strategy: the first full-finetuning (FFT) stage and another DPO (Rafailov et al. 2024) training. The learning rates are 5e-5 and 3e-6 for FFT and DPO, respectively. To accelerate training, follow the training recipe (Tunstall et al. 2023: Diao et al. 2024 Dong et al.) 2024) to pack the sam- ples and use a block size of 8192. We compare Hymba- H.ymba-l.5-IInstiee 1.5B-Instruct with competitive lightweight instruction-tuned models, i.c., Llama-3.2-1 1--.2----Innssee Instructt (Al (024c), OpenELM-1-IBB ----LL-MM---B--strtree (Mehta al., 2024).
1-.5255-11.BBInstruru (Team] 2024). and SmoILM- 1.7.. Instruct (Allal et al. 2024b). We test the instruction-tuneel models on MMLU (5-shot), IFEval, GSM8K (5-shot), GPQA (O-shot), and Berkeley Funetion-Calling Leaderboard (BFCLv2) (Yan et al. 2024). More details about experimental settings, baseline models, and evaluation tasks are shown in Append. H
Evaluation results. The evaluation results are shown in Tab. In general, Hymba- 1yymba-L.5B-Instruul Instrutt achieves the highest performance model. Qwen2.5-Instruce by around 2%. demonstrates great ability in math, reasoning, and function calling, with the best-in-class performance.
average of all tasks, outperforming the previous SoTA
In addition to full finetuning, we evaluate whether Hymba compatible with parameter- finetuning methods by finetuning the post-trained Hymba RoleBench (Wang et al.) 2023) using DoRA (Liu et al.) 2024d). Results are provided in Append. where we find that DoRA-ftmetuned Hymba significantly outperforms larger models.

---

Published as a conference paper at ICLR 2025
Table 4: The comparison between lightweight instruction tuned models. The best and second-
best results are highlighted in bold and underline, respectively. "OpenELM and SmolLM cannot understand function calling, leading to zero accuracy most categories.
Model
#Params
MMLU f
IFEval f
GSM8K f GPQA 1
BFCLv21 f
Avg.
SmolLM
1.7B
27.80
25.16
1.36
25.67
20.00
OpenELMM
1.1B
25.65
6.25
56.03
21.62
27.39
Llama-3,2
1.2B
44.41
58.92
42.99
24.11
20.27
38.14
Qwen2.5
1.5B
59.73
46.78
56,03
30.13
43.85
47.30
SmolLM2
1.7B
49.11
55.06
47.68
29.24
22.83
40.78
Hymba-1.33 1.5B
1.5B
52.79
57.14
58.76
31.03
46.40
49.22
RELATED WORKS
Large language models. Prior to the rise of LLMs, transformerrbased models (Vaswani 2017
Devlin 2018, Raffel al. 2020 Roberts et al., 2022) proved highly effective at captur-
ing relationships between tokens in complex sequences through the use of the attention mecha-
nism (Vaswani 2017). These models also demonstrated considerable scalability (Qin et al. 2023
Kaplan et al. 2020 Biderman et al. 2023) in terms of both model size and the volume of pretraining
data. This scalability paved the way for the development of LLMs, such as GLM (Du et al. 202
OPT (Zhang et al., 2022). Mistral (Jiang et al., 2023), the Llama series (Touvron et al., 2023 A]
2024b), Gemma (Team et al. 2024), and GPT-4 (Achiam et al. 2023). which showcase remarkable zero-shot and few-shot in-context learning abilities.
Efficient language model architectures. The quadratic computational complexity and the linearly increasing KV cache size of attention modules with longer sequences limit their processing effi- ciency. To address this, efficient LMs featuring sub-quadratic complexity in sequence length and
strong scaling properties have emerged (Peng al.] | [023al [Snn al. 22233] Gu & Dao 2023 Dao
& Gu [2024] Yang et al. 2023; Katharopoulos et al. (2020). As pointed out by (Gu & Dao 2023
popular efficient LM architectures such as RWKV (Peng et al. 2023a) and RetNet (Sun al.) 2023)
can be viewed as variants of SSMs (Gu et al. 2021a bb. Mamba (Gu & Dao 2023), one of the most widely used SSMs, improves upon previous SSMs by selectively propagating or forgetting information along the sequence length in an input-dependent manner. Follow-up works such as
Mamba2 (Dao Gu 2024) and GL.A (Yang et al., (003) introduce more hardware-friendl. gating mechanisms to enhance training throughput over Mamba.
Hybri language models. To combine the processing efficiency of SSMs with the recall capabilities of transformers, emerging trend is the creation of bybrid models that incorporate both types of operators. Specifically, (Park al. 2024) Waleffe et al. 2024) propose hybrid models that interleave Mamba and attention modules improve commonsense reasoning and in-context learning capabil- ities. Jamba (Lieber et al. 2024) and Zamba (Glorioso et al. 2024) develop sequentially stacked Mamba-Attem hybrid models. Samba (Ren et al. (024) introduces structure that sequentially stacks Mamba, SWA, and MLP layers by repeating the Mamba-MI.P WWW.MM----AALLLLL. structure, achiev- ing constant throughput
sequence lengths increase. Other recent work has also explored hybrid models that mix either linear RNNs or convolutions with attention (De et al. 2024 Pilault et al., 2024
Saon et al. 2023
Yang al. 2024).
5 CONCLUSION
In this work, we present Hymba, new family of small LMs featuring a hybrid-head architecture that combines the high-resolution recall capabilities of attention heads with the efficient context summarization of SSM heads. To further optimize the performance of Hymba, we introduce learn- able meta tokens to enhance the model's focus on salient information and propose a series KV cache optimization techniques. Through the roadmap of Hymba, comprehensive evaluations, and ablation studies, we demonstrate that Hymba sets new SOTA performance across wide range of tasks, achieving superior results in both accuracy and efficiency. Additionally, work provides detailed insights into the advantages of hybrid- head architectures, offering a promising direction for future research in efficient language and mudti-modal models.
10

---

Published as a conference paper at ICLR 2025
6 ACKNOWLEDGMENTS
This work would not have been possible without additional contributions from many people at NVIDIA, including Hanah Zhang, Maksim Khadkevich, Mohammad Shoeybi, Mostofa Patwary, Nikolaus Binder, Chenhan Yu, Meredith Price, and Oluwatobi Olabiyi.
REFERENCES
Josh Achiam, Steven Adler, Sandhini Agarwal, Lama Ahmad, Ilge Akkaya, Florencia Leoni Ale- man, Diogo Almeida, Janko Altenschmidt, Sam Altman, Shyamal Anadkat, al. Gpt-4 technical report. arXiv preprint arXiv:2303.088: 08774, 2023.
Cartesia AI. The On-Device Intelligence Update. 2024a. URL htt.at:s//arrrsiiiaa//////// on-device
Meta AI. Introducing Meta Llama The most capable openly available L.LM to date. 2024b. URL. https://ai. meta. com blog/ meta- -llp/mmttaallaaaa333
Meta AL Llama 3.2: Revolutionizing edge AI and vision with open, cus- URL https://ai. meta. cwm/blomcl..ocoooo blog/
tomizable models.
2024c.
1lama- 3-2 connectn 2024 --c.2-vviiii www.:iaaaddee-mmieeeeeeeeeeeeeee
Ameen Ali, Itamar Zimerman, and Lior Wolf. The hidden attention of mamba models. arXiv preprint arXiv:2403.01590, 2024.
Loubna Ben Allal, Anton Lozhkov, Elie Bakouch, Gabriel Martín Blázquez, Lewis Tunstall, Agustín Piqueres, Andres Marafioti, Cyril Zakka, Leandro von Werra, and Thomas Wolf. Smollm2 with great data, comes great performance, 2024a.
Loubna Ben Allal, Anton Lozhkov, Elie Bakouch, Leandro von Werra, and Thomas Wolf. Smollm blazingly fast and remarkably powerful, 2024b.
Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré. Simple linear attention language models balance the recall-throughuut tradeoff. arXiv preprint arXiv:2402. 18668, 2024a.
Simran Arora, Sabri Eyuboglu, Michael Zhang, Aman Timalsina, Silas Alberti, Dylan Zinsley, James Zou, Atri Rudra, and Christopher Ré. Simple linear attention language models balance the recall- -hrellthrooghhutt tradeoff, 2024b.
Yushi Bai, Xin Lv, Jiajie Zhang, Hongchang Lyu, Jiankai Tang, Zhidian Huang, Zhengxiao Du, Xiao Liu, Aohan Zeng. Lei Hou, et al. Longbench: A bilingual, multitask benchmark for long context understanding. arXiv preprint arXiv:2308. 14508, 2023.
Marco Bellagente, Jonathan Tow, Dakota Mahan, Duy Phung, Maksym Zhuravinskyi, Reshinth Adithyan, James Baicoianu, Ben Brooks, Nathan Cooper, Ashish Datta, al. Stable Im 2 1.6 technical report. arXiv preprint arXiv: 17834 2024.
Iz Beltagy, Matthew E Peters, and Arman Cohan. Longformer: The long-document transformer. arXiv preprint arXiv:2004.05155 2020.
Loubna Ben Allal, Anton Lozhkov. Guilherme Penedo. Thomas Wolf. and Leandro von URI.
Werra. Smollm-corpus, 2024.
HueginngFaceT///mmlll 1m-corpus
Assaf Ben-Kish, Itamar Zimerman, Shady Abu-Hussein, Nadav Cohen, Amir Globerson, Lior Wolf, and Raja Giryes. Decimamba: Exploring the length extrapolation potential of mamba, 2024, URI.
https::/azxxv.ozz//aae /20////...2222..... 14528
Stella Biderman, Hailey Schoelkopf, Quentin Gregory Anthony, Herbie Bradley, Kyle O'Brien, Eric Hallahan, Mohammad Aflah Khan, Shivanshu Purohit, USVSN Sai Prashanth, Edward Raff, al. Pythia: A suite for analyzing large language models across training and scaling. In International Conference Machine Learning, pp. 2397-2430. PMLR, 2023.

---

Published as a conference paper at ICLR 2025
bloc97. Dynamically scaled rope further increases performance of long context llama with zero finetuning, July 2023. URL htthsts::///w/s reddit 4mrgpr/dyynmicallly scaled rope further increases/
wtt.aa//////cc....//////. comment
Yelysei Bondarenko, Markus Nagel, and Tijmen Blankevoort. Quantizable transformers: Remov- ing outliers by helping attention heads do nothing, Advances in Neural Information Processing Systems, 36:75067-7506,, 2023.
William Brandon, Mayank Mishra, Aniruddha Nrusimha, Rameswar Panda, and Jonathan Ragan Kelly. Reducing transformer key-value cache size with cross-layer attention. arXiv preprini arXiv:2405.1298/ 2024.
Mikhail Burtsev, Yuri Kuratov, Anton Peganov, and Grigory Sapunov. Memory transformer. arXiv preprint arXiv:2006. 11527, 2020.
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, Christopher Hesse, and John Schulman. Training verifiers to solve math word problems. arXiv preprint urXiv:21 14168, 2021.
CodeParrot. Codeparnotggithu--code -ataee datasets at hugging face. URL htlp:ttst//hugggggfeee Co/ ooo//atteets///ooeearro/////bbb/
Tri Dao and Albert Gu. Transformers are ssms: Generalized models and efficient algorithms through structured state space duality, arXiv preprint arXiv:2405.21066 2024.
Timothée Darcet, Maxime Oquab, Julien Mairal, and Piotr Bojanowski. Vision transformers need registers. arXiv preprint 17X1v:2300.16588,, 2023.
Soham De, Samuel L Smith, Anushan Fernando, Aleksandar Botev, George Cristian-Muraru, Al- bert Gu, Ruba Haroun, Leonard Berrada, Yutian Chen, Srivatsan Srinivasan, et al. Griffin: Mix- ing gated linear recurrences with local attention for efficient language models. arXiv preprint arXiv:2402. 19427, 2024.
Jacob Devlin, Ming- Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: 04805, 2018.
Shizhe Diao, Rui Pan, Hanze Dong, Kashun Shum, Jipeng Zhang, Wei Xiong, and Tong Zhang. Lm- flow: An extensible toolkit for finetuning and inference of large foundation models. In Proceed- ings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume 3: System Demonstrations), pp. 116-127, 2024.
Charles Dickens. The Adventures of Oliver Twist. Ticknor and Fields, 1868.
Hanze Dong, Wei Xiong, Bo Pang, Haoxiang Wang, Han Zhao, Yingbo Zhou, Nan Jiang, Doyen Sahoo, Caiming Xiong, and Tong Zhang. Rlhf workflow: From reward modeling to online rlhf. arXiv preprint arXiv:2405.07863, 2024.
Alexey Dosovitskiy. An image is worth 16x16 words: Transformers for image recognition scale. arXiv preprint arXiv:2010.1 1929, 2020.
Zhengxiao Du, Yujie Qian, Xiao Liu, Ming Ding, Jiezhong Qiu, Zhilin Yang, and Jie Tang. Glm: General language model pretraining with autoregressive blank infilling. arXiv preprint arXiv:2103.10360, 2021.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, et al. The llama 3 herd of models. arXiv preprint arXiv:2407.2178. 2024.
12

---

Published as a conference paper at ICLR 2025
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Fos- ter, Laurence Golding, Jeffrey Hsu, Alain Le Noac'h, h, Haonan Li, Kyle McDonell, Niklas Muen- nighoff, Chris Ociepa. Jason Phang. Laria Reynolds. Hailey Schoelkopf, Aviya Skowron, Lin- tang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 12 2023. URL 10256836
org/ records/
Leo Gao, Jonathan Tow, Baber Abbasi, Stella Biderman, Sid Black, Anthony DiPofi, Charles Fos- ter; Laurence Golding, Jeffrey Hsu, Alain Le Noac"h, h, Haonan Li, Kyle McDonell, Niklas Muen- nighoff, Chris Ociepa, Jason Phang. Laria Reynolds, Hailey Schoelkopf, Aviya Skowron, Lin- tang Sutawika, Eric Tang, Anish Thite, Ben Wang, Kevin Wang, and Andy Zou. A framework for few-shot language model evaluation, 07 2024. URL (tLLttts:///eoodoo www.//dod..r///eorddd... 12608602
Paolo Glorioso, Quentin Anthony, Yury Tokpanov, James Whittington, Jonathan Pilault, Adam Ibrahim, and Beren Millidge. Zamba: A compact 7b ssm hybrid model. arXiv preprint 160. 2024.
Dirk Groeneveld, Iz Beltagy, Pete Walsh, Akshita Bhagia, Rodney Kinney, Oyvind Tafjord, Ananya Harsh Jha, Hamish Ivison, Ian Magnusson, Yizhong Wang, Shane Arora, David Atkinson, Russell Authur, Khyathi Chandu, Arman Cohan, Jennifer Dumas, Yanai Elazar, Yuling Gu, Jack Hessel, Tushar Khot, William Merrill, Jacob Morrison, Niklas Muennighoff, Aakanksha Naik, Crystal Nam, Matthew F. Peters, Valentina Pyatkin, Abhilasha Ravichander, Dustin Schwenk, Saurabh Shah, Will Smith, Nishant Subramani, Mitchell Wortsman, Pradeep Dasigi, Nathan Lam- bert, Kyle Richardson, Jesse Dodge, Kyle Lo, Luca Soldaini, Noah A. Smith, and Hannanch Hajishirzi. Olmo: Accelerating the science of language models. Preprint, 2024.
Albert Gu and Tri Dao. Mamba: Linear- time sequence modeling with selective state spaces. arXiv preprint ar arX1v:2312.00752, 2023.
Albert Gu, Karan Goel, and Christopher Ré. Efficiently modeling long sequences with structured state spaces. arXiv preprint arXiv:2111.00330 202
Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. Com- bining recurrett, convolutional, and continuous-time models with linear state space layers. Ad- vances neural information processing systems, 34:572-585, 2021b.
Yuxian Gu, Xu Han, Zhiyuan Liu, and Minlie Huang. Ppt: Pre-trained prompt tuning for few-shot learning, arXiv preprint N10.::210990333... 2021c.
Chi Han, Qifan Wang, Hao Peng. Wenhan Xiong, Yu Chen, Heng Ji, and Sinong Wang. Lm-infinite: Zero-shot extreme length generalization for large language models. In Proceedings of the 2024 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies (Volume I: Long Papers), pp. 3991-4008 2024.
Shengding Hu, Yuge Tu, Xu Han, Chaoqun He, Ganqu Cui, Xiang Long, Zhi Zheng, Yewei Fang, Yuxiang Huang, Weilin Zhao, et al. Minicpm: Unveiling the potential of small language models with scalable training strategies. arXiv preprint arXiv:2404.063395 2024.
Huggingface. FlnggingFacceBB//csmmooo BluggingFaceTB/ccsata
2024.
URL
co/
Samy Jelassi, David Brandfonbrener, Sham M Kakade, and Eran Malach. Repeat after me: Trans- formers are better than state space models copying, arXiv preprint arXiv:2402.01032, 2024.
Albert Q Jiang, Alexandre Sablayrolles, Arthur Mensch, Chris Bamford, Devendra Singh Chaplot, Diego de las Casas, Florian Bressand, Gianna Lengyel, Guillaume Lample, Lucile Saulnier, al. Mistral 7b. arXiv preprint arXiv:2310.06825, 2023.
Jared Kaplan, Sam McCandlish, Tom Henighan, Tom B Brown, Benjamin Chess, Rewon Child, Scott Gray, Alec Radford, Jeffrey Wu, and Dario Amodei. Scaling laws for neural language models. arXiv preprint arX1v:2001.08361, 2020.
13

---

Published as a conference paper at ICLR 2025
Angelos Katharopoulos, Apoorv Vyas, Nikolaos Pappas, and François Fleuret. Transformers mnn:: Fast autoregressive transformers with linear attention. In International conference on ma- chine learning, pp. 5156-5165 PMLR, 2020.
Brian Lester, Rami AI-Rfou, and Noah Constant, The power of scale for parameter-efficient prompt tuning. arXiv preprint 2 104. 104.08691, 2021.
Jeffrey Li, Alex Fang, Georgios Smymis. Maor Ivgi, Matt Jordan, Samir Gadre, Hritik Bansal, Etash Guha, Sedrick Keh, Kushal Arora, Saurabh Garg, Rui Xin, Niklas Muennighoff, Reinhard Heckel, Jean Mercat, Mayee Chen, Suchin Gururangan, Mitchell Wortsman, Alon Albalak, Yonatan Bit- ton, Marianna Nezhurina, Amro Abbas, Cheng-Ya Yu Hsieh, Dhruba Ghosh, Josh Gardner, Maciej Kilian, Hanlin Zhang, Rulin Shao, Sarah Pratt, Sunny Sanyal, Gabriel Ilharco, Giannis Daras, Kalyani Marathe, Aaron Gokaslan, Jieyu Zhang, Khyathi Chandu, Thao Nguyen, Igor Vasiljevic, Sham Kakade, Shuran Song, Sujay Sanghavi, Fartash Faghri, Sewoong Oh, Luke Zettlemoyer Kyle Lo, Alaaeldin El-Nouby, Hadi Pouransari, Alexander Toshev, Stephanie Wang, Dirk Groen- eveld, Luca Soldaini. Pang Wei Koh, Jenia Jitsev, Thomas Kollar, Alexandros G. Dimakis, Yair Carmon, Achal Dave, Ludwig Schmidt, and Vaishaal Shankar. Datacomp-Im: In search of the generation of training sets for language models, 2024.
Yuanzhi Li, Sébastien Bubeck, Ronen Eldan, Allie Del Giorno, Suriya Gunasekar, and Yin Tat Lee. Textbooks are all you need ii: phi-1.5 technical report. arXiv preprint (XXX:233990046633 2023.
Opher Lieber, Barak Lenz, Hofit Bata, Gal Cohen, Jhonathan Osin, Itay Dalmedigos, Erez Safahi, Shaked Meirom, Yonatan Belinkov, Shai Shalev-Shwartz, al. Jamba: hybrid transformer- mamba language model. arXiv preprint arXiv:2403. 19887, 2024.
Akide Liu, Jing Liu, Zizheng Pan, Yefei He, Gholamreza Haffari, and Bohan Zhuang. Mini- cache: Kv cache compression in depth dimension for large language models. arXiv preprint arXiv:2405. 14366, 2024a.
Jiang Liu, Jialian Wu, Prakamya Mishra, Zicheng Liu, Sudhanshu Ranjan, Pratik Prabhanjan Brahma, Yusheng Su, Gowtham Ramesh, Peng Sun, Zhe Li, Dong Li, Lu Tian, and Emad Bar- soum. Amd-olmo: series of lb language models trained from scratch by amd on amd instinct mi250 gpus., October 2024b. URL heattps://hugiiiii
co amd AMD OLMM
Nelson F Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts. Transactions of the Association for Computational Linguistics, 12:157-173, 2024c.
Shih-yang Liu, Chien-Yi Yi Wang, Hongxu Yin, Pavlo Molchanov, Yu-Chiang Frank Wang, Kwang- Ting Cheng, and Min-Hung Chen. Dora: Weigho-decomposed low-rank adaptation. Forty-first International Conference Machine Learning, 2024d.
Xiaoran Liu, Hang Yan, Shuo Zhang, Chenxin An, Xipeng Qiu, and Dahua Lin. Scaling laws of rope-based extrapolation. arXiv preprint arXiv:2310.05200 2023.
Colin Lockard, Prashant Shiralkar, and Xin Luna Dong. OpenCeres:
When open information extrac- tion meets the semi-structured web. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Com- putational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pp. Minneapolis, Minnesota, June 2019. Association for Computational Linguistics. doi: 10.18653/v 1/N19-1
URL //t./ttpsss///oaaaooooooooooo/.
Ang Lv, Kaiyi Zhang, Yuhan Chen, Yulong Wang, Lifeng Liu, Ji- Ji-Rong Wen, Jian Xie, and Rui Yan. Interpreting key mechanisms of factual recall in transformer-basee language models. arXiv preprint urXiv:2403.1995 2024.
Sachin Mehta, Mohammad Hossein Sekhavat, Qingqing Cao, Maxwell Horton, Yanzi Jin, Chenfan Sun, Seyed Iman Mirzadeh, Mahyar Najibi. Dmitry Belenko, Peter Zatloukal, al. Openelm: An efficient language model family with open training and inference framework. Workshop on Efficient Systems for Foundation Models I1@ ICMI.2024, 2024.
14

---

Published as a conference paper at ICLR 2025
Jack Merullo, Carsten Eickhoff, and Ellie Pavlick. Talking heads: Understanding inter-layer com-
munication in transformer language models. arXiv preprint arXiv:2406.099519 2024.
Evan Miller. Attention is off by one.
URL https://ww.eevamilleer.rr
attention off by one htmll
Evan Miller. Attention if off by one, 2023.
URL httpps//ww.evarmmllll
Jongho Park, Jacseung Park, Zheyang Xiong. Nayoung Lee, Jacwoong Cho. Samet Oymak, Kang- wook Lee, and Dimitris Papailiopoulos. Can mamba learn how to Icarn? a comparative study on in-context learning tasks. arXiv preprint urXiv:2402.0423 2024.
Guilherme Penedo, Hynek Kydlīček, Loubna Ben allal, Anton Lozhkov, Margaret Mitchell, Colin Raffel, Leandro Von Werra, and Thomas Wolf. The fineweb datasets: Decanting the web for the finest text data scale, 2024. URL https:::: / /arxiv.vvvv..////i///aa////aaiia org/ /ttp//bb./22..55555
Bo Peng, Eric Alcaide, Quentin Anthony, Alon Albalak, Samuel Arcadinho, Stella Biderman, Huanqi Cao, Xin Cheng, Michael Chung, Matteo Grella, et al. Rwkv: Reinventing mns for the transformer era. arXiv preprint arXiv:2305.13048 2023a.
Bowen Peng, Jeffrey Quesnelle, Honglu Fan, and Enrico Shippole. Yarn: Efficient context window extension of large language models. arXiv preprin! arXiv:2309.0007/ 2023b.
Jonathan Pilault, Mahan Fathi, Orhan Firat, Chris Pal, Pierre-Luc Bacon, and Ross Goroshin. Block- state transformers. Advances in Neural Information Processing Systems, 36, 2024.
Zhen Qin, Dong Li, Weigao Sun, Weixuan Sun, Xuyang Shen, Xiaodong Han, Yunshen Wei, Bao- hong Lv, Fei Yuan, Xiao Luo, et al. Scaling transnormer to 175 billion parameters. arXiv preprint 207.992307.14995, 2023.
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn. Direct preference optimization: Your language model is secretly a reward model. Advances in Neural Information Processing Systems, 36, 2024.
Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, and Peter J. Liu. Exploring the limits of transfer learning with unified text-to- text transformer. Journal of Machine Learning Research, 21(140):1-67, 2020. URL http:
w/w.i..or//papprs////// settapers/v21// w2lv22120000 2020- 074 html
Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. SQuAD: 100,000+ questions for machine comprehension of text. In Jian Su, Kevin Duh, and Xavier Carreras (eds.), Pro- ceedings of the 2016 Conference on Empirical Methods Natural Language Processing, pp. 2383-2392, Austin, Texas, November 2016. Association for Computational Linguistics. doi: 10.18653/v1/DIL 1264. URL
httplaattooooooooo//o--------- org/D16-1264
Pranav Rajpurkar, Robin Jia, and Percy Liang. Know what you don't 't know: Unanswerable questions for squad, 2018.
Liliang Ren, Yang Liu, Yadong Lu, Yclong Shen, Chen Liang, and Weizhu Chen. Samba: Sim- ple hybrid state space models for efficient unlimited context language modeling. arXiv preprint arXiv:2406.00 07522222..22.... 2024.
Adam Roberts, Hyung Won Chung, Anselm Levskaya, Gaurav Mishra, James Bradbury, Daniel Andor. Sharan Narang, Brian Lester, Colin Gaffney, Afroz Mohiuddin, Curtis Hawthorne, Aitor Lewkowyez, Alex Salcianu, Mare van Zee, Jacob Austin, Sebastian Goodman, Livio Baldini Soares, Haitang Hu, Sasha Tsvyashchenko, Aakanksha Chowdhery, Jasmijn Bastings, Jannis Bu- lian, Xavier Garcia, Jianmo Ni, Andrew Chen, Kathleen Kenealy, Jonathan H. Clark, Stephan Lee, Dan Garrette, James Lee-Thorp. Colin Raffel, Noam Shazeer, Marvin Ritter, Maarten Bosma, Alexandre Passos, Jeremy Mlaitin-Shepare Noah Fiedel, Mark Omernick, Brennan Saeta, Ryan Sepassi, Alexander Spiridonov, Joshua Newlan, and Andrea Gesmundo. Scaling up models and data with t5x and seqio. arXiv preprint arXiv:2203.17189 2022. URL https:://rrxt www.rriiv.oog.g/oooooo owg/ /203 17189
15

---

Published as a conference paper at ICLR 2025
George Saon, Ankit Gupta, and Xiaodong Cui. Diagonal state space augmented transformers for speech recognition. ICASSP 2023-2023 IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pp. 1-5. IEEE, 2023.
Yikang Shen, Zhen Guo, Tianle Cai, and Zengyi Qin. Jetmoe: Reaching Ilama2 performance with 0.1 m dollars. arXiv preprint arXiv:2404.0744 07411, 2024.
Philipp Singer, Pascal Pfeiffer, Yauhen Babakhin, Maximilian Jeblick, Nischay Dhankhar, Ga- bor Fodor, and Sri Satish Ambati. H2o-danube-L.8 b technical report. arXiv preprint arXiv:2401.16818. 2024.
Yutao Sun, Li Dong, Shaohan Huang, Shuming Ma, Yuging Xia, Jilong Xue, Jianyong Wang, and Furu Wei. Retentive network: A successor transformer for large language models. arXiv preprint arXiv:2307.08622 2023.
Gemma Team, Morgane Riviere, Shreya Pathak, Pier Giuseppe Sessa, Cassidy Hardin, Surya Bhu- patiraju. Léonard Hussenot, Thomas Mesnard, Bobak Shahriari, Alexandre Ramé, et al. Gemma 2: Improving open language models at a practical size. arXiv preprint a8Xiv:2408.011 2024.
Qwen Team. Qwen2.5: A party of foundation models, September 2024. URL Utthts:::///qennl github. ww.bii////oo//oo qwwp//.eennn...mm
Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Niko- lay Bashlykov, Soumya Batra, Prajiwal Bhargava, Shruti Bhosale, et al. Llama 2: Open founda- tion and fine-tuned chat models. arXiv preprint arXiv:2307.0922 00702997,,,8,, 2023.
Lewis Tunstall, Edward Beeching, Nathan Lambert, Nazneen Rajani, Kashif Rasul, Younes Belkada, Shengyi Huang, Leandro von Werra, Clémentine Fourrier, Nathan Habib, et al. Zephyr: Direct distillation of Im alignment. arXiv preprint arXiv:2310.1694. 2023.
A Vaswani. Attention is all you need. Advances in Neural Information Processing Systems, 2017.
Roger Waleffe, Wonmin Byeon, Duncan Riach, Brandon Norick, Vijay Korthikanti, Tri Dao, Albert Gu, Ali Hatamizadeh, Sudhakar Singh, Deepak Narayanan, al. An empirical study of mamba- based language models. arXiv preprint arXiv:2406.07887, 2024.
Zekun Moore Wang. Zhongyuan Peng. Haoran Que, Jiaheng Liu, Wangchunshu Zhou, Yuhan Wu, Hongcheng Guo, Ruitong Gan, Zehao Ni, Man Zhang, et al. Rolellm: Benchmarking, eliciting, and enhancing role-playing abilities of large language models. arXiv preprint 1rXiv:231000046. 2023.
Guangxuan Xiao, Yuandong Tian, Beidi Chen, Song Han, and Mike Lewis. Efficient streaming language models with attention sinks. arXiv preprint arXiv:2309. 17453, 2023.
Fanjia Yan, Huanzhi Mao, Charlie Cheng-Jie Ji, Tianjun Zhang, Shishir G. Patil, lon Stoica, and Joseph E. Gonzalez. Berkeley function calling leaderboard. berke ey www.eelllllssoee www.eeeyeeeeaeeeeeeeeeeeeereeeee fwnctiinoooo ion_ caw.aallnnng.. leaderboard. 2024.
Songlin Yang, Bailin Wang, Yikang Shen, Rameswar Panda, and Yoon Kim. Gated linear attention transformers with hardware-efffcieet training. arXiv preprint arXiv:2312.06633, 2023.
Songlin Yang, Bailin Wang, Yu Zhang, Yikang Shen, and Yoon Kim. Parallelizing linear transform- ers with the delta rule over sequence length. arXiv preprint arXiv:2406. 064844 2024.
Susan Zhang, Stephen Roller, Naman Goyal, Mikel Artetxe, Moya Chen, Shuohui Chen, Christo- pher Dewan, Mona Diab, Xian Li, Xi Victoria Lin, et al. Opt: Open pre-trained transformer language models. arXiv preprint 22Xiv:200500066 2022.
Xinrong Zhang, Yingfa Chen, Shengding Hu, Zihang Xu, Junhao Chen, Moo Hao, Xu Han, Zhen Thai, Shuo Wang, Zhiyuan Liu, et al. Inf bench: Extending long context evaluation beyond 100k tokens. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pp. 15262-15277,200
Itamar Zimerman, Ameen Ali, and Lior Wolf. A unified implicit attention formulation for gated- linear recurrent sequence models. arXiv preprint arXiv:2405. 16504, 2024.
16

---

Published as a conference paper at ICLR 2025
A EXTENSIVE BENCHMARK FOR MORE HYMBA MODEL VARIANTS
A.1 COMPARISON WITH SOTA TINY LMs AT 350M AND 125M SCALES
Besides our 1.5B model, we also evaluate the 350M and 125M Hymba models on a diverse set of benchmarks in Tab. and Tab. respectively. Consistent with the results of our 1.5B model, Hlymba-350//12555 models outperform the SOTA tiny LMs across most of tasks and achieve the best average score. This indicates that Hymba scales effectively across different model sizes.
Table 5: Benchmark Hymba with SOTA tiny LMs, all of which have fewer than 200M parameters. All results are obtained through HUGGINGRACE/LLGHTEYAL. following (Allal et al. 2024b).
MMLU
ARC
Model
#Params. (cloze) t (c++e) f
PIQA
Hella.
OBQA f Wino. 1 Avg.
Mamba- 130m-
130M
27.41
33.01
63.33
33.86
30.40
51.54
42,43
Cerebras-GPP
111M
25.56
27.75
58.16
26.32
25.40
50.28
37.58
GPT-neo
125M
27.25
31.30
62.35
29.68
29.20
51.54
40.81
LaMini-GPP
124M
26.47
33.26
62.89
30.05
27.80
50.75
40.95
Opt
125M
25.67
31.25
61.97
31.04
29.00
53.20
41.29
GPT2
137M
26.29
31.09
62.51
29.76
29.40
49.72
40.50
Pythia
160M
26.68
31.92
61.64
29.55
27.80
49.49
40.08
MobileLM
125M
35.51
65.30
38.90
39.50
53.10
46.46
SmoILM
135M
30.23
43.99
69.60
42.30
33.60
52.70
48.44
Hymba 125M
31.12
44.95
68.50
45.54
35.52
52.25
49.35
Table 6: Benchmark Hymba with SOTA tiny LMs, all of which have fewer than 400M parameters. All results are obtained through HUGGINGFACC/LLIITEVVL. following (Allal et al. 2024b).
MMLU ARC
Model
#Params. (cloze)
PIQA
Hella. f OBQA Wino. 1 Avg.
(c+e)
Bloom
560M
27.49
32.86
65.13
35.98
28.80
51.70
42.89
Cerebras-GPT-2566
256M
25.91
29.69
61.37
28.44
28.00
51.62
39.82
Cerebras-GPT-5500
590M
26.93
32.40
62.84
31.99
28.40
50.12
41.15
Opt
350M
26.57
31.94
64.36
36.09
27.80
52.57
42.55
Pythia
410M
28.94
35.05
66.92
39.21
28.40
52.80
44.48
GPT2-medinee
380M
27.77
34.30
66.38
37.06
31.20
49.49
43.69
MobileLM
350M
43.65
68.60
49.60
40.00
57.60
51.89
SmolLM
360M
34.17
51.10
72.00
53.80
37.20
53.70
53.56
Hymba 350M
34.54
52.46
72.91
55.08
38.40
57.85
55.34
A.2 EVALUATING HYMBA-1.5B TRAINED ON PUBLIC DATA ONLY
We have also trained our Hymba-1.5B model exclusively on public data and evaluated its perfor- mance. Specifically, following the training settings in Sec. [.4 we train Hymba-1.55 on DCLM- Baseline-1.0 (Li et al. 2024) for 1T tokens in the first phase and on M-ooL.M-Corpus (Ben Allal et al. 2024) for 500B tokens in the second phase, keeping all other settings the same. The results are summarized in Tab. 7 where only the most competitive baselines from Tab. 2 are included. We observe that (1) Hymba-1.5D trained exclusively on public data only still surpasses all base- line small LMs in terms of average accuracy: and (2) Hymba- 1.5B trained on public data primarily suffers from performance drops 5-shot MMLU compared to the version trained all data, in- cluding our proprictary dataset. This suggests that the public data used may lack sufficient factual knowledge, which is supplemented by proprietary one.
17

---

Published as a conference paper at ICLR 2025
Benchmark Hymba-1.5B trained with all data and public data only against SOTA small LMs. All models have fewer than 2B parameters, except for Llama-3.2-3B, which is marked in gray- The settings follow Tab. i in our main paper and we only include the most competitive baselines here. Hymba (Public Data) refers our model trained exclusively on public datasets, without using our proprietary high-quality dataset.
Table 7:
Model
#Purams.
Train tokens
Token/s
Cuche (MB)
MMLU 5-shot
ARC-E O-shot
ARC-C
PIQA O-shot
Wino.
Hella. O-shot
SQuAD-C 1-shot
Avg.
O-shot
0-shot
Phi-1.5
1.3B
0.15T
241
1573
42.56
76.18
44.71
76.56
72.85
48.00
30.09
55.85
h2o-danube2
1.8B
2T
271
492
40.05
70.66
33.19
76.01
66.93
53.70
49.03
55.65
Qwen2.5
1.5B
18T
469
229
60.92
75.51
41.21
75.79
63.38
50.20
49.53
59.51
SmolLM2
11T
238
1573
50.29
77.78
44.71
77.09
66.38
53.55
50,50
60.04
Llama-3.2-38
3.0B
91
191
918
56.03
74.54
42.32
76.66
69.85
55.29
43.46
59.74
Hymba
1.5B
1.5T
664
79
$1.19
76.94
45.90
77.31
66.61
53.55
$5.93
61.06
Hymba (Public Data)
1.5T
664
44.31
78.58
47.01
77.53
64.56
53.89
59.82
60.81
Table 8: Apple-to-apple comparison of our Hymba, pure Mamba (Gu & Dao 2023). Mamba with FFN, Llama3 (Dubey al., 2024 style, and Samba-style -aN-ta-F---ttt-
(Ren et al.
2024) architectures. All models have 300M parameters and are trained for 100B tokens from FineWeb dataset (Penedo et al. 2024) with exactly the same training recipes. results are ob- tained through I.-EVALUATION EALUUATTOO-HANNSSS HARNESS (Gao et al. 2023). The best and second best results are highlighted in bold and underline, respectively.
Arch. Style (300M)
Mamba w/ FFN
Task Type
Mamba
Llama3
Samba
SMA
Hymba
Wiki. ppl. LMB. ppl. +
30.78
33.41
30.04
31.41
29.75
28.53
Language
19.95
23.64
20.53
19.75
20.85
15.45
SQuAD-C1
21.31
17.56
22.10
39.88
44.44
45.24
Recall
SWDE
17.14
13.10
57.86
22.14
55.48
58.33
Intensive
Avg.
19.23
15.33
39.98
31.01
49.96
51.79
Lambda t
38.95
36.37
40.15
40.59
40.40
44.67
PIQA f
69.64
69.26
70.29
69.86
69.80
70.73
Common- sense
ARC-C
24.91
25.00
24.83
25.76
25.96
26.28
Reasoning
ARC-E f
50.67
50.34
50.24
49.79
49.62
53.20
and Question- answering
Hella. 1
44.95
44.08
45.69
46.45
46.42
48.23
Wino.
51.70
51.78
52.64
52.49
52.72
53.35
TruthfulQA
23.86
26.23
28.97
27.27
26.47
27.87
SIQA f
39.20
39.53
39.66
39.92
41.25
39.92
Avg.
42.98
42.82
44.08
44.02
44.08
45.53
A.3 APPLE-TO-APPLE COMPARISON WITH OTHER ARCHITECTURES AT 300M SCALE
In addition to the apple-to-apple architecture comparison under the same settings with a 1B model size in Sec. 3.3 of our main paper, we also validate the superiority of our architecture at the 300M size. Specifically, we train different 300M model architectures on 100B tokens from FineWeb (Penedo et al. 2024). We set peak learning rates 5e-4 and use warmup and cosine decay scheduler. The training sequence length is set to 1K. For models with sliding window attention, we set the sliding window size as 256. Since Samba (Ren et al., 2024) only has local attention, further build a variant of Samba where we replace its first, last, and middle local attention with global attention to ensure a fair comparison with our model. We call this variant Sequential-Mix attention (SMA). As shown in Tab. 8 Hymba achieves the best performance in almost all tasks (with a second- best result in one task), yielding an average accuracy boost of + +1.45% compared to the strongest baseline.
18

---

Published as a conference paper at ICLR 2025
Table 9: Benchmark Hymba-1.5B and other models real-world long-context tasks from Long-
Bench (Bai et 2023).
GovReport (Rouge-L)
MultiNews (Rouge-L)
QMSum (Rouge-L)
TriviaQA (F1)
SAMSum (Rouge-L)
TREC LSHT
Model
(Acc)
(Acc)
SinoLLM-1.78
4.77
12.79
8.55
1.97
3,23
1.00
0.00
h2o-danube- 1.8B
12.41
14.28
17.01
68.24
11.46
56.00
10.50
Hymba-1.5B
13.95
19.24
17.29
76.82
35.21
56.22
11.00
B PERFORMANCE ON REAL-WORLD LONG-CONTEXT TASKS
We evaluate Hymba-1.5B on a broader range of long-context tasks, including summarization and few-shot learning tasks from LongBench (Bai et al. 2023 Specifically, finetune Hymba-1.5B on an 8k context length using 50B tokens from the SmolLM corpus and benchmark against the best-performing models: h2o-danube2-11 (trained on a 2k context length). We evaluate all models on three English summarization tasks and four few-shot learning tasks from LongBench.
(trained
16k context length) and SmoILM-1.7BB
As shown in Tab, 9 Hymba achieves the best performance across both task types, even surpassing h2o-danube2-1.88 which has much larger KV cache size and was trained on longer context length. Additionally, we note that Hymba's long-context performance can be further improved by finetuning on longer sequences, which will be a focus future releases.
Table 10:
The comparison between A-fintuudd Hymba and baselines on RoleBench. All base- line results are from Wang al. )2023).
Instruction Generalization
Role Generalization
Model
#Params
Llama-7B
7B
19.2
19.3
Aplaca-7B
7B
25.6
24.5
Vicuna- 133
13B
25.0
24.3
Llama2-7B-chat
7B
18.8
20.5
RoleLlama-7B
7B
35.5
33.5
Hymba-DoRA
1.5B
40.0
37.9
C DORA-FINETUNING OF HYMBA ON ROLE-PLAY TASKS
We conduct experiments to evaluate whether Hymba is compatible with DoRA (Liu al. 2024d), parameter-efficient finetuning method that updates pretrained models using minimal set of parame- ters. This approach is especially well-suited for on-device finetuning scenarios where computational resources are constrained. Additionally, DoRA significantly reduces storage requirements for saving multiple downstream models, only requires storing the finetuned DoRA parameters, which con- stitute less than 10% of the original model's total parameters. Specifically, we further finetune the instruction-tuned Hymba on RoleBench (Wang et al.) (002) using DoR A enhance role-playing capabilities. The training set of RoleBench used for training, and the model is evaluated on two sub-tasks: instruction generalization (Inst. Gene.) and role generalization (Role. Gene.).
As shown in the Tab. 10 our Hymba-DoRA significantly outperforms larger models. For instance, MoRA--innetned Hymba achieves scores of 40.0% 37, 7.9% on instruction generaliizatonnrrll gener- alization, outperforming RoleLlama-7B (Wang et al. 2023) by 4.5%, and 4.4% respectively. This indicates the strong generalization our model and the effectiveness of using parameter-effice finetuning techniques to further enhance its performance.
19

---

Published as a conference paper at ICLR 2025
Table 11: Ablation study of the design choices of Hymba. The design finally adopted by Hymba is
highlighted in bold. Specifically, the task lists are the same as those Tab. 3 The throughput is measured with 8k sequence length and 128 batch size measured with a 8k sequence length, assuming the FP16 format.
NVIDIA A100 GPU. The cache size
Design Factor
Configuration
Param, Ratio
Avg.
Avg. (Recall)
Throughput (Token/s)
Cache
Attn:Mamba
(General)
(MB)
Mambu Heads Only
19.23
4720.8
Mamba 4 Attn Heads
1:8.48
44.20
44.65
3278.1
99.09
Attn/Mamha Ratio
3) Mamha 8 Altn Heads
1:4.24
44.95
52.53
1816.5
197.39
Mamba + 16 Attn Heads
1:2.12
45.08
56.46
656.6
394.00
4) GQA
1:3.64
45.19
49,90
876.7
148.24
Atn Heads Only (IJama)
1:0
44.08
39.98
721.1
414.72
7) 5) All SWA's
1:3.64
44.42
29.78
4485.09
5.51
Sliding
8)
5) + SWA's + Full Attn
1:3.64
44.56
48.79
2399.7
41.19
Window
8) Cross-layer KV sharing
1:5.23
45.16
2756.5
39.42
10) 6) Same KV compression
1:0
43.60
28.18
3710.0
28.98
Fusion
11) 9) Replace Mean by Concat
1: 5.82
44.56
48.94
1413.9
39.42
Meta
12)
Meta Tokens
0:1
44.01
19.34
4712.8
L887
Tokens
13) 9)+ Meta Tokens
1:5.23
45.53
51.79
2695.8
40.01
D ABLATION STUDIES OF OUR HYMBA ARCHITECTURE
We perform further ablation studies and analyses of the design factors in our Hymba.
Parallel sequential fusion. We compare the hybrid-head module with sequential counterpart, which interleaves local attention and Mamba layers as adopted by (Ren et al. (2024), by calculating the models' effective receptive field (ERF) and their overall cache size. All the compared models have the same parameter size and are training from scratch using exactly the same training recipe. ERF is an empirical measure of the averaged distance among tokens that allows effective information propagation (Ben-Kish) Kish et al. 2024: Dosovitskiy 2020) defined as the following,
2M (S,s) (S 5) (N + 1)
ERF ΣΣΣ VENASH
(5)
HN (N +1)
where S' is index of the last token in the sequence, N is index of the last layer in the model, and (* 8) is the normalized attention score between token s and the last token in head h.
As shown in Fig. we observe that (1) in line with common intuitions, Llama3 exhibits notably larger ERF compared to Mamba due to its higher recall resolution, al-
Parallel (Ours)
Effactive Racaptive Field
Llama3
13a
beit at the cost of a larger cache size; (2) our hybrid-head
structure demonstrates the best ERF across the four de- signs, with an order of magnitude larger ERF while main-
Sequential (Samba)
taining cache size comparable to the sequential struc- ture. This suggests that the parallel structure can better leverage the limited cache size to capture longer and more complex relationships among tokens compared to the se- quential one. The differences in ERF are also reflected in task accuracy: According to Tab. | the hybrid-head de- sign (Tab. [] (B) improves commonsense reasoning and recall accuracy by +1.08% and 4.74%, respectively, over the sequential design (Tab. (A)). Based on this bench- marking and analysis, adopt the hybrid-head module our basic building block.
Mamba
0.
050
5/7
1.0
1.0
1.5 209
Cache Size (MB)
Figure 8: Visualize the ERF and cache size trade-off.
The ratio of SSMs and attention hybrid heads. To determine the proper number of attention heads, we start with Mamba model and gradually replace Mamba's 's hidden dimensions with atten- tion heads, maintaining the same overall model size. As shown in Tab. II >~ (4), we observe that model performance improves the ratio of attention parameters increases and gradually saturates when the parameter ratio of attention to Mamba reaches 1:2.12. 12. We stop introducing more attention heads, considering that adding more would bring increased memory overhead.
20

---

Published as a conference paper at ICLR 2025
Heads Magnitude Before Normalization
Gates Magnitude
1250-
Attention lleads
125
Attention lleads
Magnitude
1000
SSM
Magnitude
100
Attention Heads
75-
500
50-
25-
250
0
Layer
Layer
Figure 9: Left: Visualization of output magnitudes of attention and SSM heads. SSM heads consis- tently have higher output magnitude than attention heads due to their structure. Right: Visualization of attention and SSM heads' gate magnitudes. Through model learning, the relative magnitudes of attention and SSM gates vary across different layers.
There two interesting observations: (1) Although the attention-only model outperforms the Mamba-only model, the hybrid model with both attention and Mamba heads achieves the best per- formance; (2) with further KV cache optimization, the ratio of attention heads decreases further. In our final model, attention heads occupy no more than 1/5 of the Mamba heads, yet significantly boost both recall and commonsense reasoning compared the vanilla Mamba. This suggests that the hybrid model leverages the strengths and diversity of both attention and SSM heads, achieving better trade-off between efficiency and performance.
The hybrid-head fusion strategy. We have explored two straightforward methods to fuse the out- puts of attention and SSM heads: concatenation and mean. For concatenation, we combine the outputs of all heads and use linear layer to project the concatenated output to the final output di- mension. However, the parameter size of the linear layer increases with both the number of heads and the head dimensions. Additionally, based the empirical comparison between Tab. [] (9) and (11), the performance of concatenation fusion not better than the simple mean fusion. Therefore, adopt the mean fusion strategy in our final design.
Impact of KV cache optimization. After applying a series of KV cache optimization techniques, moving from Tab. IT (5) to Tab. II (9), we observe that our Hymba maintains comparable recall and commonsense reasoning accuracy while being 2. 74 faster. contrast, applying the same KV cache optimization to a pure Transformer, as seen in the comparison between Tab. 11 (6) and (10), results in a recall accuracy drop of 10% or more and degraded commonsense reasoning accuracy. This supports our analysis in Sec. [.2], showing that the presence of SSM heads in our hybrid-head module has already summarized the global context, allowing us to more aggressively replace global full attention with local attention in our hybrid model.
E HEAD IMPORTANCE ANALYSIS
Setup. To understand how hybrid heads contribute to the final task performance, we zero out the attention or SSM heads in each layer by setting 31 A3 in Eq. to 0 and record the final accuracy. We consider four datasets, which are presented in Fig. [0] and the task performance is measured using 1000 samples from each task, evaluated with zero-shot setting.
iraluuatin-harnee (Gao al.] 2023)
Observations. As shown in Fig. 10 we observe that (1) the relative importance of attention/SSM heads in the same layer, indicated by the change in task performance before and after being removed, may vary across different tasks. In other words, the relative importance of attention/SS.M heads in the same layer is input-adaptive, indicating that different types of heads learn to serve different roles and undertake different responsibilities when handling various inputs; (2) The SSM head in the first layer is critical for language modeling and removing it causes substantial increase in PPL substantial drop in accuracy (to random guess levels). Generally, removing one attention/SSM head results in 0.46%/1. 11.2% reduction in accuracy averaged across all layers and tasks, respectively.
21

---

Published as a conference paper at ICLR 2025
Hellaswag Acc Diff, (%6
0.3
-054
-104
Orig Model Remove Attn Remove SSM
14 15 is Layer be Removed
26
Winogrande Acc Diff.
Orig Model Remove Attn
Remove SSM
3 d à 3
A $
Layer to be Removed
-h-hlllllnggeee Acc Diff.
Orig Model
Remove Attn
Remove SSM
Layer to be Removed
Orig Model
WikiText Diff.
1.2
Remove Attn
Remove SSM
0.84
0.24
0
Layer to be Removed (d)
Figure 10: Visualize the task performance difference across three tasks after removing the Attention or SSM heads in each layer. The task performance is measured using 1000 samples from each task. Note that removing critical modules specific layers causes a significant gap compared others, making their bars fall outside the box. For such layers, we annotate the task performance with text.
F HYMBA ATTENTION MAP VISUALIZATION
In this section, we visualize the actual attention map of Hymba and compare it with those of the Llama and Jamba (Lieber al. 2024) models. Specifically, categorize elements in the attention
22

---

Published as a conference paper at ICLR 2025
Hymba-1.58 (parallel fusion)
Sliding Window Att. head
State-Space Head
(lama3.2-3B
Jamba1.5-500 (sequential fusion)
head
State-Space
Selry
Meta
Distribution of attention scores
Hymba sliding window attention:
Llama3.2-38:
33%
0%
67%
56%
40%
4%
Hymba state-space heads:
Jamba state-space model:
6%
25%
69%
71%
29%
Jamba attention:
0.6
0.8
Meta
BOS
Cross
Self
62%
37%
1%
Figure 11: Sum of attention score from different categories (i.e., 'Meta', 'BOS', 'Self", 'Cross") in Llama-3.2-3B, Jamba, and Hymba-1.5B.
map into four types: (1) 'Meta': attention scores from all real tokens to meta tokens. This category reflects the model's preference for attending to meta tokens. attention map. they are usually located in the first few columns (c.g., 128 for Hymba) if model has meta tokens. (2) 'BOS': attention scores from all real tokens to the ooginning-of-sequeeee are usually located in the first column right after the meta tokens. (3) 'Self": attention scores from all real tokens to themselves. In the attention map, they are usually located in the diagonal line, (4) "Cross": attention scores from all real tokens to other real tokens. In the attention map, they are usually located in the off-diagonal area.
token. the attention map, they
In Fig. I1 visualize the real attention maps from Llama-3.2-3B and Hymba-1.5B texts from Oliver Twist Chapter 29 (Dickens) 1868) and sum up the attention scores from different categories. The summed scores normalized by the context length. For SSM heads, follow (Ben-Kish et al. 2024) and (Zimerman et al. 2024 to calculate their attention maps and normalize the attention maps ensure each row sums I.
We observe that the attention pattern of Hymba significantly different from the vanilla Transform- In vanilla Transformers, attention scores concentrated on 'BOS', which is consistent
23

---

Published as a conference paper at ICLR 2025
Article
Average Attn. Score
0.020
0.015
0.010
0.005
0.000
30
60
90
120
Math
Average Attn. Score
0.020
0.015
0.010
0.005
0.000
30
60
90
120
Code
Average Attn. Score
0.020
0.015
0.010
0.005
0.000
30
60
90
120
Meta Token Index
Averaged attention scores received by the meta tokens in the last layer of Hymba- 1.5B model. Prompts of 'Article', 'Math' and "Code" are from SQuAD (Rajpurkar et al. (016)) GSM8K (Cobbe et al. 2021), and GitHub-Code (CodeParrot) datasets, respectively.
Figure 12:
with the findings in (iiao et al. 2023) In addition, vanilla Transformers also have a higher pro- portion of 'Self" attention scores. In Hymba, meta tokens, attention heads and SSM heads work complimentary to each other, leading to more balanced distribution of attention scores across dif- ferent types of tokens. Specifically, meta tokens offload the attention scores from 'BOS' allowing the model to focus more on the real tokens. SSM heads summarize the global context, which fo- cus more on current tokens (i.e., 'Self" attention scores). Attention heads, on the other hand, pay less attention 'Self' and 'BOS' tokens, and more attention to other tokens (i.e., 'Cross" attention scores). This suggests that the hybrid-head design Hymba can effectively balance the attention distribution across different types of tokens, potentially leading to better performance.
G META TOKENS: MORE ANALYSIS AND VISUALIZATION
Interpretation from the memory aspect. Similar to the analogy in Sec. 2.1 the meta tokens participate in the attention and SSM calculations of all subsequent tokens, analogous metamemory
24

---

Published as a conference paper at ICLR 2025
Attention Heads
SSM Heads
Meta Teklens
Tocans Ours w/ Tokens
Entropy
Entropy
8
Layer
Layer
(a)
(b)
Figure 13: Visualize the layer-wise attention map entropy of (a) attention heads, and (b) SSM heads with and without meta tokens.
in the human brain, which helps recognize where to locate needed information in other memories. To see this, we visualize the averaged attention scores received by the meta tokens in Fig. 12 for a Hymba-1.55 model. We observe that when the prompts are from different domains (e.g., article, math, and codes), different meta tokens are activated. This suggests that different meta tokens encapsulate different world knowledge, which can be leveraged to guide the attention mechanism to focus on relevant information.
Relationship with prior works. Learnable tokens have also been leveraged in previous transformer- based models. Previous prompt tuning works (Lester et al. 2021 Gu et al. (202cc) prepend learn- able prompts while keeping the model weights frozen during the task-specific tuning stage, aiming to adapt pretrained LM to downstream tasks in parameter-efficient manner. (Burtsev et al. 2020) introduces both learnable tokens and corresponding memory update modules to augment the memory mechanism in transformers. (Darcet et al.) 2023) appends a set of learnable tokens called registers to the image patches of vision transformers (Dosovitskiy 2020 to store global information and improve visual recognition. Our method integrates ideas from all these works in more flexible manner. It jointly optimizes meta tokens with model weights during pretraining, remains compat- ible with SWA heads and other attention types SSMs, and converts meta tokens into KV-cache initialization during inference without modifying the architecture.
Meta tokens reduce attention map entropy. We visualize the entropy of the attention map for both the attention and SSM heads (Ali et al. 20024 Benn-Kish et al. (2024) before and after intro- ducing meta tokens. As introduced Sec. 2.3 of our main paper, the attention map entropy reflects the distribution of attention scores across tokens, where lower entropy indicates stronger retrieval effects (Ren et al. 2024), as the attention scores are concentrated around a smaller subset of tokens.
As shown in Fig. 13 we observe that after introducing meta tokens, both the attention and SSM heads exhibit an overall reduction entropy. Specifically, entropy is significantly reduced in all attention heads and in 10 out of 12 layers of the SSM heads. This suggests that meta tokens can reduce attention map entropy, potentially helping both the attention and SSM heads focus more on subset of important tokens that contribute most to task performance, as indicated by the boosted performance in Tab.
Table 12:
Ablation study on the number of meta tokens.
Model (300M)
Wiki. ppl.
LMB. ppl.
Lambda PIQA
ARC-C
ARC-E Hella.
Wino.
TruthfulQA Avg.
SIQA
Hymba w/o meta tokens
28.99
18.68
41.26
71.55
24.66
$1.43
47.48
55.17
29.21
40.53
45.16
Hymba w/ 128 meta tokens
28.53
15.45
44.67
70.73
26.28
$3.20
48.24
53.35
27.88
39.92
45.53
Hymba w/ 256 meta tokens
28.85
16.20
43.43
72.47
26.37
51.68
48.33
53.75
28.42
40.07
45.57
Number of meta tokens. To better understand the relationship between the number of meta tokens and the performance of the model, we further compare the performance 128, and 256 meta tokens, trained on 100B tokens from Fineweb (Penedo al. 2024), following the apple-to-apple comparison in Tab. g As shown in the Tab. 12 we observe that (1) compared to Hymba without meta tokens, adding meta tokens consistently boosts the average accuracy and reduces the language model perplexity: (2) increasing the number of meta tokens from 128 to 256 does not result in notable boost in average accuracy. As such, we adopt 128 meta tokens in Hymba.
flymba-300M with 0,
25

---

Published as a conference paper at ICLR 2025
Table 13: Architecture details of Hymba models of different size.
Attribute
125M
350M
1.5B
Blocks
24
32
32
Hidden Size
512
768
1600
SSM State
16
16
16
Attn. Heads
8
12
25
Query Groups
4
4
5
Num. Full Attn
3
3
3
Window Size
1024
1024
1024
MLP Hidden
1664
2432
5504
Tie Embedding
True
True
True
Parameters
125M
350M
1.52B
H PRETRAINING AND POST-TRAINING IMPLEMENTATION DETAILS
Pretraining settings. We train Hymba--155/3355 125M/350M/1.5B models on 1.5T tokens, using a mix of 1.0 (Li al. 2024), SmolL (Ben Allal al. 2024), and internal high-quality dataset for IT, 250B, and 50B tokens, respectively. We adopt the WSD learning rate scheduler (Hu et 2024) with three phases: (1) warmup steps set to 1% of the total steps, (2) stable phase maintaining the peak learning rate of 3e-3, and (3) a decay phase reducing the learning rate le-5 over 20% of the total steps, while gradually annealing to smaller, higher-quality datasets like AmolLM-Corpuu and the internal dataset. We use sequence length of 2K and batch size of 2M tokens throughout the training process, which conducted on 128 NVIDIA A100 GPUs. Details of Hymba-15 1aiiaa-125M///3 125M/350M/1.55 models are shown in Tab. 13 We also show the training curves of Hymba- 1.5B in Fig. 14
Implementation details of post-training. We post-train our 1.5B base model with two-stage strategy: the first full-finetuning (FFT) stage and another DPO (Rafailov et al.] 2024) training. The learning rates are 5e-5, and 3e-6 for FFT and DPO, respectively. Both FFT and DPO training carried out for one epoch with cosine scheduler. The global batch size is set to 1024. To accelerate training, we follow the training recipe (Tunstall al. 2023) Diao et al., 2024 Dong et al. 2024) to pack the samples and use a block size of 2048. We implement the finetuning and DPO training with the LMFlow toolkit (Diao al. 2024). In addition to full-finetuning, also leverage Dora (Liu et al. 2024d) to do parameter-elficient finetuning.
Baselines and downstream tasks. We compare Hymba-1.5B-Imme Instruct with competitive lightweight instruction-tuned models, including Llama-3.2-IB-Inntruue (Al 2024c). OpenELM- I-IB-Instruct (Mehta al. 2024), Qwen2.55-1.BBIInttuu
(Team 2024 and Smo1LM-1.7BB
Instruct Alla et al.) 2024bb. We evaluate these instruction-tunee models on MMLU (5-shot), IFEval, GSM8K (5-shot), GPQA (0-shot), and the Berkeley Canetion-Callingg Leaderboard v2 (BF- CLv2) (Yan et al. 2024). For BFCLv2, we use the official code from the Gorilla project (Yan al.) 2024] and assess the BFCLv2-live category, which includes live sivele, live multiple, live parallel, live per-parrallllmmutppe and live_relevance. We exclude live_ireeevance since we found that some
Val. Loss (Phase
Val. Loss (Phase
Learning Rate
2.500
1.85
3.0
2.475
1.80
2,450
1.75
1.70
2.400
1.5-
1.65
2.375
Aym
1.60
1.0
2.350
0.5+
Phase
Phase 2 (2K context)
Phase 12K context)
1.50
context)
2,300
0.00
0.25
0.50
0.75
1.00
1.0
1.1
1.2
1.3
0.0
0.5
1.0
15
Figure 14: Training curves of Hymba-1.5B.
26

---

Published as a conference paper at ICLR 2025
baseline models without function-calling capabilities achieved high scores in this category (where function calling is not required) but performed poorly on other tasks. As result, these models at- tained high overall accuracy despite being ineffective for function calling. For the remaining tasks, we use the Im-evaluation-harme (Gao et al. 2024) for evaluation.
Implementation of Hymba's forward process. We provide illustrative pseudocode for Hymba's forward process in Alg. I
Algorithm 1: Forward Process of Hymba-1.5B
Input: X .......,,,,,,, where X E ((n.d) are text input tokens.
Model Configurations:
Number of blocks: 32
Block indices with global attention: [1, 16, // Three global attention
KV reusing groups: [2. 3). [4, 5), [6, 7]. [8.9). [10. [12, 13). [14, 15], [17, 18, 19), [20, 21], [22, 23], [24, 25]. [26, 27], [28, [30, (30 // share KV per group
Model Forward:
X
X = [1i, , I1y. // Prepend n meta tokens Re R(m.d)
for block-i in [I,... [2] do
block-i in [1, 16, 32) then
X HYMBABLOCK-GA(X
// global attention
else
if
block-i is the first block its KV reusing group then X1, KV HYYBABBLOCK-SWA -1 // sliding window attention
else
Retrieve KV cache from the previous layer: KVi-1
X HYYMBABLOCK-SW((T ,KVi-l) // reuse KV
end
end
end
27


Reviews:
[{"review_id": "vk9cUSxnME", "reviewer": "Reviewer_HFsT", "summary": "The paper introduces **Hymba**, a new family of small language models designed with a hybrid-head architecture that merges attention mechanisms with state space models (SSMs) for improved memory functions. Hymba utilizes attention heads for precise recall of high-resolution information and SSM heads to summarize broader context efficiently, mirroring aspects of human memory. A significant innovation is the introduction of learnable meta tokens, which act as a dynamic cache initialization, enhancing focus on key information during inference. \n\nThe authors outline a systematic approach to developing Hymba, from creating fused hybrid modules to scaling model and data sizes. Experimental results show that Hymba sets new benchmarks for small language models, achieving an improved accuracy-efficiency balance. Notably, *Hymba-1.5B* matches the commonsense reasoning performance of larger models, such as *LLaMA 3.2 3B*, while operating more efficiently. The meta tokens also reduce attention map entropy, potentially aiding the model in identifying and focusing on salient tokens. Hymba’s design offers promising advances in both the performance and efficiency of compact language models.", "questions": "1. Equation 1 does not mention a scaling factor. Is it included in the actual implementation?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "1. this is a solid and well-written paper.\n2. the hybrid-head design, which combines attention and state space models, is an innovative approach that provides Hymba with both fine-grained recall and efficient long-range summarization. The introduction of learnable meta tokens as a dynamic cache initialization mechanism is also novel, drawing a parallel to human metamemory functions.\n3. the experiments are extensive and well-documented, including ablation studies that thoroughly evaluate the impact of each component, such as the hybrid heads and meta tokens. The benchmarks are comprehensive and competitive, providing a robust demonstration of Hymba's capabilities.", "weaknesses": "1. It would be even better if the effectiveness of the Hymba could be validated on image or speech modalities.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "ETs3mbMXQ3", "reviewer": "Reviewer_EGZg", "summary": "This paper introduces Hymba,  a small language model that combines attention mechanisms with SSMs in a hybrid-head architecture. Authors did the following, \n\n1. A hybrid-head architecture that processes inputs through parallel attention and SSM heads in each layer, leveraging attention's high-resolution recall and SSM's efficient context summarization\n\n2. Learnable meta tokens prepended to input sequences that act as learned cache initialization to modulate token processing\n\n3. Optimization techniques including local/global attention combination and cross-layer KV sharing to improve efficiency\n\nThe authors validate their approach through extensive experiments showing Hymba1.5B achieves comparable performance to larger models while being 3x faster and using 15x less cache yielding memory gain.", "questions": "1. The interplay between SSM and Attention, see weakness \n\n2. Consider a more fair comparison for efficiency and performance gain? see weakness\n\n3. How does the model perform for long context tasks as this is the where the gain of Hymba goes significant?", "limitations": "", "rating": 8, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. Empirical study is comprehensive and clear, particularly the ablation studies;\n\n2. Consistent gains for models with different sizes under 1.5B\n\n3. The paper is easy to follow", "weaknesses": "In general I think this is a strong paper. I have the following comments and questions. \n\n- Some implementation details can be added\n1. I am a bit lost while reading the cache optimization, and meta-token, maybe worth more explanations or pseudocode?\n\n- How many meta tokens are needed, and how they are related to the performance in downstream tasks?\n\n- The ratio between SSM and Attention is not clear. And I understand that recent papers demonstrate that it is important to integrate attention for linear RNN models, but attention layer still added overhead, though coupled with all those techniques. A fair comparison could be a pure attention model with all methods proposed in this paper, comparing their efficiency gain and performance curve.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "EcAxFIBowN", "reviewer": "Reviewer_nHjx", "summary": "This paper introduces Hymba, a hybrid architecture that combines Transformer and SSM within a single layer. The authors also propose several additional techniques to enhance further efficiency and performance, such as combining global and local attention, cross-layer KV cache sharing, and introducing meta-tokens, which act as a learned prefix. Through extensive evaluation, the authors show that Hymba performs best among small LMs while being significantly more efficient.", "questions": "1. Does the Samba baseline in Figure 3 also use the same number of global attention layers? If not, how can we tell if the performance gain comes from the parallel design or the introduction of global attention layers?\n2. How does the parallel design impact the throughput? Empirically, are the SSM heads and Attention heads computed in parallel or sequentially? (e.g., if you forward the input through the SSM heads, then forward the input through the attn heads, and then aggregate them, then the implementation is done sequentially, even if the design is conceptually ‘parallel’) Would ‘true’ parallel computation require a specialized GPU kernel?\n3. Is the concept of meta-tokens useful for SSMs only, or would general Transformer models also benefit from the technique?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The paper introduces several useful designs that could be empirically used for training hybrid Transformer + SSM models.\n2. The proposed method shows high performance, outperforming most small LMs while achieving better computation efficiency.\n3. The authors perform extensive evaluations across diverse tasks and setups.a", "weaknesses": "1. Limited novelty. The paper seems to suggest a combination of implementation tricks rather than proposing a significant idea.\n2. The hybrid head design seems to be the most significant component of the proposed method, but the evaluation justifying its efficacy is confusing. In Figure 3, the authors compare their method with Samba and claim they achieved a larger ERF. However, it is unclear if the gain comes from the parallel design or the introduction of global attention heads (not present in Samba).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "VzmCCbhfJw", "reviewer": "Reviewer_wExb", "summary": "This paper introduces a new hybrid model named Hymba that integrates attention mechanisms with SSMs in a hybrid-head manner. The main difference between existing models like Samba is that Hymba enables hybrids to operate in parallel rather than sequentially.\n\nThe authors progressively propose several augmentations to the hybrid-head framework, including local/global attention, KV cache sharing, and meta tokens. They found that Hymba performs strongly against existing pure attention-based models and other hybrid models. By training the small-scale Hymba on trillions of tokens, Hymba performs well on common benchmarks and achieves near-perfect results on NIAH tests.", "questions": "* In Table 2, why are the ARC-C scores reported for 25 shots? I believe the common choice is zero shot.\n* I am curious about how the throughputs in Table 2 are measured. Given that RWKV6 is reported to be much faster than others, this does not match my impressions. What is the input to the model? Can A100 GPUs with 80GB of memory handle an input size of 128 * 8K?", "limitations": "", "rating": 8, "confidence": 5, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "* This is a very solid work, with numerous experimental results and ablation studies verifying the effectiveness of Hymba, which is quite convincing.\n* The results on NIAH tests are very impressive, especially given the model scales.", "weaknesses": "* I suggest the authors to add discussions with InfiniteTransformer, which fuses LA and attn in similar manners (Eq. 10)\n* Why not the authors conduct experiments on Mamba2 rather than Mamba?\n* If possible, I suggest the authors to add some discussions with more existing linear attention works like RetNet/GLA/HGRN2/YOCO\n\nInfiniteTransformer: Efficient Infinite Context Transformers with Infini-attention", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Spotlight)