Summary (Overview)

  • Priming is a method for initializing Hybrid State-Space Models (SSMs) from pre-trained Transformers, requiring less than 0.5% of the source model's pre-training token budget. It enables exploring the Hybrid architecture design space without costly from-scratch training, turning it into a knowledge transfer problem.
  • The method is grounded in realization theory, which provides a principled framework for determining when and why an SSM layer can approximate an Attention layer's behavior, guiding layer selection and the initialization and alignment stages.
  • The paper establishes an empirical performance hierarchy among SSMs: Gated KalmaNet (GKA) > Gated DeltaNet (GDN) > Mamba-2, consistent with their theoretical expressiveness. All SSM-based Hybrids also outperform Hybrids using Sliding Window Attention (SWA), even with a 4× larger KV cache.
  • At scale, a GKA-based Primed Hybrid (32B) achieves a +3.8 average percentage-point gain over its source Transformer on reasoning benchmarks while being up to ~2.3× faster in decode throughput and roughly halving the memory footprint, enabling ~2× more concurrent sequences.
  • The paper also introduces a training-free context extension method (state composition) that extends native context from 128K to 256K tokens.

Introduction and Theoretical Foundation

The paper addresses the prohibitive cost of exploring Hybrid architecture design space. Hybrid State-Space Models combine Attention layers with recurrent SSM layers to balance two forms of memory:

  • Eidetic memory from Attention, which stores every token verbatim in a Key-Value (KV) cache that grows linearly with context length, at quadratic cost.
  • Fading memory from SSMs, which compresses the entire history into a fixed-dimensional latent state at linear cost.

Conventionally, exploring the design space for Hybrid models (choosing layer ratios, placement, SSM types) requires training each configuration from scratch, costing trillions of tokens and GPU-hours. The paper introduces Priming as a solution: initializing a Hybrid model directly from existing pre-trained Transformer weights.

Its theoretical foundation is classical realization theory from system identification. The paper formalizes both Attention and SSMs as applying a T×TT \times T mixing matrix to a token sequence. For an SSM, this matrix is semi-separable with an order equal to its state dimension. The central theoretical result (Theorem 3.1) states that a causal Attention map can be exactly realized by a state-space model with a state dimension related to the Hankel rank of its mixing matrix.

This theory has two crucial implications:

  1. It guides layer selection: high-Hankel-rank Attention layers (which require complex, global retrieval) are poor candidates for SSM realization and are retained as Attention. Low-Hankel-rank layers are good candidates and are replaced with SSMs.
  2. It motivates a multi-stage pipeline: Priming does not construct a realization per input; it amortizes the realization across the input distribution through a short alignment phase and further end-to-end adaptation.

Methodology

The Priming pipeline consists of several distinct stages, each a pragmatic approximation to the realization-theoretic ideal:

  1. Layer Selection: To identify which layers to keep as Attention, it uses an empirical proxy for Hankel rank: each layer is individually replaced with a Sliding Window Attention (SWA) layer of small window size, and the degradation on a set of long-context benchmarks (e.g., Synthetic Recall, RAG, ICL) is measured. Layers with high degradation (high importance) are those with high Hankel rank and are retained as Attention. This importance-based layer selection outperforms a uniform (evenly-spaced) pattern by 30% (relative) on long-context aggregate performance.

  2. Stage 0: Knowledge Transfer (Weight Initialization). This phase exploits the structural correspondence between Attention and SSM projections. The Query, Key, Value, and Output projections from the source Attention layers are transferred directly to the SSM layers. Other SSM-specific parameters (e.g., gating, transition dynamics) are initialized randomly. All MLPs, layer norms, embeddings, and LM head are copied. This provides a well-aligned starting point, not the exact realization.

  3. Stage 1: Layerwise Alignment. This phase aligns each new SSM layer to the Attention layer it replaces, amortizing the realization across the data distribution. The student Hybrid is trained on web text to replicate the teacher Transformer's final hidden states using a Mean-Squared Error (MSE) loss:

    Lmin=XHybridXTransformer2.\mathcal{L}_{\mathrm{min}} = \left\| \mathbf{X}_\ell^{\mathrm{Hybrid}} - \mathbf{X}_\ell^{\mathrm{Transformer}} \right\|^2 .

    This is done with a memory-efficient fused architecture (teacher and student share all non-SSM layers).

  4. Stage 2: Task Adaptation Following alignment, the Hybrid model is fine-tuned end-to-end with a standard next-token prediction loss on task-specific data (e.g., instruction-following, chain-of-thought reasoning).

  5. State Expansion (Adaptive GQA): A key trick to restore SSM capacity when inheriting Grouped Query Attention (GQA) structure, which reduces per-head state capacity. AGQA adds a cheap, low-rank projection to map shared heads to slightly more distinct ones to recover performance.

The paper experimentally validates that End-to-end MSE outperforms Layerwise MSE for alignment, and Adaptive GQA outperforms standard GQA at 32B scale.

Methodology

The core experimental methodology is driven by the controlled comparison made possible by Priming. The authors meticulously compare different SSM families, Hybrid ratio, layer placement, and layer types under identical conditions (same source model, data, hyperparameters, and Hybrid ratio of 50%). This uniformity allows for a precise, head-to-head evaluation.

The main research focuses on:

  • Three pure SSM families: Mamba-2, Gated DeltaNet (GDN), and Gated KalmaNet (GKA), which differ in their transition/gating operators.
  • One Hybrid layer: B’MOJO-F, which fuses an SSM with a Sliding Window Attention.
  • Two model scales: 8B and 32B, using the Qwen3 family as the source Transformer.
  • Two target task families: instruction-tuned (IT) models for standard retrieval, QA, and generation; and Reasoner models for chain-of-thought, math, science, and coding tasks.

Training recipes: 8B and 32B models are trained with targeted objectives. Stage 1 uses ~40B tokens of mixed behavior data. Stage 2 for IT models uses long-context continued pre-training (28B tokens) and SFT (1.6B tokens). The reasoning recipe uses a multi-stage approach (SFT → context extension → instruction alignment), consuming a total of ~150B tokens.

Empirical Validation / Results

The paper provides extensive comparisons and ablations.

  • Controlled SSM comparison (At 8B scale, 50% ratio, same data, same recipe). The performance ranking is GKA > GDN for long-context and reasoning, and GKA > GDN > Mamba2 for IT. All models using SSMs are within 3.5% (relative) of the Transformer baseline, while Mamba2 is consistently worse. This establishes the empirical hierarchy.

  • SSM vs. SWA: The all SSM-based Hybrids outperform the SWA-based Hybrid model by at least 5% (relative) on aggregate long-context performance, even though the SWA window has a KV cache 2-4 times larger than the SSM state. This demonstrates the advantage of compressed fading memory over windowed verbatim memory.

  • Scaling to 8B and 32B reasoning models:

    • GKA-Primed-HQwen3-Reasoner (8B) matches or outperforms the GDN-based counterpart uniformly across all tasks (by up to +156% relative on SciCode).
    • At 1B-3B scale, the GKA-based Hybrid Reasoner trails the Transformer Reasoner-SFT (following the same recipe alignment) by only 1% (relative) in aggregate, while also achieving higher performance on some tasks.
    • At 32B scale, GKA-Primed-HQwen3-Reasoner achieves a +3.8 average percentage-point gain over its source Transformer, remaining within 1% of the same Transformer post-trained with the same recipe.
  • Inference Efficiency: A key advantage comes from the smaller KV cache (SSM layers require none) and fixed-size recurrent state. At 128K context, the GKA-Primed-HQwen3-32B delivers up to 2.06× decode throughput (Table 13, 14 of the paper). It also enables ~2× more concurrent sequences on the same hardware. On AIME 2025, the GKA model reaches the same target accuracy threshold 1.6× faster than the Transformer, which is critical for RL post-training speed.

  • GKA's Runtime Compute-Quality Knob: Unlike Mamba-2 and GDN, GKA solves an online ridge-regression problem via an iterative Chebyshev solver. By trading accuracy for compute at serving time (reducing iteration count from 30 to 10), a single trained model can adjust its speed and accuracy without retraining, a significant practical advantage.

  • '. Decode throughput': Up to 2.3× for GDN Primed models at 128K context. The advantage grows with context length.

theoretical and practical implications

  • Theoretical: This work provides strong experimental evidence that the expressiveness of the state update is a crucial design dimension for SSMs when combined with Attention, and directly correlates with downstream task performance. It provides an empirical, large-scale validation of the theoretical framework that Hybrids are a better design to complement eidetic memory with a compressed fading memory than with a larger window of verbatim tokens, confirming this with a relationship and unifying the perspective of the design space.

  • Practical: Priming is a general, cost-effective recipe that can be ported to any open-weights Transformer (Qwen, Llama, Mistral) and any SSM layer, turning architectural design from a pre-training problem into a knowledge transfer problem. It is a critical leap forward for the efficient deployment of long-context LLMs.

Conclusion

The paper presents Priming, a foundational framework for building Hybrid State-Space models at scale. It definitively validates that SSM-Attention hybrids are a viable and even superior alternative to pure Transformers for reasoning and long-context tasks, with large gains in inference efficiency.

The work has three main future directions:

  • Scaling beyond 32B: To test efficiency and quality at 100B+ scale.
  • Higher Hybrid ratios: To explore the design trade-off between qfizz and SSM layers.
  • Integration in RL loops and dynamic iteration scheduling: To test the throughput advantage directly improves reinforcement learning efficiency, and to further use the GKA iteration knob adaptively.

Related papers