Hymba: A Hybrid-head Architecture for Small Language Models | Summary

This article explains the key points of Hymba: A Hybrid-head Architecture for Small Language Models.

• 2024-11-20 (arXiv)
• Dong, Xin, Fu, Yonggan, Diao, Shizhe, Byeon, Wonmin, Chen, Zijia, Mahabaleshwarkar, Ameya Sunil, Liu, Shih-Yang, Van Keirsbilck, Matthijs, Chen, Min-Hung, Suhara, Yoshi, et al.
• NVIDIA, Georgia Institute of Technology, USA, HKUST
• Paper
• Github

Read this article in Korean

Summary

• Small language models must retrieve earlier input accurately while limiting the memory and computation needed to retain it. Hymba-1.5B-Base scores 61.06% versus Llama-3.2-3B’s 59.74% on the paper’s seven-task average, with a reported 11.67× smaller cache and 3.49× throughput on an NVIDIA A100 at 8k sequence length, starting with a batch size of 128 and halving it for models that exhaust memory.
• Hymba processes each layer’s input through attention heads, which retrieve detailed information from earlier tokens, and state space model (SSM) heads, which recurrently summarize context in a fixed-size state. It combines these parallel branches with 128 learned prefix embeddings called meta tokens, local attention restricted to a recent token window in most layers, and cross-layer key-value (KV) sharing, which reuses cached attention keys and values between consecutive layers.
• In controlled experiments with 1B parameters and the same 100B training tokens, where B denotes billion, Hymba achieves 54.57% on the eight-task general average versus 52.83% for Samba, a model that stacks attention and SSM operators sequentially. Hymba does not lead every individual task, and conflicting training specifications in the PDF leave details of reproduction unresolved.

1. Introduction

Transformers retain detailed token information through attention caches, but full attention has quadratic sequence-processing cost and cache requirements that grow with input length. Mamba-style SSMs summarize history recurrently with fixed-size memory, at the expense of detailed recall; the paper argues that sequentially alternating these operators can create bottlenecks when a layer cannot preserve information needed later.

• Hymba places both operators in the same layer so they process the same input and contribute complementary representations.
• The released model identifiers are Hymba-1.5B-Base and Hymba-1.5B-Instruct, available at https://huggingface.co/nvidia/Hymba-1.5B-Base and https://huggingface.co/nvidia/Hymba-1.5B-Instruct.

2. Hymba: The Proposed Hybrid-Head Architecture

The design roadmap separates architecture development from scaling. Initial ablations, experiments that change individual design components, use 300M parameters and 100B training tokens, where M denotes million, to test parallel fusion, cache optimization, and meta tokens; later experiments scale the model and training data and extend the context from 2K to 8K tokens.

• Table 1 reports 45.19% general accuracy and 49.90% recall accuracy for parallel fusion, versus 44.07% and 45.16% for sequential fusion.
• The final roadmap configuration reports 60.64% general accuracy and 68.79% recall accuracy. These averages use different task groupings from Table 2’s seven-task public-model comparison.

2.1. A Fused Hybrid-Head Module

A shared input projection produces the attention queries, keys, and values together with the SSM features and output gate, which modulates the SSM contribution. Because the measured SSM outputs have consistently larger magnitudes, Hymba normalizes and rescales both branches with learned channel-wise vectors before fusion and a final output projection.

• Equation (3) expresses the fusion as Y = W_out_proj(β₁ norm(M_attn X̃) + β₂ norm(M_ssm X̃)), where M_attn and M_ssm denote the attention and SSM operators on the prefix-augmented input X̃.
• The paper describes the selected fusion as a mean; its displayed equation writes a weighted sum without an explicit averaging factor.
• Removing a branch from an individual layer affects tasks differently. On 1000 HellaSwag samples, a sentence-completion task, the reported average accuracy reductions are 0.24% for attention removal and 1.1% for SSM removal, with especially severe degradation after removing the first-layer SSM branch.

Both branches receive features derived from the same prefix-augmented input and are normalized before fusion. The memory analogy distinguishes detailed attention retrieval, recurrent context summaries, and learned prefix information; it is a conceptual interpretation of the operators.

Parallel fusion improves the general and recall averages over sequential fusion, while local attention and cache sharing deliver most of the throughput and cache improvements. The table averages eight general tasks and two recall tasks, which differ from the seven-task public-model comparison.

2.2. KV Cache Optimization

Sliding Window Attention (SWA), which restricts attention to a recent token window, replaces full attention in most layers. Hymba retains global attention in the first, middle, and last layers and shares keys and values between consecutive layers to reduce cache storage and redundant projections.

• Replacing every attention layer with SWA reduces recall accuracy from 49.90% to 29.78% in Table 10. Restoring three global-attention layers recovers 48.79%, with throughput rising from 876.7 to 2399.7 tokens/s and cache falling from 148.24 to 41.19 MB.
• Cross-layer sharing then reaches 2756.5 tokens/s with a 39.42 MB cache, while general accuracy increases from 44.56% to 45.16% and recall changes from 48.79% to 48.04%.
• The SSM branch permits substantial cache reduction, but the all-local configuration shows that it does not remove the need for global attention in these experiments.

The stack preserves global attention at three positions while repeating local-attention blocks elsewhere. Sharing cached keys and values between consecutive local blocks reduces storage while retaining the parallel SSM path.

2.3. Meta Tokens

Meta tokens are 128 input-independent embeddings learned jointly with the model and prepended to every sequence. Their fixed prefix permits offline computation at inference, making them a learned initialization for the attention cache and SSM state; one shared set is used for all tasks.

• The authors hypothesize that these tokens absorb attention sinks, positions that receive substantial attention despite limited semantic relevance, and provide an outlet for softmax attention’s requirement to distribute its probability mass across available positions.
• Last-layer attention profiles differ for article, math, and code prompts. This shows varying use of the learned prefix across the sampled domains; the interpretation that individual tokens encode particular world knowledge remains a hypothesis.
• In Table 10, adding meta tokens to the compressed hybrid raises general accuracy from 45.16% to 45.53% and recall from 48.04% to 51.79%, while throughput changes from 2756.5 to 2695.8 tokens/s and cache from 39.42 to 40.01 MB. Table 1 instead lists 45.59% for the corresponding general-accuracy result.

Article, math, and code prompts produce different attention profiles over the 128 learned prefix positions. This shows varying use of meta tokens across the sampled domains, without identifying their precise semantic contents.

2.4. Hymba Attention Map

The attention-map analysis divides token contributions into Meta, attention to learned prefix tokens; BOS, attention to the beginning-of-sequence token; Self, attention to the current token; and Cross, attention to other real tokens. For SSMs, the authors compute implicit attention maps, token-to-token contributions derived from the recurrent operator, and normalize each row to sum to one.

• On text from Oliver Twist Chapter 29, Hymba’s local attention allocates approximately 33% to Meta and 67% to Cross. Its SSM branch allocates approximately 6% to Meta, 25% to Cross, and 69% to Self.
• Llama-3.2-3B allocates approximately 56% to BOS, 40% to Cross, and 4% to Self in the illustrated comparison. These patterns show complementary branch behavior on the analyzed text; the paper does not establish that these proportions hold across tasks.

Hymba’s attention branch distributes its illustrated contribution primarily between meta tokens and other real tokens, while its SSM branch emphasizes the current token. The comparison with Llama-3.2-3B and Jamba shows complementary branch behavior in the analyzed example; SSM contributions are represented through normalized implicit attention maps.

2.5. Hymba Model Family

The Hymba family contains 125M, 350M, and nominally 1.5B parameter models. Section 2.5 describes a mixture of 1T DCLM-Baseline-1.0 tokens, 250B SmolLM-Corpus tokens, and 50B proprietary high-quality tokens, where T denotes trillion, with a Warmup-Stable-Decay schedule that increases, holds, and then decreases the learning rate between 3e-3 and 1e-5.

• The main text specifies a 2M-token batch and 2k training sequences until the final 100B tokens, when sequences increase to 8k and the rotary positional encoding (RoPE) base, a parameter controlling position encoding, changes.
• Figure 8 presents 1T tokens of general pretraining followed by 500B tokens of learning-rate annealing, consistent with the 1.5T total reported in the headline results. The listed data-mixture counts total 1.3T, leaving an unresolved discrepancy.
• Table 11 gives 24, 32, and 32 blocks and hidden sizes of 512, 768, and 1600 for the three model sizes. All use three full-attention layers, a 1024-token local window, and SSM state size 16.

3. Model Evaluations

Evaluation covers released base models, controlled architecture comparisons, and instruction-tuned models, which receive additional training to follow requests. The task groups test language modeling, commonsense reasoning, question answering, retrieval from supplied context, mathematical reasoning, function calling, and role-playing.

3.1. Experiment Settings

The public-model benchmark compares Hymba with existing small transformers and hybrid models whose training data and token budgets differ. The controlled benchmark trains architectures from scratch with matched parameter counts, layer counts, data, hyperparameters, and codebase to assess architectural effects more directly.

• Public base-model evaluation uses lm-evaluation-harness, a language-model benchmark runner. Table 2 uses five demonstrations for MMLU, a knowledge benchmark, one for SQuAD-C, a question-answering completion variant, and zero for the other listed tasks.
• Throughput is measured using PyTorch on an NVIDIA A100 with 8k sequences and an initial batch size of 128. If a model exhausts memory, the batch size is repeatedly halved until measurement succeeds; cache comparisons assume FP16, a 16-bit floating-point format.
• The throughput comparison therefore reflects feasible batch sizes under this procedure. The paper does not report corresponding measurements on a mobile device.

3.2. Benchmark with SOTA Small LMs

Hymba-1.5B achieves the highest Table 2 average, 61.06%, among the evaluated sub-2B models, compared with 60.04% for SmolLM2-1.7B and 59.74% for Llama-3.2-3B. The paper reports 19.91× cache reduction and 2.79× throughput versus SmolLM2, and 11.67× cache reduction and 3.49× throughput versus Llama-3.2-3B.

• Table 2 lists 664 tokens/s and a rounded 79 MB cache for Hymba, 238 tokens/s and 1573 MB for SmolLM2, and 191 tokens/s and 918 MB for Llama-3.2-3B. These rounded entries do not reproduce every reported ratio exactly.
• The 1.02 and 1.32 differences in average accuracy are percentage-point differences. Hymba’s reported 1.5T training tokens are fewer than SmolLM2’s 11T and Llama-3.2-3B’s 9T, but differing data and training recipes prevent attributing the public-model gap solely to architecture.
• The ranking concerns the average: Qwen2.5 scores higher on MMLU, SmolLM2 leads ARC-E, the science-question subset labeled Easy, Phi-1.5 leads Winogrande, a pronoun-resolution task, and h2o-danube2 leads HellaSwag among the sub-2B models. Hymba leads ARC-C, the science-question subset labeled Challenge, PIQA, a physical-commonsense task, and SQuAD-C.

Hymba leads the seven-task average at 61.06% and combines 664 tokens/s with a rounded 79 MB cache. Individual columns qualify the aggregate ranking, including Qwen2.5’s stronger MMLU score and Llama-3.2-3B’s stronger Winogrande and HellaSwag scores.

The pipeline separates base pretraining and learning-rate annealing from supervised adaptation and preference optimization. Its 1T-plus-500B pretraining schedule agrees with the reported 1.5T total but differs from the 1.3T data-mixture total described elsewhere.

The plots place Hymba near the highest illustrated commonsense accuracy while using less cache than SmolLM2 and achieving higher throughput. The six-task commonsense average differs from Table 2’s seven-task average, and the measurements use the stated A100, 8k-sequence conditions.

3.3. Benchmark Different Architectures Under The Same Setting

At 1B parameters, all controlled models are trained for 100B SmolLM-Corpus tokens using the same recipe and evaluated without demonstrations. Hymba obtains 54.57% on the eight-task general average versus Samba’s 52.83%, and 49.50% on the two-task recall average versus Llama3’s 47.33%.

• The recall tasks are SWDE, extraction of semi-structured relations from raw HTML websites, and SQuAD-C, answer completion using supplied context. Hymba scores 54.29% and 44.71%; Llama3 scores 75.95% and 18.70%, revealing substantial task-specific differences beneath the average.
• Hymba also has the lowest reported perplexity, a measure of uncertainty in predicting text, on both language-modeling evaluations labeled Wiki. and LMB.: 18.62 and 10.38, respectively.
• Needle-in-the-haystack testing inserts an informative sentence into a longer document and asks the model to retrieve it. The matched 1B models are pretrained at 1k sequence length, fine-tuned at 4k, and tested up to 16k; Hymba’s heatmap shows broader successful retrieval across positions and lengths than Mamba2 and Llama3, although the figure supplies no single aggregate score.

Matched 1B models trained on the same 100B tokens provide stronger evidence for an architectural advantage than the public-model comparison. Hymba leads both aggregate accuracy groups, but Llama3’s 75.95% SWDE result exceeds Hymba’s 54.29%, showing substantial variation within the recall average.

The heatmaps compare retrieval across document lengths and insertion positions after fine-tuning at 4k. Hymba shows broader successful retrieval, while Mamba2 struggles with earlier insertions and Llama3 shows position-dependent failures beyond its fine-tuning range.

3.4. Instruction-tuned Model

Post-training combines supervised full fine-tuning, which updates the model using example responses, and Direct Preference Optimization (DPO), which trains from preferred and dispreferred responses. The specified learning rates are 5e-5 and 3e-6, respectively; Hymba-1.5B-Instruct reaches a 49.22% five-benchmark average versus Qwen2.5’s 47.30%.

• Hymba scores 58.76% on GSM8K, a math-word-problem benchmark, 31.03% on GPQA, a question-answering benchmark, and 46.40% on Berkeley Function-Calling Leaderboard v2 (BFCLv2), which tests function-call generation. It trails Qwen2.5 on MMLU and Llama-3.2 on IFEval, an instruction-following evaluation.
• BFCLv2 uses live_simple, live_multiple, live_parallel, live_parallel_multiple, and live_relevance, excluding live_irrelevance, the category in which no function call is required. The reported score is specific to this selection.
• With DoRA, a parameter-efficient adaptation method that trains and stores a small set of additional parameters, Hymba reaches 40.0% instruction generalization and 37.9% role generalization on RoleBench, a role-playing evaluation. Table 5 reports 35.5% and 33.5% for RoleLlama-7B, with baseline results taken from prior work.

Hymba’s 49.22% average exceeds Qwen2.5’s 47.30%, with higher math, GPQA, and selected function-calling scores. Its lower MMLU score than Qwen2.5 and lower IFEval score than Llama-3.2 show that the aggregate lead does not extend to every capability.

Hymba-DoRA scores 40.0% on instruction generalization and 37.9% on role generalization, exceeding the listed RoleLlama-7B results by 4.5 and 4.4 percentage points. The baseline results come from prior work, giving this comparison weaker experimental control than Table 3.

4. Related Works

The related-work discussion connects Hymba to transformer language models, efficient recurrent architectures, and hybrids that combine attention with SSMs. Jamba, Zamba, and Samba combine these operators through sequential structures; Hymba processes them in parallel within the same fused module.

• The learned prefix also connects to prompt tuning, which adapts a model through learned prompt embeddings, and prior learned memory tokens. Hymba learns its prefix jointly during pretraining and converts its fixed contribution into inference-time initialization.

5. Conclusion

The paper concludes that parallel attention and SSM processing, learned meta tokens, and cache optimization jointly improve accuracy and efficiency in the reported evaluations. Controlled comparisons support the architectural contribution, while cache ablations identify configurations that preserve much of the recall accuracy and others that substantially reduce it.

Appendix

• A. Extensive Benchmark for More Hymba Model Variants and A.1. Comparison with SOTA Tiny LMs at 350M and 125M Scales report LightEval, a benchmark runner, averages of 49.35% for Hymba-125M and 55.34% for Hymba-350M, versus 48.44% and 53.56% for the corresponding SmolLM baselines. These results extend the reported average-performance advantage to smaller models, using an evaluation setup that differs from Table 2.
• A.2. Evaluating Hymba-1.5B Trained on Public Data Only uses 1T DCLM-Baseline-1.0 tokens followed by 500B SmolLM-Corpus tokens. The public-data model averages 60.81%, versus 61.06% for the main model and 60.04% for SmolLM2; its MMLU score falls from 51.19% to 44.31%, while several other task scores improve.
• A.3. Apple-to-Apple Comparison with Other Architectures at 300M Scale trains models on 100B FineWeb tokens with a peak learning rate of 5e-4, warmup and cosine decay, 1K sequences, and a 256-token window where local attention applies. Hymba achieves 45.53% general accuracy versus Llama3’s 44.08%, and 51.79% recall accuracy versus Llama3’s 39.98%.
• B. Ablation Studies of Our Hymba Architecture, including Parallel vs. Sequential fusion, compares effective receptive field, an empirical measure of the distance over which tokens propagate information, against cache size. The parallel design has a larger measured receptive field than the sequential alternative at a comparable cache size. Table 10 further shows that increasing attention capacity improves recall while increasing memory cost, and concatenation does not improve the general-task average over the selected mean fusion.
• C. Head Importance Analysis removes attention or SSM branches one layer at a time and evaluates 1000 samples per task without demonstrations. The reported accuracy reductions averaged across layers and tasks are 0.46% for attention removal and 1.2% for SSM removal, with the first-layer SSM particularly important. These interventions establish branch importance for the evaluated models and tasks.
• D. Meta Tokens: More Analysis and Visualization relates the learned prefix to prompt tuning, memory tokens, and vision-transformer registers, learned tokens used to retain global information. Figure 15 shows lower attention-map entropy, meaning more concentrated token contributions, in all illustrated attention layers and 10 of 12 SSM layers after adding meta tokens; the proposed link between concentration and improved retrieval remains interpretive.
• E. Pretraining and Post-training Implementation Details specifies 128 NVIDIA A100 GPUs, 1% warmup, a 20% learning-rate decay phase, and a 2M-token training batch. It also specifies one epoch each for full fine-tuning and DPO, a global batch size of 1024, and LMFlow, a fine-tuning toolkit. Reproduction requires resolving differences between this appendix and the main text: 1.3T versus 1.5T pretraining tokens, the appendix’s 2K sequence description versus the main text’s final 8k extension, and post-training block sizes of 2048 versus 8192.

Brief Thoughts

The controlled comparisons support parallel fusion as an architectural contribution: its advantage persists with matched model size, data, and training recipe, and the cache ablations expose the cost of removing global attention. The memory analogy and attention visualizations help interpret the design, but do not identify world knowledge stored in particular meta tokens. Deployment assessment requires measurements on the intended hardware and workload, while reproduction requires clarification of the conflicting training specifications.