PaperlayerView exact revision on arXiv

Source-linked AI summary

On-Policy Attention Linearization

Arian Raje, Anupam Nayak, Anthony Fei, Akaash Parthasarathy, Mohamed Abdelfattah, Gauri Joshi

arXiv:2609.31947v1cs.LG

TL;DR

Distilled linear and hybrid attention models can lag on long-context retrieval and reasoning, the settings where their memory efficiency matters. OPAL trains a hybrid student on its own trajectories with teacher supervision. It recovers teacher-level retrieval and substantial teacher performance on reasoning and commonsense benchmarks.

  • Problem

    Distilled hybrid models can fall short of full-attention teachers on long-context retrieval and multi-step reasoning, despite strong likelihood-benchmark performance.

    Source: §1 · p. 2

  • Method

    OPAL adds prolonged on-policy distillation, training the hybrid student on its own sampled continuations with the frozen teacher scoring tokens.

    Source: §1 · p. 2

  • Results

    OPAL matches teacher retrieval, retains 83% and 93% of teacher reasoning performance, and retains 87% and 94% of teacher commonsense performance.

    Source: §1 · p. 2, §4.2 · p. 7

  • Takeaways & Limitations

    A trained transformer’s long-context retrieval and reasoning abilities can survive late architectural conversion, with performance recovered through distillation.

    Source: §5 · p. 9

    Fixed-corpus off-policy distillation can provide a calibrated warm start but does not directly address compounding state drift during long student rollouts.

    Source: §3 · p. 5

Abstract

from arXiv · show

Hybrid transformer architectures that replace most softmax attention layers with linear attention offer transformer-level quality at a fraction of the memory cost. Rather than pretraining such models, a growing body of work distills them from already trained full-attention transformers. However, these distilled models often collapse on long-context retrieval and reasoning tasks, particularly when operating in thinking mode, where the efficiency gains of hybrid architectures matter most. Since linear attention layers must compress context into a fixed-size state, their errors compound over long sequences. As off-policy distillation never teaches the student model to recover from this drift, tasks that necessitate longer sequence lengths become especially challenging. We introduce On-Policy Attention Linearization (OPAL) in which the hybrid attention student samples its own long-context trajectories and receives dense supervision from the frozen full-attention teacher. Applying OPAL to Qwen3-4B and MiMo-7B-RL-0530, we recover $87$--$94\%$ of full-attention performance on commonsense reasoning, $100\%$ on needle-in-a-haystack (NIAH) retrieval, and $83$--$93\%$ on mathematical reasoning with only 3B training tokens. We achieve these results without supervised fine-tuning (SFT) or reinforcement learning with verifiable rewards (RLVR). Compared with the strongest prior linearization method, which recovers $68\%$ of its teacher's retrieval performance and $21.6\%$ absolute average mathematical reasoning accuracy, OPAL fully recovers retrieval and achieves $67.6$--$72.2\%$ on math reasoning.

1 INTRODUCTION

Hybrid attention offers a way to reduce long-context memory costs, but distilled models remain weak on the retrieval and reasoning tasks that motivate linearization. OPAL adds student-generated-trajectory distillation to the conversion pipeline, recovering strong performance across these tasks.

1 INTRODUCTION

  • Linear attention replaces a growing KV cache with a constant-sized recurrent state, targeting long-context retrieval and multi-step reasoning where efficiency gains matter.Prior distilled models can score well on likelihood benchmarks yet fall short on these long-input tasks.Source: §1 · pp. 1–2
  • Prior linearization distillation typically aligns converted-layer hidden states, then trains the whole student to match the teacher’s next-token distribution on pretraining-like data.Some pipelines first select which layers to convert.Source: §1 · pp. 1–2
  • OPAL converts a pretrained full-attention model into a hybrid student through hidden-state alignment, fixed-corpus distillation, and on-policy distillation on student-generated continuations scored by the teacher.The approach appends on-policy distillation to the standard two-stage recipe.Source: §1 · p. 2, Figure 1 · p. 2
  • OPD uses progressively longer student rollouts, a continual-learning-style schedule extended from 200M to 2B tokens, and a curriculum that introduces harder reasoning prompts as easier-task accuracy saturates.Rollout length doubles when teacher-student KL plateaus.Source: §1 · p. 2
  • OPAL recovers 100% of teacher retrieval and 83% and 93% of teacher reasoning performance; its math average reaches 67.6–72.2% versus 21.6% for the strongest prior linearized model.Across six prior methods, the best baseline retains 68% of teacher retrieval accuracy and 66.4% of its reasoning average.Source: §1 · p. 2

2 PROBLEM SETUP AND MOTIVATION

Full attention retains token history in a growing KV cache, whereas linear attention compresses the prefix into a fixed-size recurrent state. Existing off-policy distillation trains on fixed prefixes, leaving students without direct supervision on their own trajectories, where state drift can compound.

2 PROBLEM SETUP AND MOTIVATION

  • Linear attention compresses the prefix into recurrent state S_t, keeping per-token compute and state memory O(1) in sequence length rather than attending to exact prior keys and values.Its state update can be additive or, in GDN, use a delta rule.Source: §2 · p. 3
  • GDN combines data-dependent decay and update strength to provide global forgetting and targeted overwriting in its recurrent state.The passage says this lets GDN perform comparably to full attention.Source: §2 · p. 3
  • Linearization distillation first aligns converted-layer outputs with the teacher, then fine-tunes the student to match the teacher’s next-token distribution on a fixed corpus.The described objective is forward KL divergence.Source: §2 · p. 3
  • Off-policy distillation supervises fixed-corpus prefixes rather than student-generated traces, so it gives no direct signal for recovering from student-induced states.Prior methods including RADLADS, Zebra-Llama, and HALO optimize this objective or a close variant.Source: §2 · p. 3
  • A hybrid model’s memory combines a context-length-growing full-attention term with a fixed recurrent-state term, making savings increasingly significant as context grows.The fixed-size state is maintained by linear attention layers.Source: §2 · pp. 3–4
  • On long rollouts, a changed prefix alters the compressed recurrent state and subsequent predictions, allowing deviations to compound into repetitive or degenerate trajectories.Strong short-context performance does not guarantee long-context robustness.Source: §2 · p. 4, Figure 2 · p. 4

3 PROPOSED METHOD

OPAL converts a full-attention teacher into a hybrid student through representation matching, off-policy distillation, and prolonged on-policy distillation. Its final stage trains on student-generated trajectories under a frozen teacher.

Architecture

  • OPAL retains every fourth full-attention layer and replaces the others with GDN, yielding 9 full-attention and 27 GDN layers in the 36-layer models.Embeddings, MLPs, normalization layers, and the output head are inherited from the teacher.Source: §3 · p. 4

Stage 1: Hidden State Matching

  • Stage 1 independently aligns each converted GDN layer to its corresponding teacher layer, updating only GDN parameters while freezing the retained full-attention layers and other model parameters.This isolates the architectural mismatch but produces coherent text that can remain inaccurate and unrelated.Source: §3 · pp. 4–5

Stage 2: Off-Policy Distillation

  • Stage 2 matches student and teacher predictive distributions on fixed corpus sequences, moving from 4K-token to 32K-token contexts while training all student weights.It serves as a warm start: fixed-corpus prefixes do not expose the student to states induced by its own predictions or directly address compounding state drift.Source: §3 · p. 5

Stage 3: Prolonged On-Policy Distillation

  • Prolonged on-policy distillation trains the student on its own sampled continuations, with the frozen teacher supervising those trajectories using token-level reverse KL.The student is substantially weaker than its teacher on long contexts after Stage 2, making sustained on-policy training necessary as its trajectory distribution evolves.Source: §3 · pp. 5–6

KL-Guided Rollout Horizon Doubling

  • Rollouts start at 512 tokens and double only after reverse-KL improvement plateaus, up to a maximum horizon of 32K.The adaptive schedule progressively trains the student to remain stable over longer sequences of recurrent state updates.Source: §3 · p. 5

Sample Curriculum

  • Once the rollout horizon reaches 1K, the sample curriculum adds retrieval examples spanning 2K–32K tokens, then replaces initial reasoning data with harder math problems at later checkpoints.Difficulty is controlled at the dataset level rather than estimated per sample.Source: §3 · p. 5

Prolonged OPD with WSD

  • A warmup-stable-decay learning-rate schedule lets OPD extend incrementally to 2B tokens while preserving intermediate checkpoints for evaluating retrieval and reasoning recovery.Training continues from pre-decay checkpoints, allowing capability recovery to be measured as a function of OPD compute.Source: §3 · p. 5

Joint Training

  • Retained full-attention layers remain trainable because they alone directly access the full token history, and they train jointly with GDN layers during Stage 3.Stage 2 uses a smaller learning rate for retained full-attention layers; Stage 3 switches to a shared flat rate.Source: §3 · p. 6

4 RESULTS

OPAL students recover strong retrieval and mathematical reasoning performance while retaining the efficiency benefits of hybrid attention. Results cover training-budget trends, matched off-policy controls, and alternative architectures and conversions.

4.1 EXPERIMENTAL SETUP

  • The evaluation covers short-context commonsense, retrieval at 4K–32K, and mathematical reasoning; students receive about 3B total training tokens, including 2B student-generated OPD tokens.The study compares against released linearized models and each full-attention teacher, emphasizing teacher-relative recovery.Source: §4.1 · p. 6

4.2 MAIN RESULTS

  • 100.0% average NIAH accuracy lets OPAL-Qwen3 match its teacher; OPAL-MiMo retains 93% of its teacher’s average math reasoning and fully recovers AIME’24 pass@8.OPAL-Qwen3 and OPAL-MiMo retain 87% and 94% of their teachers’ commonsense-suite performance, respectively.Source: §4.2 · pp. 6–7, Table 1 · p. 7, Table 2 · p. 7
  • At 3B training tokens, OPAL-Qwen3 and OPAL-MiMo use no separate post-training phase beyond distillation, while accepting a commonsense-performance tradeoff.The reported commonsense retention is 87% for OPAL-Qwen3 and 94% for OPAL-MiMo.Source: §4.1 · p. 6, §4.2 · p. 7
  • MATH-500 improves rapidly and saturates early, whereas AIME continues improving through the final checkpoints near the 2B-token OPD budget.The authors suggest that further OPD scaling may yield additional AIME gains because performance has not clearly saturated.Source: §4.2 · p. 7, Figure 3 · p. 8
  • The matched off-policy controls use the Stage 3 data mixture and post-Stage 2 hybrid initialization, but train with off-policy KD rather than student-generated trajectories.The supplied passages describe the control design but do not state its comparative outcome.Source: §4.2 · p. 8, Table 3 · p. 8

4.3 ADDITIONAL EXPERIMENTS

  • At 128K context, cache and recurrent-state memory is approximately 4× lower than full attention, and throughput is approximately 2× higher at highest concurrency.The comparison uses the 1:3 full-attention-to-GDN ratio and vLLM serving.Source: §4.3 · p. 8, Figure 4 · p. 8
  • At matched OPD budgets, Mamba-3 hybrids achieve reasoning performance competitive with GDN, with reported MATH-500 and AIME scores for both models.Qwen3 with Mamba-3 reaches 91.0%, 50.0%, and 35.0% on MATH-500, AIME’24, and AIME’25 at 1B OPD tokens; MiMo reaches 92.4%, 55.0%, and 41.2% at 800M tokens.Source: §4.3 · p. 9, Figure 5 · p. 9
  • OPAL hybrids outperform the non-distillation SWA and TransMLA conversions on AIME and 32K-context NIAH, where both conversion baselines collapse.On MATH-500, OPAL-Qwen3 and OPAL-MiMo reach 92.0% and 93.8%, versus 17.4% and 51.0% for SWA and 27.0% for MiMo TransMLA.Source: §4.3 · p. 9, Table 4 · p. 9

5 CONCLUSION

OPAL converts pretrained full-attention transformers into hybrid models while maintaining long-context retrieval and reasoning abilities. Prior work spans several post-training conversions, but existing linearization methods often lag on long-context tasks.

5 CONCLUSION

  • OPAL fully recovers long-context retrieval and closes the mathematical-reasoning gap while retaining only a quarter of the full-attention layers.The method adds prolonged on-policy distillation so the hybrid student learns from states encountered during generation.Source: §5 · p. 9

A RELATED WORKS

  • Hybrid architectures retain some full attention because pure fixed-state models have recall and retrieval limitations, despite enabling subquadratic sequence processing.Design choices include the ratio and placement of full and linear attention layers, the attention mixer, and whether mechanisms are interleaved or combined.Source: A RELATED WORKS · p. 17
  • Prior conversion methods range from adding modules to modifying pretrained backbones, but benchmarks vary and some approaches do not recover long-context performance.Capability recovery has generally relied on significant continued pretraining.Source: A RELATED WORKS · p. 17
  • Existing linearization methods typically combine layerwise alignment, off-policy distillation, or continued pretraining, yet long-context reasoning and retrieval often lag.Priming is an exception, but uses up to 128K context lengths, supervised fine-tuning on chain-of-thought traces, instruction alignment, and 150B tokens.Source: A RELATED WORKS · p. 17
  • On-policy distillation replaces sparse outcome rewards with dense token-level supervision on student generations and has transferred reasoning capabilities to smaller models.It has also been used in multi-teacher post-training and in a 550B-parameter hybrid model's pipeline.Source: A RELATED WORKS · pp. 17–18

B.1 MODELS AND HYBRID CONSTRUCTION

The converted hybrids retain the same full-attention layers and use Gated DeltaNet for linear attention.

B.1 MODELS AND HYBRID CONSTRUCTION

  • Both teacher models and converted hybrids retain the same nine full-attention layers and use Gated DeltaNet as the linear attention layer.Source: Table 5 · p. 18

B.2 TRAINING PROCEDURE AND HYPERPARAMETERS

The appendix reports optimizer settings, stage configurations, rollout sampling conditions, and the hardware used for training.

B.2 TRAINING PROCEDURE AND HYPERPARAMETERS

  • Runs use 3–4 NVIDIA H100 or H200 GPUs, depending on availability.Source: §B.2 · p. 18
  • All stages use AdamW with β1=0.9, β2=0.95, weight decay 0, and gradient clipping at norm 1.0.“Inherited” denotes all non-GDN weights.Source: Table 6 · p. 18
  • The reported training configurations use 100M, 600M, 294M, or 2B tokens and sequence lengths of 4K or 32K.Source: §B.2 · p. 18
  • The final two listed configurations use WSD learning-rate schedules, while the earlier listed configurations use cosine schedules.Source: §B.2 · p. 18
  • Stage-3 rollouts are sampled with vLLM at T=1.0 and topp=1.0 under the thinking-mode chat template.Source: Table 7 · p. 18

B.3 OPD SCALING LADDER

The OPD scaling ladder continues each rung from the prior rung's pre-decay snapshot before linearly decaying the learning rate.

B.3 OPD SCALING LADDER

  • Each OPD scaling-ladder rung continues from the previous rung’s pre-decay snapshot at constant learning rate, then decays linearly to the final learning rate over its decay window.Source: Table 8 · p. 19

C ADDITIONAL EXPERIMENTS

Table 9 reports commonsense reasoning performance, with teacher and student models identified.

C.1 COMMONSENSE REASONING PERFORMANCE

  • Table 9 presents commonsense reasoning performance (%) and identifies the teacher and student models.Source: Table 9 · p. 19

C.2 WSD TRAINING CURVES

The WSD training curves cover both models across OPAL scaling rungs, while the decay horizons lengthen and then plateau.

C.2 WSD TRAINING CURVES

  • Decay horizons grow at successive OPAL scaling rungs before plateauing in the final rungs.Source: Figure 6 · p. 20
  • Qwen3-4B training curves are shown for each rung using the WSD schedule.Source: Figure 7 · p. 20
  • MiMo-7B-RL-0530 training curves are shown for each rung using the WSD schedule.Source: Figure 8 · p. 20
  • Position-wise reverse-KL curves are provided for the first Qwen3-4B rung and its 1B–1.2B OPD-token run.Source: Figure 9 · p. 21, Figure 10 · p. 21
  • Position-wise reverse-KL curves are provided for the first MiMo-7B-RL-0530 rung and its 1B–1.2B OPD-token run.Source: Figure 11 · p. 22, Figure 12 · p. 22

C.3 SAMPLE RESPONSES

The sample responses include long-context number retrieval and mathematical reasoning examples from hybrid models. The math examples give final answers of 110 and 104, while the retrieval examples reproduce the requested hidden numbers.

AIME 2024 I, Problem 13 — Qwen3-4B GDN Hybrid

  • The smallest prime p for which some n makes n^4 + 1 divisible by p^2 is 17; for 17^2 = 289, the least positive m is 110.Source: §C.3 · p. 26

NIAH Multi-Key, 32K Context — Qwen3-4B GDN Hybrid Prompt

  • In a 32K-context multi-key retrieval example, Qwen3-4B returns 1144632 for luminous-ember.Source: §C.3 · pp. 26–27

NIAH Multi-Query, 32K Context — MiMo-7B-RL-0530 GDN Hybrid Prompt

  • In a 32K-context multi-query retrieval example, MiMo-7B-RL-0530 returns all four requested key-number pairs.The pairs are vivid-walnut: 8616278, placid-lantern: 7006521, mellow-bramble: 8084209, and rugged-ember: 5916232.Source: §C.3 · p. 29

Related papers

Loading 2609.31947v1…