Apriel-H1: Distilling Efficient Reasoning Models via Mamba Hybrids

ServiceNow AI has developed Apriel-H1, a family of 15B reasoning models that convert a full-attention teacher model into a Mamba hybrid architecture. This approach achieves up to 2.1x throughput increases with minimal loss in reasoning quality, demonstrating that efficiency can be retrofitted into existing strong models through targeted distillation.

The Core Insight: Capability-Matched Distillation

Distilling a reasoning model requires using data that matches the specific capability being preserved, rather than general-purpose pretraining data. Initial attempts to distill using pretraining data or a mix of pretraining and SFT data resulted in dramatic losses in reasoning quality.

The successful approach utilized high-quality reasoning traces from the teacher's SFT dataset, including:

  • Multi-step math proofs
  • Coding tasks with clear logical dependencies
  • Scientific analysis with detailed explanatory chains

This is necessary because replacing attention mechanisms (such as retrieval and induction heads) with Mamba's linear recurrence disrupts the computational paths used for reasoning. High-quality reasoning traces provide the explicit structure needed for the hybrid model to discover new paths to the same outcomes. Additionally, the team used reverse KL divergence (temperature 1) for distillation, as its mode-seeking behavior encourages the student model to commit to the high-confidence, structured predictions of the teacher.

Apriel-H1 Performance and Benchmarks

The flagship model, Apriel-H1-15b-Thinker-SFT, achieves 2.1x throughput compared to the full-attention teacher. The impact on reasoning benchmarks is as follows:

Benchmark Teacher Score Apriel-H1-15b-Thinker-SFT Score
MATH500 0.90 $\rightarrow$ 0.92 Improvement
MTBench 8.30 $\rightarrow$ 8.58 Improvement
GSM8k 0.97 $\rightarrow$ 0.95 Slight Regression
GPQA 0.59 $\rightarrow$ 0.55 Slight Regression
AIME24 0.70 $\rightarrow$ 0.65 Slight Regression

Across the Apriel-H1 family, checkpoints range from 25 to 40 Mamba layers (out of 50 total). While the H-30 variant offers the best balance, the H-40 variant maximizes throughput (up to 3.4x) for latency-critical workloads.

Staged Distillation Process

Converting a full-attention model to a hybrid cannot be done in a single step. ServiceNow AI employed a three-stage procedure:

Stage 1: Initial Layer Replacement

Using Leave-One-Out (LOO) analysis on MMLU, the team identified the 25 least-important layers. These were replaced with Mamba-in-Llama (MIL) initialized mixers and distilled end-to-end.

Stage 2: Progressive Conversion

Because LOO analysis becomes unreliable beyond 25 layers, the team implemented MIL-Mamba-Replacement (MMR). This dynamic heuristic initializes a Mamba mixer, runs 100 training steps, and records the distillation loss. Layers that converge to a lower loss are prioritized for replacement. This process was performed incrementally (25 $\rightarrow$ 27 $\rightarrow$ 30 $\rightarrow$ 34 $\rightarrow$ 37 $\rightarrow$ 40 layers).

Stage 3: Final SFT

After reaching the target Mamba layer count, a final SFT pass was conducted until reasoning performance stabilized. The final Apriel-H1-15b-Thinker-SFT model was produced after 55.9B distillation tokens and 20.9B SFT tokens.

Implementation and Reproducibility

Apriel-H1 was built using Fast-LLM, an open-source (Apache 2.0) training framework that treats attention and Mamba as modular mixing interfaces. This allows the layer order to be specified via a pattern field in the configuration, enabling the free swapping of block types.

Production Deployment

Apriel-H1 is integrated with Hugging Face Transformers and vLLM. The vLLM integration utilizes Mamba cache operations for continuous batching, prefix caching, and chunked prefill. The authors note that while throughput gains are significant, deploying hybrid models currently requires custom code and careful numerical validation due to maturing tooling.

Technical Specifications

  • Mamba Hyperparameters: State size 16, DT rank 16, inner dimension 4096.
  • Architecture: Hybrid of Mamba-1 and Attention layers.
  • Training Data: 76.8B total tokens for the flagship model.

Sources