기술 요약: 자동회귀적 추측적 디코딩에서의 어텐션 드리프트

문제 제기

추측적 디코딩은 경량 "드래프터" 모델을 사용하여 미래 토큰을 예측하고, 이를 더 큰 "타겟" 모델이 검증함으로써 대규모 언어 모델 (LLM) 추론을 가속화합니다. 효과적이기는 하지만, 현재 드래프터 모델 (특히 EAGLE-3 및 멀티 토큰 예측 헤드) 은 템플릿 교란 및 긴 컨텍스트 입력과 같은 배포 환경의 변화 하에서 심각한 취약성을 보입니다. 실무자들은 종종 성능을 유지하기 위해 특정 추론 엔진이나 프롬프트에 맞춰 드래프터를 재학습해야 하며, 이는 근본적인 구조적 불안정성을 시사합니다.

본 논문은 어텐션 드리프트 (Attention Drift) 라는 이전에 보고되지 않은 현상을 규명합니다. 드래프터가 추측 토큰 시퀀스를 생성함에 따라, 그 어텐션 메커니즘은 점차 "어텐션 싱크 (attention sink, 일반적으로 프롬프트 시작부 근처에 있는 안정적인 앵커 토큰)"에서 멀어지고, 최근 생성된 자신의 토큰에 집중하게 됩니다. 이러한 드리프트는 특히 긴 컨텍스트 시나리오와 입력 템플릿이 학습 분포에서 벗어날 때 수용률 (acceptance rates) 의 저하와 상관관계가 있습니다.

방법론 및 분석

저자들은 Llama 3.1, Qwen 3/3.5, GPT-oss 등 여러 모델 계열에 걸쳐 경험적 관찰과 구조적 분석을 결합하여 어텐션 드리프트의 근본 원인을 규명했습니다.

1. 현상학적 관찰

어텐션 히트맵과 정량적 지표를 활용하여 저자들은 다음을 입증했습니다:

• 싱크 이동 (Sink Migration): 강력한 어텐션 싱크를 가진 모델 (예: Llama, Qwen) 에서 드래프터의 싱크 토큰에 대한 어텐션은 추측 체인이 깊어짐에 따라 감소하는 반면, 가장 최근에 작성된 토큰에 대한 어텐션은 증가합니다.
• 보편성: 이 드리프트는 사후에 고정된 타겟 모델로 학습된 EAGLE-3 드래프터와 타겟 모델과 공동 학습된 MTP 헤드 모두에서 관찰되며, 이는 자신의 출력을 소비하는 자동회귀적 설계의 고유한 특성임을 시사합니다.
• 싱크 부재: 싱크를 억제하도록 설계된 모델 (게이트드 어텐션 또는 학습 가능한 편향을 갖춘 GPT-oss 등) 에서조차, 프롬프트의 시작이 아닌 반복되는 템플릿 토큰에 어텐션이 집중되는 형태의 드리프트가 발생합니다.

2. 근본 원인 분석: 은닉 상태 크기

본 논문은 표준 프리-노름 (pre-norm) 아키텍처 내 추측 단계 간의 비정규화 잔차 경로 (unnormalized residual path) 에서 어텐션 드리프트의 원인을 추적합니다:

• 단조 증가: 프리-노름 트랜스포머에서 잔차 연결은 정규화 없이 크기를 누적합니다. 저자들은 드래프터의 은닉 상태 크기 ($\|h^*_k\|$) 가 추측 깊이 ($k$) 에 따라 단조 증가함을 보여줍니다.
• 레이어 적층 vs 자동회귀: 이러한 증가는 드래프터가 독립적인 자동회귀 예측기로 작동하기보다 타겟 모델 위에 적층된 추가 트랜스포머 레이어의 시퀀스처럼 행동하게 만듭니다. $k$ 단계에서 드래프터는 안정적인 분포에 기반하여 다음 토큰을 예측하는 대신, 깊이 $N+k$ 의 타겟 모델을 근사하게 됩니다.
• 불균형 입력: 타겟 모델에서 융합된 은닉 상태 ($h_{FC}$) 는 종종 더 큰 크기를 가진 깊은 레이어에 의해 지배되며, 이는 드래프터로 입력될 때 크기 불일치를 더욱 악화시킵니다.

3. 제안된 개입

크기 누적을 상쇄하고 드래프팅 과정을 안정화하기 위해 저자들은 두 가지 구조적 수정을 제안합니다:

1. 드래프터 은닉 상태에 대한 포스트-노름 (Post-Norm): 잔차 연결 (추측 단계 간) 이후 드래프터의 은닉 상태에 RMSNorm 을 적용하여 크기의 단조 증가를 방지합니다. 이는 드래프터가 깊이별 정제 (depth-specific refinements) 가 아닌 깊이 불변의 자동회귀 규칙을 학습하도록 강제합니다.
2. 은닉 상태별 RMSNorm: 타겟의 은닉 상태 ($h_{low}, h_{mid}, h_{high}$) 를 드래프터 입력으로 융합하기 전에 RMSNorm 레이어를 추가하여 유입 신호의 크기를 균형 있게 맞춥니다.

저자들은 또한 (고정된 싱크에 대한 의존성을 억제하는) 게이트드 어텐션을 별도의 개입으로 평가했습니다. 게이트는 싱크를 제거하지만 크기 증가를 방지하지는 못하며, 게이트와 포스트-노름을 결합하면 과도한 정규화 (엔트로피 붕괴) 로 이어질 수 있음을 발견했습니다.

주요 결과

제안된 포스트-노름 아키텍처는 다양한 벤치마크 및 조건 하에서 현재 최첨단인 프리-노름 EAGLE-3 와 비교 평가되었습니다:

• 템플릿 교란: 템플릿 교란 (예: 채팅 템플릿 또는 시퀀스 시작 토큰 제거) 하에서 포스트-노름 드래프터는 프리-노름 대비 수용 길이를 최대 2 배까지 개선했습니다. 프리-노름 모델은 치명적인 하락 (최대 52% 손실) 을 겪은 반면, 포스트-노름 모델은 탄력성을 유지했습니다 (최대 5% 손실).
• 긴 컨텍스트 작업: 긴 컨텍스트 작업 (LongBench, 슬라이딩 윈도우 어텐션을 사용한 멀티 턴 채팅) 에서 포스트-노름은 수용 길이를 1.18 배 개선했습니다. 풀-어텐션 긴 컨텍스트 설정에서 프리-노름 모델은 수용률이 거의 0 에 수렴하여 붕괴된 반면, 포스트-노름은 사용 가능한 성능을 유지했습니다.
• 표준 벤치마크: 일곱 가지 표준 벤치마크 (멀티 턴 채팅, 수학, 코딩) 에서 포스트-노름은 수용 길이를 일관되게 1.10 배 개선했습니다.
• 학습 효율성: 포스트-노름 드래프터는 학습 시계열을 넘어선 추론 깊이에도 일반화됩니다. 이를 통해 더 짧은 "학습 시간 - 테스트" (TTT) 깊이 (예: 8 단계에서 4 단계로 감소) 로 학습이 가능해져 성능을 희생하지 않으면서 학습 시간을 약 3 분의 1 로 단축할 수 있습니다.
• 노이즈 강건성: 포스트-노름 아키텍처는 은닉 상태의 노이즈에 대해 훨씬 더 강건합니다 ($\alpha=0.5$ 노이즈에서 58% 수용 유지 vs 프리-노름의 5%), 이는 양자화 및 분포 변화에 대한 더 나은 강건성을 시사합니다.

중요성 및 주장

본 논문은 어텐션 드리프트가 비정규화 잔차 경로와 자동회귀적 자기 피드백 간의 상호작용에서 비롯된, 추측적 디코딩의 치명적이고 이전에 보고되지 않은 실패 모드라고 주장합니다.

• 진단: 저자들은 어텐션 드리프트가 더 깊은 문제의 증상이라고 주장합니다. 즉, 은닉 상태 크기의 누적으로 인해 드래프터가 안정적인 토큰 예측이 아닌 깊이 의존적 변환을 학습하게 된다는 것입니다.
• 해결책: 그들은 단순히 어텐션 싱크를 수정하는 것 (예: 게이팅을 통한) 은 불충분하며, 근본적인 크기 역학을 해결해야 한다고 주장합니다. 제안된 포스트-노름 아키텍처는 드래프팅 과정을 안정화하는 간단하고 효과적인 해결책으로 제시됩니다.
• 일반화: 이 발견들은 포스트-노름이 다양한 프롬프트와 긴 컨텍스트에 대한 강건성이 필요한 프로덕션 시스템을 위해 추측적 디코딩 드래프터의 기본 구조 선택으로 고려되어야 함을 시사합니다.

본 논문은 "어텐션 드리프트"를 명명하고 특성화하는 것이 드래프터 강건성을 위한 유용한 진단 도구를 제공하며, 크기 누적을 대상으로 하는 구조적 개입이 신뢰할 수 있는 추측적 디코딩에 필수적이라고 결론지었습니다.