Aufmerksamkeitsamnesie in hybriden LLMs: Wenn CoT-Feinabstimmung den Langzeitabruf beeinträchtigt, und wie man ihn behebt

Zusammenfassung

Chain-of-Thought (CoT) überwachtes Feintuning (SFT) wird häufig eingesetzt, um die Reasoning-Fähigkeit zu verbessern. Wir stellen jedoch fest, dass es bei hybriden Lineare-Attention-Modellen systematisch den Langkontext-Abruf verschlechtert. Über Architekturen wie HypeNet und Jet-Nemotron hinweg verschlechtert sich die Retrieval-Leistung bei der Nadel-im-Heuhaufen-Aufgabe (NIAH) nach CoT-SFT erheblich, und die Verschlechterung wird unter schwierigeren Retrieval-Bedingungen und bei längeren Kontextfenstern noch gravierender. Beispielsweise sinkt HypeNet-9B bei NIAH-S2@256K von 67,2 % auf 9,4 %. Wir führen dies darauf zurück, dass CoT-SFT die Aufmerksamkeitsgradienten hin zu kurzfristigen Mustern verzerrt und dabei die Query-Key-Projektionen (W_Q, W_K) stört, die für das Langstrecken-Routing verantwortlich sind. Ausgehend von dieser Beobachtung schlagen wir QK-Restore vor, eine trainingsfreie Methode, die nur W_Q und W_K aus dem Pre-SFT-Checkpoint wiederherstellt, während alle anderen Post-SFT-Parameter erhalten bleiben. Darüber hinaus führen wir eine Procrustes-Variante ein, um eine Balance zwischen Routing-Erhaltung und Reasoning-Anpassung zu erreichen. Über Architekturen hinweg stellt QK-Restore konsistent die Langkontext-Fähigkeit ohne Trainingskosten wieder her, während die Reasoning-Leistung erhalten bleibt; beispielsweise verbessert es bei HypeNet-5B die S3@256K von 65,4 % auf 76,4 %, während eine starke Reasoning-Leistung beibehalten wird.

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.