AWS Priming Method Converts Transformers into Hybrid SSMs Without Retraining

The attention mechanism at the core of Transformers requires a growing KV cache for every token processed, making memory costs scale with context length — a growing problem as AI agents run longer, more autonomous sessions. Hybrid architectures that mix attention with more efficient mechanisms like sparse attention, sliding-window attention, or state space models (SSMs) offer a potential fix, but exploring this design space has traditionally required pretraining each configuration from scratch, costing trillions of tokens and hundreds of thousands of GPU-hours.

A new paper from an AWS research team introduces 'priming,' a method that reuses pretrained Transformer weights as a starting point, cutting the pretraining cost of hybrid architectures by more than 99%. The technique exploits the fact that a Transformer is itself a special case of a state space realization: when an SSM's mixing matrix can approximate an attention layer's output closely enough, the query, key, and value projections from the source Transformer can be transferred directly, leaving only SSM-specific parameters to be learned through a short distillation and fine-tuning process.

The team compared three SSM variants — Gated KalmaNet, Gated DeltaNet, and Mamba-2 — finding Gated KalmaNet performs best and uniquely allows developers to tune the speed-accuracy trade-off after training. A 32-billion-parameter hybrid model built this way matches an equivalently sized Transformer's accuracy while reaching correct answers faster, decoding up to twice as fast at 128,000 tokens of context, and arriving at correct answers in two-thirds the wall-clock time. The team has released primed hybrid models at 8B and 32B scales, along with training code, optimized kernels, and a vLLM serving plugin, under the Apache 2.0 license.