Amazon Researchers Cut Hybrid Model Pretraining Costs by 99% via Priming
A novel pipeline converts Transformer attention layers into state space models without expensive retraining.
1Transformer Memory Scalability Problem
Modern autonomous AI agents require extended contexts and deeper reasoning chains, yet standard Transformers impose an expanding memory footprint. The key-value cache scales directly with sequence length, making long-context inference compute-intensive and costly for high-concurrency production deployments across enterprise infrastructure environments.
2Cost Barrier of Hybrid Exploration
Hybrid architectures combine attention mechanisms with state space models to balance memory usage, throughput, and generation quality. Historically, evaluating novel hybrid configurations required complete pretraining runs from scratch, consuming trillions of tokens and hundreds of thousands of GPU-hours to test layer replacements.
3Priming Architecture Conversion Pipeline
Amazon Web Services researchers introduced Priming, a method transforming pretrained Transformers into hybrid SSM models while saving over 99% of pretraining compute. Priming reuses existing model projection weights, requiring less than 0.5% of the original pretraining token budget to adapt dense or MoE architectures.
4Eidetic Versus Fading Memory Layers
Standard attention maintains eidetic memory, recalling every token verbatim through a dynamically growing key-value cache. Conversely, state space layers utilize fading memory, condensing past context into a fixed-size mathematical state suitable for local processing where global full-context retrieval across every layer is not required.
5Empirical Layer Selection Proxy
The Priming framework identifies viable layers by replacing individual attention blocks with sliding-window attention to quantify degradation. Layers tolerating local substitution exhibit lower Hankel rank in their mixing matrices, proving that a fixed-size state bottleneck can accurately replicate their sequence-mixing behavior without structural fidelity loss.
6Three Stage Conversion Workflow
The Priming pipeline executes across three phases: selecting compressible layers, transferring query-key-value projection weights into target SSM layers with local distillation alignment, and conducting end-to-end adaptation fine-tuning. This workflow generalizes across major foundational families, including Qwen, Llama, and Mistral open-weight architectures.
7State Space Architecture Hierarchy
Controlled benchmarks revealed clear performance disparities among fading-memory layers: Gated KalmaNet surpassed Gated DeltaNet, while Mamba-2 ranked third. Mamba-2 relies on uniform scalar decay, DeltaNet incorporates selective key erasure, and KalmaNet utilizes history-aware update vectors derived from cumulative historical sequence statistics.
8SSM Outperforms Sliding Window Attention
Evaluations demonstrated that hybrid models using compressed SSM layers consistently outperformed sliding-window attention configurations on long-context benchmarks. Lossy global history compression proved functionally superior to verbatim local retention, even when sliding-window baselines were allocated two to four times greater total memory capacity.
9Dynamic Runtime Computation Control
Gated KalmaNet introduces test-time variable computation by solving its internal regression mechanism iteratively without requiring model retraining. Operators can dynamically reduce serving iterations from 30 down to 10 at inference time, achieving 10% faster token decoding while conceding approximately 1% accuracy on AIME.
10Throughput Gains on Reasoning Benchmarks
On the AIME 2025 benchmark, a primed 32-billion GKA hybrid matched the reasoning accuracy of its source Qwen3-32B baseline while cutting memory usage in half. The hybrid reached target solution thresholds 1.2 to 1.6 times faster in wall-clock time due to larger batched execution.
11Long Context Throughput Multipliers
Throughput acceleration scales aggressively with sequence length under GPU memory saturation constraints. The 32-billion parameter GKA hybrid achieved a 1.50x decoding speedup at 16,000 tokens, reaching 2.06x at 128,000 tokens, while the eight-billion model scale achieved up to 2.29x throughput gains at maximum context.
12Post Training Acceleration and Open Artifacts
Beyond serving efficiencies, halving key-value memory enables 20% to 100% higher trajectory generation throughput during reinforcement learning post-training. AWS open-sourced primed eight-billion and 32-billion parameter model checkpoints, optimized kernel implementations, and dedicated vLLM serving plugins under an Apache 2.0 software license.