Amnésie attentionnelle dans les LLM hybrides : lorsque le fine-tuning par CoT brise le rappel à longue distance, et comment y remédier

Résumé

Le fine-tuning supervisé par chaîne de pensée (CoT-SFT) est largement utilisé pour améliorer la capacité de raisonnement, mais nous constatons qu'il dégrade systématiquement le rappel en contexte long dans les modèles hybrides à attention linéaire. Sur différentes architectures, notamment HypeNet et Jet-Nemotron, les performances de recherche sur le benchmark Needle-In-A-Haystack (NIAH) se détériorent considérablement après CoT-SFT, et cette dégradation s'aggrave dans des conditions de recherche plus difficiles et avec des fenêtres de contexte plus longues. Par exemple, HypeNet-9B sur NIAH-S2@256K passe de 67,2% à 9,4%. Nous attribuons cela au fait que CoT-SFT biaise les gradients d'attention vers des motifs à courte portée, perturbant les projections requête-clé (W_Q, W_K) responsables du routage à longue portée. Motivés par cette observation, nous proposons QK-Restore, une méthode sans entraînement qui restaure uniquement W_Q et W_K à partir du point de contrôle antérieur à CoT-SFT, tout en conservant tous les autres paramètres post-SFT. Nous introduisons également une variante de Procruste pour équilibrer la préservation du routage et l'adaptation au raisonnement. Sur plusieurs architectures, QK-Restore rétablit systématiquement la capacité en contexte long sans coût d'entraînement, tout en préservant les performances de raisonnement ; par exemple, sur HypeNet-5B, il améliore S3@256K de 65,4% à 76,4% tout en maintenant des performances solides en raisonnement.

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.