Attention Drift 발견

Speculative decoding drafter가 다음 토큰을 이어 갈수록 attention이 프롬프트에서 자기 생성 토큰 쪽으로 이동하는 attention drift 현상을 확인했다.

이 현상은 EAGLE3 drafter와 MTP heads 모두에서 나타났고, template perturbation과 long-context 입력에서 acceptance가 급격히 떨어지는 원인으로 제시됐다.

원인은 chain step마다 비정규화된 residual path가 hidden state magnitude를 계속 키우는 구조였다. 저자들은 이를 target 위에 pre-norm transformer layer를 추가한 것과 유사한 동역학으로 해석했다.

개선안은 두 가지다.

• drafter hidden state에 Post-norm 적용
• target hidden state 캡처 후 각 hidden state에 RMSNorm 적용

효과도 수치로 제시됐다.

• template perturbation에서 acceptance length 최대 2배
• long-context 작업에서 1.18배
• 7개 표준 벤치마크에서 1.10배

또 train-time test depth를 짧게 학습해도 더 긴 drafting sequence로 일반화되는 경향을 보였다.