PaperlayerView exact revision on arXiv

Source-linked AI summary

Attention Amnesia in Hybrid LLMs: When CoT Fine-Tuning Breaks Long-Range Recall, and How to Fix It

Xinyu Zhou, Boyu Zhu, Yi Xu, Zhiwei Li, Yingfa Chen, Huiming Wang, Zhijiang Guo

arXiv:2606.11052v1cs.CL

TL;DR

CoT-SFT can substantially weaken long-context recall in distilled hybrid models, creating a gap between reasoning gains and retrieval stability. This paper analyzes the routing changes behind that failure and introduces QK-Restore, which restores query-key projections without retraining; across architectures, it recovers long-context capability while largely preserving reasoning performance.

  • Problem

    CoT-SFT improves reasoning but can substantially degrade long-context recall in distilled hybrid models, especially for challenging retrieval settings and extended contexts.

  • Method

    QK-Restore restores query and key projections from the pre-SFT checkpoint while preserving all other post-SFT parameters, with a Procrustes variant balancing routing preservation and reasoning adaptation.

  • Results

    Across multiple architectures, QK-Restore substantially recovers long-context capability while largely preserving SFT reasoning gains; HypeNet-9B NIAH-S2@256K decreases from 67.2% to 9.4% after SFT.

  • Takeaways & Limitations

    Routing stability in retained softmax-attention layers is a key factor for efficient long-context model adaptation after reasoning-oriented SFT.

  • Takeaways & Limitations

    The theoretical analysis assumes CoT reasoning traces follow a latent Markov structure with first-order Markov transitions.

Abstract

from arXiv · show

Chain-of-thought (CoT) supervised fine-tuning (SFT) is widely adopted to improve reasoning ability, yet we find that it systematically degrades long-context recall in hybrid linear-attention models. Across architectures including HypeNet and Jet-Nemotron, retrieval performance on Needle-In-A-Haystack (NIAH) deteriorates substantially after CoT-SFT, and the degradation becomes more severe under harder retrieval settings and longer context windows. For example, HypeNet-9B on NIAH-S2@256K decreases from $67.2\%$ to $9.4\%$. We attribute this to CoT-SFT biasing attention gradients toward short-range patterns, disrupting query-key projections ($W_Q, W_K$) that are responsible for long-range routing. Motivated by this observation, we propose QK-Restore, a training-free method that restores only $W_Q$ and $W_K$ from the pre-SFT checkpoint while preserving all other post-SFT parameters. We further introduce a Procrustes variant to balance routing preservation and reasoning adaptation. Across architectures, QK-Restore consistently restores long-context capability at zero training cost while preserving reasoning performance; for instance, on HypeNet-5B it improves S3@256K from $65.4\%$ to $76.4\%$ while maintaining strong reasoning performance.

1. Introduction

The introduction identifies CoT-SFT as a post-training failure mode that improves local reasoning while substantially degrading long-context recall in distilled hybrid models. It attributes this degradation to locality-biased changes in query-key routing and introduces QK-Restore, which restores long-range recall while preserving post-SFT reasoning adaptations.

  • Problem: CoT-SFT improves mathematical and multi-step reasoning but can substantially degrade the long-context recall acquired during pretraining or architectural conditioning.This degradation is highlighted in distilled hybrid models used as efficient base models.
  • Mechanism: The CoT-Markov analysis predicts exponentially decaying attention-logit gradients with token distance, favoring local interactions over long-range routing.Token autocorrelation and attention-logit gradient measurements empirically validate this prediction.
  • Mechanism: CoT-SFT induces locality-biased drift in query-key projections W_Q and W_K, which determine information retrieval, while value-side extraction can still benefit from adaptation.Long-context recall depends heavily on a small number of retained softmax-attention layers.
  • Method: QK-Restore restores only query and key projections from the pre-SFT checkpoint while leaving value-side and other post-SFT parameters unchanged.The method is training-free and targets the retained softmax-attention layers in hybrid models.
  • Outcome: The proposed restoration recovers long-context recall while largely preserving reasoning performance after CoT-SFT.The contribution is framed as a structured post-training failure analysis paired with a targeted repair.

2. Related Work

Related work establishes hybrid attention models as efficient long-context systems and highlights unresolved risks from reasoning-oriented post-training. It also frames CoT as reinforcing local dependencies, potentially undermining the query-key routing needed for long-range recall.

  • Efficient Long-Context Models and Hybrid Attention: Hybrid attention models interleave softmax-attention with recurrent or linear-attention layers to support reasoning, long-context recall, and efficient inference.The passage notes that pretraining strong hybrid models from scratch at large scale remains prohibitively expensive.
  • Distilling Transformers into Hybrid Models: Transformer-to-hybrid distillation converts selected softmax-attention layers into linear or recurrent mixers, with layer selection critically influencing long-context retrieval.Strong recall after conversion does not necessarily imply stability after reasoning-oriented post-training.
  • CoT Dynamics and Attention Routing: A Markovian view of CoT models reasoning as local transitions between latent states, suggesting that CoT supervision reinforces short-range step-to-step dependencies.Recent work treats CoT reasoning as a structured generation process over intermediate reasoning states.
  • CoT Dynamics and Attention Routing: Locality-biased CoT-SFT can erode long-range attention routing required for recall in efficient long-context models.Attention analyses distinguish query-key interactions, which define routing logits selecting positions, from value projections, which determine extracted content.

3. Background and Problem Setup

Hybrid models combine a small number of softmax-attention layers with efficient linear-attention layers, making the routing behavior of the softmax layers central to long-context recall. Because routing is determined by W_Q and W_K, distance-dependent attention gradients explain how training can shift these matrices toward local patterns.

  • Hybrid architectures interleave a small set of softmax-attention layers with efficient linear-attention layers to improve long-context efficiency while retaining softmax attention’s expressive power.
  • Long-range recall depends disproportionately on designated softmax-attention layers, so preserving their routing behavior is critical in hybrid models.
  • Attention routing is determined solely by W_Q and W_K, whose changes directly alter long-range retrieval behavior.
  • The distance dependence of ∂L/∂e_t_s is central to understanding how CoT-SFT affects retrieval, through its gradient updates to QK projections.
  • When attention-logit gradients emphasize small distances and neglect large ones, gradient descent pushes W_Q and W_K toward local patterns regardless of context length.

4. Why CoT-SFT Disrupts Routing?

The section models CoT data as a latent Markov process and shows that its training gradients become increasingly local with distance. Empirical decay patterns support the mechanism: CoT-SFT strengthens short-range reasoning while weakening long-range retrieval.

  • CoT-Markov structure: CoT data is modeled with latent reasoning states following Markov transitions, ergodicity, reversibility, and a spectral gap.The assumption’s validity is discussed in Appendix A.
  • Routing mechanism: Softmax-attention layers are the sole locus of long-range recall, so the analysis targets how CoT-SFT erodes their routing capacity through gradient locality.The focus is on attention layers in the hybrid model rather than other components.
  • Gradient locality: Theorem 4.6 predicts that expected attention-gradient magnitude decays exponentially with distance under the CoT-Markov model.The decay rate is determined by the latent transition matrix’s spectral gap, while the multiplicative constant depends on model norm bounds.
  • Implication: CoT-SFT improves short-range reasoning ability in hybrid models at the cost of long-range retrieval.This summarizes the section’s observed tradeoff between local reasoning and distant recall.
  • Empirical validation: 172 tok versus 115 tok: CoT text has a longer token-autocorrelation scale than general prose, while both token and gradient correlations follow exponential decay fits.The reported scales are W_corr ≈172 tok for mathematical reasoning and W_corr ≈115 tok for general prose; Figure 3 validates the decay patterns on HypeNet-2B.
  • Empirical validation: When W_grad < W_corr, distances between those scales require long-range attention but receive no reinforcing training gradient.This creates a band where data demands long-range attention while optimization no longer reinforces it.

5. Method

The method identifies an asymmetry in CoT-SFT: W_Q and W_K lose long-range routing capacity, while W_V and W_O retain extraction improvements. QK-Restore repairs routing by transplanting pre-SFT W_Q and W_K while preserving post-SFT W_V and W_O.

  • Gradient asymmetry: Under CoT-SFT, gradients to W_Q and W_K decay as ρ^τ with pair distance, concentrating routing updates on short-range interactions.The result assumes token generation follows a stationary reversible ergodic Markov chain with spectral gap 1 − ρ > 0.
  • Gradient asymmetry: Gradients reaching W_V remain uniformly bounded below across positions, enabling value-processing improvements independent of context length.The extraction-side gradient behavior contrasts with the short-range concentration affecting routing parameters.
  • Method motivation: CoT-SFT corrupts W_Q and W_K while improving W_V and W_O, separating routing damage from extraction gains across disjoint parameter sets.This separation motivates repairing only the routing parameters rather than reverting the entire post-SFT model.
  • QK-Restore: QK-Restore transplants W_Q and W_K from the pre-SFT checkpoint while retaining post-SFT W_V and W_O.The procedure initializes the repaired model from post-SFT weights and applies the transplant in softmax-attention layers.

6. Experiments

Experiments evaluate CoT-SFT and QK-Restore across HypeNet and Jet-Nemotron on long-context recall and mathematical reasoning. CoT-SFT substantially harms difficult, long-context retrieval, whereas QK-Restore recovers retrieval while largely preserving reasoning performance.

  • Experimental setup: Experiments use HypeNet models from 2B to 9B and Jet-Nemotron-2B1, evaluating NIAH recall alongside GSM8K and MATH500 reasoning.HypeNet is trained in pre-training and SFT stages, while Jet-Nemotron receives SFT from pre-trained checkpoints.
  • CoT-SFT degradation: CoT-SFT weakens long-context retrieval most on harder NIAH-Single-2 and NIAH-Single-3 settings and at longer context lengths.NIAH-Single-1 remains nearly saturated, while retrieval complexity amplifies degradation by demanding more precise long-range token interactions.
  • CoT-SFT degradation: HypeNet-2B on NIAH-Single-3 at 128K drops from 30.4 to 7.60 after SFT, while HypeNet-9B on NIAH-Single-2 at 256K decreases from 67.2 to 9.40.The degradation consistently becomes more severe as context length increases.
  • QK-Restore retrieval: QK-Restore consistently mitigates SFT-induced long-context losses, with the largest gains occurring where degradation is most severe.It can also surpass the original pre-training baseline in some configurations.
  • QK-Restore retrieval: HypeNet-2B NIAH-Single-3 at 256K improves from 10.6 to 30.2 (+19.6), while HypeNet-9B improves from 22.8 to 42.6 (+19.8).Under OpenThoughts-3 CoT-SFT, HypeNet-9B improves from 22.4 to 45.0 on NIAH-Single-2 at 256K and from 40.6 to 61.2 on NIAH-Single-3 at 256K.
  • Reasoning preservation: QK-Restore largely preserves CoT-SFT reasoning gains, with HypeNet-5B changing by only −0.6 points on MATH500 and −0.4 points on GSM8K.Similar trends appear across model scales while long-context retrieval improves.

7. Analysis

The analysis shows that long-context degradation is specific to CoT-style supervision rather than post-training generally, and extends from mathematical reasoning to coding. Attention visualizations and ablations further support QK-Restore’s routing-based mechanism and motivate the QK-Pro variant.

  • CoT versus instruction-following SFT: Instruction-following SFT preserves or improves long-context retrieval, unlike CoT-SFT, showing the degradation is not inherent to post-training.HypeNet-5B improves from 75.2 to 83.8 on NIAH-Single-3 at 256K, while HypeNet-9B improves from 52.0 to 86.8.
  • Coding-domain generalization: LiveCodeBench-V5 rises from 0.27 to 6.37 on HypeNet-5B and from 1.66 to 13.11 on HypeNet-9B after CoT-SFT, alongside long-context retrieval degradation.The same retrieval–reasoning trade-off therefore persists beyond mathematical reasoning in the coding domain.
  • Coding-domain generalization: QK-Restore improves HypeNet-5B coding from 6.37 to 6.92 while recovering NIAH-Single-3 at 256K from 75.2 to 79.4.On HypeNet-9B, NIAH-Single-3 at 256K improves from 40.6 to 61.2 while coding changes from 13.11 to 12.91.
  • Attention-map analysis: Layer 33 is most affected after CoT-SFT, with 25 of 32 heads shifting toward shorter attention distances.The attention-map analysis identifies localized routing changes in HypeNet-5B’s softmax-attention layers.
  • Attention-map analysis: QK-Restore partially recovers suppressed long-range routing, with distant-key attention stronger than post-SFT attention for most query positions.The visualization shows this recovery in the off-diagonal lower triangle of Layer-33 Head-8.
  • QK restoration analysis: Full QK restoration recovers long-range recall but can erase beneficial math adaptation because harmful routing drift and math ability are entangled in W_Q.This motivates modeling the correction as a Procrustes problem and introducing QK-Pro.
  • QK restoration analysis: V-Restore fails to recover NIAH, while isolated Q-Restore and K-Restore provide only partial recovery; K-Restore reaches 13.2 versus 30.8 on NIAH-Single-3 at 128K.The ablation attributes the mismatch to isolated restoration and supports joint restoration for coherent routing geometry.

8. Conclusion … G. Proof for Token Autocorrelation Decay

The paper concludes that CoT-SFT can disrupt long-context recall by perturbing query-key routing, motivating training-free QK-Restore, while its theoretical analysis explains gradient locality through Markov-chain correlation decay. The appendices formalize assumptions, spectral arguments, gradient bounds, and token autocorrelation decay underlying this explanation.

  • 8. Conclusion: CoT-SFT-induced long-context recall degradation is identified as a critical failure mode caused by short-range optimization that perturbs query-key routing geometry.The proposed QK-Restore restores only query and key projections from the pre-SFT checkpoint without training.
  • A. Discussion of the CoT-Markov Structure Assumption: A k-step Markov chain can be represented as a 1-step chain on a K^k-state product space, preserving exponential decay with a larger reasoning horizon W.The 1-step Markov condition is therefore presented as an analytical idealisation.
  • A. Discussion of the CoT-Markov Structure Assumption: Reversibility makes the Markov operator P self-adjoint on L2(π), but causally generated CoT chains need not satisfy detailed balance.The discussion characterizes reversibility as sufficient rather than necessary for the bound.
  • B. Justification for Norm Bounds Assumption: Layer normalization and the fixed value projection provide a finite model-dependent bound on the value vector norm.The bound depends on learned scale and bias parameters and the operator norm of W_V.
  • C. Justification for Score-function Identity Assumption: At the optimum p_θ = p_data, the cross-entropy gradient with respect to attention output o_t vanishes conditionally and therefore has zero expectation.The result follows by first conditioning on x_1:t and then applying the law of total expectation over x_1:t.
  • D. Preliminary Definition and Lemma: For a finite reversible Markov chain, detailed balance makes P symmetric in L2(π), enabling an orthonormal eigenbasis with spectral gap 1 − ρ.Mean-zero functions have no component along the stationary eigenfunction and expand only over nonstationary modes.
  • E. Proof for Spectral Correlation Decay Lemma: The spectral-correlation proof uses stationarity, the τ-step prediction P^τf, eigenbasis expansion, orthonormality, and Cauchy-Schwarz to derive decay bounds.These steps connect Markov transition powers to correlations of mean-zero functions.
  • G. Proof for Token Autocorrelation Decay: Token autocorrelation decays geometrically because conditionally independent emissions reduce the covariance to latent mean-zero functions governed by the second-largest eigenvalue ρ.The conclusion states that observed ¯ρ(τ) ∼ Cρ̂^τ with ρ̂ < 1 implies spectral gap 1 − ρ̂ > 0.

H. Discussion of Attention Gradient Decay in Pre-training Stage

Attention-gradient locality applies beyond CoT to general pre-training data, where routing degrades more slowly but through the same mechanism. In hybrid architectures, this pre-training erosion reduces recall capacity before fine-tuning and is amplified by CoT fine-tuning.

  • H. Discussion of Attention Gradient Decay in Pre-training Stage: g(τ) ≤ C′e−τ/W applies to any training corpus, including general pre-training data.The bound follows from assumptions about ergodic Markov chains with a spectral gap, which are also observed in general pre-training corpora.
  • H. Discussion of Attention Gradient Decay in Pre-training Stage: A smaller ΔW = Wcorr − Wgrad mismatch in general pre-training makes routing degrade more slowly than under CoT, through the same mechanism.The mismatch measures how far the data requires attention versus how far gradients reinforce it.
  • H. Discussion of Attention Gradient Decay in Pre-training Stage: Hybrid architectures lose meaningful recall capacity during pre-training, leaving them closer to the threshold before CoT fine-tuning amplifies the erosion.Their vulnerability therefore begins in pre-training and is further intensified by CoT fine-tuning.

I. Proof for the Routing-Extraction Gradient Decoupling Theorem · J. Training Data and Configuration Details · K. Evaluation Details

The paper proves bounded gradient behavior for routing-related terms, specifies HypeNet’s pre-training and CoT-SFT setup, and details evaluation protocols for recall, reasoning, and coding tasks.

  • I. Proof for the Routing-Extraction Gradient Decoupling Theorem: The proof invokes Theorem 4.6 together with bounded norms ∥k_s∥≤B_K and ∥h_t∥≤B_h.These bounds are used in the routing-extraction gradient decoupling argument.
  • I. Proof for the Routing-Extraction Gradient Decoupling Theorem: The derivation accumulates downstream dependencies and separately tracks extraction parameters through v_s=W_Vh_s entering the loss via o_t.The supplied proof passages distinguish downstream accumulation from value-vector extraction pathways.
  • I. Proof for the Routing-Extraction Gradient Decoupling Theorem: The theorem’s assumptions require A_ss≥δ_A>0 and E[∥G_s∥]≥c_G>0 for every s.These conditions establish positive lower bounds for the attention and gradient quantities used in the proof.
  • J. Training Data and Configuration Details: HypeNet uses FineWeb-Edu and UltraData-Math during pre-training, with an 80% general-web and 20% math-focused mixture.The mixture is intended to combine broad linguistic and factual competence with mathematical and STEM reasoning support.
  • J. Training Data and Configuration Details: HypeNet follows a three-stage pre-training setup consistent with Chen et al. (2026), with configurations listed for 2B, 5B, and 9B models.The detailed pre-training configurations appear in Tables 5–7.
  • J. Training Data and Configuration Details: CoT-SFT uses the MiroMind-M1 dataset, assembled from OpenR1, Open-thoughts, Light-R1, and Synthetic-1.SFT configurations are provided separately for 2B and 5B/9B models in Tables 8 and 9.
  • K. Evaluation Details: NIAH evaluation uses LM Evaluation Harness for the official test.The supplied passage identifies the harness used for NIAH evaluation.
  • K. Evaluation Details: MATH500 and GSM8K report average pass@1 over 16 generations and Maj@16, while LiveCodeBench reports average pass@1 over 8 generations.Generation lengths are 8,192 for MATH500, 2,048 for GSM8K, and 16,384 for LiveCodeBench.

L. Analysis on Pure Softmax-Attention Model · M. Method of Identifying Retrieval Layers in Pure Softmax-Attention Models · N. Details of deriving QK-Pro

In pure softmax-attention models, SFT can severely impair long-context retrieval in some models but has little effect in others. Retrieval-layer identification uses leave-one-out ablations, while QK-Pro derives its constrained restoration through linearization and Lagrange multipliers.

  • L. Analysis on Pure Softmax-Attention Model: The pure softmax-attention analysis covers Qwen2.5-3B, Qwen2.5-7B, and Misrtral-7B-Instruct-v0.3, with retrieval-layer identification detailed in Appendix M.Attention-gradient decay on Qwen2.5-3B differs from HypeNet’s exponential-decay pattern.
  • L. Analysis on Pure Softmax-Attention Model: At 64K, Mistral-7B’s NIAH-Single-3 performance drops from 81.0 to 4.80 after SFT, indicating severe long-range retrieval degradation.The degradation suggests a breakdown of long-range retrieval.
  • L. Analysis on Pure Softmax-Attention Model: QK-Restore only limitedly recovers Mistral-7B’s degradation, indicating weaker long-context representations in the pre-trained model.The limited recovery contrasts with cases where restoration fully recovers lost capability.
  • L. Analysis on Pure Softmax-Attention Model: Qwen2.5-7B drops from 95.8 to 92.6 at 64K, while Qwen2.5-7B and Qwen2.5-3B remain highly stable across NIAH-Single-2 and Single-3.These models appear to learn robust long-context retrieval during pre-training, so QK-Restore provides only marginal gains.
  • M. Method of Identifying Retrieval Layers in Pure Softmax-Attention Models: For each softmax-attention layer ℓ, leave-one-out ablation sets W^ℓ_K = 0 and evaluates the resulting model on Ruler-QA-SQuAD, FDA, and SWDE.The ablation study is conducted over the attention-layer set L_attn.
  • M. Method of Identifying Retrieval Layers in Pure Softmax-Attention Models: Layers are ranked by descending δRecall(ℓ), with large scores marking critical routing weights; the top-k layers are restored using k = L_attn/4.S(·) denotes average retrieval accuracy across the benchmarks, and near-zero scores indicate redundancy.
  • N. Details of deriving QK-Pro: Because K = R_pre has no closed-form joint solution, QK-Pro fixes one factor to its pre-SFT value and solves the remaining constrained least-squares problem.The constrained system is solved via Lagrange multipliers, followed by the stationarity condition ∂L/∂X = 0.
Loading 2606.11052v1…