Zamba Model: Hybrid SSM-Transformer

• Zamba Model is a hybrid architecture that combines state-space model (SSM) layers with a shared, minimal Transformer self-attention mechanism for efficient language processing.
• It employs an 80-block SSM stack interleaved with shared attention every six layers, reducing compute overhead, memory footprint, and inference latency compared to pure Transformer models.
• The models leverage open pretraining corpora and a two-phase curriculum to optimize performance, making them competitive on benchmarks and valuable for research exploration.

Zamba denotes a family of compact, open-weight hybrid models that combine state-space model (SSM) principles (primarily the Mamba/Mamba2 architecture) with minimal Transformer self-attention. The initial 7B parameter Zamba model and its successors in the Zamba2 suite represent a paradigm for efficient, high-performing language modeling at moderate parameter scales, outperforming pure-transfomer architectures of similar and larger sizes in throughput, inference memory footprint, and performance per training token. Developed by the Zyphra team, Zamba and Zamba2 models have been openly released with full training checkpoints and, for Zamba2, their pretraining corpora.

1. Architecture: Hybrid SSM-Transformer Design

Zamba pioneered an 80-block stack based primarily on Mamba SSM layers interleaved with a single shared Transformer-style self-attention module applied recurrently every six blocks, totaling thirteen such applications in the model. Each standard block consists of:

• Residual-channel dimension $d=3712$
• Mamba "memory" state dimension $p=16$
• 1D convolution mixing dimension $c=4$

The predominant forward pass is governed by SSM recurrence:

$$\begin{aligned} h_{t+1} &= e^{A\,\delta_t}\,h_t + B_t\,x_t \ y_t &= C_t\,h_{t+1} \end{aligned}$$

where $A\in\mathbb{R}^{p\times p}$ is fixed per-layer; $(\delta_t, B_t, C_t)$ are token-dependent and computed from input through small dense and convolutional projections.

A global shared attention (GSA) module is inserted every $N=6$ Mamba blocks, using a single multi-head self-attention+MLP parameter set across all 13 slots. At each invocation:

• The current residual $x_l$ is concatenated with the input embedding $x_0$, giving $[x_l \| x_0]\in\mathbb{R}^{2d}$
• Standard multi-head attention and an MLP with inner width $p=16$0 are applied, followed by a linear projection and a skip connection.

This architectural pattern minimizes the parameter and compute overhead of the attention mechanism, as each GSA instance reuses weights and stores only a single set of keys/values per sequence for generation, as opposed to per-layer caching in conventional transformers (Glorioso et al., 2024).

Zamba2 models generalize this approach by adopting Mamba2 blocks (for increased throughput) and, in larger variants, using two alternating shared transformer blocks, each equipped with rotary position embeddings (RoPE). The use of attention LoRA adapters is reserved for the 1.2B parameter variant, targeting optimal expressivity vs. overhead at small scale (Glorioso et al., 2024).

Parameterization

Model	Layers (D)	Width (W)	Params (B)	Shared Attention Blocks
Zamba-7B	80	3712	7.0	1 (every 6 blocks)
Zamba2-1.2B	32	1792	1.2	1 (with LoRA)
Zamba2-2.7B	40	2560	2.7	2
Zamba2-7.4B	48	3840	7.4	2

The total parameter count is approximated by $p=16$1 for $p=16$2 layers and channel width $p=16$3, accounting for the backbone, transformer modules, embeddings, and output heads (Glorioso et al., 2024).

2. Training Corpus and Curriculum

Zamba's pretraining corpus is drawn from open datasets including RefinedWeb, Pile, C4, peS2o, Cosmopedia, and arXiv, yielding 950B tokens. Filtering comprises LSH-minhash deduplication, removing 10% of Pile and 0.5% of C4. The pretraining sequence context is 4096 tokens, using the Mistral vocabulary and a batch size of 512 on 128 H100 GPUs in BF16 precision.

A two-phase curriculum is employed:

• Phase 1: Cosine learning rate (LR) decay from $p=16$4 to $p=16$5, with 1% warmup.
• Phase 2 (Annealing): ~50B tokens on high-quality and LM-generated synthetic data (OpenOrca, code/math benchmarks, etc), 60% web corpus replay, 40% new data. Learning rate re-warm to $p=16$6, then exponential decay to $p=16$7 over $p=16$8 steps, $p=16$9 (Glorioso et al., 2024).

Zamba2 extends this regimen using the Zyda-2 dataset, a cross-deduplicated, model-filtered corpus aggregating Zyda-1, DCLM, FineWeb-Edu3, and Dolma. Zyda-2 provides ~5T post-filtering tokens for training. Zamba2-7.4B undergoes 2T phase-1 tokens, followed by 100B tokens of annealing (two epochs of a 60:40 mix, as above). The curriculum incorporates strict minibatch replay of phase-1 samples during annealing to mitigate catastrophic forgetting. All Zyda-2 data uses normalization, document and sentence-level deduplication, and alignment to Mistral-compatible tokenization (Glorioso et al., 2024).

3. Optimization and Implementation Details

Zamba is optimized with AdamW ($c=4$0, weight decay $c=4$1), 1.0 gradient clipping, activation checkpointing, and ZeRO-1 sharding across multiple GPUs.

Zamba2 introduces architectural performance enhancements:

• Replacement of Mamba1 by Mamba2 blocks for approximately $c=4$2 higher throughput.
• Two shared transformer blocks (except for the 1.2B model), greatly reducing transformer FLOPs vs full attention at each layer.
• Rotary position embeddings (RoPE) integrated into all shared attention blocks (except Zamba2-2.7B interim checkpoints).
• Post-hoc LoRA adapters for MLP subblocks in larger models.

Learning rates in Zamba2 follow a standard cosine decay for phase-1 ($c=4$3), and a warmed-restart rapid cosine decay in phase-2. Precise learning rates per variant are ablation-derived (Glorioso et al., 2024).

For deployment, weights and checkpoints are available via Hugging Face under Apache-2.0. Inference is accelerated via optimized Mamba kernels (BF16), FlashAttention v2, and NVIDIA Transformer Engine RMSNorm. Recommended hardware is a single H100 or A100 with $c=4$440GB VRAM to support extended sequence lengths.

4. Empirical Evaluation

Zamba demonstrates competitive or superior performance to leading 7B open-weight transformer models across multiple zero-shot and few-shot benchmarks. Zamba2 not only outperforms earlier Zamba variants but also surpasses contemporary models on major tasks per training token.

Selected benchmark results (1T tokens for Zamba, 2-3T for Zamba2):

Model	PIQA	ARC-E	ARC-C	BoolQ	Wino	Hella	OpenB	MMLU
Llama2 7B	79.05	74.58	46.16	77.68	69.06	75.97	44.2	45.88
Mistral 7B	82.26	79.59	53.92	83.64	73.88	81.07	44.2	62.30
Gemma 7B	81.12	80.77	54.27	83.12	73.80	80.46	45.2	62.85
Zamba (Phase 1)	79.00	70.74	43.90	77.50	70.80	77.14	44.2	50.82
Zamba (Annealed)	81.37	74.50	46.48	83.60	76.40	80.24	44.6	57.72
Zamba2-7.4B	81.00	82.00	68.90	84.10	77.30	81.50	48.20	67.20

Ablation studies showed (i) the addition of a single shared attention block every $c=4$56 layers consistently outperformed both pure Mamba and parameter-matched “Mambaformer” hybrids with per-layer attention; (ii) the annealing phase was critical for boosting MMLU and reasoning metrics (+7 points MMLU post-annealing). Some benchmarks (e.g., ARC-challenge, BoolQ) realized significant gains from annealing; others, such as OpenBookQA, showed limited sensitivity to this phase (Glorioso et al., 2024).

On efficiency, Zamba achieves a forward inference latency of $c=4$6 ms/token (sequence length $c=4$7) on dual H100 GPUs, outperforming Llama2-7B ($c=4$8 ms) and Mistral-7B ($c=4$9 ms). The memory footprint for KV-caching is $$\begin{aligned} h_{t+1} &= e^{A\,\delta_t}\,h_t + B_t\,x_t \ y_t &= C_t\,h_{t+1} \end{aligned}$$08GB at $$\begin{aligned} h_{t+1} &= e^{A\,\delta_t}\,h_t + B_t\,x_t \ y_t &= C_t\,h_{t+1} \end{aligned}$$1, compared to 18–20GB for baseline transformer models, a direct consequence of the minimal number of cached attention instances (Glorioso et al., 2024). Zamba2-7.4B further reduces time-to-first-token by up to 30–50% and achieves $$\begin{aligned} h_{t+1} &= e^{A\,\delta_t}\,h_t + B_t\,x_t \ y_t &= C_t\,h_{t+1} \end{aligned}$$26\times smaller KV-cache than pure transformer baselines, enabling $$\begin{aligned} h_{t+1} &= e^{A\,\delta_t}\,h_t + B_t\,x_t \ y_t &= C_t\,h_{t+1} \end{aligned}$$3K token contexts on commodity GPUs (Glorioso et al., 2024).

5. Instruction Tuning and Downstream Capabilities

Instruction-tuned versions of Zamba2 ("Instruct" variants) are released across all scales, trained on approximately 4M supervized samples via prompt templates such as ChatML. Further refinement is conducted via Direct Preference Optimization (DPO) using about 200k preference pairs. On MT-Bench and IFEval, Zamba2-Instruct models achieve parity or better robustness on unseen prompt formats compared to Llama-3 and Mistral models of similar scale (Glorioso et al., 2024).

6. Open Source Artifacts and Research Utility

All Zamba v1 and v2 model weights (pretrain, annealed, instruct-tuned), tokenizer files, and the full Zyda-2 corpus (for Zamba2) are released on Hugging Face under permissive licensing. Zamba’s open checkpoints at every phase enable detailed investigation of large-scale SSM vs. transformer learning dynamics, phase-wise curriculum effects, and scaling laws for hybrid architectures.

Optimization and evaluation scripts, as well as full training configuration, are provided in the Zyphra Cookbook resource. The availability of full pretraining datasets and extensive checkpoints establishes Zamba as an empirical anchor in open research on efficient, hybrid LLMs (Glorioso et al., 2024, Glorioso et al., 2024).

7. Research Outlook

Zamba2 extends the original Zamba’s principles by leveraging improved SSM kernels (Mamba2), architectural search balancing loss-per-FLOP and parameter, and a scalable, openly released data curriculum. Future directions enumerated by the Zyphra team include:

• Optimized SSM/attention fusion
• Expanded synthetic-data pretraining curricula
• Mixture-of-Experts (MoE) SSM hybrid models
• Distillation from larger models
• Extended context-window scaling

A plausible implication is that the Zamba model line will serve both as a production-ready open-weight series and as a platform for algorithmic and theoretical advances in SSM-transformer hybrid architectures (Glorioso et al., 2024).