Adaptive Mean Estimation by In-Context Learning: A Gradient-Flow Analysis

Martin Eppert; Krishna Balasubramanian; Subhro Ghosh; Jason Klusowski; Yan Shuo Tan

cs.LGOct 6, 2026

Adaptive Mean Estimation by In-Context Learning: A Gradient-Flow Analysis

Authors: Martin Eppert, Krishna Balasubramanian, Subhro Ghosh, Jason Klusowski, Yan Shuo Tan

Organizations: National University of Singapore · University of California, Davis · Princeton University

Abstract

Prior Fitted Networks (PFNs) such as TabPFN now rival established statistical procedures across prediction and estimation tasks. A natural explanation is that PFNs have the property of statistical adaptivity, that is, they perform nearly as well as a method tailored to the true data-generating model for a heterogeneous set of models, while not being told which model the data comes from. We study how such adaptivity is learned in a controlled location-estimation problem. Each task is an unlabeled sample whose family is hidden: Gaussian data call for averaging, with error of order $n^{-1}$ , whereas uniform data are best estimated from their extremes, at the faster rate $n^{-2}$ . We also provide the example of a symmetric Gaussian mixture, for which a rate of $σ^2_n/n$ can be attained. On scalar inputs, softmax attention computes the derivative of the empirical cumulant-generating function. A single primitive therefore both supplies features that distinguish the families and forms estimators interpolating between the sample mean and the mid-range. We combine attention experts through either a softmax mixture of experts or a gated linear unit (GLU), and analyze stagewise gradient flow. With $\widetildeΩ(n^{1+ε})$ pretraining tasks, the learned estimator is asymptotically efficient on Gaussian tasks, within a factor $n^ε$ of the minimax rate on uniform tasks, and order-optimal on mixtures in a shrinking-variance regime. These guarantees extend to new locations and longer contexts. A risk decomposition separates expert error, routing error and normalization error, which clarifies the architectural contrast. Softmax gating enforces normalization and exact translation equivariance, whereas the GLU must learn it: its dynamics separate into fast bias removal followed by slow expert selection. End-to-end experiments recover the predicted specialization.

Figures & tables

Panel of Figure 1 — Figure 1: Illustration of our studied architecture. We also highlight the parts which were fixed for simplicity.

Alternative family	terminal $\|\lambda_{2}\|$	Gaussian risk	Alternative risk	$R_{\mathrm{A},n}^{\star}$
Uniform	$n^{1-\epsilon}$	$(1+o(1))/n$	$\mathcal{O}(n^{-2+\epsilon})$	$\asymp n^{-2}$
Gaussian mixture	$\frac{\log(1/\sigma_{n}^{2})}{4\sqrt{1-\sigma_{n}^{2}}}$	$(1+o(1))/n$	$\mathcal{O}(\sigma_{n}^{2}/n)$	$\geq\sigma_{n}^{2}/n$

Table 1: The terminal expert parameters and fresh-context risks certified by the main theorem.

Figure 2: Softmax MoE, joint training with the uniform alternative. Left: $(\bar{w}_{1\mathrm{G}},\bar{w}_{1\mathrm{A}})$ across runs, colored by iteration. Middle: median $|\lambda_{1}|$ and $|\lambda_{2}|$ , where thin lines are individual runs. Right: median validation error of $\widehat{\mu}$ on Gaussian and uniform contexts, with the sample mean and sample mid-range as references. Shading spans all ten runs.

Figure 3: GLU, joint training with the uniform alternative, with panels and conventions as in Figure 2 .

Appendix figures & tables11 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Softmax MoE, joint training with the Gaussian-mixture alternative, as in Figure 2 . The dashed line in the middle panel marks $\log(1/\sigma_{n}^{2})/(4\sqrt{1-\sigma_{n}^{2}})$ .

Figure 5: GLU, joint training with the Gaussian-mixture alternative, as in Figure 3 . The dashed line in the middle panel marks $\log(1/\sigma_{n}^{2})/(4\sqrt{1-\sigma_{n}^{2}})$ .

Figure 6: GLU, joint training with the uniform alternative: population gates

\bar{w}_{kj}

(left) and normalization error

|\bar{w}_{1j}+\bar{w}_{2j}-1|

(right). Curves show medians, and shading spans all ten runs.

Figure 7: GLU, joint training with the Gaussian-mixture alternative, with panels and conventions as in Figure 6 .

Figure 8: First stage with the uniform alternative: $(\bar{w}_{1j},\bar{w}_{2j})$ for Gaussian (left) and uniform (right) contexts across runs, colored by iteration. The dashed line is the critical manifold $\bar{w}_{1j}+\bar{w}_{2j}=1$ .

Figure 9: Uniform alternative: validation error of each expert

g_{k}

and of

\widehat{\mu}

on Gaussian (left) and uniform (right) contexts, with the sample mean and sample mid-range as references. The vertical line marks

t_{\theta}

Figure 10: Uniform alternative: validation error of each expert

g_{k}

and of

\widehat{\mu}

on Gaussian (left) and uniform (right) contexts, repeated for comparison with Figure 12 . The vertical line marks

t_{\theta}

Figure 11: First stage with the Gaussian-mixture alternative: $(\bar{w}_{1j},\bar{w}_{2j})$ for Gaussian (left) and Gaussian-mixture (right) contexts across runs, colored by iteration. The dashed line is the critical manifold $\bar{w}_{1j}+\bar{w}_{2j}=1$ .

Figure 12: Gaussian-mixture alternative: validation error of each expert

g_{k}

and of

\widehat{\mu}

on Gaussian (left) and Gaussian-mixture (right) contexts. The vertical line marks

t_{\theta}

Figure 13: Less constrained architecture, end-to-end training with the uniform alternative. Left: training loss. Middle and right: validation error on Gaussian and uniform contexts, with the sample mean and sample mid-range as references.

Figure 14: Stream plots of the reduced flow studied in Lemma D.8 . Left: population features with LayerNorm. Right: uncentered population features at

\mu=0

without LayerNorm.

Explore similar work

Context-Adaptive Inference: A Unified Statistical and Foundation-Model View
Jul 25, 2026Yue Yao, Caleb N. Ellington, Jingyun Jia +9Parameter-Efficient Adaptation In-Context Learning
Department of Statistics, University of Wisconsin-Madison · Computational Biology Department, Carnegie Mellon University · Department of Computer Science, Yale University +5
Adaptive Kernel Density Estimation with Pre-training
May 13, 2026Ruitong Zhang, Ke DengDensity Ratio Estimation Kernel Method
Department of Statistics and Data Science Tsinghua University Beijing, China
Adaptive Bayesian Online Learning via Expert Aggregation
Jul 22, 2026Jungbin Jun, Ilsang OhnLearning-Augmented Algorithms Bayesian
Department of Statistics Inha University Incheon, 22212, South Korea

Jul 25, 2026, stat.MLJ/K move · Enter open · S save

Context-Adaptive Inference: A Unified Statistical and Foundation-Model View

Yue Yao, Caleb N. Ellington, Jingyun Jia, Baiheng Chen+8

Department of Statistics, University of Wisconsin-Madison · Computational Biology Department, Carnegie Mellon University · Department of Computer Science, Yale University+5

Parameter-Efficient Adaptation In-Context Learning Foundation Model Multimorbidity

Open paper AskPDF arXiv

Modern predictive systems are expected to adapt their behavior to the specific situation they are facing. A clinical model should not treat every patient the same; a retrieval-augmented model should change its answer when given different evidence; a mixture-of-experts model should route different inputs to different experts. We call this capability context-adaptive inference: before predicting, the system uses information about the current context to specialize its parameters or computation for that instance. This article provides a unified view of context-adaptive inference across three traditions that are usually treated separately: (i) explicit adaptation in statistics (e.g. varying-coefficient models, local regression, hierarchical sharing), (ii) rapid task-specific adaptation in meta-learning and transfer, and (iii) implicit adaptation in large foundation models via prompting, retrieval, and expert routing. We formalize these approaches under a common objective: to map context

c

to adapted parameters

θ(c)

, then to predict via

f(x; θ(c))

. Under squared loss, linear prediction heads, and fixed features, we prove that explicit parameter adaptation and implicit routing are mathematically equivalent to kernel ridge regression on joint features of inputs and context. Building on this bridge, we propose practical design principles and evaluation metrics including adaptation-efficiency, routing stability, and context-specific robustness to guide when to specialize, how to constrain that specialization, and how to audit context-adaptive models in deployment. Finally, we identify open problems in identifiability, robustness under distribution shift, and efficient large-scale adaptation, outlining design principles for methods that are scalable, reliable, and transparent in real-world settings.

Abstract

Explore similar work

Context-Adaptive Inference: A Unified Statistical and Foundation-Model View

Adaptive Kernel Density Estimation with Pre-training

Adaptive Bayesian Online Learning via Expert Aggregation