Skip to main content
Back to timeline
Amazon Science homepageSource publication:

Hybrid architectures skip pretraining from scratch: reusing Transformer weights cuts cost by over 99%

Synopsis

Amazon Web Services introduces priming, which initializes hybrid state space models from pretrained Transformer weights and cuts the pretraining cost of hybrid architectures by more than 99%, enabling controlled comparisons of three state space layer types.

Source-provided article image: Can we change an AI model’s architecture without retraining it from scratch?
Amazon Science

Interpretation

The pretraining cost of hybrid architectures drops to less than 0.5% of the source model's pretraining token budget, making architecture comparisons feasible at a scale that was previously too expensive. Exploring this design space previously required training every new architecture from scratch, using trillions of tokens and hundreds of thousands of GPU-hours, which made controlled comparisons extremely expensive. The entire pipeline consumes less than 0.5% of the source model's pretraining token budget and works across Transformer families (Qwen, Llama, Mistral), model classes (dense or MoE), and scales from eight billion to 32 billion parameters and more.

Under the same data, the same training budget, and the same ratio of eidetic to fading layers, the three state space layer types show a predictable ordering: Gated KalmaNet beats Gated DeltaNet, which beats Mamba-2. Mamba-2 applies a uniform scalar decay to the entire state at each step, Gated DeltaNet adds targeted updates on top of uniform decay, and Gated KalmaNet's erase direction also depends on statistics of all past keys. The theoretical expressiveness hierarchy predicts the empirical one, with the gap widening on long-context reasoning and narrowing on short-context tasks; at the 32-billion-parameter scale, hybrid GKA is within 1% of its source Transformer trained on the same data while requiring half the KV cache at inference.

Hybrid models built on state space layers consistently outperform corresponding sliding-window attention hybrids on long-context benchmarks, even when the window gets two to four times as much memory. Sliding-window attention retains a fixed window of recent tokens verbatim, so no information is lost within the window's range, and stacking multiple such layers increases the receptive field. Every SSM-based hybrid tested outperformed the corresponding SWA hybrid on long-context benchmarks, even when the window received two to four times as much memory as the state.

Gated KalmaNet offers variable computation at inference time, so a single trained model can serve multiple points on the speed-accuracy frontier. Mamba-2 and Gated DeltaNet have the same computation budget for every input token, regardless of its content. On AIME 2025, serving GKA at 10 iterations instead of its usual 30 sacrifices roughly 1% of accuracy while delivering 10% faster decoding; the 32-billion-parameter GKA hybrid decodes up to twice as quickly with 128,000 tokens of context and reaches correct answers in as little as two-thirds of the wall-clock time.

Perspective

The report is aimed at researchers and engineers working on long-context reasoning and agent workloads. The demonstrated scope includes priming across the Qwen, Llama, and Mistral families and dense or MoE model classes, from eight billion to 32 billion parameters and more; the tested hybrids natively support 128,000 tokens, and partitioning a long input into chunks with a shared prefix, concatenating KV caches, and merging SSM states via the weighted-averaging scheme from PICASO extends context to 256,000 tokens without training while preserving substantial accuracy on RULER NIAH retrieval and BABILong reasoning. Released artifacts include primed hybrids at eight-billion and 32-billion parameters, priming training code, point-to-point sequence parallelism for SSM layers, a universal sequence parallelism module, Triton kernels for GKA decoding and training, and a vLLM serving plugin with continuous batching, chunked prefill, and tensor parallelism, all under the Apache 2.0 License.

The material describes priming as an approximation: the source Transformer's projections can be reused only when a fixed-dimensional-state SSM can produce a mixing matrix close enough to the attention layer's, and whether that holds depends on whether the layer performs complex long-range retrieval or simple local operations. Layer selection relies on a cheap empirical proxy, replacing each layer with sliding-window attention and measuring the degradation, which connects to the Hankel rank of the layer's mixing matrix from dynamical-systems theory. Serving GKA at 10 iterations instead of 30 sacrifices roughly 1% of accuracy, so the speed-accuracy trade-off is adjustable but not free. The context-extension method preserves substantial accuracy from 128K to 256K, and the material states no conclusion for longer extrapolation.