[논문 리뷰] Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection

Yu Bai, Fan Chen외 3명|arXiv (Cornell University)|2023. 06. 07.

Adversarial Robustness in Machine Learning인용 수 9

한 줄 요약

본 논문은 변환기가 광범위한 인-컨텍스트 학습 알고리즘(예: 최소제곱, 릿지, 라쏘, GLMs)을 구현하고 다양한 작업에 적응하기 위한 인-컨텍스트 알고리즘 선택을 수행할 수 있음을 이론적으로 입증하는 포괄적 이론을 제시하며, 이론적 보장과 실험적 검증을 제공한다.

ABSTRACT

Neural sequence models based on the transformer architecture have demonstrated remarkable \emph{in-context learning} (ICL) abilities, where they can perform new tasks when prompted with training and test examples, without any parameter update to the model. This work first provides a comprehensive statistical theory for transformers to perform ICL. Concretely, we show that transformers can implement a broad class of standard machine learning algorithms in context, such as least squares, ridge regression, Lasso, learning generalized linear models, and gradient descent on two-layer neural networks, with near-optimal predictive power on various in-context data distributions. Using an efficient implementation of in-context gradient descent as the underlying mechanism, our transformer constructions admit mild size bounds, and can be learned with polynomially many pretraining sequences. Building on these ``base'' ICL algorithms, intriguingly, we show that transformers can implement more complex ICL procedures involving \emph{in-context algorithm selection}, akin to what a statistician can do in real life -- A \emph{single} transformer can adaptively select different base ICL algorithms -- or even perform qualitatively different tasks -- on different input sequences, without any explicit prompting of the right algorithm or task. We both establish this in theory by explicit constructions, and also observe this phenomenon experimentally. In theory, we construct two general mechanisms for algorithm selection with concrete examples: pre-ICL testing, and post-ICL validation. As an example, we use the post-ICL validation mechanism to construct a transformer that can perform nearly Bayes-optimal ICL on a challenging task -- noisy linear models with mixed noise levels. Experimentally, we demonstrate the strong in-context algorithm selection capabilities of standard transformer architectures.

연구 동기 및 목표

• 파라미터 업데이트 없이도 변환기가 표준 ML 알고리즘을 인-컨텍스트에서 구현할 수 있음을 입증한다.
• 표현력, 인-컨텍스트 예측, 그리고 프리트레이닝 샘플 복잡도에 대한 엔드-투-엔드 이론을 제공한다.
• 인-컨텍스트 알고리즘 선택 메커니즘(사전-ICL 테스트와 사후-ICL 검증)을 도입하고 분석한다.
• 단일 트랜스포머가 태스크와 데이터 분포에 걸쳐 기본 ICL 알고리즘을 적응적으로 선택할 수 있음을 보인다.

제안 방법

• 인-컨텍스트 릿지 회귀와 최소제곱에 대한 명시적 트랜스포머 기반 구현을 구성한다.
• 인-컨텍스트 일반화 선형 모형(GLM) 및 볼록 리스크 최소화로 확장한다.
• 트랜스포머 내에서 효율적인 인-컨텍스트 경사 하강 메커니즘을 개발한다.
• 사후-ICL 검증과 사전-ICL 테스트의 두 가지 알고리즘 선택 메커니즘을 증명한다.
• 다항적으로 많은 시퀀스들로부터 다수의 ICL 태스크를 학습할 수 있음을 보여주는 프리트레이닝 결과를 제공한다.
• 노이즈가 있는 상황에서 단일 트랜스포머가 Bayes-optimal ICL에 근접할 수 있음을 이론적으로 및 경험적으로 검증한다.

실험 결과

연구 질문

• RQ1 트랜스포머가 인-컨텍스트에서 광범위한 표준 ML 알고리즘(예: 최소제곱, 릿지, 라쏘, GLMs)을 구현할 수 있는가?
• RQ2 하나의 트랜스포머가 시퀀스 전반에 걸쳐 서로 다른 기본 ICL 알고리즘 중에서 선택하는 메커니즘은 무엇인가?
• RQ3 정확한 인-컨텍스트 학습에 필요한 레이어/헤드 복잡도와 가중치 노름을 어떻게 상한화할 수 있는가?
• RQ4 다양한 태스크에서 학습된 트랜스포머가 도전적인 노이즈 환경에서 거의 Bayes-최적의 ICL을 보이는가?
• RQ5 프리트레이닝이 트랜스포머의 인-컨텍스트 학습 및 알고리즘 선택 능력에 어떤 영향을 미치는가?

주요 결과

• 트랜스포머는 입증 가능한 오차 보장을 갖고 인-컨텍스트 릿지 회귀와 최소제곱을 근사할 수 있다.
• 트랜스포머는 GLM에 대한 볼록 리스크 최소화와 2층 네트워크에 대한 인-컨텍스트 경사하강을 구현할 수 있다.
• 두 가지 일반적인 알고리즘 선택 메커니즘이 가능하다: 사전-ICL 테스트와 사후-ICL 검증으로 적응형 태스크/알고리즘 선택이 가능하다.
• 포스트-ICL 검증을 사용하여 혼합 노이즈 레벨의 노이즈 선형 모델에서 거의 Bayes-최적의 ICL을 달성할 수 있다.
• 프리트레이닝은 다항적으로 많은 학습 시퀀스로부터 여러 ICL 태스크를 학습하는 것을 지원한다.
• 대표적인 태스크에서 강한 인-컨텍스트 알고리즘 선택 능력을 실험적으로 확인한다.