[논문 리뷰] Uncovering mesa-optimization algorithms in Transformers

Johannes von Oswald, Eyvind Niklasson외 10명|arXiv (Cornell University)|2023. 09. 11.

Topic Modeling인용 수 5

한 줄 요약

이 논문은 자동회귀적으로 훈련된 트랜스포머가 추론 중에 내부 목적 함수를 문맥에서 구성하고 이를 최적화하여 예측을 생성하는 경량 기반 메사 최적화를 암묵적으로 구현한다는 것을 보여준다. 저자들은 이 과정을 역공학적으로 분석하여, 문맥 내에서 최적화 문제를 명시적으로 해결하는 새로운 메사 레이어를 제안하고, 이로 인해 few-shot 문맥 학습 및 시퀀스 모델링 작업에서의 성능 향상이 이루어짐을 보였다.

ABSTRACT

Transformers have become the dominant model in deep learning, but the reason for their superior performance is poorly understood. Here, we hypothesize that the strong performance of Transformers stems from an architectural bias towards mesa-optimization, a learned process running within the forward pass of a model consisting of the following two steps: (i) the construction of an internal learning objective, and (ii) its corresponding solution found through optimization. To test this hypothesis, we reverse-engineer a series of autoregressive Transformers trained on simple sequence modeling tasks, uncovering underlying gradient-based mesa-optimization algorithms driving the generation of predictions. Moreover, we show that the learned forward-pass optimization algorithm can be immediately repurposed to solve supervised few-shot tasks, suggesting that mesa-optimization might underlie the in-context learning capabilities of large language models. Finally, we propose a novel self-attention layer, the mesa-layer, that explicitly and efficiently solves optimization problems specified in context. We find that this layer can lead to improved performance in synthetic and preliminary language modeling experiments, adding weight to our hypothesis that mesa-optimization is an important operation hidden within the weights of trained Transformers.

연구 동기 및 목표

• 자기회귀적으로 훈련된 트랜스포머가 메타학습을 통해 few-shot 학습을 위해 설계되지 않았음에도 불구하고, 추론 중에 내부 최적화 과정을 수행하는지 여부를 조사하기 위해.
• 간단한 시퀀스 모델링 작업을 역공학적으로 분석하여 어텐션 메커니즘이 어떻게 내부 훈련 세트를 구성하고 목적 함수를 최적화하는지 밝혀내기 위해.
• 이러한 메사 최적화가 모델 크기와 무관하게 표준 자기회귀 트랜스포머에서 문맥 기반 few-shot 학습 능력을 제공하는지 평가하기 위해.
• 내부 최적화 문제를 명시적이고 효율적으로 해결하는 새로운 어텐션 메커니즘인 메사 레이어를 설계하고 평가하기 위해.
• 메사 레이어의 성능을 합성 및 초도 언어 모델링 작업에서 테스트하여, 문맥 기반 학습 능력을 향상시킬 잠재력을 평가하기 위해.

제안 방법

• 간단한 시퀀스 작업에서 훈련된 자기회귀 트랜스포머를 역공학적으로 분석하여 내부 메커니즘을 규명하고, 훈련 데이터를 구성하고 목적 함수를 최적화하는 방식을 파악하기 위해.
• 초기 어텐션 레이어가 입력 토큰을 그룹화하고 복사함으로써 내부 데이터셋을 형성하고, 암묵적으로 목적 함수를 정의함을 확인함.
• 더 깊은 레이어를 사용하여 이러한 내부 목적 함수를 기울기 기반 최적화로 최적화하여 예측을 생성함.
• 메사 레이어를 제안함 — 이는 기울기 스텝 대신 닫힌 형태의 해를 통해 최소 제곱 최적화 문제를 직접 해결하는 새로운 자기어텐션 메커니즘.
• 메사 레이어의 훈련 안정성을 향상시키기 위해 키 및 쿼리 정규화를 적용하여 수렴성과 성능 향상.
• GPT-2 스타일 아키텍처의 표준 자기어텐션 레이어를 메사 레이어로 교체하고, 표준 자기회귀 언어 모델링 설정을 사용하여 The Pile 데이터셋에서 평가함.

실험 결과

연구 질문

• RQ1 자기회귀적으로 훈련된 트랜스포머는 메타학습 없이도 추론 중에 기울기 기반 메사 최적화를 암묵적으로 수행하는가?
• RQ2 트랜스포머 내부의 최적화 과정을 역공학적으로 분석하여, 문맥에서 내부 목적 함수를 어떻게 구성하고 해결하는지 밝힐 수 있는가?
• RQ3 표준 자기회귀 트랜스포머가 메사 최적화 덕분에 얼마나 강력한 문맥 기반 few-shot 학습 능력을 보이는가?
• RQ4 메사 레이터와 같은 새로운 어텐션 메커니즘이 내부 최적화 문제를 명시적으로 해결함으로써 표준 자기어텐션보다 시퀀스 모델링에서 더 우수한 성능을 낼 수 있는가?
• RQ5 표준 어텐션을 메사 레이어로 교체하면 언어 모델링 작업에서 문맥 기반 학습 성능이 향상되는가?

주요 결과

• 자기회귀적으로 훈련된 트랜스포머는 두 단계의 메사 최적화 과정을 수행한다: (i) 어텐션 메커니즘을 통한 내부 훈련 데이터 구성, (ii) 예측 생성을 위한 내부 목적 함수의 기울기 기반 최적화.
• 훈련된 트랜스포머의 순방향 전파 과정에는 문맥에 기반한 목적 함수를 운영하는 암묵적인 기울기 하강 알고리즘이 포함되어 있으며, 이는 표준 LLM에서의 메사 최적화를 보여준다.
• 표준 자기회귀 손실로 훈련된 작고 단순한 트랜스포머도 모델 크기와 무관하게 강력한 문맥 기반 few-shot 학습 능력을 보이며, 이는 메사 최적화의 결과이다.
• 프롬프트 튜닝은 이러한 모델에서 문맥 기반 학습 성능을 향상시키며, 이는 대규모 언어 모델에서의 기능과 유사하다.
• 최소 제곱 최적화 문제를 직접 해결하는 것으로 제안된 메사 레이어는 합성 시퀀스 작업에서 표준 소프트맥스 및 선형 자기어텐션 트랜스포머보다 우수한 성능을 보였다.
• 초기 언어 모델링 실험 결과, 표준 어텐션을 메사 레이어로 교체하면 The Pile 데이터셋에서 성능 향상과 강력한 문맥 기반 학습 능력 향상이 이루어졌다.