Amazon's priming method turns Transformers into hybrid models for under 0.5% of pretraining cost

Why hybrid architectures have been hard to study

The Transformer architecture carries a memory cost that grows with the number of tokens it processes. As AI agents run longer sessions and longer reasoning chains, that cost grows with them, according to the Amazon Science post.

Hybrid architectures try to balance memory, speed and quality by combining several methods. These include sparse attention, which picks a subset of past tokens to consider, sliding-window attention (SWA), which keeps only the most recent tokens, and state space models (SSMs), which compress the full history into a fixed-size state.

Until now, exploring this design space has meant training each new architecture from scratch, using trillions of tokens and hundreds of thousands of GPU-hours. That has made controlled comparisons between hybrids extremely expensive, the team writes.

How priming works

Attention gives Transformers what the authors call eidetic memory: verbatim recall of every past token, stored in a growing KV cache. Not every attention layer needs that. Some do long-range retrieval across tens of thousands of tokens, while others work locally and can get by with a compact, fixed-size summary, which the authors call fading memory. Hybrid models use attention layers where global recall matters and state space layers for the rest.

Priming rests on the observation that a Transformer is a special case of a state space realization. If an SSM can produce a mixing matrix close enough to that of an attention layer, the query, key and value projections of the source Transformer can be reused as the SSM's input projections. Only the SSM-specific parameters then need to be learned.

To decide which layers can be converted, the team replaces each layer with SWA and measures the degradation. Layers that tolerate the substitution are candidates for SSM conversion. The pipeline has three stages:

• Layer selection: identify layers that can be compressed into a fixed-size state and keep the rest as attention.
• Weight transfer and alignment: initialize SSM layers from the corresponding attention weights, then run a short distillation process to align each SSM with the layer it replaces.
• End-to-end adaptation: fine-tune the full hybrid model on next-token prediction.

According to the authors, the whole pipeline uses less than 0.5% of the source model's pretraining token budget. They say it works across Transformer families (Qwen, Llama, Mistral), model classes (dense or MoE) and scales from eight billion to 32 billion parameters and more.

Findings from the comparison

The paper compares hybrids built with three state space models: Gated KalmaNet (GKA), Gated DeltaNet (GDN) and Mamba-2. Under controlled conditions, the team reports that GKA performs best, followed by GDN and then Mamba-2. According to the authors, GKA also offers variable-length inference computation without retraining, giving users direct control over their position on the speed-accuracy frontier.

The researchers also report that SSMs consistently outperform SWA, even when SWA gets twice the memory budget. Their summary: a good lossy summary of everything beats a perfect recording of the most recent past.

On speed, a 32-billion-parameter GKA-based hybrid matches its source Transformer on reasoning benchmarks while decoding up to twice as quickly with 128,000 tokens of context. It reaches correct answers in as little as two-thirds of the wall-clock time, the authors say.

Open release and what it means

To support further research, the team has released a zoo of primed hybrid models at the eight-billion and 32-billion-parameter scales, for long-context reasoning and instruction following. The release also includes training code, optimized kernels and a vLLM-serving plugin, all under the Apache 2.0 License.

For model developers, the practical change is that questions once too costly to test at scale become approachable: which attention layers can be replaced, how many can go without hurting quality, which state space layer works best, and what the speed-accuracy trade-off looks like. Because every variant starts from the same pretrained initialization and needs minimal training, the authors say layer type, hybrid ratio and layer placement can now be explored cheaply and systematically.