Amnésia de Atenção em LLMs Híbridos: Quando o Fine-Tuning de Cadeia de Pensamento Quebra a Recordação de Longo Alcance, e Como Corrigi-la

Resumo

O ajuste fino supervisionado por cadeia de pensamento (CoT-SFT) é amplamente adotado para melhorar a capacidade de raciocínio, mas descobrimos que ele degrada sistematicamente a recuperação de contexto longo em modelos híbridos de atenção linear. Em diversas arquiteturas, incluindo HypeNet e Jet-Nemotron, o desempenho de recuperação no teste Needle-In-A-Haystack (NIAH) deteriora-se substancialmente após o CoT-SFT, e a degradação torna-se mais severa em configurações de recuperação mais difíceis e janelas de contexto mais longas. Por exemplo, o HypeNet-9B no NIAH-S2@256K cai de 67,2% para 9,4%. Atribuímos isso ao fato de o CoT-SFT enviesar os gradientes de atenção para padrões de curto alcance, perturbando as projeções consulta-chave (W_Q, W_K) que são responsáveis pelo roteamento de longo alcance. Motivados por essa observação, propomos o QK-Restore, um método sem treinamento que restaura apenas W_Q e W_K a partir do checkpoint pré-SFT, preservando todos os outros parâmetros pós-SFT. Introduzimos ainda uma variante de Procrustes para equilibrar a preservação do roteamento e a adaptação do raciocínio. Em diversas arquiteturas, o QK-Restore restaura consistentemente a capacidade de contexto longo com custo zero de treinamento, preservando o desempenho de raciocínio; por exemplo, no HypeNet-5B, melhora o S3@256K de 65,4% para 76,4%, mantendo um desempenho de raciocínio robusto.

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.