[논문 리뷰] In-Context Convergence of Transformers

Huang Yu, Yuan Cheng외 1명|arXiv (Cornell University)|2023. 10. 08.

Domain Adaptation and Few-Shot LearningComputer Science인용 수 3

한 줄 요약

이 논문은 선형 함수의 인컨텍스트 학습을 위한 소프트맥스 어텐션을 갖는 일층 트랜스포머에서 경사 하강법 훈련 동역학에 대한 최초의 이론적 분석을 제공한다. 균형 잡힌 및 균형 잡히지 않은 특징 분포 모두에 대해 유한 시간 수렴성을 입증하며, 주로 우세한 특징이 먼저 학습되고, 그 후에 부족한 특징들이 학습되는 단계별 학습 과정을 드러낸다. 이는 경쟁하는 어텐션 가중치 동역학에 대한 새로운 분석을 통해 이루어진다.

ABSTRACT

Transformers have recently revolutionized many domains in modern machine learning and one salient discovery is their remarkable in-context learning capability, where models can solve an unseen task by utilizing task-specific prompts without further parameters fine-tuning. This also inspired recent theoretical studies aiming to understand the in-context learning mechanism of transformers, which however focused only on linear transformers. In this work, we take the first step toward studying the learning dynamics of a one-layer transformer with softmax attention trained via gradient descent in order to in-context learn linear function classes. We consider a structured data model, where each token is randomly sampled from a set of feature vectors in either balanced or imbalanced fashion. For data with balanced features, we establish the finite-time convergence guarantee with near-zero prediction error by navigating our analysis over two phases of the training dynamics of the attention map. More notably, for data with imbalanced features, we show that the learning dynamics take a stage-wise convergence process, where the transformer first converges to a near-zero prediction error for the query tokens of dominant features, and then converges later to a near-zero prediction error for the query tokens of under-represented features, respectively via one and four training phases. Our proof features new techniques for analyzing the competing strengths of two types of attention weights, the change of which determines different training phases.

연구 동기 및 목표

• 경사 하강법으로 훈련된 소프트맥스 기반 트랜스포머의 인컨텍스트 학습에서의 학습 동역학을 이해하는 것.
• 일부 특징이 부족하게 나타나는 불균형 특징 분포가 존재할 경우의 수렴 행동을 분석하는 것.
• 경쟁하는 어텐션 가중치에 의해 이끌어지는 단계 전이에 의해 유도되는 훈련 과정을 기술하는 것.
• 균형 잡힌 및 균형 잡히지 않은 설정 모두에 대해 유한 시간 수렴 보장을 확립하는 것.

제안 방법

• K개의 특징 벡터에서 확률 {pk}k=1K로 샘플링된 토큰을 갖는 구조화된 데이터 모델에서 소프트맥스 어텐션을 갖는 일층 트랜스포머를 경사 하강법으로 훈련하는 분석.
• 프롬프트 시퀀스에 입력-라벨 쌍과 쿼리 토큰을 사용하여 인컨텍스트 학습을 선형 회귀 문제로 모델링.
• 균형 잡힌 특징에 대해 두 단계 분석을 도입: 빠른 어텐션 정렬 이후 손실 최소화.
• ‘목표’ 및 ‘비목표’ 어텐션 가중치 간의 상호작용을 추적하는 데 기반한 새로운 증명 기법을 제안하여 훈련 단계를 식별.
• 주의 맵의 진화를 경사 하강법 하에서 제어하기 위해 농도 불등식과 재귀적 경계를 사용.
• 확률적 추론을 활용하여 다항식(K)의 표본 크기와 오차 허용 오차에 대해 log(1/ε) 의존성에서 고확률 수렴을 입증.

실험 결과

연구 질문

• RQ1 소프트맥스 기반 트랜스포머가 경사 하강법 훈련 동역학을 거쳐 인컨텍스트 학습을 수행할 때 수렴 방식은 어떻게 되는가?
• RQ2 목표 특징와 비목표 특징 간의 어텐션 가중치 경쟁은 훈련 단계를 어떻게 형성하는가?
• RQ3 왜 불균형 데이터 분포에서 부족한 특징이 존재하더라도 인컨텍스트 학습이 성공하는가?
• RQ4 균형 잡힌 및 균형 잡히지 않은 특징 분포 모두에 대해 유한 시간 수렴을 보장할 수 있는가?
• RQ5 특징이 불균형할 경우 학습 동역학에서 나타나는 구체적인 훈련 단계는 무엇인가?

주요 결과

• 균형 잡힌 특징에 대해, 모델는 두 단계로 near-zero 예측 오차를 달성한다: 빠른 어텐션 정렬 이후 손실 최소화.
• 불균형 특징에 대해, 모델는 단계별 수렴을 보이며, 먼저 우세한 특징이 신속히 학습되고, 그 후에 네 개의 구분 가능한 훈련 단계를 거쳐 부족한 특징들이 학습된다.
• 우세한 특징의 예측 오차는 O(log(ε^(-1/2)) / (ηε)) 경사 하강법 스텝 내에서 수렴하며, 오차는 ε 이내로 제한된다.
• 정확한 특징에 대한 어텐션 스코어는 충분한 반복 후 고확률 1 - e^(-Ω(poly(K))) 에서 near-1로 집중된다.
• 분석 결과, ‘목표’ 대비 ‘비목표’ 어텐션 가중치의 지배력이 시간이 지남에 따라 변화하며, 이는 훈련 단계를 정의한다.
• 경쟁 어텐션 가중치 동역학에 기반한 이 증명 기법은 일반화 가능하며, 다른 트랜스포머 관련 문제에도 적용 가능할 수 있다.