Амнезия внимания в гибридных LLM: когда тонкая настройка CoT нарушает долгосрочное запоминание, и как это исправить

Аннотация

Цепочка рассуждений (CoT) под контролируемой тонкой настройкой (SFT) широко применяется для улучшения способности к рассуждению, однако мы обнаруживаем, что она систематически ухудшает извлечение информации из длинного контекста в гибридных моделях линейного внимания. Для архитектур, включая HypeNet и Jet-Nemotron, производительность извлечения на тесте «Иголка в стоге сена» (NIAH) существенно падает после CoT-SFT, причём ухудшение становится более выраженным в более сложных условиях извлечения и для более длинных окон контекста. Например, у HypeNet-9B показатель NIAH-S2@256K снижается с 67,2% до 9,4%. Мы объясняем это тем, что CoT-SFT смещает градиенты внимания в сторону короткодиапазонных шаблонов, нарушая проекции «запрос–ключ» (W_Q, W_K), отвечающие за маршрутизацию на дальние расстояния. Основываясь на этом наблюдении, мы предлагаем QK-Restore — метод, не требующий обучения, который восстанавливает только W_Q и W_K из контрольной точки до SFT, сохраняя все остальные параметры после SFT. Мы также вводим вариант с преобразованием Прокруста для баланса между сохранением маршрутизации и адаптацией рассуждения. На различных архитектурах QK-Restore последовательно восстанавливает способность работы с длинным контекстом при нулевых затратах на обучение, сохраняя при этом производительность рассуждений; например, для HypeNet-5B он улучшает показатель S3@256K с 65,4% до 76,4% при сохранении высокого качества рассуждений.

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.