AWS introduces priming to convert pre-trained Transformers into hybrid state space models with 99% less pretraining cost

AWS researchers introduced priming, a method that converts pre-trained Transformer weights into hybrid state space models (SSMs) to reduce pretraining costs by more than 99%. This approach allows developers to explore the design space of hybrid architectures—mixing attention layers with SSM layers—without the prohibitive expense of training from scratch.

How Priming Works

Priming leverages the structural similarity between Transformers and SSMs, treating Transformers as a special case of state space realization. The process involves three stages:

• Layer selection: Identifying attention layers that can be compressed into fixed-size states using a sliding-window attention proxy.
• Weight transfer: Initializing SSM layers from Transformer weights and aligning them via distillation.
• End-to-end adaptation: Fine-tuning the hybrid model on next-token prediction.

The pipeline consumes less than 0.5% of the source model's pretraining token budget and works across families like Qwen, Llama, and Mistral.

Performance and Findings

The study compared three SSM types: Gated KalmaNet (GKA), Gated DeltaNet (GDN), and Mamba-2. Key findings include:

• GKA performed best, followed by GDN and Mamba-2, with the gap widening on long-context reasoning tasks.
• A 32-billion-parameter GKA-based hybrid matched its source Transformer on reasoning benchmarks while decoding up to twice as quickly with 128,000 tokens of context.
• SSM-based hybrids consistently outperformed sliding-window attention (SWA) hybrids, even when SWA had twice the memory budget.

Variable Computation and Context Extension

GKA offers variable-length inference computation, allowing users to adjust the speed-accuracy trade-off at serving time without retraining. For example, reducing iterations from 30 to 10 sacrifices roughly 1% of accuracy on AIME 2025 but delivers 10% faster decoding. Additionally, the team released a method to extend context length from 128K to 256K tokens without further training by composing states across chunks.

Open Source Release

AWS released a zoo of primed hybrid models (8B and 32B parameters), training code, optimized kernels, and a vLLM serving plugin under the Apache 2.0 License on GitHub and Hugging Face.