Transformers Learn Robust In-Context Regression under Distributional Uncertainty

TL;DR

This paper investigates whether Transformers can perform robust in-context linear regression under realistic distributional uncertainties, including non-Gaussian and dependent data, and finds they often outperform classical methods.

Contribution

It demonstrates that Transformers can effectively learn in-context regression under broad distributional shifts, extending beyond traditional assumptions.

Findings

Transformers match or outperform classical estimators under various distributional shifts.

Robust in-context learning is achievable with Transformers even with non-Gaussian and dependent data.

Transformers demonstrate adaptability beyond classical maximum-likelihood-based methods.

Abstract

Recent work has shown that Transformers can perform in-context learning for linear regression under restrictive assumptions, including i.i.d. data, Gaussian noise, and Gaussian regression coefficients. However, real-world data often violate these assumptions: the distributions of inputs, noise, and coefficients are typically unknown, non-Gaussian, and may exhibit dependency across the prompt. This raises a fundamental question: can Transformers learn effectively in-context under realistic distributional uncertainty? We study in-context learning for noisy linear regression under a broad range of distributional shifts, including non-Gaussian coefficients, heavy-tailed noise, and non-i.i.d. prompts. We compare Transformers against classical baselines that are optimal or suboptimal under the corresponding maximum-likelihood criteria. Across all settings, Transformers consistently match or outperform these baselines, demonstrating robust in-context adaptation beyond classical estimators.