Aandachtsamnesie in hybride LLM's: Wanneer CoT-fijnafstemming de langeafstandsherinnering verstoort, en hoe dit te verhelpen

Samenvatting

Ketens-van-gedachten (Chain-of-Thought, CoT) begeleid fine-tunen (supervised fine-tuning, SFT) wordt op grote schaal toegepast om het redeneervermogen te verbeteren. Toch stellen wij vast dat het systematisch het lange-contextherinneringsvermogen in hybride lineaire-aandachtsmodellen (linear-attention models) aantast. Bij verschillende architecturen, waaronder HypeNet en Jet-Nemotron, verslechtert de retrievalprestatie op Needle-In-A-Haystack (NIAH) aanzienlijk na CoT-SFT, en deze verslechtering wordt ernstiger naarmate de retrievalomstandigheden moeilijker worden en de contextvensters langer zijn. Zo daalt HypeNet-9B op NIAH-S2@256K van 67,2% naar 9,4%. Wij schrijven dit toe aan het feit dat CoT-SFT de aandachtgradiënten naar kortepatronen verschuift, waardoor de projecties van query en key (W_Q, W_K) – die verantwoordelijk zijn voor routering over lange afstanden – worden verstoord. Naar aanleiding van deze waarneming stellen wij QK-Restore voor, een trainingsvrije methode die alleen W_Q en W_K uit het pre-SFT-controlepunt herstelt, terwijl alle andere post-SFT-parameters behouden blijven. Verder introduceren wij een Procrustes-variant om een balans te vinden tussen het behoud van routering en de aanpassing aan redeneren. Bij meerdere architecturen herstelt QK-Restore op consistente wijze het lange-contextvermogen zonder enige trainingskosten, terwijl de redeneerprestatie behouden blijft; bijvoorbeeld op HypeNet-5B verbetert het S3@256K van 65,4% naar 76,4% met behoud van een sterke redeneerprestatie.

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.