[논문 리뷰] On the Optimization and Generalization of Multi-head Attention

Puneesh Deora, Rouzbeh Ghaderi외 2명|arXiv (Cornell University)|2023. 10. 19.

Neural Networks and Applications인용 수 4

한 줄 요약

이 논문은 로지스틱 손실을 사용한 단일층 다중헤드 자기주의 모델의 경사하강법 학습에 대해 처음으로 수렴성과 일반화 보장을 제공한다. 다항로그 수준의 헤드 수(H = Ω(log⁶n))와 약한 실현 가능성 및 NTK 분리 조건 하에서, 모델은 ˜O(1/n) 수준의 훈련 손실과 일반화 갭을 달성하며, 오버파rameterized MLP 이론과의 연결을 통해 알고리즘 안정성과 자기구속 손실 분석을 활용한다.

ABSTRACT

The training and generalization dynamics of the Transformer's core mechanism, namely the Attention mechanism, remain under-explored. Besides, existing analyses primarily focus on single-head attention. Inspired by the demonstrated benefits of overparameterization when training fully-connected networks, we investigate the potential optimization and generalization advantages of using multiple attention heads. Towards this goal, we derive convergence and generalization guarantees for gradient-descent training of a single-layer multi-head self-attention model, under a suitable realizability condition on the data. We then establish primitive conditions on the initialization that ensure realizability holds. Finally, we demonstrate that these conditions are satisfied for a simple tokenized-mixture model. We expect the analysis can be extended to various data-model and architecture variations.

연구 동기 및 목표

• 유한한 너비 모델에 대한 다중헤드 어텐션 학습 역학에 대한 이론적 이해의 격차를 메우기 위해.
• 이전에 단일헤드 어텐션에 국한된 분석을 실용적인 헤드 수를 갖는 다중헤드 구조로 확장하기 위해.
• 오버파arameterized 신경망 이론의 도구를 활용하여 다중헤드 어텐션에 대한 유한시간 수렴성 및 일반화 경계를 확립하기 위해.
• 실현 가능성 보장과 날카운 일반화 경계 달성을 가능하게 하는 초기화 조건을 식별하고, 이를 검증 가능한 형태로 제시하기 위해.
• 토큰화된 믹스처 데이터 모델에 프레임워크의 적용 가능성을 시연하며, 초기화 이후 한 단계의 경사하강 스텝으로 데이터가 빠르게 분리됨을 보여주기 위해.

제안 방법

• 모델 파라미터에 대해 약간의 의존성만을 가지는 매개변수 κ에 의해 결정되는 헤시안 곡률을 정량화하는 방식으로, 경험적 손실 ̂L(θ)의 자기구속성 및 약한 볼록성 성질을 유도한다.
• 오버파arameterized MLP에서 유래된 알고리즘 안정성 도구(Taheri & Thrampoulidis, 2023)를 활용하여 일반화 갭을 훈련 손실과 초기화로부터의 거리로 경계한다.
• 다중헤드 어텐션을 오버파arameterized 은닉층과 유사한 병렬화된 구조로 간주하는 새로운 분석 프레임워크를 도입하여, 기존의 일반화 기법을 이전 모델로 이식할 수 있도록 한다.
• 모델 출력이 훈련 세트 크기 n에 대해 로그 수준이 되도록 하는 초기화 조건을 정식화하여, NTK 특징 기반 데이터 실현 가능성 보장한다.
• 초기화 이후 한 번의 랜덤화된 경사하강 스텝을 거친 후, NTK 특징이 상수 마진 분리 가능성을 확보하는 토큰화된 믹스처 데이터 모델을 분석한다.
• 로지스틱 손실의 자기구속성과 곡률 경계를 활용하여, 표준 경사하강법(스텝 크기 η = ˜O(1)) 하에서 유한시간 최적화 및 일반화 속도를 유도한다.

실험 결과

연구 질문

• RQ1 오버파arameterized 완전연결 네트워크에 대해 개발된 프레임워크를 다중헤드 어텐션 분석에 적용할 수 있는가?
• RQ2 어떤 초기화 조건이 NTK 특징 기반 데이터 실현 가능성 보장을 보장하며, 이를 통해 수렴성과 일반화를 가능하게 하는가?
• RQ3 주어진 오차율을 달성하기 위해 필요한 최소한의 어텐션 헤드 수 H는 얼마인가?
• RQ4 다중헤드 어텐션에서 경사하강법 하에서 일반화 갭은 훈련 세트 크기 n에 따라 어떻게 스케일링되는가?
• RQ5 제안된 프레임워크는 토큰화된 믹스처 모델과 같은 구체적인 데이터 모델에 적용 가능하며, 최소한의 최적화 스텝 이후 데이터 분리가 빠르게 이루어지는지를 예측할 수 있는가?

주요 결과

• H = Ω(log⁶n) 헤드를 갖는 모델은 상수 마진 NTK 분리 조건과 η = ˜O(1) 스텝 크기 하에서 ˜O(1/n) 수준의 훈련 손실과 일반화 갭을 달성한다.
• 경험적 손실 ̂L(θ)는 자기구속성 약한 볼록성 조건을 만족한다: λmin(∇²̂L(θ)) ≳ −κ/√H ⋅ ̂L(θ), 여기서 κ는 파라미터 벡터에 대해 약간의 의존성을 가진다.
• 간단한 초기화 조건이 실현 가능성 보장을 보장한다: 초기화 시 모델 출력이 O(log n)이면, 데이터가 NTK 분리 가능성 마진 γ를 갖는다면 경계가 성립한다.
• 영초기화 상태에서 한 번의 랜덤화된 경사하강 스텝을 거친 후, MHA 모델의 NTK 특징은 토큰화된 믹스처 데이터를 마진 γ⋆로 분리한다.
• 일반화 갭은 ˜O(1/n) 비례하는 항으로 경계되며, 이는 n에 비해 작은 수의 헤드를 가진 다중헤드 어텐션도 잘 일반화될 수 있음을 보여준다.
• 이 분석 프레임워크는 다양한 데이터-모델 및 아키텍처 변형으로 확장 가능하며, 현재 설정을 초월해 광범위하게 적용 가능함을 시사한다.