Papers·어제
CoT-SFT 가 hybrid linear-attention 모델의 long-context recall 을 67.2%→9.4% 로 떨어뜨리는 문제와 QK-Restore 해결책
HKUST LARK Lab 연구진이 CoT SFT 가 hybrid linear-attention 모델의 long-context recall 을 심각하게 저하시킨다는 점을 발견했습니다. HypeNet-9B 의 경우 NIAH-S2@256K 성능이 67.2% 에서 9.4% 로 급감합니다. 원인은 CoT-SFT 가 attention gradient 를 단거리 패턴에 편향시켜 장거리 라우팅을 담당하는 W_Q, W_K 를 망가뜨리기 때문. 제안된 QK-Restore 는 SFT 전 checkpoint 의 W_Q, W_K 만 복원하는 training-free 방법으로, HypeNet-5B 에서 S3@256K 를 65.4%→76.4% 로 회복시키면서 reasoning 성능도 유지합니다.
CoT SFT 가 hybrid linear-attention 모델의 long-context recall 을 급격히 떨어뜨린다는 분석과, training-free 복원 방법 QK-Restore 를 제안한 논문입니다.
핵심 결론
- 문제 — CoT SFT 후 HypeNet-9B 의 NIAH-S2@256K 가 67.2% → 9.4% 로 하락.
- 원인 — CoT-SFT 가 attention gradient 를 단거리 패턴에 편향시켜 W_Q, W_K 의 장거리 라우팅 능력을 손상시킴.
- 해결 — QK-Restore: SFT 전 checkpoint 의 W_Q, W_K 만 복원하는 training-free 방법. HypeNet-5B 에서 S3@256K 65.4%→76.4%.
방법
- QK-Restore — SFT 후 모델에서 W_Q, W_K 만 pre-SFT 값으로 교체하고 나머지 파라미터는 그대로 유지.
- Procrustes 변형 — W_Q, W_K 를 복원할 때 orthogonal 변환을 적용해 reasoning adaptation 과 routing preservation 사이 균형을 맞춤.
- 추가 학습이 전혀 필요 없으며, 다양한 hybrid attention 모델(HypeNet, Jet-Nemotron)에서 일관된 회복 효과를 보임.
한계·조건
- 대상 — Hybrid linear-attention 모델에 한정됨; full attention 모델에서는 동일 현상이 관찰되지 않음.
- 벤치 — Needle-In-A-Haystack (NIAH) 및 LongBench 등 long-context retrieval 태스크 중심; reasoning 성능은 GSM8K, MATH 등으로 측정.
- 코드 — 논문에 코드 공개 여부 명시되지 않음; 실험은 8B~9B 스케일 모델에서 수행.
편집자 한 줄
CoT SFT 의 부작용을 구체적으로 밝힌 점이 흥미롭습니다. training-free 복원 방법이라 실용성이 높아 보이네요.
- #long-context
- #cot
- #sft
- #hybrid-attention
- #hku
LARK Lab@HKUST (GZ)