[논문 리뷰] Transformers learn in-context by gradient descent

Johannes von Oswald, Eyvind Niklasson외 5명|arXiv (Cornell University)|2022. 12. 15.

Domain Adaptation and Few-Shot Learning인용 수 89

한 줄 요약

본 논문은 트랜스포머의 컨텍스트 내 학습이 기계적으로는 그래디언트 하강 업데이트로 이해될 수 있음을 보이며, 셀프 어텐션 레이어가 컨텍스트 내 데이터에 대해 GD 스텝을 구현할 수 있음을 증명하고, MLP를 추가하면 딥 표현에서의 그래디언트 기반 학습을 통해 비선형 회귀가 가능하다는 것을 보여준다.

ABSTRACT

At present, the mechanisms of in-context learning in Transformers are not well understood and remain mostly an intuition. In this paper, we suggest that training Transformers on auto-regressive objectives is closely related to gradient-based meta-learning formulations. We start by providing a simple weight construction that shows the equivalence of data transformations induced by 1) a single linear self-attention layer and by 2) gradient-descent (GD) on a regression loss. Motivated by that construction, we show empirically that when training self-attention-only Transformers on simple regression tasks either the models learned by GD and Transformers show great similarity or, remarkably, the weights found by optimization match the construction. Thus we show how trained Transformers become mesa-optimizers i.e. learn models by gradient descent in their forward pass. This allows us, at least in the domain of regression problems, to mechanistically understand the inner workings of in-context learning in optimized Transformers. Building on this insight, we furthermore identify how Transformers surpass the performance of plain gradient descent by learning an iterative curvature correction and learn linear models on deep data representations to solve non-linear regression tasks. Finally, we discuss intriguing parallels to a mechanism identified to be crucial for in-context learning termed induction-head (Olsson et al., 2022) and show how it could be understood as a specific case of in-context learning by gradient descent learning within Transformers. Code to reproduce the experiments can be found at https://github.com/google-research/self-organising-systems/tree/master/transformers_learn_icl_by_gd .

연구 동기 및 목표

• 트랜스포머에서 컨텍스트 내 학습 메커니즘에 대한 이해를 촉진한다.
• 선형 자기-어텐션 업데이트와 선형 회귀에 대한 한 단계의 그래디언트 하강 간의 동등성을 보인다.
• 어텐션 레이어를 쌓으면 반복적인 GD와 곡률 보정(GD++)에 해당하는 업데이트가 가능함을 시연한다.
• MLP가 딥 표현에서의 그래디언트 하강을 통해 비선형 회귀를 가능하게 하는 방법을 설명한다.
• 메타학습, 빠른 가중치, 그리고 induction-head 메커니즘과의 연관성을 논의한다.

제안 방법

• 선형 회귀 손실에 대한 그래디언트 하강 업데이트와 동등하도록 단일 선형 자기-어텐션 단계가 되도록 가중치 구성(weight construction)을 유도한다.
• 선형 회귀 과제에서 학습된 선형 자기-어텐션 레이어를 GD 구성과 실증적으로 비교하여 정렬 정도를 평가한다.
• 다층 자기-어텐션으로 확장하고 반복적 데이터 변환에서 GD-유사 동작(GD++)과 잔차 곡률 보정을 보인다.
• 트랜스포머에 MLP를 포함시키면 딥 표현에서의 그래디언트 하강을 통해 비선형 회귀를 해결할 수 있음을 보이며(커널 회귀 관점).
• 포지션 구성 및 토큰 구성 및 데이터 변환을 조사하여 트랜스포머가 순전파 내의 그래디언트 기반 업데이트를 통해 컨텍스트 내 학습을 구현할 수 있음을 보여준다.

실험 결과

연구 질문

• RQ1 단일 선형 자기-어텐션 레이어가 선형 회귀 과제에서 그래디언트 하강 단계를 구현할 수 있는가?
• RQ2 자기-어텐션 레이어를 가진 학습된 트랜스포머가 선형 회귀 데이터에서 GD-유사 해에 수렴하는가?
• RQ3 다중 어텐션 레이어와 MLP가 트랜스포머의 GD++ 및 비선형 과제 수행 능력에 어떤 영향을 미치는가?
• RQ4 트랜스포머의 컨텍스트 내 학습을 순전파 내에서 알고리즘(메사-최적화)을 학습하는 것으로 이해할 수 있는가?
• RQ5 포지션 구성 및 데이터 변환이 순전파 외부의 명시적 가중치 업데이트 없이 컨텍스트 내 학습을 가능하게 하는 역할은 무엇인가?

주요 결과

• 단일 헤드 선형 자기-어텐션 레이어가 선형 회귀를 위한 학습 데이터에 대해 그래디언트-하강 유사 업데이트를 수행할 수 있다.
• 학습된 선형 자기-어텐션 레이어는 구성된 GD 업데이트와 예측 및 민감도까지 포함하여 밀접하게 일치한다.
• 다층 자기-어텐션은 반복적 곡률 보정(GD++)을 구현하고 선형 과제에서 일반 GD보다 우수하다.
• MLP를 포함하면 트랜스포머가 딥 표현에서의 그래디언트 하강을 수행하여 비선형 회귀를 해결하고, 사실상 커널 유사 회귀를 가능하게 한다.
• 트랜스포머는 학습된 데이터 변환과 태스크 특이적 표현을 통해 컨텍스트 내 학습을 얻을 수 있으며, 이는 메사-최적화 및 빠른 가중치의 개념과 일치한다.
• 해당 아키텍처는 분포 내/분포 외 태스크에서 그래디언트 기반 학습 역학을 재현하거나 근사할 수 있다.