[Paper] Attention logits의 변화를 제어하기

Source: arXiv - 2511.21377v1

개요

대형 트랜스포머 모델을 학습하는 과정은 생각보다 매우 불안정할 수 있습니다. 어텐션 메커니즘에서 쿼리와 키 행렬이 크기가 급격히 커지면서 학습이 불안정해지고 성능이 저하됩니다. Anson과 Aitchison의 이 논문은 그 근본 원인이 attention logits(softmax 이전의 원시 점수)의 제어되지 않은 드리프트임을 보여줍니다. 저자들은 쿼리와 키 가중치에 대해 간단한 파라미터‑종속 학습률 스킴을 적용해 로그잇이 변할 수 있는 양을 제한함으로써, QK norm이라는 무거운 트릭 없이도 안정성을 회복합니다. 특히 QK norm을 사용할 수 없는 상황에서도 효과를 발휘합니다.

주요 기여

• logits drift를 트랜스포머 학습 불안정성의 주요 원인으로 규명함.
• 경량화된 개입 제안: 로그잇 업데이트 크기를 직접 제한하는 파라미터별 학습률 스케일링.
• **Multi‑Latent Attention (MLA)**와의 호환성을 입증함. MLA는 전체 쿼리/키 정규화를 사용할 수 없는 메모리 효율 어텐션 변형입니다.
• 실험적 이득을 보여줌: 제안 방법으로 더 높은 기본 학습률을 사용할 수 있으며, MLA에서는 기존 안정화 기법을 능가하고, 표준 multi‑head attention에서는 QK norm과 동등한 성능을 달성함.
• 실용적인 레시피 제공: 몇 줄의 코드만 추가하면 되며, 추가적인 forward‑pass 연산이 필요 없음.

방법론

1. 문제 정의 – 저자들은 어텐션 logits (L = QK^\top / \sqrt{d})가 학습 단계마다 급격히 변할 수 있어 softmax 분포가 지나치게 날카롭거나 평평해진다는 관찰에서 시작합니다.
2. 파라미터‑종속 학습률 – 모든 가중치에 동일한 학습률 (\eta)를 적용하는 대신, 쿼리와 키 행렬에 스케일된 학습률 (\eta_{Q,K} = \alpha \cdot \eta)를 할당합니다. 여기서 (\alpha)는 작은 상수(예: 0.1)이며, 이는 한 번의 업데이트에서 로그잇이 이동할 수 있는 양을 직접 제한합니다.
3. 구현 세부사항 – 스케일링은 옵티마이저 수준에서 적용됩니다(예: Adam의 커스텀 파라미터 그룹). 추가적인 forward 또는 backward pass가 필요 없으며, 파라미터별 LR을 지원하는 모든 옵티마이저와 호환됩니다.
4. 실험 설정 – 두 가지 관점에서 평가합니다:
   • (a) 언어 모델링 벤치마크에서의 표준 multi‑head attention (MHA)
   • (b) 추론 시 전체 쿼리/키 텐서를 물리화하지 않는 Multi‑Latent Attention (MLA)
     베이스라인으로는 vanilla 학습, QK norm, 그리고 기타 최신 안정화 기법을 포함합니다.

결과 및 발견

• 파라미터별 LR 스킴은 attention logits의 분산을 낮게 유지하여 softmax 포화 현상을 방지합니다.
• QK norm을 적용할 수 없는 MLA에서는 새로운 방법이 기존 모든 안정화 기법을 능가하며, 더 빠른 수렴을 가능하게 합니다.
• 두 설정 모두에서 최종 모델 품질은 저하되지 않으며, 단지 더 큰 학습률을 안전하게 사용할 수 있게 됩니다.

실용적 함의

• 학습 속도 향상 – 개발자는 학습률을 2–3배 올려도 발산 위험이 적어, 대형 트랜스포머 사전학습의 실제 시간(벽시계 시간)을 크게 단축할 수 있습니다.
• 메모리 효율 어텐션 – MLA나 스트리밍/디바이스 추론 등 메모리 절약 어텐션 기법을 사용하는 모델에 대해, QK norm이 제공하지 못하는 안정성 수정을 제공합니다.
• 드롭인 교체 – 학습률 조정만으로 구현되므로 기존 코드베이스(Pytorch, TensorFlow, JAX)에 최소한의 리팩터링만으로 적용할 수 있습니다.
• 하이퍼파라미터 강건성 향상 – 학습률 스케줄 튜닝에 드는 노력을 크게 줄여, 비용이 많이 드는 프로덕션 파이프라인에서 특히 유용합니다.

제한점 및 향후 연구

• 이 방법은 **수동으로 선택한 스케일링 팩터 (\alpha)**에 의존합니다. 저자들은 비교적 견고한 기본값을 제시하지만, 작업이나 모델 규모에 따라 최적값이 달라질 수 있습니다.
• 다른 불안정성 원인(예: 피드‑포워드 레이어의 그래디언트 폭발, layer‑norm 스케일링 문제)에는 직접적인 해결책을 제공하지 않습니다.
• 논문은 언어 모델링 및 비전 분류에 초점을 맞추었으며, 멀티모달이나 강화학습 트랜스포머에 대한 확장은 아직 남아 있습니다.
• 향후 연구에서는 logit drift를 실시간으로 관찰해 자동으로 (\alpha)를 조정하는 적응형 스킴이나, 다른 정규화 기법과 결합해 더욱 강력한 안정성을 추구할 수 있습니다.

저자

• Ben Anson
• Laurence Aitchison

논문 정보

• arXiv ID: 2511.21377v1
• Categories: cs.LG
• Published: November 26, 2025
• PDF: Download PDF