[논문 리뷰] Transformers learn to implement preconditioned gradient descent for in-context learning

Kwangjun Ahn, Xiang Cheng외 2명|arXiv (Cornell University)|2023. 06. 01.

Stochastic Gradient Optimization Techniques인용 수 6

한 줄 요약

이 논문은 무작위 선형 회귀 예제를 학습시킨 트랜스포머가 손실 곡면 분석을 통해 조정된 경사하강법을 학습함을 보여준다. 단일 어텐션 레이어의 경우, 전역 최적점은 데이터 분포와 데이터 부족으로 인한 분산에 적응하는 한 단계의 조정된 경사하강법에 해당하며, 더 깊은 트랜스포머는 다중 반복을 구현하며, 임계점은 GD++를 포함한 적응형 최적화 알고리즘과 일치한다.

ABSTRACT

Several recent works demonstrate that transformers can implement algorithms like gradient descent. By a careful construction of weights, these works show that multiple layers of transformers are expressive enough to simulate iterations of gradient descent. Going beyond the question of expressivity, we ask: Can transformers learn to implement such algorithms by training over random problem instances? To our knowledge, we make the first theoretical progress on this question via an analysis of the loss landscape for linear transformers trained over random instances of linear regression. For a single attention layer, we prove the global minimum of the training objective implements a single iteration of preconditioned gradient descent. Notably, the preconditioning matrix not only adapts to the input distribution but also to the variance induced by data inadequacy. For a transformer with $L$ attention layers, we prove certain critical points of the training objective implement $L$ iterations of preconditioned gradient descent. Our results call for future theoretical studies on learning algorithms by training transformers.

연구 동기 및 목표

• 트랜스포머가 수작업으로 설계된 가중치에 의존하지 않고, 무작위 문제 예제를 학습시켜 기반 경사 최적화 알고리즘을 학습할 수 있는지 조사하기 위해.
• 무작위 선형 회귀 예제를 학습시킨 선형 트랜스포머의 손실 곡면을 분석하여 학습 과정에서 최적화 알고리즘이 어떻게 유도되는지 이해하기 위해.
• 트랜스포머의 파rameter 공간 내 전역 최소점과 임계점의 구조를 특성화하여 어떤 최적화 알고리즘을 실제로 실행하는지 규명하기 위해.
• 트랜스포머의 이론적 표현 능력과 인라인 학습에서의 실질적 행동 간 격차를 메우기 위해, 특히 기반 경사 방법에 대해.
• 이론적 발견을 실험적으로 검증하여 학습된 임계점이 GD++와 같은 알려진 적응형 최적화 알고리즘과 일치하는지 확인하기 위해.

제안 방법

• 비소프트맥스 어텐션 메커니즘을 사용하여, 무작위 등방성 선형 회귀 예제를 학습시킨 단일 레이어 선형 트랜스포머의 손실 곡면을 분석한다.
• 학습 목표의 전역 최소점이 데이터 분포와 데이터 부족으로 인한 분산에 적응하는 조정된 경사하강법의 한 단계에 해당함을 증명한다.
• 파rameter 공간에 희소성 조건을 도입하여 k단계 적응형 기반 경사 최적화 알고리즘으로의 검색을 제한하고, 다중 레이어 트랜스포머의 임계점 특성화를 가능하게 한다.
• 더 깊은 트랜스포머(L 레이어)의 경우, 특정 임계점이 데이터 의존적 조정을 갖는 L단계의 조정된 경사하강법을 실행함을 보여준다.
• 희소성 조건을 완화하여 전체 파rameter 공간을 연구하고, 조건을 개선하기 위해 선형 변환을 통한 기울기 하강법과 조합된 새로운 기반 최적화 알고리즘을 구현하는 임계점을 발견한다.
• 이론적 발견을 실험적으로 검증하기 위해 학습된 가중치를 시각화하고, 트랜스포머 예측기의 테스트 손실을 표준 최적화 기준(경사하강법, 조정된 경사하강법, OLS)과 비교한다.

(a) $\operatorname{Dist}(\Sigma^{1/2}A_{0}\Sigma^{1/2},I)$

실험 결과

연구 질문

• RQ1 무작위 선형 회귀 예제를 학습시킨 트랜스포머가 비凸 최적화를 통해 기반 경사하강법을 학습할 수 있는가?
• RQ2 무작위 회귀 문제를 학습시킨 단일 레이어 선형 트랜스포머의 전역 최소점이 어떤 최적화 알고리즘을 실행하는가?
• RQ3 더 깊은 트랜스포머(L 레이어)의 임계점은 반복 최적화 알고리즘과 어떻게 관련이 있는가?
• RQ4 파rameter의 희소성 제약 조건을 제거했을 경우 학습된 알고리즘은 어떻게 되는가?
• RQ5 이론적으로 유도된 임계점은 훈련된 트랜스포머에서 관측된 경험적 행동과 일치하는가?

주요 결과

• 단일 레이어 선형 트랜스포머의 전역 최소점은 입력 데이터 분포와 데이터 부족으로 인한 분산에 적응하는 한 단계의 조정된 경사하강법을 실행한다.
• 희소성 조건 하에서 이중 레이어 트랜스포머의 경우, 전역 최소점은 적응형 스텝 사이즈를 갖는 경사하강법에 해당하며, 효과적으로 k단계 적응형 최적화 과정을 수행한다.
• 더 깊은 트랜스포머(L 레이어)의 경우, 학습 목표의 특정 임계점은 데이터 의존적 조정을 갖는 L단계의 조정된 경사하강법을 실행한다.
• 희소성 제약 조건을 제거하면, 기울기 하강법과 선형 변환을 조합하여 조건을 개선하는 새로운 알고리즘을 구현하는 임계점이 나타난다. 데이터 공분산이 등방성일 경우 이는 GD++ 알고리즘과 정확히 일치한다.
• 경험적 검증 결과, 삼중 레이어 트랜스포머의 학습된 가중치는 이론적 정적점과 일치하며, 이 점에서의 목표 값은 거의 0에 가까워 전역 최적점임을 시사한다.
• 테스트 손실 비교 결과, 트랜스포머가 학습한 예측기는 세 단계의 조정된 경사하강법과 유사한 성능을 보이며, 이는 이론적 일치와 실질적 최적화 행동 간의 일치를 검증한다.

(b) $\operatorname{Dist}(\Sigma^{1/2}A_{1}\Sigma^{1/2},I)$