Amnesia atencional en LLMs híbridos: cuando el ajuste fino de CoT rompe el recuerdo a largo plazo, y cómo solucionarlo

Resumen

La supervisión fina con razonamiento encadenado (CoT-SFT) se adopta ampliamente para mejorar la capacidad de razonamiento, pero encontramos que degrada sistemáticamente la recuperación de contexto largo en modelos híbridos de atención lineal. En arquitecturas como HypeNet y Jet-Nemotron, el rendimiento de recuperación en la prueba Needle-In-A-Haystack (NIAH) se deteriora sustancialmente después de CoT-SFT, y la degradación se vuelve más severa bajo condiciones de recuperación más difíciles y ventanas de contexto más largas. Por ejemplo, HypeNet-9B en NIAH-S2@256K disminuye del 67,2% al 9,4%. Atribuimos esto a que CoT-SFT sesga los gradientes de atención hacia patrones de corto alcance, alterando las proyecciones consulta-clave (W_Q, W_K) responsables del enrutamiento de largo alcance. Motivados por esta observación, proponemos QK-Restore, un método sin entrenamiento que restaura únicamente W_Q y W_K del punto de control anterior a SFT, preservando todos los demás parámetros posteriores a SFT. Además, introducimos una variante Procrustes para equilibrar la preservación del enrutamiento y la adaptación al razonamiento. En distintas arquitecturas, QK-Restore restaura consistentemente la capacidad de contexto largo sin costo de entrenamiento, manteniendo el rendimiento de razonamiento; por ejemplo, en HypeNet-5B mejora S3@256K del 65,4% al 76,4% mientras mantiene un sólido rendimiento de razonamiento.

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.