Attention Drift: Autoregressive Speculative Decoding 모델은 무엇을 학습하는가

요약

본 논문은 Speculative Decoding 과정에서 발생하는 'attention drift'라는 새로운 현상을 식별했습니다. 이는 drafter 모델이 추측 체인 내에서 연속적인 토큰을 생성할 때, 어텐션 메커니즘이 입력 프롬프트로부터 점진적으로 자신이 최근에 생성한 토큰들로 이동하는 현상입니다. 연구진은 이 원인을 잔차 경로의 비정규화된 누적 증가와 연결 지었으며, 이를 해결하기 위해 drafter hidden state에 Post-norm 적용 및 RMSNorm을 제안했습니다. 이러한 구조적 개선을 통해 다양한 벤치마크에서 수락 길이 및 전반적인 성능 향상을 입증했습니다.

핵심 포인트

• Speculative Decoding의 초안 작성 모델(drafter)이 긴 문맥이나 템플릿 섭동 환경에서 성능 저하를 보이는 'attention drift' 현상이 발견됨.
• Attention Drift는 drafter가 추측 체인 내에서 토큰을 생성할 때, 어텐션이 프롬프트에서 최근 생성된 토큰으로 점진적으로 이동하는 현상임.
• 원인은 drafter의 hidden state 크기가 체인 깊이에 따라 비정규화된 잔차 경로를 통해 단조 증가하기 때문이며, 이는 독립적인 예측기라기보다 레이어 누적 구조와 유사함.
• 이 문제를 해결하기 위해 drafter hidden states에 Post-norm을 적용하고 각 hidden state에 RMSNorm을 적용하는 두 가지 구조적 변경 사항이 제안됨.
• 제안된 개선책은 다양한 벤치마크에서 수락 길이 및 전반적인 성능을 크게 향상시켰으며, 더 짧은 훈련 시간으로도 긴 초안 시퀀스에 대한 일반화가 가능함을 보임.