Late-Stage Patching Full-Attention Transformers to Gated DeltaNets

Hybrid Models

Recent frontier models increasingly rely on mixing GDNs and full-attention transformers than choosing between them. Examples include Kimi K3[1] (Moonshot AI, 2.8T parameters, 1M-token context) and Qwen3.8[2] (Alibaba, 2.4T MoE).

A hybrid of full-attention and GDN layers can be substantially more efficient to serve than a pure attention model, particularly at long-context lengths, because most of its layers replace a KV cache that grows with every token and an attention computation that revisits the entire history at every step with a fixed-size memory and a constant per-token update.

The intuition behind the hybrid design is that attention and GDN layers provide different forms of memory. GDN layers can carry information over long spans through a compact recurrent state, making them efficient at preserving and updating context without attending over the entire history. Full-attention layers, meanwhile, retain direct access to individual past token positions, which is useful when the model needs to recover a specific detail or model an exact interaction across a long distance. Combining the two therefore offers efficient long-range memory through GDNs while preserving the precise retrieval capabilities of full-attention.

Morphing is the first stage, aimed at converting an off-the-shelf pretrained full-attention model into a hybrid architecture, before any training begins. Every 4^th layer retains its original softmax full-attention unchanged, while the attention blocks in the remaining layers are replaced with Gated DeltaNet[3] modules matching the Qwen3-Next[4] design. We apply this transformation to two base architectures: Qwen3-4B[5] and MiMo-7B-RL[6].

This is the first stage of our recovery pipeline. Throughout all stages, we use the original base model as the teacher and the morphed hybrid model as the student. The goal of this stage is to provide a strong initialization for subsequent training. Concretely, we train each student GDN layer to match the output of the corresponding teacher full-attention layer, while feeding both layers the hidden representation produced by the teacher’s preceding layer (hence called off-policy). This encourages the morphed model GDN layers to mimic the teacher full-attention layer's outputs by minimizing distribution shift between the teacher and student representations.

• Objective: A single teacher forward per batch yields every hidden state, where h^{T}_{l} is the input to layer l and h^{T}_{l+1} is the output.
  h^{T}_{0} = \mathrm{Embed}(x), \qquad h^{T}_{l+1} = f^{T}_{l}\!\left(h^{T}_{l}\right)
  Let f^{S}_{l} denote the functional map for the layer l of the student model we minimize
  \mathcal{L} \;=\; \sum_{l \,\in\, \mathrm{GDN}} \;\left\lVert\, f^{S}_{l}\!\left(h^{T}_{l}\right) \;-\; h^{T}_{l+1} \,\right\rVert_{2}^{2}
  Every layer's input is the teacher's h^{T}_{l}, not the student's output from the preceding layer, so \partial\mathcal{L}/\partial\theta_{l} depends on layer l alone. The objective decomposes into 27 independent layer-wise feature-matching problems, each minimizing a squared reconstruction error. The gradient flows through the inherited q, k, v, output projection, and the parameters with no counterpart in attention: the forget gate \alpha, the update gate \beta, the output gate z, the short convolution and the decay terms — which is 42.1M per layer on Qwen3-4B and 67.4M on MiMo-7B, or 1.14B and 1.82B across the 27. Everything outside those blocks is frozen. The reason for using a layerwise loss here is to avoid compounding approximation errors caused from distribution shift at the input across the depth of the network
• Data and Optimizer: 102M tokens of training

  Data: Both arms use the same 102M-token mix: 70% FineWeb-Edu, 20% Python code, 10% open-web-math. Optimizer: AdamW, no weight decay, betas 0.9 and 0.95, gradients clipped at global norm 1.0. Learning rate warmup over 3% of steps then cosine decay to 10% with peak LR 3e-3 on Qwen3-4B, 1e-3 on MiMo-7B with a global batch of 32K tokens on Qwen3-4B (~3,050 steps) and 16K on MiMo-7B (~6,100 steps). Weight decay is off because there is a fixed right answer here. Decay nudges every weight a little toward zero at each step — it is a regulariser, meant to keep a model from fitting noise in its training data too closely. But these blocks are being fitted to a target function that already exists, starting from weights copied out of a working model. There is no noise to guard against, so shrinking them only drags the fit away from the answer.

Stage 1 calibrates each GDN block independently, using hidden-state inputs generated by the teacher model. While this provides a strong layerwise initialization, it does not account for interactions between the newly introduced GDN blocks: each block is optimized on teacher-generated representations and not the representations it will receive from preceding student layers at inference time. Consequently, approximation errors and representation shifts may accumulate when the student is composed end-to-end.

Stage 2 is the first stage that trains the language model end-to-end distilling the teacher’s output distribution into the complete student model, thereby encouraging globally consistent behavior across the full network. It runs in two substages: first at short context, then at the teacher's full window.

• Objective: Both student and teacher models use same text as the input, and at every position the student's output distribution is trained to match the teacher distribution by minimizing the forward KL Loss:
  \mathcal{L}_{2} \;=\; \mathbb{E}_{x \sim \mathcal{D}}\left[\, \sum_{t} \mathrm{KL}\!\left(p^{T}(\cdot \mid x_{<t}) \;\middle\|\; p^{S}(\cdot \mid x_{<t})\right) \right]
  Forward KL, with the teacher as the reference distribution, is mode-covering: it penalizes the student for assigning too little probability wherever the teacher places mass, encouraging the student to cover the full range of outputs considered plausible by the teacher. The training data consists of ordinary pretraining-style packed corpus text rather than prompt–response pairs or teacher-generated samples. The teacher does not generate any tokens; instead, it provides its predictive distribution at each position in the fixed corpus. Because the trajectories are drawn from an external data distribution rather than sampled from the teacher, this stage is off-policy.

Stage 2a: Short-Context Distillation

• Data and Optimizer: 600M tokens of training at 4K context

  Data: 600M tokens of pretraining-style text — web, code and maths — packed into 4K-token sequences, one pass, no repetition. Optimizer: AdamW, betas 0.9 and 0.95, gradients clipped at 1.0. Warmup over 3% of steps, then cosine decay to a 10%. Global batch 2M tokens. Both the full-attention and GDN layers use a cosine learning rate schedule decayed to 10% over the training horizon, while the peak learning rate of the full-attention layers is fixed to 10% of that of the GDN layers. The peak learning rate used for the GDN layers is 1.5e-4 for Qwen3-4B and 5e-5 for MiMo-7B.

Stage 2b: Long-Context Distillation

• Data and Optimizer: 294M tokens of training at 32K context

  Data: 294M tokens packed into 32K-token sequences, chosen so that those sequences are genuinely long documents — including books, long web pages, and mathematics content and not short documents concatenated together. The mix is approximately 70% long documents and 30% short-context replay, with the replay included to prevent short-context behavior from regressing while the recurrent state learns to retain information over longer inputs. The 32K context length is set by the teacher's own context window — the teacher cannot provide supervision beyond what it can itself read. Optimizer: AdamW, betas 0.9 and 0.95, clip 1.0, 3% warmup, cosine decay to 10%, global batch 2M tokens. Peak LR 1e-4 for the GDN group on Qwen3-4B and 2.5e-5 on MiMo-7B, and the peak LR of transformer layers is 10% of the GDN layers.

Many prior approaches keep the retained full-attention layers frozen throughout training while the GDN layers learn how to update and use the fixed size memory state. However, we freeze them only during stage 1, while the newly introduced blocks are being initialized. In stage 3, which is the longest stage of our training pipeline, the retained attention layers are updated jointly with the rest of the model.

This is important because the full-attention layer in pure transformers and hybrid models differ in the role those layers must play. Before conversion, every attention layer contributes directly to long-range interaction at its own depth. After conversion, the retained full-attention layers become the only points in the network where information can be accessed beyond the compressed state, and must therefore support long-range retrieval on behalf of the converted layers as well. Their pretrained weights were not optimized for this altered role, so they must adapt during recovery training. We observe that the best results are obtained when Stage 2 serves as a gradual bridge between stages 1 and 3, jointly updating the full model while using a lower learning rate for the retained full-attention layers so that they adapt on a slower timescale than the newly introduced GDN blocks.

Stage 3 is the final and most important stage of the recovery pipeline. As in the earlier stages, the original base model serves as the teacher and the morphed hybrid as the student, with the objective of recovering the capabilities lost during conversion.

This is also the stage where long-context reasoning is restored. After Stage 2, the student can retrieve information from long inputs, but it still struggles to reason effectively over long contexts. Its AIME accuracy remains at 0, while MATH-500 scores under a 32K-token budget are only around 2.0%, substantially worse than performance with shorter 1–4K token budgets and thinking disabled. In other words, giving the model more room to reason actually hurts performance without stage 3.

This gap between long-context retrieval and long-context generation comes from a difference in distribution shift. During retrieval, nearly all of the tokens the model conditions on are provided in the input, so teacher-forced training covers the relevant distribution reasonably well. During long-form generation, however, the model repeatedly conditions on its own outputs. After only a few thousand generated tokens, it can enter regions of the sequence distribution that are rarely, if ever, represented in static training corpora. Off-policy training would need to anticipate and cover this space in advance, which is impractical for open-ended generation. Training directly on the student's own generated trajectories instead exposes the model to exactly the states it will encounter at inference time. Stage 3 additionally also improves long form retrieval performance.

• Objective: The student generates a continuation from a prompt. The frozen teacher then scores those exact tokens, and the loss is the divergence between the two distributions at each generated position:
  \mathcal{L}_{3} \;=\; \mathbb{E}_{y \sim \pi^{S}}\left[\, \sum_{t} \mathrm{KL}\!\left(p^{S}(\cdot \mid y_{<t}) \;\middle\|\; p^{T}(\cdot \mid y_{<t})\right) \right]
  Reverse KL is mode-seeking, and it punishes the student for assigning high probability where the teacher does not. The prompt is context only; loss is taken on generated positions, so generated tokens and loss tokens are the same count. Generation length starts at 512 and doubles whenever the smoothed divergence stops improving, out to 16K and later 32K, so the model is supervised at depths it will actually be used at.
• Why does Stage 3 have to be prolonged: On-policy distillation is most effective when the teacher and student induce reasonably similar output distributions[8] — that is, the student must produce good trajectories often enough that there is some positive signal for the teacher to reinforce. A well trained model such as off the shelf Qwen3-4B checkpoint used as a student starts inside that regime and converges quickly. Our model post stage 2 does not. The student is very weak at long-context generation when this stage 3 begins, so good trajectories are rare, and the teacher is the base model itself rather than a larger one as is the case typically in standard on policy distillation. The same-size teacher is a deliberate constraint rather than a limitation we worked around. A stronger teacher could introduce capabilities that the original pre-Stage-0 model never possessed, turning the process into knowledge transfer rather than capability recovery. Hence as a result of this weak student - weak teacher combination, Stage 3 requires prolonged training for effective recovery.
• Pretraining-like step: Each optimizer step consumes a fixed budget of 128K generated tokens — the same for Qwen3-4B and MiMo-7B — split across the trainer processes. An update is defined by a token count, as in pretraining, and not by a number of prompts as in post-training distillation/RL. How many sequences that budget buys is whatever the current generation length allows. The learning rate is scheduled on consumed tokens rather than on steps, so the horizon ladder described below can change sequence length without disturbing it.
• Generation-length curriculum: Generation length starts at 512 tokens and doubles every phase to 16K over the first 600M tokens, then to 32K over the next 400M. The generation length is increased when the loss plateaus and doesn't make any meaningful progress. Within each phase, we track the lowest smoothed KL seen so far. A 2% improvement resets the counter and sets a new record; otherwise, the counter increments. After 30 steps without sufficient improvement, the generation cap doubles and all phase statistics reset. This curriculum is more efficient than starting with long rollouts: with a fixed token budget, longer generations mean fewer samples, noisier gradients, and slower sampling. Early in training, the student also struggles to stay coherent over long horizons, making much of the extra text useless.

• Data: prompts only — 65% Math Reasoning OpenThoughts/DAPO/Deepscaler, 20% Dolly-15k, 15% retrieval

  Data: The data used for the first 1B tokens in this stage consists of prompts drawn from approximately 65% OpenThoughts[10], 20% Dolly-15k[11], and 15% synthetic long-context retrieval tasks, with all answers generated on-policy by the student model. For Stage 3 (1B–2B tokens), OPAL-MiMo replaces the mathematical reasoning data with a mixture of 43.5% DeepScaleR[23] and 16.5% DAPO-17K[24], while increasing the retrieval data proportion to 20%. OPAL-Qwen3 retains the same data mixture used during the first 1B tokens. During inference time retrieval, the model must recover a fact from a long-context in the prompt and retain it while continuing to generate its response. These tasks are introduced only once the generation cap reaches 1024 tokens, since shorter rollouts do not provide enough generation length for retrieval and subsequent use of the retrieved information to become meaningful. Sampling is performed at temperature 1.0, so the student is trained on the same distribution from which its answers are generated.
• Optimizer and learning-rate schedule: For the prolonged on-policy distillation phase, we use a Warmup-Stable-Decay (WSD) learning-rate schedule[9] (Figure 1). WSD is predominantly a pretraining LR schedule. It belongs to a family of infinite LR schedules designed for continual pretraining. We adopt it here specifically because our post-training regime is unusually long and is extended incrementally in blocks of 100–200M tokens, without knowing in advance when the loss will saturate. WSD is well suited to this setting because it defines an effectively open-ended schedule: the stable phase can be extended cheaply as additional training is added, and the final decay can be triggered only once we decide to stop. In contrast, cosine decay assumes a fixed training horizon and generally requires the schedule to be specified from the outset. This flexibility is the main reason we prefer WSD despite cosine decay being the standard choice in post-training. We parameterize the schedule by the number of tokens trained on rather than the number of generated sequences or optimization steps, since rollout lengths vary and an equal generation or step budget can correspond to substantially different amounts of training data. All GDN and full-attention layers are trained using the same learning-rate schedule. We use AdamW with a peak learning rate of 2e-5 for Qwen3-4B and 1e-5 for MiMo-7B, and fix the global batch size to 2M tokens. The learning rate is linearly warmed up for 150 steps, while the decay phase is reconfigured at each training increment and spans approximately 10% of the cumulative Stage 3 steps reached at that point. All results reported in this work use 2B tokens of Stage 3 training unless stated. We note that the loss has not yet saturated at 2B tokens (Figure 3), suggesting that extending Stage 3 further is likely to yield additional gains.

• Asynchronous sampling: Generation and training run concurrently rather than in alternation.

  Asynchronous sampling: Generation and training run concurrently rather than in alternation. We maintain multiple trainers and a single sampler copy. Sampler processes hold an inference copy of the student and do nothing but produce continuations; trainer processes hold the training copy and do only training via forward and backward pass on generations. The two are coupled by a queue of rollouts and a weight publish: after every optimizer step the trainer writes out its current weights, the samplers hot-load them. The process is asynchronous in the sense that the sampler does not wait for the trainer; multiple trainers per sampler ensure that the sampler never idles through backward passes due to a long queue.

Commonsense Reasoning Suite : Short Context

Each linearized model is compared against its corresponding teacher. The final column reports the average across all six tasks; percentages in parentheses denote performance retained relative to the teacher average. This part is easy to recover even without Stage 3. We donot beat the baselines here and pay the short context tax for long-context performance. We think that Commonsense performance can be recovered somewhat cheaply with additional SFT as done by some of the baselines.

Needle-in-a-haystack: Long Form Retrieval

Retrieval accuracy across context lengths, grouped by task type. The average is the macro-average over all twelve task–length combinations; percentages in parentheses denote performance relative to the corresponding teacher. One can unlock non trivial performance here without stage 3. Some baselines such as Zebra-Llama[14] perform well in shorter contexts (4–8K) but severely underperform at longer contexts (16-32K) and in the multi-key/query settings compared to the corresponding full-attention model, which has near perfect retrieval scores across the board. OPAL on the other hand has scores that match the corresponding full-attention model.

Mathematical reasoning : Long-Context Generation

This is the stage where Prolonged On-Policy Distillation yields the largest gains, with the improvements over baselines being most pronounced on harder benchmarks such as AIME, where most baselines fail to recover any performance. In fact the numbers reported below correspond to premature early stopped checkpoints. Extended training still shows signs of decreasing training loss.

How reasoning returns, stage by stage

The tables above compare finished models. The tables below trace the same models across the pipeline, which is where the long-context reasoning improvement is visible.

Long-form retrieval recovers faster and is essentially saturated after the first 200M tokens.

We establish prolonged on-policy distillation as an effective approach for converting trained Transformer layers into efficient architectures such as Mamba-3 and Gated DeltaNet, while recovering both short- and long-context capabilities. Our method requires no additional teacher beyond the original full-attention Transformer and can successfully hybridize off-the-shelf pretrained models, including those that have undergone extensive RL post-training, with only a modest additional training budget of roughly 3B tokens. Moreover, performance continues to improve as the on-policy stage is prolonged, with gains remaining unsaturated at the end of our training budget. This brings us to the question:

Why convert?: Organizations with sufficient compute may prefer to pretrain the target architecture directly, while compute-constrained users may simply fine-tune an existing hybrid model such as Qwen3.5[18].

In this sense, our result is primarily a demonstration that architecture conversion that retains long-context abilities from transformer layers to GDN or Mamba-3 like layers is possible even late in training, including after RL post-training, at relatively low additional cost. This may become more useful as architectural innovation outpaces pretraining cycles, allowing already-trained models to adopt newer architectures without restarting from scratch. The concept could also be particularly relevant for architectures that are substantially harder or more expensive to pretrain, such as models with deep-memory or TITANS-style[19] modules.

Looking forward, several directions are particularly promising. Beyond extending the approach to more sophisticated architectures such as looped Transformers, architectural conversion itself could be treated as a curriculum, gradually modifying the backbone rather than introducing a large architectural shift at once. More ambitiously, one could study scaling laws for architecture conversion: how the required recovery compute depends on model scale, architectural distance, context length, and the fraction of layers being replaced. Such a characterization could help turn post-hoc architecture conversion into a more predictable and general mechanism for adapting pretrained models to new computational constraints.

1. Kimi Team, Moonshot AI. Kimi K3: Open Frontier Intelligence. Technical report, 2026. 2.8T-parameter MoE built on Kimi Delta Attention and Attention Residuals, 1M-token context. arxiv.org/abs/2607.24653
2. Qwen Team, Alibaba. Qwen3.8-2.4T-A95B — 2.4T parameters, 95B active: Model card and accompanying blog post, 2026. huggingface.co/Qwen/Qwen3.8-2.4T-A95Bqwen.ai/blog?id=qwen3.8
3. Songlin Yang, Jan Kautz, Ali Hatamizadeh. Gated Delta Networks: Improving Mamba2 with Delta Rule. ICLR 2025. arxiv.org/abs/2412.06464
4. Qwen Team, Alibaba. Qwen3-Next —. Model card, 2025. huggingface.co/Qwen/Qwen3-Next-80B-A3B-Instruct
5. Qwen Team, Alibaba. Qwen3 Technical Report. 2025. Covers the Qwen3 family, including the 4B model used as a teacher here. arxiv.org/abs/2505.09388
6. Xiaomi LLM-Core Team. MiMo: Unlocking the Reasoning Potential of Language Model — From Pretraining to Posttraining. 2025. MiMo-7B-RL is the RL-tuned model from this report. arxiv.org/abs/2505.07608
7. Kaiyue Wen, Zhiyuan Li, Jason Wang, David Hall, Percy Liang, Tengyu Ma. Understanding Warmup-Stable-Decay Learning Rates: A River Valley Loss Landscape Perspective. 2024. arxiv.org/abs/2410.05192
8. Kevin Lu et al., Thinking Machines Lab. On-Policy Distillation. 2025. thinkingmachines.ai/blog/on-policy-distillation
9. Shengding Hu et al. MiniCPM: Unveiling the Potential of Small Language Models with Scalable Training Strategies. 2024. arxiv.org/abs/2404.06395
10. Etash Guha et al. OpenThoughts: Data Recipes for Reasoning Models. 2025. arxiv.org/abs/2506.04178
11. Mike Conover, Matt Hayes, Ankit Mathur, Jianwei Xie, Jun Wan, Sam Shah, Ali Ghodsi, Patrick Wendell, Matei Zaharia, Reynold Xin. Free Dolly: Introducing the World’s First Truly Open Instruction-Tuned LLM. 2023. databricks.com/blog/2023/04/12/dolly-first-open-commercially-viable-instruction-tuned-llm
12. Daniel Goldstein, Eric Alcaide, Janna Lu, Eugene Cheah. RADLADS: Rapid Attention Distillation to Linear Attention Decoders at Scale. 2025. arxiv.org/abs/2505.03005
13. Aviv Bick, Tobias Katsch, Nimit Sohoni, Arjun Desai, Albert Gu. Llamba: Scaling Distilled Recurrent Models for Efficient Language Processing. 2025. arxiv.org/abs/2502.14458
14. Mingyu Yang, Mehdi Rezagholizadeh, Guihong Li, Vikram Appia, Emad Barsoum. Zebra-Llama: Towards Extremely Efficient Hybrid Models. 2025. openreview.net/forum?id=l42UGsdrNn
15. Yingfa Chen, Zhen Leng Thai, Zihan Zhou, Zhu Zhang, Xingyu Shen, Shuo Wang, Chaojun Xiao, Xu Han, Zhiyuan Liu. Hybrid Linear Attention Done Right: Efficient Distillation and Effective Architectures for Extremely Long Contexts. 2026. arxiv.org/abs/2601.22156
16. Mingjie Liu et al. ProRL: Prolonged Reinforcement Learning Expands Reasoning Boundaries in Large Language Models. 2025. arxiv.org/abs/2505.24864
17. Aakash Lahoti, Kevin Y. Li, Berlin Chen, Caitlin Wang, Aviv Bick, J. Zico Kolter, Tri Dao, Albert Gu. Mamba-3: Improved Sequence Modeling using State Space Principles. 2026. arxiv.org/abs/2603.15569
18. Qwen Team. Qwen3.5: Towards Native Multimodal Agents. February 2026. qwen.ai/blog?id=qwen3.5
19. Ali Behrouz, Peilin Zhong, Vahab Mirrokni. Titans: Learning to Memorize at Test Time. 2024. openreview.net/forum?id=8GjSf9Rh7Z
20. Disen Lan, Jianbin Zheng, Yuxi Ren, Xin Xia, Xuanda Wang, Xuefeng Xiao, Xipeng Qiu, Yu Cheng. Morphing into Hybrid Attention Models. 2026. arxiv.org/abs/2606.30562
21. Aditya Chattopadhyay, Elvis Nunez, Prannay Kaul, Benjamin Bowman, Evan Becker, Luca Zancato, David Thomas, Wei Xia, Stefano Soatto. Priming: Hybrid State Space Models From Pre-trained Transformers. 2026. arxiv.org/abs/2605.08301
22. Yanhong Li, Songlin Yang, Shawn Tan, Mayank Mishra, Rameswar Panda, Jiawei Zhou, Yoon Kim. Distilling to Hybrid Attention Models via KL-Guided Layer Selection. 2025. arxiv.org/abs/2512.20569
23. Michael Luo, Sijun Tan, Justin Wong, Xiaoxiang Shi, William Y. Tang, Manan Roongta, Colin Cai, Jeffrey Luo, Li Erran Li, Raluca Ada Popa, Ion Stoica. DeepScaleR: Surpassing O1-Preview with a 1.5B Model by Scaling RL. 2025. notion.site/DeepScaleRhuggingface.co/datasets/agentica-org/DeepScaleR-Preview-Dataset
24. ByteDance Seed, Tsinghua AIR. DAPO: An Open-Source LLM Reinforcement Learning System at Scale. 2025. arxiv.org/abs/2503.14476huggingface.co/datasets/BytedTsinghua-SIA/DAPO-Math-17k