posts

Reusing transformer weights to build hybrid models

The memory tax of long-context inference usually locks teams into the architectures they already trained. Running full attention across 128,000 tokens means maintaining a key-value cache that expands with every step, but replacing those layers with recurrent or state space alternatives has historically required retraining from scratch on trillions of tokens. In a recent paper, AWS researchers outline a method called priming, which reuses pretrained Transformer projections to build hybrid models for less than one percent of the original training compute.

The approach relies on the observation that many attention layers only handle local context. Instead of preserving an exact record of every previous token, those layers can be approximated by a fixed-size recurrent state. The AWS team maps the original query, key, and value weights directly into state space layers, runs a short distillation step to align their outputs, and finishes with end-to-end fine-tuning. In their tests, a 32-billion-parameter hybrid using Gated KalmaNet matched its source Transformer on reasoning benchmarks while decoding twice as fast at 128,000 tokens of context.

What makes the release useful for builders is that the team published the 8-billion and 32-billion model weights alongside vLLM serving plugins under Apache 2.0. If weight transfer across architectures proves stable on other open model families, testing non-Transformer inference ceases to be a luxury reserved for labs with massive training budgets.

Source: amazon.science/blog/can-we-change-an-ai-models-architecture-w…

← all posts