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

Jul 25, 2026stat.ML

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

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.

Yue Yao, Caleb N. Ellington, Jingyun Jia +9

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

Adaptive Inference Meta-Learning Mixture-of-Experts Inference In-Context Learning

Ask PDF

Oct 5, 2026stat.ML

Learning Decision-Stump Thresholds in Context: Dynamics of Softmax Attention

Estimating a decision threshold requires locating observations near an unknown boundary. We study how gradient-based pretraining learns this statistical rule in a two-parameter softmax-attention model with a fixed feature and inequality direction. Pretraining uses labeled contexts and their true thresholds; a fresh threshold must be inferred from context alone. Under a large-resolution initialization, constant-step gradient descent on

m

tasks with

n

examples each produces a frozen estimator with error

\widetilde O((m\wedge n)^{-1}+N^{-1})

for each fixed interior threshold and every fresh-context size

N

. The two terms separate finite-pretraining accuracy from fresh-context localization. The mechanism is coordinated parameter divergence: population training calibrates the relative label and feature scores, then increases the attention scale as

t^{1/4}

, giving population threshold error

O(t^{-1/4})

. To transfer this mechanism to a fixed finite corpus, we control gradient errors relative to the shrinking directions of progress at successive parameter scales. This certifies a growing training interval without requiring long-time tracking of the population trajectory. We also identify the boundary limitation of the one-head model and explain statistically what a reflected symmetrization could achieve.

Hong Ha Le, Jackie Lok, Atsushi Nitanda +1

National University of Singapore · Princeton University · Nanyang Technological University

Language Model Pretraining Softmax Attention Parameter Estimation Gradient Descent In-Context Learning

Ask PDF

May 13, 2026stat.ML

Adaptive Kernel Density Estimation with Pre-training

Density estimation in high-dimensional settings is an important and challenging statistical problem.Traditional methods based on kernel smoothing are inefficient in high dimensions due to the difficulties in specifying appropriate location-adaptive kernels. In this work, we introduce pre-training, a key idea behind many cutting-edge AI technologies, to the context of non-parametric density estimation. By establishing a pre-trained neural network that can recommend an appropriate location-adaptive kernel for each sample point, efficient density estimation with adaptive kernels is achieved in high dimensions. A wide range of numerical experiments show that this strategy is highly effective for improving density-estimation accuracy, when the target distribution is close to the distribution family for pre-training. When the target distribution is substantially different from the pre-training distribution family, the benefit from the proposed pre-training strategy may be diluted, but can be reactivated by an additional fine-tuning procedure.

Ruitong Zhang, Ke Deng

Department of Statistics and Data Science Tsinghua University Beijing, China

Kernel Density Estimation Density Estimation

Ask PDF