하이브리드 LLM에서의 어텐션 건망증: CoT 미세 조정이 장거리 회상 능력을 저해하는 이유와 해결 방법

Chain-of-thought (CoT, 사고의 사슬) 지도 미세 조정 (Supervised Fine-Tuning, SFT)은 추론 능력을 향상시키기 위해 널리 채택되고 있지만, 우리는 이것이 하이브리드 선형 어텐션 (Linear-attention) 모델에서 장기 문맥 회상 (Long-context recall) 능력을 체계적으로 저하시킨다는 것을 발견했습니다. HypeNet 및 Jet-Nemotron을 포함한 아키텍처 전반에서, Needle-In-A-Haystack (NIAH) 테스트의 검색 성능이 CoT-SFT 이후 실질적으로 악화되었으며, 이러한 저하는 더 어려운 검색 설정과 더 긴 문맥 창 (Context window) 환경에서 더욱 심각해졌습니다. 예를 들어, NIAH-S2@256K 환경에서 HypeNet-9B의 성능은 $67.2%$에서 $9.4%$로 감소합니다. 우리는 이러한 현상의 원인이 CoT-SFT가 어텐션 그래디언트 (Attention gradients)를 단거리 패턴 쪽으로 편향시켜, 장거리 라우팅 (Long-range routing)을 담당하는 쿼리-키 투영 (Query-key projections, $W_Q, W_K$)을 방해하기 때문이라고 분석합니다. 이러한 관찰에 착안하여, 우리는 다른 모든 SFT 이후의 파라미터는 보존하면서 SFT 이전의 체크포인트로부터 $W_Q$와 $W_K$만을 복구하는 훈련이 필요 없는 방법론인 QK-Restore를 제안합니다. 나아가, 우리는 라우팅 보존과 추론 적응 사이의 균형을 맞추기 위한 Procrustes 변형 방식을 도입합니다. 여러 아키텍처에 걸쳐 QK-Restore는 추론 성능을 유지하면서도 추가적인 훈련 비용 없이 장기 문맥 능력을 일관되게 복구합니다. 예를 들어, HypeNet-5B의 경우 S3@256K 성능을 $65.4%$에서 $76.4%$로 향상시키는 동시에 강력한 추론 성능을 유지했습니다.