Source-linked AI summary
Zamba: A Compact 7B SSM Hybrid Model
Paolo Glorioso, Quentin Anthony, Yury Tokpanov, James Whittington, Jonathan Pilault, Adam Ibrahim, Beren Millidge
TL;DR
Memory costs can constrain serving language models on local devices. Zamba combines a Mamba backbone with shared attention and uses a two-phase training approach; it matches leading 7B models on many linguistic evaluations while slightly lagging on reasoning and in-context learning.
Problem
For local-device language-model serving, loading models into memory can be a greater bottleneck than forward-pass FLOPs.
Source: §I · p. 1Method
Zamba combines a Mamba backbone with a global shared attention layer and uses general pretraining followed by high-quality-data annealing with rapid learning-rate decay.
Source: A. Overview · p. 3, §II · p. 4Results
Zamba matches state-of-the-art 7B models on many linguistic evaluations but lags slightly on reasoning and in-context-learning tests.
Source: A. Overview · p. 3Takeaways & Limitations
The shared-attention design provides a way to boost performance at small and constant parameter and memory cost.
Source: A. Overview · p. 3
Abstract
from arXiv · showhide
In this technical report, we present Zamba, a novel 7B SSM-transformer hybrid model which achieves competitive performance against leading open-weight models at a comparable scale. Zamba is trained on 1T tokens from openly available datasets and is the best non-transformer model at this scale. Zamba pioneers a unique architecture combining a Mamba backbone with a single shared attention module, thus obtaining the benefits of attention at minimal parameter cost. Due to its architecture, Zamba is significantly faster at inference than comparable transformer models and requires substantially less memory for generation of long sequences. Zamba is pretrained in two phases: the first phase is based on existing web datasets, while the second one consists of annealing the model over high-quality instruct and synthetic datasets, and is characterized by a rapid learning rate decay. We open-source the weights and all checkpoints for Zamba, through both phase 1 and annealing phases.
I. INTRODUCTION
The introduction frames the search for architectures that can scale without transformers’ quadratic attention cost. It motivates state-space models and hybrids, while noting efficiency trade-offs in serving and the use of high-quality data during training.
I. INTRODUCTION
- Transformers scale across model size, data, and compute, but their core attention computation has quadratic cost, motivating alternatives that preserve performance and scaling.The introduction leaves open whether transformers are the only architecture able to leverage scale.Source: §I · p. 1
- SSMs use linear dynamical systems for sequence mixing, enabling parallel-scan training and recurrent generation with linear-time, constant-memory inference.Their fixed-size hidden state avoids inference costs quadratic in context length and a linearly growing KV cache.Source: §I · p. 1
- Recent SSMs are closing the performance gap, but reported weaknesses in expressivity and in-context learning motivate combining SSM and attention blocks.A study cited in the introduction found roughly one quarter self-attention layers optimal among the hybrid designs it tested.Source: §I · p. 1
- MoE models trade greater parameter and memory costs for fewer active FLOPs, a disadvantage when loading models on memory-limited local devices is the bottleneck.This motivates exploring performance gains from additional FLOPs rather than extra parameters.Source: §I · p. 1
- The introduction describes a two-phase curriculum: broad web-data pretraining followed by rapid learning-rate decay on higher-quality instruct and synthetic data.The proposed rationale is that rapid decay can focus training on scarce, newer high-quality data.Source: §I · p. 2
A. Overview
Zamba combines a Mamba backbone with shared attention and trains through general pretraining followed by high-quality-data annealing. The report presents its hybrid design, training approach, and open checkpoints as contributions.
A. Overview
- Zamba approaches leading ∼7B models on many linguistic evaluations, while lagging slightly on reasoning and in-context learning tests.The authors suggest the gap may relate to the substantial difference in training data volume.Source: A. Overview · p. 3
- Zamba combines a Mamba backbone with a global shared self-attention layer, aiming to retain attention’s retrieval and in-context-learning benefits with small, constant parameter and memory cost.The shared block is repeated every six Mamba blocks, with shared parameters and input embeddings concatenated to the residual stream.Source: A. Overview · p. 3, Figure 2 · p. 3
- The architecture is inspired by the cortex and hippocampus analogy, with different processing layers sending information to a single shared memory store.The authors connect shared layers that spend additional FLOPs to improve performance with recurrent-model work.Source: A. Overview · p. 3
- Zamba’s two-phase training uses general pretraining followed by high-quality-data annealing with rapid decay; rewarming before exponential decay improved some evaluations but had little effect on others.The report releases phase-1 and annealed models, as well as checkpoints throughout both phases.Source: A. Overview · p. 3
B. Contributions
- The contributions include a 7B Transformer–SSM hybrid intended to preserve SSM FLOP efficiency alongside attention’s in-context-learning ability.The report also claims shared attention preserves the modeling performance of independent attention blocks while saving memory.Source: B. Contributions · p. 3
- The authors report replicating two-phase training methods at large scale and provide phase-1 and annealing checkpoints.They present open checkpoints as a resource for studying learning dynamics.Source: B. Contributions · p. 3, A. Overview · p. 3
II. MODEL
Zamba combines a Mamba backbone with a global shared self-attention block every N Mamba blocks. The design shares attention weights and reuses initial residual activations to provide attention while reducing memory costs.
II. MODEL
- Unlike transformers, which separate sequence mixing and per-token processing, Mamba combines both in one block.Mamba’s input-dependent control and observation matrices and timestep distinguish its linear dynamical system; its gating and projections also perform token processing.Source: §II · pp. 3–4
- Zamba inserts a global shared self-attention block every N Mamba blocks into an otherwise pure Mamba backbone.The block contains self-attention and an MLP with shared weights, reducing attention-parameter and KV-cache memory costs.Source: §II · p. 4
- Each shared-attention call concatenates the layer’s residual activations with the initial residual activations before attention.The authors report this gives a consistent performance improvement; the resulting query, key, and value vectors have double dimension before projection back to residual width.Source: §II · p. 4
- Early small-model ablations found consistent improvement over pure Mamba and significant gains over parameter-matched Mambaformer (Park et al., 2024).The authors suggest self-attention parameters may be partly redundant and describe the design as trading additional inference FLOPs for fewer parameters while maintaining performance.Source: §II · p. 4
A. Phase 1
Phase 1 pretrained Zamba-v1 on open web datasets for approximately 950 billion tokens, using a deliberately gradual cosine learning-rate decay ahead of annealing.
A. Phase 1
- 950 billion tokens were used to train Zamba-v1 in phase 1 on open web datasets.Training lasted approximately 30 days on 128 H100 GPUs.Source: A. Phase 1 · p. 4
- The learning rate followed a cosine decay from 1.5 ∗10−4 to 7.5 ∗10−5, less steeply than common schedules to leave room for rapid annealing decay.The model used a sequence length of 4096 tokens and the Mistral tokenizer.Source: A. Phase 1 · pp. 4–5
- Training used data and tensor parallelism across two ranks, Zero-1 optimizer-state sharding, and activation checkpointing.Parallelism was applied across attention, MLP, and Mamba blocks.Source: A. Phase 1 · pp. 4–5
B. Datasets
The phase-1 corpus combines open web and other datasets, with filtering and fuzzy deduplication applied to selected components.
B. Datasets
- The phase-1 mixture includes The Pile, RefinedWeb, C4, PeS2o, arxiv, and Cosmopedia.Table I reports each sub-dataset’s proportion of the full 1T-token dataset, whose datasets are open and available.Source: B. Datasets · p. 5, Table I · p. 5
- Filtering and 13-gram word-based MinHash deduplication removed roughly 10% of Pile-dedup and 0.5% of C4-en.The filters applied to Pile and C4-en included document-length and character-composition thresholds; deduplication targeted 50% Jaccard similarity.Source: B. Datasets · p. 5
- Deduplication checked datasets against themselves and one another, but RefinedWeb was checked only against other datasets and Cosmopedia was unprocessed.The authors cite RefinedWeb’s existing deduplication and Cosmopedia’s synthetic nature as reasons for those choices.Source: B. Datasets · p. 5
- The full dataset contained just over 1 trillion tokens, of which phase 1 used 950 billion, with no upsampling or multiple epochs.Table II reports the percentages of documents marked as duplicates by fuzzy LSH-Minhash deduplication.Source: B. Datasets · p. 5, Table II · p. 5
C. Annealing
Annealing continued training on a mixture of high-quality and pretraining data with rapid learning-rate decay; the authors found rewarming and exponential decay preferable in their evaluations.
C. Annealing
- 60% pretraining data and 40% new annealing data formed the annealing mixture, drawn from over 100 high-quality datasets.The new data included math, code, instruction-finetuning, and synthetic examples; most datasets were used for one epoch, with select subsets used for two.Source: C. Annealing · p. 5
- Rewarming the learning rate before annealing performed better than starting at the prior run’s final learning rate, and exponential decay outperformed linear decay.The schedule used γ = 0.25, with η0 = 1.1 ∗10−4 and ηT = 10−7.Source: C. Annealing · p. 5
- The exponential annealing schedule uses ηt = Ae−t/(γT ) + B, with γ = 0.25.Here t is the iteration number and T is the total annealing iterations; A and B are set by the learning rates after warmup and at annealing’s end.Source: C. Annealing · p. 5
- Loss on the original pretraining dataset decreased during annealing, but less than reported for miniCPM.The authors suggest model size, token count, annealing datasets, or schedule differences could account for the contrast.Source: C. Annealing · p. 5
III. PERFORMANCE
Zamba performs strongly against open-checkpoint models and is competitive on several language-modeling benchmarks, while slightly trailing leading models overall. Its inference speed and KV-cache memory use compare favorably with similar-scale models.
III. PERFORMANCE
- Zamba performs well for its size, generally leads open-dataset competitors, and slightly lags leading models overall.It is often competitive with leading models on PIQA, WinoGrande, and HellaSwag, but its code performance is relatively poor.Source: §III · pp. 5–6
- Zamba’s forward pass is significantly faster than competing 7B-scale models, with the advantage increasing at longer sequence lengths.Its SSM backbone also substantially reduces KV-cache memory relative to similar-scale models, supporting more efficient generation and longer contexts on one device.Source: §III · p. 6
- One attention block is applied 13 times in Zamba-7B, with independent activations and KV-cache entries at each invocation.The report identifies Zamba as the best-performing open-checkpoint model available at the time of writing.Source: §III · p. 6
- Zamba’s phase 1 and annealed models significantly outperform competing open-checkpoint models on the aggregate evaluation.The aggregate is the mean of PIQA, ARC-Easy, ARC-Challenge, HellaSwag, WinoGrande, BoolQ, and OpenBookQA.Source: Figure 3 · p. 6, §III · p. 6
IV. RELATED WORK
State-space models offer efficient sequence processing but have fixed-size memory and restricted sequence mixing. Zamba’s one-attention-layer hybrid addresses reported SSM weaknesses in in-context learning and algorithmic tasks.
A. State-space-models
- OLMo 1.0 MMLU and all OLMo 1.7 evaluation results were unavailable because of incompatibilities with evaluation-harness v0.4.0.Source: A. State-space-models · p. 6
- Figure 4 inference and generation measurements use two H100 GPUs with tensor parallelism and batch size one; competitors were evaluated using vLLM (Kwon et al., 2023).Source: Figure 4 · p. 7
- SSMs use linear dynamics and fixed-size memory, enabling convolution or parallel-scan computation but restricting sequence mixing and expressive power.The passage contrasts their fixed-memory approach with transformer attention’s sequence-length cost.Source: A. State-space-models · pp. 6–7
- Recent work reports potential pure-SSM deficiencies in in-context learning and algorithmic tasks; Zamba proposes that one attention layer may suffice in a hybrid.The passage notes that many prior hybrids add multiple attention layers.Source: A. State-space-models · p. 7
B. Annealing
Zamba uses an annealing phase after initial training, following a two-phase schedule discussed by earlier released models. The report also situates its open checkpoints alongside efforts to make pretraining models and training details available.
B. Annealing
- Zamba follows a two-phase approach with high-quality data in a rapid exponential decay schedule and learning-rate rewarmup from zero.Its dataset mix differs from earlier approaches discussed in the passage.Source: B. Annealing · p. 7
C. Open checkpoints
- Pythia, OLMo, LLM360, and Zyphra are cited as efforts releasing training checkpoints or details, with Zyphra’s open checkpoints covering a 370M pure Mamba model.The passage distinguishes Pythia’s scientific purpose and performance from later pretraining-democratization efforts.Source: C. Open checkpoints · p. 7
V. DISCUSSION
The discussion presents Zamba as a competitive 7B SSM with shared attention. It reports that annealing narrows the gap to leading models while leaving the gap’s causes and the interpretation of annealing gains unresolved.
V. DISCUSSION
- Zamba approaches leading models with a relatively small budget—about $200k, seven researchers, and one month—and uses a shared global attention block.The authors present this as evidence that approaching state-of-the-art LLM pretraining need not require vast budgets or teams.Source: §V · p. 8
- After phase 1, Zamba reaches Llama2-level performance on 1T tokens versus Llama2’s at least 2T, and its base model is comparable on MMLU.The authors offer dataset differences or architectural advantages as possible explanations, not established causes.Source: §V · p. 8
- The remaining gap may reflect token volume and data quality: Zamba used 1T tokens and straightforward filtering and deduplication, while performance was still increasing near training’s end.The authors note that competitors used substantially more tokens and that more advanced dataset-preparation techniques exist.Source: §V · p. 8
- Annealing significantly improves many evaluation scores and closes about half or more of the Llama2-to-state-of-the-art gap, especially on MMLU and ARC.The authors say the gap may narrow further with an improved annealing phase on more synthetic data or pretraining on more tokens.Source: §V · p. 8
- Whether annealing gains reflect true capability improvements or more evaluation-friendly responses remains unresolved, as do the optimal schedule, replay fraction, and dataset composition.The authors release the base model and annealing checkpoints to support study of these questions.Source: §V · p. 8
- As small-scale data-scaling benefits become marginal, the authors argue that architectural innovation matters; they release pretraining and annealing checkpoints for scientific study.The released checkpoints cover both training phases, with one checkpoint every 2,500 steps.Source: §V · p. 8