On-Policy Attention Linearization

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

cs.LGSep 25, 2026

On-Policy Attention Linearization

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

Organizations: Carnegie Mellon University · Cornell University

Abstract

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.

Figures & tables

Figure 1: Overview of Opal . We convert a full attention teacher into a hybrid student by replacing $75\%$ of attention layers with linear attention. We recover the teacher’s capabilities through hidden-state alignment, short- and long-context off-policy distillation, and prolonged on-policy distillation.

Figure 2: Degenerate rollouts during student decoding. Off-policy distillation aligns the student on fixed prefixes, but provides no signal on the student’s own trajectories. Student rollouts can collapse into degenerate rollouts such as repetitive text or unrelated prose. Open circles indicate

0\%

	Single Needle				Multi-Key				Multi-Query
Model	4K	8K	16K	32K	4K	8K	16K	32K	4K	8K	16K	32K	Avg.
Qwen2.5-7B-Instruct	$100.0$	$100.0$	$100.0$	$100.0$	$100.0$	$100.0$	$100.0$	$100.0$	$100.0$	$100.0$	$100.0$	$100.0$	$100.0$
QRWKV6-7B-Instruct	$0.0$	$0.0$	$0.0$	$0.0$	$1.0$	$0.0$	$0.0$	$0.0$	$0.0$	$0.0$	$0.0$	$0.0$	$0.1$ ( $0.1\%$ )
Llama-3.2-1B-Instruct	$100.0$	$100.0$	$100.0$	$100.0$	$94.0$	$93.0$	$71.0$	$67.0$	$98.5$	$94.2$	$93.2$	$96.5$	$92.3$
Llamba-1B	$0.0$	$0.0$	$0.0$	$0.0$	$0.0$	$0.0$	$0.0$	$0.0$	$0.0$	$0.0$	$0.0$	$0.0$	$0.0$ ( $0.0\%$ )
Llama-3.2-3B-Instruct	$100.0$	$100.0$	$100.0$	$100.0$	$98.0$	$95.0$	$94.0$	$87.0$	$100.0$	$99.2$	$87.5$	$68.5$	$94.1$

Table 1: Needle-in-a-Haystack retrieval accuracy (%). Teacher Model ; Student Model .

					AIME’24		AIME’25
Model	GSM8K	GSM-Plus	GSM-Sym	MATH-500	pass@1	pass@8	pass@1	pass@8	Avg.
Qwen2.5-7B-Instruct	$91.9$	$72.7$	$83.4$	$76.4$	$10.4$	$26.7$	$7.5$	$23.3$	$49.0$
QRWKV6-7B-Instruct	$27.4$	$16.8$	$9.3$	$24.0$	$0.4$	$3.3$	$0.4$	$3.3$	$10.6$ ( $21.6\%$ )
Llama-3.2-1B-Instruct	$38.5$	$24.9$	$26.2$	$25.8$	$0.8$	$6.7$	$0.0$	$0.0$	$15.4$
Llamba-1B	$28.1$	$12.9$	$13.0$	$8.6$	$0.0$	$0.0$	$0.0$	$0.0$	$7.8$ ( $51.0\%$ )
Llama-3.2-3B-Instruct	$71.1$	$58.3$	$59.5$	$39.0$	$2.5$	$13.3$	$0.0$	$0.0$	$30.5$

Table 2: Mathematical reasoning performance (%). Teacher Model ; Student Model .

Figure 3: Math reasoning performance as a function of OPD training tokens. MATH-500 improves rapidly and saturates early, while AIME continues to benefit from additional OPD compute.

Figure 4: Memory and throughput improvements of GDN hybrids over the full attention teacher.

	MATH-500		AIME’24		AIME’25
Training	200M	400M	200M	400M	200M	400M
OPD (Student)	$\mathbf{83.6}$	$\mathbf{87.0}$	$\mathbf{23.3}$	$\mathbf{33.3}$	$\mathbf{18.8}$	$\mathbf{31.7}$
Off-Policy (Dataset)	$77.3$	$78.1$	$12.5$	$17.1$	$14.6$	$15.8$
Off-Policy (Teacher)	$75.5$	$80.0$	$12.9$	$17.5$	$13.8$	$16.7$

Table 3: OPD versus off-policy controls at matched training budgets. Parenthesis show source of training traces.

Figure 5: Comparison of how reasoning scales for various hybrid architectures.

	NIAH, 32K Context			Math Reasoning
Model / Conversion	Single Needle	Multi-Key	Multi-Query	MATH-500	AIME’24	AIME’25
Qwen3-4B	$100.0$	$99.8$	$100.0$	$96.4$	$73.8$	$63.7$
$\hookrightarrow$ SWA $(256,4)$	$0.0$	$1.0$	$0.2$	$17.4$	$0.0$	$0.0$
$\hookrightarrow$ GDN (ours, 2B OPD)	$\mathbf{100.0}$	$\mathbf{99.6}$	$\mathbf{100.0}$	$\mathbf{92.0}$	$\mathbf{52.5}$	$\mathbf{35.0}$
MiMo-7B-RL-0530	$100.0$	$92.8$	$76.2$	$97.0$	$73.3$	$67.9$
$\hookrightarrow$ SWA $(256,4)$	$0.0$	$1.0$	$0.3$	$51.0$	$1.7$	$1.7$

Table 4: Comparison to non-distillation baselines on NIAH (32K) and math reasoning (%).

Appendix figures & tables12 assets

Supplementary material from the paper’s appendix.

Appendix

	Qwen3-4B	MiMo-7B-RL-0530
Teacher	Qwen/Qwen3-4B	XiaomiMiMo/MiMo-7B-RL-0530
Teacher Parameters	4.02B	7.62B (MTP head dropped)
Layers	36, GQA 32 Q/8 KV, $d_{h}=128$	36, GQA 32 Q/8 KV, $d_{h}=128$
Context Length	32,768	65,536
Retained Full Attention Layers	9 (Layers $3,7,\dots,35$ )
Converted Layers	27 (Layers 0, 1, 2, 4, 5, 6, …, 34)

Appendix

Table 5: Teacher models and the converted hybrids. Both retain the same nine full attention layers and use Gated DeltaNet (GDN) as the linear attention layer.

	Stage 1		Stage 2A		Stage 2B		Stage 3
	Qwen	MiMo	Qwen	MiMo	Qwen	MiMo	Qwen	MiMo
Trainable	GDN	GDN	All	All	All	All	All	All
Tokens	100M	100M	600M	600M	294M	294M	2B	2B
Sequence length	4K	4K	4K	4K	32K	32K	32K	32K
Tokens per batch	32K	16K	2.1M	2.1M	2.1M	2.1M	131K	131K
Optimizer Steps	3.1K	6.1K	287	287	140	140	12K	12K

Appendix

Table 6: Optimization hyperparameters by stage. All stages use AdamW (

\beta_{1}{=}0.9

\beta_{2}{=}0.95

, weight decay 0), gradient clipping at norm 1.0. “Inherited” denotes all non-GDN weights.

	Qwen3-4B	MiMo-7B-RL-0530
Generated Tokens per Step	131,072	131,072
Prompt + Rollout Limit	40,960	65,536
Initial Horizon $H$	512	512
Horizon Doubling Rule	KL Plateau: rel. gain $<$ 2% for 30 steps
Maximum Horizon	32K	32K
Sampler Batch (Prompts)	64; 32 at $H{=}32$ k	64; 32 at $H{=}32$ k

Appendix

Table 7: Stage-3 on-policy distillation settings. Rollouts are sampled with vLLM at

T{=}1.0

, top-

p{=}1.0

under the thinking-mode chat template.

Qwen3-4B
Rung	Max Output Horizon	Decay Start	Decay Tokens
0 $\rightarrow$ 200M	512 $\rightarrow$ 16K	161M	39M
200M $\rightarrow$ 400M	16K	340M	60M
400M $\rightarrow$ 500M	16K	430M	70M
500M $\rightarrow$ 600M	16K	520M	80M
600M $\rightarrow$ 800M	16K $\rightarrow$ 32K	710M	90M

Appendix

Table 8: OPD scaling ladder configuration. Each rung continues the previous one from its pre-decay snapshot at the constant learning rate, then decays linearly to the final LR over the decay window.

Model	PIQA	HellaSwag	ARC-E	ARC-C	Winogrande	MMLU	Avg.
Qwen2.5-7B-Instruct	$80.2$	$80.5$	$80.9$	$55.0$	$70.5$	$74.3$	$73.6$
QRWKV6-7B-Instruct	$\mathbf{79.9}$	$78.9$	$\mathbf{79.3}$	$55.9$	$71.4$	$\mathbf{64.2}$	$\mathbf{71.6}$ ( $97.3\%$ )
Llama-3.2-1B-Instruct	$74.9$	$61.6$	$63.9$	$37.6$	$61.4$	$46.0$	$57.6$
Llamba-1B	$73.7$	$61.8$	$65.2$	$37.5$	$61.9$	$31.5$	$55.3$ ( $96.0\%$ )
Llama-3.2-3B-Instruct	$76.8$	$71.6$	$71.0$	$46.4$	$68.7$	$60.7$	$65.9$
Llamba-3B	$78.0$	$74.0$	$73.6$	$46.2$	$\mathbf{71.7}$	$50.1$	$65.6$ ( $99.6\%$ )

Appendix

Table 9: Commonsense reasoning performance (%). Teacher Model ; Student Model .

Figure 6: Training horizons for each rung of the OPD scaling ladder. Decay horizons gradually become longer at subsequent rungs before plateauing in the last few rungs.

Figure 7: Training curves for Qwen3-4B at each rung using WSD schedule.

Figure 8: Training curves for MiMo-7B-RL-0530 at each rung using WSD schedule.

Figure 9: Training curves for position-wise reverse-KL at the first rung of Qwen3-4B training.

Figure 10: Training curves for position-wise reverse-KL at the run 1B-1.2B OPD tokens of Qwen3-4B training.

Figure 11: Training curves for position-wise reverse-KL at the first rung of MiMo-7B-RL-0530 training.

Figure 12: Training curves for position-wise reverse-KL at the run 1B–1.2B OPD tokens of MiMo-7B-RL-0530 training.

Linear Attention

Linear attention reduces sequence-length costs by replacing pairwise token interactions with kernel features, recurrent states, or compact memory.

See all 127 papers

STILL: Selecting Tokens for Intra-Layer Hybrid Attention to Linearize LLMs
Feb 2, 2026Weikang Meng, Liangyu Huo, Yadan Luo +4LLM Inference Acceleration Linear Attention
SMULL Group, Harbin Institute of Technology, Shenzhen · Pengcheng Laboratory · UQMM Lab, University of Queensland +1
Teach Yourself Where to Look: On-Policy Attention Self-Distillation for Reasoning
Sep 27, 2026Safaeid Hossain Arib, Rabeya Akter, Ismam Nur Swapnil +3Numerical Reasoning in Language Models Attention-Guided Distillation
ACI PLC, Bangladesh · University of Dhaka, Bangladesh · BRAC University, Bangladesh +2
The Key to Going Linear: Analysis-Driven Transformer Linearization
Jul 8, 2026Anna Kuzina, Paul N. Whatmough, Babak Ehteshami BejnordiSoftmax Attention LLM Inference Acceleration
Qualcomm AI Research∗ · Qualcomm AI Research

Feb 2, 2026J/K move · Enter open · S save

STILL: Selecting Tokens for Intra-Layer Hybrid Attention to Linearize LLMs

Weikang Meng, Liangyu Huo, Yadan Luo, Jiawen Guan+3

SMULL Group, Harbin Institute of Technology, Shenzhen · Pengcheng Laboratory · UQMM Lab, University of Queensland+1

LLM Inference Acceleration Linear Attention Attention Mechanisms Hybrid Attention

Open paper AskPDF arXiv

Linearizing pretrained large language models (LLMs) primarily relies on intra-layer hybrid attention mechanisms to alleviate the quadratic complexity of standard softmax attention. Existing methods perform token routing based on sliding-window partitions, resulting in position-based selection and fails to capture token-specific global importance. Meanwhile, linear attention further suffers from distribution shift caused by learnable feature maps that distort pretrained feature magnitudes. Motivated by these limitations, we propose STILL, an intra-layer hybrid linearization framework for efficiently linearizing LLMs. STILL introduces a Self-Saliency Score with strong local-global consistency, enabling accurate token selection using sliding-window computation, and retains salient tokens for sparse softmax attention while summarizing the remaining context via linear attention. To preserve pretrained representations, we design a Norm-Preserved Feature Map (NP-Map) that decouples feature direction from magnitude and reinjects pretrained norms. We further adopt a unified training-inference architecture with chunk-wise parallelization and delayed selection to improve hardware efficiency. Experiments show that STILL matches or surpasses the original pretrained model on commonsense and general reasoning tasks, and achieves up to a 86.2% relative improvement over prior linearized attention methods on long-context benchmarks. Code is available at this URL.

Abstract

Linear Attention

Explore similar work

STILL: Selecting Tokens for Intra-Layer Hybrid Attention to Linearize LLMs

Teach Yourself Where to Look: On-Policy Attention Self-Distillation for Reasoning

The Key to Going Linear: Analysis-Driven Transformer Linearization