Transformers Learn Robust In-Context Regression under Distributional Uncertainty

• URL: http://arxiv.org/abs/2603.18564v1
• Date: Thu, 19 Mar 2026 07:24:48 GMT
• Title: Transformers Learn Robust In-Context Regression under Distributional Uncertainty
• Abstract summary: We study in-context learning for noisy linear regression under a broad range of distributional shifts.<n>We compare Transformers against classical baselines that are optimal or suboptimal under the corresponding maximum-likelihood criteria.<n>Across all settings, Transformers consistently match or outperform these baselines.
• Score: 14.348332291561086
• License: http://creativecommons.org/licenses/by-nc-sa/4.0/
• 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.

Related papers

• Variational Rectification Inference for Learning with Noisy Labels [74.85528327499662]
  We propose variational rectification inference (VRI) to formulate the adaptive rectification for loss functions.<n>VRI is constructed as a hierarchical Bayes by treating the rectifying vector as a latent variable.<n>By introducing the variational term in VRI, the conditional posterior is estimated accurately and avoids collapsing to a Dirac delta function.
  arXiv  Detail & Related papers  (2026-03-18T01:25:08Z)
• Diffusion-Inspired Reconfiguration of Transformers for Uncertainty Calibration [52.017716672255524]
  Uncertainty calibration in pre-trained transformers is critical for their reliable deployment in risk-sensitive applications.<n>We propose a diffusion-inspired reconfiguration of transformers in which each feature transformation block is modeled as a probabilistic mapping.<n>Our method achieves superior calibration and predictive accuracy compared to existing uncertainty-aware transformers.
  arXiv  Detail & Related papers  (2026-02-09T17:24:47Z)
• Learning on Transformers is Provable Low-Rank and Sparse: A One-layer Analysis [63.66763657191476]
  We show that efficient numerical training and inference algorithms as low-rank computation have impressive performance for learning Transformer-based adaption. We analyze how magnitude-based models affect generalization while improving adaption. We conclude that proper magnitude-based has a slight on the testing performance.
  arXiv  Detail & Related papers  (2024-06-24T23:00:58Z)
• Linear Transformers are Versatile In-Context Learners [19.988368693379087]
  We prove that each layer of a linear transformer maintains a weight vector for an implicit linear regression problem. We also investigate the use of linear transformers in a challenging scenario where the training data is corrupted with different levels of noise. Remarkably, we demonstrate that for this problem linear transformers discover an intricate and highly effective optimization algorithm.
  arXiv  Detail & Related papers  (2024-02-21T23:45:57Z)
• Learning a Gaussian Mixture for Sparsity Regularization in Inverse Problems [2.174919458782602]
  In inverse problems, the incorporation of a sparsity prior yields a regularization effect on the solution.<n>We propose a probabilistic sparsity prior formulated as a mixture of Gaussians, capable of modeling sparsity with respect to a generic basis.<n>We put forth both a supervised and an unsupervised training strategy to estimate the parameters of this network.
  arXiv  Detail & Related papers  (2024-01-29T22:52:57Z)
• Trained Transformers Learn Linear Models In-Context [39.56636898650966]
  Attention-based neural networks as transformers have demonstrated a remarkable ability to exhibit inattention learning (ICL) We show that when transformer training over random instances of linear regression problems, these models' predictions mimic nonlinear of ordinary squares.
  arXiv  Detail & Related papers  (2023-06-16T15:50:03Z)
• The noise level in linear regression with dependent data [36.252641692809924]
  We derive upper bounds for random design linear regression with dependent ($beta$-mixing) data. Up to constant factors, our analysis correctly recovers the variance term predicted by the Central Limit Theorem.
  arXiv  Detail & Related papers  (2023-05-18T17:55:52Z)
• Understanding Noise-Augmented Training for Randomized Smoothing [14.061680807550722]
  Randomized smoothing is a technique for providing provable robustness guarantees against adversarial attacks. We show that, without making stronger distributional assumptions, no benefit can be expected from predictors trained with noise-augmentation. Our analysis has direct implications to the practical deployment of randomized smoothing.
  arXiv  Detail & Related papers  (2023-05-08T14:46:34Z)
• Predicting Out-of-Domain Generalization with Neighborhood Invariance [59.05399533508682]
  We propose a measure of a classifier's output invariance in a local transformation neighborhood. Our measure is simple to calculate, does not depend on the test point's true label, and can be applied even in out-of-domain (OOD) settings. In experiments on benchmarks in image classification, sentiment analysis, and natural language inference, we demonstrate a strong and robust correlation between our measure and actual OOD generalization.
  arXiv  Detail & Related papers  (2022-07-05T14:55:16Z)
• Optimal Online Generalized Linear Regression with Stochastic Noise and Its Application to Heteroscedastic Bandits [88.6139446295537]
  We study the problem of online generalized linear regression in the setting of a generalized linear model with possibly unbounded additive noise. We provide a sharp analysis of the classical follow-the-regularized-leader (FTRL) algorithm to cope with the label noise. We propose an algorithm based on FTRL to achieve the first variance-aware regret bound.
  arXiv  Detail & Related papers  (2022-02-28T08:25:26Z)
• Interpolation can hurt robust generalization even when there is no noise [76.3492338989419]
  We show that avoiding generalization through ridge regularization can significantly improve generalization even in the absence of noise. We prove this phenomenon for the robust risk of both linear regression and classification and hence provide the first theoretical result on robust overfitting.
  arXiv  Detail & Related papers  (2021-08-05T23:04:15Z)
• Consistency Regularization for Certified Robustness of Smoothed Classifiers [89.72878906950208]
  A recent technique of randomized smoothing has shown that the worst-case $ell$-robustness can be transformed into the average-case robustness. We found that the trade-off between accuracy and certified robustness of smoothed classifiers can be greatly controlled by simply regularizing the prediction consistency over noise.
  arXiv  Detail & Related papers  (2020-06-07T06:57:43Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.