Can we change an AI model’s architecture without retraining it from scratch?

來源原標題

Can we change an AI model’s architecture without retraining it from scratch?

機構

Amazon

原始連結

https://www.amazon.science/blog/can-we-change-an-ai-models-architecture-without-retraining-it-from-scratch

摘要原文

AI agents are rapidly scaling: longer sessions, deeper reasoning chains, more autonomy. But the Transformer architecture underneath them still pays a memory tax that grows with the number of tokens it processes. As agents become more capable, that cost scales with them. Now, there’s a proliferation of hybrid architectures that use some combination of methods to achieve a better balance of memory, speed, and quality. Those methods include sparse attention, which selects a subset of past tokens to consider when selecting the next token; sliding-window attention, which preserves only the most recent tokens; and state space models, which compress information from the full history into a fixed-size state. These hybrid architectures can be understood as points in a sequence-model design space: to differing degrees, they mix information while traversing a sequence and passing through different states. What distinguishes them from each other is how they retain and use the past. Until now, exploring this design space has required training every new architecture from scratch, using trillions of tokens and hundreds of thousands of GPU-hours. This has made controlled comparisons of hybrid architectures extremely expensive. In a new paper, “Priming: Hybrid state space models from pre-trained Transformers,” our team at Amazon Web Services addresses this problem. We introduce priming, a method that leverages pretrained Transformer weights as a starting point to reduce the pretraining cost of hybrid architectures by more than 99%. With priming, model developers can ask questions that were previously too expensive to answer at scale: Which attention layers can be replaced by more-efficient layers? How many attention layers can be replaced without sacrificing quality? Which state space layer works best? And what does the trade-off between speed and accuracy look like? In our paper, we compare hybrid models that use three different types of state space models (Gated KalmaNet, Gated DeltaNet, and Mamba-2). Some of our findings: Not all state space models are equal. Under controlled conditions, Gated KalmaNet (GKA) performs best, followed by Gated DeltaNet (GDN) and then Mamba-2. GKA uniquely offers variable-length inference computation without retraining and gives the user direct control over its position on the speed-accuracy frontier. A good lossy summary of everything beats a perfect recording of the most recent past. State space models consistently outperform sliding-window attention (SWA), even when SWA gets twice the memory budget. Same quality, twice the throughput. A 32-billion-parameter GKA-based hybrid model matches its source Transformer on reasoning benchmarks while decoding up to twice as quickly with 128,000 tokens of context and reaching correct answers in as little as two-thirds of the wall-clock time. To foster further research on hybrid architectures, we have released a zoo of primed hybrid models at the eight-billion-parameter and 32-billion-parameter scales for long-context reasoning and instruction following, together with training code, optimized kernels, and a vLLM-serving plugin under the Apache 2.0 License. The attention mechanism A Transformer’s attention mechanism — which determines how relevant previously seen tokens are to the token currently being processed — relies on three different vectors, which take their names from database design: the query, the key, and the value. For each token, the query captures what that token is “looking for,” the key captures what it “offers,” and the value carries the actual information it will pass along. Relevance is computed by comparing the current token’s query against every other token’s key; the closer the match, the more that token’s value contributes to the output of attention. Each of these vectors is produced by multiplying the token by a learned weight projection matrix. The Transformer outputs one token at a time, and each output token becomes a new input, so the context grows one token at a time as well. To save computation, the Transformer thus stores the key and value projections of each new input in a “KV cache”. Two kinds of memory layers: Eidetic versus fading The attention mechanism gives Transformers what we call eidetic memory: verbatim recall of every past token. The possibility of recalling over an LLM’s entire context window is what made in-context learning, long-document reasoning, and retrieval-augmented generation possible. But it comes at a cost that scales with the number of input tokens. However, not all attention layers do the same job. Some perform long-range retrieval across tens of thousands of tokens. Others operate more locally, tracking the current clause, copying nearby tokens, and the like. The local ones don't need a growing cache: a compact, fixed-size summary of the past — what we call fading memory — suffices. Hybrid models mix both: attention layers (eidetic, growing cache) for the layers that need global recall and state space layers (fading, fixed size) for the rest. The result is a model with a reduced memory footprint that still has access to verbatim recall where it matters. Priming: Knowledge isn't locked to the architecture Exploring hybrids’ design space has been impractical because each hybrid configuration would normally require pretraining from scratch. Priming sidesteps this by reusing the stored knowledge in a pretrained Transformer's weights. Priming is possible because a Transformer is itself a special case of a state space realization. The outputs of both attention and SSM layers can be expressed as the result of multiplying the value sequence by a “mixing” matrix that decides how much each past token contributes to each current token. Attention computes the mixing matrix by multiplying the query and key matrices; an SSM produces a similar output through a recurrent state of fixed dimension that all past information must pass through. Priming takes advantage of this structural similarity: if an SSM with a fixed-dimensional state can produce a mixing matrix close enough to the attention layer's, the query, key, and value projections from the source Transformer can be reused as the SSM's own input projections, with only the SSM-specific parameters needing to be learned. Whether the approximation holds depends on whether the attention layer performs complex, long-range retrievals or simple, local operations. We distinguish between the two using a cheap empirical proxy: replace each layer with an SWA — which preserves only a fraction of the context — and measure the degradation. Layers that tolerate this substitution are natural candidates for SSM conversion: their mixing matrices don’t need to capture long-range relationships, so the state of an SSM can faithfully represent what they do. This proxy connects to a formal quantity from dynamical-systems theory, the Hankel rank of the layer's mixing matrix, which measures how much information must pass through the state bottleneck and therefore the minimum SSM state dimension required to realize the layer exactly. Priming proceeds in three stages: Layer selection: identify which layers can be faithfully compressed into a fixed-size state while retaining the rest as attention; Weight transfer and alignment: initialize SSM layers from the corresponding attention weights, then run a short distillation process aligning each SSM to the behavior of the layer it replaces; End-to-end adaptation: fine-tune the full hybrid model on next-token prediction, letting all layers coadapt. The entire pipeline consumes less than 0.5% of the source model's pretraining token budget, and it works across Transformer families (Qwen, Llama, Mistral), model classes (dense or MoE), and scales (eight billion to 32 billion and more).