Transformers Learn Robust In-Context Regression under Distributional Uncertainty
Authors: Hoang T. H. Cao, Hai D. V. Trinh, Tho Quan, Lan V. Truong
Organizations: Faculty of Computer Science and Engineering, Ho Chi Minh University of Technology (HCMUT), Vietnam National University Ho Chi Minh City (VNU-HCM), Vietnam
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.
Figures & tables
Figure 1: In-context prediction error under different coefficient priors in the noiseless setting. All methods are trained with ℓ2 loss. Top row: exponential (left) and Laplace (right) coefficient distributions. Bottom row: coefficients sampled uniformly from the unit hypersphere. We compare Transformer-based in-context learning against OLS, Ridge, and ℓ1 -based solvers (LP and ADMM). Errors are reported as normalized excess loss relative to a curriculum-dependent baseline.
Figure 2: Effect of feature distribution shifts on in-context learning. Gamma features induce skewed, non-negative marginals while preserving independence across context steps. VAR(1) features introduce temporal dependence across the prompt through an autoregressive structure.
Figure 3: Robustness of in-context learning under non-Gaussian and heavy-tailed noise distributions. All models in this figure are trained and evaluated using ℓ1 loss. We compare Transformer-based in-context learning with classical estimators, including least squares, Ridge regression, and ℓ1 solvers (LP and ADMM). For Bernoulli, exponential, Gamma, and Poisson noise, the ℓ1 objective corresponds to the maximum-likelihood estimator. The Student- t case ( ν=2 ) is shown separately, as it admits no finite-variance likelihood and falls outside the classical ML framework.
Appendix figures & tables7 assets
Supplementary material from the paper’s appendix.
Appendix
Figure 4: Extension of Figure 2 (Gamma-distributed features). In-context learning performance across different Gamma shape and rate parameters.
Figure 5: Extension of Figure 2 (VAR(1) features). Effect of increasing temporal correlation strength across the prompt.
Figure 6: Extension of Figure 3 (noise parameter sweeps). Performance under varying Bernoulli, Gamma, Poisson, and Student- t noise parameters. The qualitative trends observed in the main paper persist across a broad range of noise intensities.
Figure 7: Extension of Figure 3 (Exponential noise). Performance under exponential noise with varying rate parameter λ .
Figure 8: Extension of Figure 3 (Gamma noise). In-context learning under Gamma-distributed noise with varying shape and rate parameters.
Figure 9: Extension of Figure 3 (Poisson noise). Performance under Poisson noise with varying rate parameter λ .
Figure 10: Extension of Figure 3 (Student- t noise). Performance under heavy-tailed noise beyond the regimes covered by standard ℓ1 or ℓ2 optimality guarantees.
Transformers have demonstrated remarkable in-context learning (ICL) capabilities. The strong ICL performance of transformers is commonly believed to arise from their ability to implicitly execute certain algorithms on the context, thereby enhancing prediction and generation. In this work, we investigate how transformers with softmax attention perform in-context learning on linear classification data. We first construct a class of multi-layer transformers that can perform in-context logistic regression, with each layer exactly performing one step of normalized gradient descent on an in-context loss. Then, we show that our constructed transformer can be obtained through (i) training a single self-attention layer supervised by one-step gradient descent, and (ii) recurrently applying the trained layer to obtain a looped model. Training convergence guarantees of the self-attention layer and out-of-distribution generalization guarantees of the looped model are provided. Our results advance the theoretical understanding of ICL mechanism by showcasing how softmax transformers can effectively act as in-context learners.
Chenyang Zhang, Yuan Cao
School of Computing & Data Science, The University of Hong Kong
Pre-trained transformers are able to learn from examples provided as part of the prompt without any weight updates, a remarkable ability known as in-context learning (ICL). Despite its demonstrated efficacy across various domains, the theoretical understanding of ICL is still developing. Whereas most existing theory has focused on linear models, we study ICL in the nonlinear regression setting. Through the interaction mechanism in attention, we explicitly construct transformer networks to realize nonlinear features, such as polynomial or spline bases, which span a wide class of functions. Based on this construction, we establish a framework to analyze end-to-end in-context nonlinear regression with the constructed features. Our theory provides finite-sample generalization error bounds in terms of context length and training set size. We numerically validate the theory on synthetic regression tasks.
Alexander Hsu, Zhaiming Shen, Wenjing Liao +1
Department of Mathematics, Purdue University · School of Mathematics, Georgia Institute of Technology
In-context learning is a remarkable property of transformers and has recently received a lot of interest. In many studies of in-context learning, it has been shown that transformers are capable of implementing solver for linear and non-linear regression problems, in which the most of them implement gradient descent algorithm. However, it is still unclear whether those implementations have actually been acquired through training. In this paper, we construct a transformer with linear self-attention, which in-context learns the least squares estimate in a simple regression task. The point here is that the closed form (analytical) solution is approximately obtained by using layer normalization rather than an approximate solution based on gradient descent algorithm. Then, we show an experimental example, in which our implementation is mainly used in the transformer trained with l1 regularization when the target output is the least squares estimate.
Katsuyuki Hagiwara
Faculty of Education, Mie University · Kurima-Machiya-cho, Tsu, 514-8507, Japan