混合型大語言模型中的注意力遺忘：思維鏈微調如何破壞長程召回及其修復方法

摘要

思維鏈監督微調（CoT-SFT）被廣泛用於提升推理能力，但我們發現它會系統性地降低混合線性注意力模型中的長上下文召回效能。在包括HypeNet和Jet-Nemotron在內的多種架構中，CoT-SFT導致「大海撈針」（NIAH）檢索任務的表現大幅下降，且在更嚴苛的檢索設定與更長的上下文視窗下，效能退化更加嚴重。例如，HypeNet-9B在NIAH-S2@256K上的準確率從67.2%降至9.4%。我們將此歸因於CoT-SFT使注意力梯度偏向短程模式，從而擾亂了負責長程路由的查詢-鍵投影（W_Q, W_K）。基於此發現，我們提出QK-Restore，一種無需額外訓練的方法，僅從SFT前的檢查點恢復W_Q與W_K，同時保留所有其他SFT後的參數。我們進一步引入普氏分析（Procrustes）變體，以平衡路由保留與推理適應。在各種架構中，QK-Restore在零訓練成本下持續恢復長上下文能力，同時保持推理效能；例如，在HypeNet-5B上，它將S3@256K從65.4%提升至76.4%，並維持強大的推理表現。

English

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.