Transformers, with their attention-based architecture, have become the dominant choice for language models (LMs) due to their strong performance, parallelization capabilities, and long-term recall through key-value (KV) caches. However, their quadratic computational cost and high memory demands pose efficiency challenges. In contrast, state space models (SSMs) like Mamba and Mamba-2 offer constant complexity and efficient hardware optimization but struggle with memory recall tasks, affecting their performance on general benchmarks.
NVIDIA researchers recently proposed
Hymba
, a family of small language models (SLMs) featuring a hybrid-head parallel architecture that integrates transformer attention mechanisms with SSMs to achieve both enhanced efficiency and improved performance. In Hymba, attention heads provide high-resolution recall, while SSM heads enable efficient context summarization.
The novel architecture of Hymba reveals several insights:
Overhead in attention:
Over 50% of attention computation can be replaced by cheaper SSM computation.
Local attention dominance:
Most global attention can be replaced by local attention without sacrificing performance on general and recall-intensive tasks, thanks to the global information summarized by SSM heads.
KV cache redundancy:
Key-value cache is highly correlated across heads and layers, so it can be shared across heads (group query attention) and layers (cross-layer KV cache sharing).
Softmax attention limitation:
Attention mechanisms are constrained to sum to one, limiting sparsity, and flexibility. We introduce learnable meta-tokens that are prepended to prompts, storing critical information and alleviating the “forced-to-attend” burden associated with attention mechanisms.
This post shows that Hymba 1.5B performs favorably against state-of-the-art open-source models of similar size, including Llama 3.2 1B, OpenELM 1B, Phi 1.5, SmolLM2 1.7B, Danube2 1.8B, and Qwen2.5 1.5B. Compared to Transformer models of similar size, Hymba also achieves higher throughput and requires 10x less memory to store cache.
Hymba 1.5B is released to the
Hugging Face
collection and
GitHub
.
Hymba 1.5B performance
Figure 1 compares Hymba 1.5B against sub-2B models (Llama 3.2 1B, OpenELM 1B, Phi 1.5, SmolLM2 1.7B, Danube2 1.8B, Qwen2.5 1.5B) in terms of average task accuracy, cache size (MB) relative to sequence length, and throughput (tok/sec).
Figure 1. Performance comparison of Hymba 1.5B Base against sub-2B models
In this set of experiments, the tasks include MMLU, ARC-C, ARC-E, PIQA, Hellaswag, Winogrande, and SQuAD-C. The throughput is measured on an NVIDIA A100 GPU with a sequence length of 8K and a batch size of 128 using PyTorch. For models encountering out of memory (OOM) issues during throughput measurement, the batch size was halved until the OOM is resolved to measure the maximal achievable throughput without OOM.
Hymba model design
SSMs such as Mamba were introduced to address the quadratic complexity and large inference-time KV cache issues of transformers. However, due to their low-resolution memory, SSMs struggle with memory recall and performance. To overcome these limitations, we propose a road map for developing efficient and high-performing small LMs in Table 1.
Configuration
Commonsense reasoning (%) ↑
Recall (%) ↑
Throughput (token/sec) ↑
Cache size (MB) ↓
Design reason
Ablations on 300M model size and 100B training tokens
Transformer (Llama)
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
A. + Attention heads (sequential)
44.07
45.16
776.3
156.3
Enhance recall capabilities
B. + Multi-head heads (parallel)
45.19
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
D. + KV cache sharing
45.16
48.04
2756.5
39.4
Cache efficiency
E. + Meta-tokens
45.59
51.79
2695.8
40.0
Learned memory initialization
Scaling to 1.5B model size and 1.5T training tokens
F. + Size / data
60.56
64.15
664.1
78.6
Further boost task performance
G. + Extended context length (2K→8K)
60.64
68.79
664.1
78.6
Improve multishot and recall tasks
Table 1. Design road map of the Hymba model
Fused hybrid modules
Fusing attention and SSM heads in parallel within a hybrid-head module outperforms sequential stacking, according to the ablation study. Hymba fuses attention and SSM heads in parallel within a hybrid head module, enabling both heads to process the same information simultaneously. This architecture improves reasoning and recall accuracy.
Figure 2. The hybrid-head module in Hymba
Efficiency and KV cache optimization
While attention heads improve task performance, they increase KV cache requirements and reduce throughput. To mitigate this, Hymba optimizes the hybrid-head module by combining local and global attention and employing cross-layer KV cache sharing. This improves throughput by 3x and reduces cache by almost 4x without sacrificing performance.
Figure 3. Hymba model architecture
Meta-tokens
A set of 128 pretrained embeddings prepended to inputs, functioning as learned cache initialization to enhance focus on relevant information. These tokens serve a dual purpose:
Mitigating attention drain by acting as backstop tokens, redistributing attention effectively
Encapsulating compressed world knowledge
Figure 4. Interpretation of Hymba from the memory aspect
Model analysis
This section presents an apples-to-apples comparison across different architectures under the same training settings. We then visualize the attention maps of SSM and Attention in different pretrained models. Finally, we perform head importance analysis for Hymba through pruning. All the analyses in this section help to illustrate how and why the design choices for Hymba are effective.
Apples-to-apples comparison
We performed an apples-to-apples comparison of Hymba, pure Mamba2, Mamba2 with FFN, Llama3 style, and Samba style (Mamba-FFN-Attn-FFN) architectures. All models have 1 billion parameters and are trained from scratch for 100 billion tokens from SmolLM-Corpus with exactly the same training recipe. All results are obtained through lm-evaluation-harness using a zero-shot setting on Hugging Face models. Hymba performs the best on commonsense reasoning as well as question answering and recall-intensive tasks.
Table 2 compares various model architectures on language modeling and recall-intensive and commonsense reasoning tasks, with Hymba achieving strong performance across metrics. Hymba demonstrates the lowest perplexity in language tasks (18.62 for Wiki and 10.38 for LMB) and solid results in recall-intensive tasks, particularly in SWDE (54.29) and SQuAD-C (44.71), leading to the highest average score in this category (49.50).
Model
Language (PPL) ↓
Recall intensive (%) ↑
Commonsense reasoning (%) ↑
Mamba2
15.88
43.34
52.52
Mamba2 w/ FFN
17.43
28.92
51.14
Llama3
16.19
47.33
52.82
Samba
16.28
36.17
52.83
Hymba
14.5
49.5
54.57
Table 2. Comparison of architectures trained on 100 billion tokens under the same settings
In commonsense reasoning and question answering, Hymba outperforms other models in most tasks, such as SIQA (31.76) and TruthfulQA (31.64), with an average score of 54.57, slightly above Llama3 and Mamba2. Overall, Hymba stands out as a balanced model, excelling in both efficiency and task performance across diverse categories.
Attention map visualization
We further categorized elements in the attention map into four types:
Meta:
Attention scores from all real tokens to meta-tokens. This category reflects the model’s preference for attending to meta-tokens. In attention maps, they are usually located in the first few columns (for example, 128 for Hymba) if a model has meta-tokens.
BOS:
Attention scores from all real tokens to the beginning-of-sequence token. In the attention map, they are usually located in the first column right after the meta-tokens.
Self:
Attention scores from all real tokens to themselves. In the attention map, they are usually located in the diagonal line.
Cross:
Attention scores from all real tokens to other real tokens. In the attention map, they are usually located in the off-diagonal area.
The attention pattern of Hymba is significantly different from that of vanilla Transformers. In vanilla Transformers, attention scores are more concentrated on BOS, which is consistent with the findings in Attention Sink. In addition, vanilla Transformers also have a higher proportion of Self attention scores. In Hymba, meta-tokens, attention heads, and SSM heads work complementary to each other, leading to a more balanced distribution of attention scores across different types of tokens.
Specifically, meta-tokens offload the attention scores from BOS, enabling the model to focus more on the real tokens. SSM heads summarize the global context, which focuses more on current tokens (Self attention scores). Attention heads, on the other hand, pay less attention to Self and BOS tokens, and more attention to other tokens (that is, Cross attention scores). This suggests that the hybrid-head design of Hymba can effectively balance the attention distribution across different types of tokens, potentially leading to better performance.
Figure 5. Schematics of the attention map of Hymba as a combination of meta-tokens, sliding window attention, and Mamba contributions
Figure 6. Sum of the attention score from different categories in Llama 3.2 3B and Hymba 1.5B
Heads importance analysis
We analyzed the relative importance of attention and SSM heads in each layer by removing them and recording the final accuracy. Our analysis reveals the following:
The relative importance of attention/SSM heads in the same layer is input-adaptive and varies across tasks, suggesting that they can serve different roles when handling various inputs.
The SSM head in the first layer is critical for language modeling, and removing it causes a substantial accuracy drop to random guess levels.
Generally, removing one attention/SSM head results in an average accuracy drop of 0.24%/1.1% on Hellaswag, respectively.
Figure 7. The achieved accuracy, measured using 1K samples from Hellaswag, after removing the Attention or SSM heads in each layer
Model architecture and training best practices
This section outlines key architectural decisions and training methodologies for Hymba 1.5B Base and Hymba 1.5B Instruct.
Model architecture
Hybrid architecture:
Mamba is great at summarization and usually closer focuses on the current token, while attention is more precise and acts as snapshot memory. Combining them in parallel merges these benefits, but standard sequential fusion does not. We chose a 5:1 parameter ratio between SSM and attention heads.
Sliding window attention:
Full attention heads are preserved in three layers (first, last, and middle), with sliding window attention heads used in the remaining 90% layers.
Cross-layer KV cache sharing:
Implemented between every two consecutive attention layers. It is done in addition to GQA KV cache sharing between heads.
Meta-tokens:
These 128 tokens are learnable with no supervision, helping to avoid entropy collapse problems in large language models (LLMs) and mitigate the attention sink phenomenon. Additionally, the model stores general knowledge in these tokens.
Training best practices
Pretraining:
We opted for two-stage base model training. Stage 1 maintained a constant large learning rate and used less filtered large corpus data. Continuous learning rate decay was then performed to 1e-5 using high-quality data. This approach enables continuous training and resuming of Stage 1.
Instruction fine-tuning:
Instruct model tuning is performed in three stages. First, SFT-1 provides the model with strong reasoning abilities by training on code, math, function calling, role play, and other task-specific data. Second, SFT-2 teaches the model to follow human instructions. Finally, DPO is leveraged to align the model with human preferences and improve the model’s safety.
Figure 8. Training pipeline adapted for the Hymba model family
Performance and efficiency evaluation
With only 1.5T pretraining tokens, the Hymba 1.5B model performs the best among all small LMs and achieves better throughput and cache efficiency than all transformer-based LMs.
For example, when benchmarking against the strongest baseline, Qwen2.5, which is pretrained on 13x more tokens, Hymba 1.5B achieves a 1.55% average accuracy improvement, 1.41x throughput, and 2.90x cache efficiency. Compared to the strongest small LM trained on fewer than 2T tokens, namely h2o-danube2, our method achieves a 5.41% average accuracy improvement, 2.45x throughput, and 6.23x cache efficiency.
Model
# Para-ms
Train tokens
Token
per sec
Cache
(MB)
MMLU  5-
shot
ARC-E 0-shot
ARC-C 0-shot
PIQA 0-shot
Wino.  0-shot
Hella. 0-shot
SQuAD -C
1-shot
Avg
Open
ELM-1
1.1B
1.5T
246
346
27.06
62.37
19.54
74.76
61.8
48.37
45.38
48.57
Rene
v0.1
1.3B
1.5T
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
30.09
55.85
Smol
LM
1.7B
1T
238
1573
27.06
76.47
43.43
75.79
60.93
49.58
45.81
54.15
Cosmo
1.8B
.2T
244
1573
26.1
62.42
32.94
71.76
55.8
42.9
38.51
47.2
h20
dan-ube2
1.8B
2T
271
492
40.05
70.66
33.19
76.01
66.93
53.7
49.03
55.65
Llama 3.2 1B
1.2B
9T
535
262
32.12
65.53
31.39
74.43
60.69
47.72
40.18
50.29
Qwen
2.5
1.5B
18T
469
229
60.92
75.51
41.21
75.79
63.38
50.2
49.53
59.51
AMD
OLMo
1.2B
1.3T
387
1049
26.93
65.91
31.57
74.92
61.64
47.3
33.71
48.85
Smol
LM2
1.7B
11T
238
1573
50.29
77.78
44.71
77.09
66.38
53.55
50.5
60.04
Llama
3.2 3B
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.5T
664
79
51.19
76.94
45.9
77.31
66.61
53.55
55.93
61.06
Table 2. Hymba 1.5B Base model results
Instructed models
The Hymba 1.5B Instruct model achieves the highest performance on an average of all tasks, outperforming the previous state-of-the-art model, Qwen 2.5 Instruct, by around 2%. Specifically, Hymba 1.5B surpasses all other models in GSM8K/GPQA/BFCLv2 with a score of 58.76/31.03/46.40, respectively. These results indicate the superiority of Hymba 1.5B, particularly in areas requiring complex reasoning capabilities.
Model
# Params
MMLU ↑
IFEval ↑
GSM8K ↑
GPQA ↑
BFCLv2 ↑
Avg. ↑
SmolLM
1.7B
27.80
25.16
1.36
25.67
-*
20.00
OpenELM
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.5B
1.5B
52.79
57.14
58.76
31.03
46.40
49.22
Table 3. Hymba 1.5B Instruct model results
Conclusion
The new Hymba family of small LMs features 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, learnable meta-tokens are introduced to act as a learned cache for both attention and SSM heads, enhancing the model’s focus on salient information. Through the road map of Hymba, comprehensive evaluations, and ablation studies, Hymba sets new state-of-the-art performance across a wide range of tasks, achieving superior results in both accuracy and efficiency. Additionally, this work provides valuable insights into the advantages of hybrid-head architectures, offering a promising direction for future research in efficient LMs.
Learn more about
Hybma 1.5B Base
and
Hymba 1.5B Instruct
.
Acknowledgments
This work would not have been possible without contributions from many people at NVIDIA, including Wonmin Byeon, Zijia Chen, Ameya Sunil Mahabaleshwarkar, Shih-Yang Liu, Matthijs Van Keirsbilck, Min-Hung Chen, Yoshi Suhara, Nikolaus Binder, Hanah Zhang, Maksim Khadkevich, Yingyan Celine Lin, Jan Kautz, Pavlo Molchanov, and Nathan Horrocks.
Reading Time:
4
minutes
Transformer は、その Attention ベースのアーキテクチャによる、強力なパフォーマンス、並列化能力、および KV (Key-Value) キャッシュを通じた長期記憶のおかげで、言語モデル (LM) の主流となっています。しかし、二次計算コストと高いメモリ要求により、効率性に課題が生じています。これに対し、Mamba や Mamba-2 のような状態空間モデル (SSMs) は、複雑さを一定にして効率的なハードウェア最適化を提供しますが、メモリ想起タスクが苦手でそれは一般的なベンチマークでのパフォーマンスに影響を与えています。
NVIDIA の研究者は最近、効率性とパフォーマンスの両方を向上させるために、Transformer の Attention メカニズムを SSM と統合したハイブリッド ヘッド並列アーキテクチャを特徴とする小規模言語モデル (SLM) ファミリである
Hymba
を提案しました。Hymba では、Attention ヘッドが高解像度の記憶能力を提供し、SSM ヘッドが効率的なコンテキストの要約を可能にします。
Hymba の新たなアーキテクチャは、いくつかの洞察を明らかにしています。
Attention のオーバーヘッド:
Attention 計算の 50% 以上を、より安価な SSM 計算に置き換えることができます。
ローカル Attention の優位性:
SSM ヘッドにより要約されたグローバル情報のおかげで、一般的なタスクやメモリ想起に集中するタスクのパフォーマンスを犠牲にすることなく、ほとんどのグローバル Attention をローカル Attention に置き換えることができます。
KV キャッシュ冗長性:
Key-value キャッシュは、ヘッド間とレイヤー間で高い相関性があるため、ヘッド間 (GQA: Group Query Attention) およびレイヤー間 (Cross-layer KV キャッシュ共有) で共有できます。
Softmax の Attention の制限:
Attention メカニズムは、合計が 1 になるように制限されており、疎性と柔軟性に制限があります。NVIDIA は、プロンプトの先頭に学習可能なメタトークンを導入し、重要な情報を格納し、Attention メカニズムに関連する「強制的に Attention を行う」負担を軽減します。
この記事では、Hymba 1.5B が同様の規模である最先端のオープンソース モデル、Llama 3.2 1B、OpenELM 1B、Phi 1.5、SmolLM2 1.7B、Danube2 1.8B、Qwen2.5 1.5B などと比較して、良好なパフォーマンスを発揮することが示されています。同等のサイズの Transformer モデルと比較すると、Hymba はより高いスループットを発揮し、キャッシュを保存するために必要なメモリが 10 分の 1 で済みます。
Hymba 1.5B は
Hugging Face
コレクションと
GitHub
で公開されています。
Hymba 1.5B のパフォーマンス
図 1 は、Hymba 1.5B と 2B 未満のモデル (Llama 3.2 1B、OpenELM 1B、Phi 1.5、SmolLM2 1.7B、Danube2 1.8B、Qwen2.5 1.5B) を、平均タスク精度、シーケンス長に対するキャッシュ サイズ (MB)、スループット (tok/sec) で比較したものです。
図 1. Hymba 1.5B Base と 2B 未満のモデルのパフォーマンス比較
この一連の実験には、MMLU、ARC-C、ARC-E、PIQA、Hellaswag、Winogrande、SQuAD-C などのタスクが含まれています。スループットは、シーケンス長 8K、バッチ サイズ 128 で PyTorch を使用して NVIDIA A100 GPU で測定します。スループット測定中にメモリ不足 (OOM: Out of Memory) 問題が発生したモデルでは、OOM が解決されるまでバッチ サイズを半分にして、OOM なしで達成可能な最大スループットを測定しました。
Hymba モデルのデザイン
Mamba のような SSM は、Transformer の二次的な複雑性と推論時の KV キャッシュが大きい問題に対処するために導入されました。しかし、メモリ解像度が低いために、SSM は記憶想起とパフォーマンスの点で苦戦しています。これらの制限を克服するために、表 1 で効率的で高性能な小規模言語モデルを開発するためのロードマップを提案します。
構成
常識推論 (%) ↑
リコール (%) ↑
スループット (token/sec) ↑
キャッシュ サイズ (MB) ↓
設計理由
300M モデル サイズと 100B トレーニング トークンのアブレーション
Transformer (Llama)
44.08
39.98
721.1
414.7
非効率的ながら正確な記憶
状態空間モデル (Mamba)
42.98
19.23
4720.8
1.9
効率的だが不正確な記憶
A. + Attention ヘッド (連続)
44.07
45.16
776.3
156.3
記憶能力を強化
B. + 複数ヘッド (並列)
45.19
49.90
876.7
148.2
2 つのモジュールのバランスの改善
C. + ローカル / グローバル Attention
44.56
48.79
2399.7
41.2
演算 / キャッシュの効率を向上
D. + KV キャッシュ共有
45.16
48.04
2756.5
39.4
キャッシュ効率化
E. + メタトークン
45.59
51.79
2695.8
40.0
学習した記憶の初期化
1.5B モデル サイズと 1.5T トレーニング トークンへのスケーリング
F. + サイズ / データ
60.56
64.15
664.1
78.6
タスク パフォーマンスのさらなる向上
G. + コンテキスト長の拡張 (2K→8K)
60.64
68.79
664.1
78.6
マルチショットとリコール タスクの改善
表 1. Hymba モデルのデザイン ロードマップ
融合型ハイブリッド モジュール
アブレーション研究によると、ハイブリッド ヘッド モジュール内で Attention と SSM ヘッドを並列にして融合するほうが、シーケンシャルにスタッキングするより優れていることが分かっています。Hymba は、ハイブリッド ヘッド モジュール内で Attention と SSM ヘッドを並列に融合させ、両ヘッドが同時に同じ情報を処理できるようにします。このアーキテクチャは、推論と記憶の正確さを高めます。
図 2. Hymba のハイブリッド ヘッド モジュール
効率性と KV キャッシュの最適化
Attention ヘッドはタスクのパフォーマンスを向上させますが、KV キャッシュの要求を増大させ、スループットを低下させます。これを緩和するために、Hymba はローカルおよびグローバルの Attention を組み合わせ、 Cross-layer KV キャッシュ共有を採用することで、ハイブリッド ヘッド モジュールを最適化します。これにより、パフォーマンスを犠牲にすることなくスループットが 3 倍向上し、キャッシュがほぼ 4 分の 1 に削減されます。
図 3. Hymba モデルのアーキテクチャ
メタトークン
入力の先頭に置かれる 128 の事前学習済みの埋め込みのセットであり、学習済みキャッシュの初期化として機能し、関連情報への注意を強化します。このようなトークンには 2 つの目的があります。
バックストップ トークンとして機能し、Attention を効果的に再分配することで Attention の流出を軽減する
圧縮された世界知識をカプセル化する
図 4. メモリの側面から見た Hymba の解釈
モデル解析
このセクションでは、同一のトレーニング設定における異なるアーキテクチャを比較する方法を紹介します。それから、SSM と Attention の Attention マップを異なる学習済みモデルで可視化し、最後に、剪定 (pruning) を通じて Hymba のヘッド重要度分析を行います。このセクションのすべての分析は、Hymba のデザインにおける選択の仕組みと、それが効果的な理由を説明するのに役立ちます。
同一条件での比較
Hymba、純粋な Mamba2、Mamba2 と FFN、Llama3 スタイル、Samba スタイル (Mamba-FFN-Attn-FFN) のアーキテクチャを同一条件で比較しました。すべてのモデルが 10 億のパラメーターで、まったく同じトレーニング レシピで SmolLM-Corpus から 1,000 億トークンをゼロから学習しています。すべての結果は、Hugging Face モデルでゼロショット設定を使用して lm-evaluation-harness を通じて取得されています。Hymba は、常識推論だけでなく、質問応答タスクや記憶想起タスクでも最高のパフォーマンスを発揮します。
表 2 は、言語モデリングタスクと記憶想起タスクおよび常識推論タスクに関するさまざまなモデル アーキテクチャを比較しており、Hymba はすべての評価基準で卓越したパフォーマンスを達成しています。Hymba は、言語モデリングタスクで最も低い Perplexity を示し (Wiki で 18.62、LMB で 10.38)、特に SWDE (54.29) と SQuAD-C (44.71) の記憶想起タスクにおいて堅実な結果を示し、このカテゴリで最高の平均スコア (49.50) を達成しました。
モデル
言語モデリング (PPL) ↓
記憶想起型 (%) ↑
常識推論 (%) ↑
Mamba2
15.88
43.34
52.52
Mamba2 と FFN
17.43
28.92
51.14
Llama3
16.19
47.33
52.82
Samba
16.28
36.17
52.83
Hymba
14.5
49.5
54.57
表 2. 同じ設定で 1,000 億トークンで学習されたアーキテクチャの比較
常識推論と質問応答において、Hymba は平均スコア 54.57 で、 SIQA (31.76) や TruthfulQA (31.64) などのほとんどのタスクで、Llama3 や Mamba2 をやや上回っています。全体的に、Hymba はバランスの取れたモデルとして際立っており、多様なカテゴリで効率性とタスク パフォーマンスの両方で優れています。
Attention マップの可視化
さらに、Attention マップの要素を 4 つのタイプに分類しました。
Meta:
すべての実トークンからメタトークンへの Attention スコア。このカテゴリは、モデルがメタトークンに Attention を向ける傾向を反映するものです。Attention マップでは、通常、モデルにメタトークンがある場合、最初の数列 (例えば Hymba の場合は 128) に位置しています。
BOS:
すべての実トークンからセンテンスの開始トークンまでの Attention スコア。Attention マップでは、通常、メタトークンの直後の最初の列に位置します。
Self:
すべての実トークンからそれ自身への Attention スコア。Attention マップでは、通常、対角線上に位置しています。
Cross:
すべての実トークンから他の実トークンへの Attention スコア。Attention マップでは、通常、対角線外の領域に位置しています。
Hymba の Attention パターンは、vanilla (加工されていない) Transformer のそれとは大きく異なります。vanilla Transformer の Attention スコアは BOS に集中しており、Attention Sink の結果と一致しています。さらに、vanilla Transformer は、Self-Attention スコアの比率も高くなっています。Hymba では、メタトークン、Attention ヘッド、SSM ヘッドが互いに補完し合うように機能し、異なるタイプのトークン間で、よりバランスの取れた Attention スコアの分布を実現しています。
具体的には、メタトークンが BOS からの Attention スコアをオフロードすることで、モデルがより実際のトークンに集中できるようになります。SSM ヘッドはグローバルなコンテキストを要約し、現在のトークン (Self-Attention スコア) により重点を置きます。一方、Attention ヘッドは、Self と BOS トークンに対する注意が低く、他のトークン (すなわち、Cross Attention スコア) への注意が高くなります。これは、Hymba のハイブリッド ヘッド デザインが、異なるタイプのトークン間の Attention 分布のバランスを効果的に取ることができ、パフォーマンスの向上につながる可能性があることを示唆しています。
図 5. メタトークン、Sliding Window Attention、Mamba 貢献の組み合わせによる Hymba の Attention マップの概略図
図 6. Llama 3.2 3B と Hymba 1.5B の異なるカテゴリからの Attention スコアの合計
ヘッド重要度分析
各レイヤーのAttention と SSM ヘッドの相対的な重要性を分析するために、それぞれを削除して最終的な精度を記録しました。分析の結果、以下のことが明らかになりました。
同じレイヤーの  Attention / SSM ヘッドの相対的な重要性は入力適応であり、タスクによって異なります。これは、さまざまな入力の処理において、異なる役割を果たす可能性があることを示唆しています。
最初のレイヤーの SSM ヘッドは言語モデリングタスクに不可欠で、これを削除すると、ランダム推測レベルにまで大幅に精度が低下します。
一般的に、Attention / SSM ヘッドを 1 つ削除すると、Hellaswag ではそれぞれ平均 0.24%/1.1% 精度が低下します。
図 7. Hellaswag の 1K サンプルを使用して測定した、各レイヤーの Attention または SSM ヘッドを削除した後の達成精度
モデル アーキテクチャと学習のベスト プラクティス
このセクションでは、Hymba 1.5B Base と Hymba 1.5B Instruct の主要アーキテクチャ上の決定事項と学習方法の概要について説明します。
モデル アーキテクチャ
ハイブリッド アーキテクチャ:
Mamba は要約に優れ、通常は現在のトークンにより重点を置きます。Attention はより正確でスナップショット メモリとして機能します。標準的なシーケンシャル融合ではなく、並列に組み合わせることで利点を統合することができます。SSM と Attention ヘッド間のパラメーター比は 5:1 を選択しました。
Sliding Window Attention:
完全な Attention ヘッドは 3 つのレイヤー (最初、最後、中間) に維持され、残りの 90% のレイヤーで Sliding Window Attention ヘッドが使用されます。
Cross-layer KV キャッシュ共有:
連続する 2 つの Attention レイヤー間に実装されます。これは、ヘッド間の GQA KV キャッシュ共有に加えて行われます。
メタトークン:
これらの 128 トークンは教師なし学習が可能であり、大規模言語モデル (LLM) におけるエントロピー崩壊の問題を回避し、Attention Sink 現象を緩和するのに役立ちます。さらに、モデルはこれらのトークンに一般的な知識を格納します。
学習のベスト プラクティス
事前学習:
2 段階のベースモデル学習を選択しました。ステージ 1 では、一定の高い学習率を維持し、フィルタリングされていない大規模なコーパス データの使用しました。続いて、高品質のデータを用いて 1e-5 まで継続的に学習率を減衰させました。このアプローチにより、ステージ 1 の継続的な学習と再開が可能になります。
指示ファインチューニング:
指示モデルの調整は 3 つの段階で行われます。まず、SFT-1 は、コード、数学、関数呼び出し、ロール プレイ、その他のタスク固有のデータで学習を実施し、強力な推論能力をモデルに付与します。次に、SFT-2 はモデルに人間の指示に従うことを教えます。最後に、DPO を活用して、モデルを人間の好みに合わせ、モデルの安全性を高めます。
図 8. Hymba モデル ファミリに適応した学習パイプライン
パフォーマンスと効率性の評価
1.5T の事前学習トークンだけで、Hymba 1.5B モデルはすべての小規模言語モデルの中で最高の性能を発揮し、Transformer ベースの LM よりも優れたスループットとキャッシュ効率を実現します。
例えば、13 倍以上のトークン数で事前学習された最も強力なベースラインである Qwen2.5 に対してベンチマークした場合、Hymba 1.5B は平均精度が 1.55%、スループットが 1.41 倍、キャッシュ効率が 2.90 倍に向上します。2T 未満のトークンで学習された最も強力な小規模言語モデル、すなわち h2o-danube2 と比較すると、この方法は平均精度が 5.41%、スループットが 2.45 倍、キャッシュ効率が 6.23 倍に向上しています。
モデル
パラメーター数
学習トークン
トークン
(1 秒あたり)
キャッシュ
(MB)
MMLU 5-
shot
ARC-E 0-shot
ARC-C 0-shot
PIQA 0-shot
Wino. 0-shot
Hella. 0-shot
SQuAD -C
1-shot
平均
OpenELM-1
1.1B
1.5T
246
346
27.06
62.37
19.54
74.76
61.8
48.37
45.38
48.57
Renev0.1
1.3B
1.5T
800
113
32.94
67.05
31.06
76.49
62.75
51.16
48.36
52.83
Phi1.5
1.3B
0.15T
241
1573
42.56
76.18
44.71
76.56
72.85
48
30.09
55.85
SmolLM
1.7B
1T
238
1573
27.06
76.47
43.43
75.79
60.93
49.58
45.81
54.15
Cosmo
1.8B
.2T
244
1573
26.1
62.42
32.94
71.76
55.8
42.9
38.51
47.2
h20dan-ube2
1.8B
2T
271
492
40.05
70.66
33.19
76.01
66.93
53.7
49.03
55.65
Llama 3.2 1B
1.2B
9T
535
262
32.12
65.53
31.39
74.43
60.69
47.72
40.18
50.29
Qwen2.5
1.5B
18T
469
229
60.92
75.51
41.21
75.79
63.38
50.2
49.53
59.51
AMDOLMo
1.2B
1.3T
387
1049
26.93
65.91
31.57
74.92
61.64
47.3
33.71
48.85
SmolLM2
1.7B
11T
238
1573
50.29
77.78
44.71
77.09
66.38
53.55
50.5
60.04
Llama3.2 3B
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.5T
664
79
51.19
76.94
45.9
77.31
66.61
53.55
55.93
61.06
表 2. Hymba 1.5B ベース モデルの結果
指示モデル
Hymba 1.5B Instruct モデルは、全タスク平均で最高のパフォーマンスを達成し、直近の最高性能モデルである Qwen 2.5 Instruct を約 2% 上回りました。特に、Hymba 1.5B は GSM8K/GPQA/BFCLv2 で、それぞれ 58.76/31.03/46.40 のスコアで他のすべてのモデルを上回っています。これらの結果は、特に複雑な推論能力を必要とする分野において、Hymba 1.5B の優位性を示しています。
モデル
パラメーター数
MMLU ↑
IFEval ↑
GSM8K ↑
GPQA ↑
BFCLv2 ↑
平均↑
SmolLM
1.7B
27.80
25.16
1.36
25.67
-*
20.00
OpenELM
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.5B
1.5B
52.79
57.14
58.76
31.03
46.40
49.22
表 3. Hymba 1.5B Instruct モデルの結果
まとめ
新しい Hymba ファミリの小規模言語モデルは、ハイブリッド ヘッド アーキテクチャを採用し、Attention ヘッドの高解像な記憶能力と SSM ヘッドの効率的なコンテキストの要約を組み合わせています。Hymba のパフォーマンスをさらに最適化するために、学習可能なメタトークンが導入され、Attention ヘッドと SSM ヘッドの両方で学習済みキャッシュとして機能し、顕著な情報に注目するモデルの精度を強化しました。Hymba のロードマップ、包括的な評価、アブレーション研究を通じて、Hymba は幅広いタスクにわたって新たな最先端のパフォーマンスを確立し、正確さと効率性の両面で優れた結果を達成しました。さらに、この研究は、ハイブリッド ヘッド アーキテクチャの利点に関する貴重な洞察をもたらし、効率的な言語モデルの今後の研究に有望な方向性を示しています。
Hybma 1.5B Base
と
Hymba 1.5B Instruct
の詳細はこちらをご覧ください。
謝辞
この成果は、Wonmin Byeon、Zijia Chen、Ameya Sunil Mahabaleshwarkar、Shih-Yang Liu、Matthijs Van Keirsbilck、Min-Hung Chen、Yoshi Suhara、Nikolaus Binder、Hanah Zhang、Maksim Khadkevich、Yingyan Celine Lin、Jan Kautz、Pavlo Molchanov、Nathan Horrocks など、NVIDIA の多くのメンバーの貢献なくしては実現しませんでした。
関連情報
GTC セッション:
Optimizing Large Language Models: An Experimental Approach to Pruning and Fine-Tuning LLama2 7B (大規模言語モデルの最適化: LLama2 7B の剪定とファインチューニングの実験的アプローチ)
GTC セッション:
Accelerating End-to-End Large Language Models System using a Unified Inference Architecture and FP8 (統一推論アーキテクチャと FP8 を用いたエンドツーエンドの大規模言語モデル システムの高速化)
NGC コンテナー:
Llama-3.1-Nemotron-70B-Ins
truct
NGC コンテナー:
Llama-3-Swallow-70B-Instruct-v0.1
SDK:
NeMo Megatron