PaperlayerView exact revision on arXiv

Source-linked AI summary

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

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

arXiv:2610.07804v1cs.LGstat.ML

TL;DR

The paper asks how gradient-based pretraining can learn statistical adaptivity when a location-estimation task’s distribution family is hidden. It studies attention experts and gating, and reports family-specific adaptive guarantees for Gaussian, uniform, and shrinking-variance mixture tasks.

  • Problem

    How gradient-based pretraining learns statistical adaptivity remains an open question, studied here in location estimation where the distribution family is hidden.

    Source: §1 · p. 2, §2 · p. 4

  • Method

    The paper analyzes stagewise gradient flow for gated attention experts, which supply both distribution-sensitive features and candidate location estimators.

    Source: §1 · p. 2, §7 · p. 10

  • Results

    The learned estimator achieves Gaussian asymptotic efficiency, nearly optimal uniform rates, and order-optimal rates for Gaussian mixtures with shrinking component variance.

    Source: §1 · p. 3, §7 · p. 10

  • Takeaways & Limitations

    Softmax gating enforces normalization and exact translation equivariance, while the GLU learns approximate normalization before selecting experts.

    Source: §1 · p. 3, §7 · p. 10

    The analysis relies on simplifying assumptions about training and architecture, including fixed parameters; removing those constraints makes the problem harder to study despite being easier to learn in practice.

    Source: §6 · p. 9

Abstract

from arXiv · show

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 $\sigma^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{\Omega}(n^{1+\epsilon})$ pretraining tasks, the learned estimator is asymptotically efficient on Gaussian tasks, within a factor $n^{\epsilon}$ 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.

1 INTRODUCTION

The introduction frames statistical adaptivity as learning to perform well across hidden data-generating families, then studies its emergence through gradient-based pretraining in a one-dimensional location-estimation task. It presents attention as both a source of distribution-sensitive features and a way to construct estimators, with gating combining the resulting experts.

1 INTRODUCTION

  • Gaussian samples favor the sample average at rate n^-1, while uniform samples permit the faster n^-2 rate through the sample mid-range, but the family is hidden.The learner receives neither the family label nor a prescribed estimator choice during pretraining.Source: §1 · p. 2
  • Scalar attention evaluates the derivative of the empirical cumulant-generating function, providing family-distinguishing features and estimators from the sample mean to the mid-range.Setting λ = 0 gives the mean; taking λ toward either infinity approaches the sample extremes.Source: §1 · p. 2
  • A Softmax MoE or GLU gates a bank of attention experts, and stagewise gradient flow first trains gates with experts fixed, then trains experts with gates fixed.Experiments indicate that this staged constraint is not strictly necessary.Source: §1 · p. 2
  • The learned estimator achieves Gaussian asymptotic efficiency, nearly optimal uniform rates, and order-optimal Gaussian-mixture rates, with guarantees extending to new locations and longer contexts.The uniform guarantee is within a factor n^ε of the minimax rate.Source: §1 · p. 3
  • Softmax gating enforces normalized weights and exact translation equivariance, whereas the GLU must learn approximate normalization; the risk analysis separates expert variance, estimator bias, and gating error.Distribution-sensitive weighting preserves the accuracy of the appropriate expert.Source: §1 · p. 3
  • Bias and variance evolve on different timescales: bias is tied to the sum of weights and prior variance, while variance accumulates across experts.The analysis tracks these components separately as training learns the gating and expert parameters.Source: §1 · p. 3

2 PROBLEM AND ARCHITECTURES

The section defines unlabeled location-estimation tasks where the hidden distribution family determines which estimator is appropriate, then presents a two-layer attention architecture and a stagewise training protocol. Its attention features distinguish families while its experts estimate location, with softmax gating providing normalization and exact translation equivariance.

2 PROBLEM AND ARCHITECTURES

  • Tasks provide observations from a distribution whose location is to be estimated, while the family is hidden and distributions differ only in shape, not variance.The context contains n observations generated from a shifted centered law.Source: §2 · p. 4

2.1 ARCHITECTURE

  • The two-layer transformer separates attention heads into location-estimating experts and heads that identify the distribution family, then combines expert outputs using either Softmax MoE or GLU gating.Each expert symmetrizes exponentially tilted averages at λ_k and −λ_k.Source: §2.1 · pp. 4–5
  • The location-invariant feature z converges to normalized CGF-derivative features that distinguish families; softmax gating guarantees exact translation equivariance, while GLU must learn normalization.With the population routing feature, unbiased experts give an unbiased Softmax MoE output, whereas GLU must also learn normalization.Source: §2.1 · p. 5

2.2 TRAINING PROTOCOL

  • Stagewise gradient flow first trains gating parameters θ with λ fixed, then trains expert parameters λ with θ fixed; the population version uses population loss in both stages.The empirical protocol samples m independent tasks and optimizes the empirical loss in the two stages.Source: §2.2 · p. 5

3 MAIN RESULTS

The section gives risk guarantees for the trained estimator across Gaussian, uniform, and Gaussian-mixture tasks, then analyzes how gating and expert parameters are learned in stages. It separates expert, routing, and normalization error and contrasts softmax normalization with the GLU’s learned normalization.

3 MAIN RESULTS

  • Softmax MoE has a stronger guarantee for large |µ| because it is exactly translation equivariant, while the GLU must learn approximate normalization.For the GLU, the stated risks are invariant to location.Source: §3 · pp. 5–6

3.1 PRETRAINING DYNAMICS

  • Pretraining analysis first trains gating parameters with fixed experts, then specializes expert parameters with fixed gates.The analysis defines population attention responses and family features, with barred quantities evaluated at population moments rather than as expectations of nonlinear statistics.Source: §3.1 · p. 6

3.2 STAGE 1: GATING DYNAMICS

  • The first-stage flow approaches a manifold of unbiased predictors while optimizing variance; GLU unbiasedness is not automatic, and larger prior variance can amplify normalization error.For some initializations, normalization inaccuracies can lead the gating layer toward a single expert.Source: §3.2 · p. 7

3.3 STAGE 2: EXPERT DYNAMICS

  • With gates frozen, the Gaussian expert parameter moves toward zero, the uniform expert toward endpoint estimation, and the mixture expert toward a limit.The tasks used to train expert parameters may also have been used to train the gates.Source: §3.3 · p. 7

3.4 EFFICIENCY OF THE FINAL ESTIMATOR

  • Final-estimator error separates expert, routing, and normalization terms; softmax normalization error vanishes, and the risk bounds preserve the Gaussian leading constant.The bounds control routing and normalization errors through the population gate’s distance from its assigned expert.Source: §3.4 · p. 8

4 EXPERIMENTS

Joint training learns family-specific routing and expert parameters in both architectures: Gaussian error stays near the sample-mean error, while uniform error moves toward the sample-mid-range error.

4 EXPERIMENTS

  • Gaussian error stays near the sample-mean error, while uniform error falls below it toward the sample-mid-range error.Both architectures show family-specific selection as |λ1| decreases toward zero and |λ2| grows; specialization begins with more variation across GLU runs.Source: §4 · p. 8, Figure 2 · p. 8, Figure 3 · p. 9
  • The figures track routing weights, absolute attention parameters, and median validation error, with shading spanning ten runs.Figure 2 reports these panels for Softmax MoE; Figure 3 uses the same conventions for GLU.Source: Figure 2 · p. 8, Figure 3 · p. 9

5 RELATED WORK

Related work addresses in-context learning, adaptive prediction guarantees, and attention training dynamics; this paper instead follows how pretraining learns family-dependent estimation rates.

5 RELATED WORK

  • Prior pretrained models address prediction or posterior approximation, while this work studies how pretraining attains estimation rates that depend on an unknown family.The cited model families include PFNs and tabular foundation models, as well as pretrained estimators of treatment effects, survival curves, mutual information, and Poisson means.Source: §5 · p. 9
  • Transformer theory covers in-context algorithm selection, minimax rates over known classes, and Bayesian predictors adapting to unknown structure.Those guarantees concern empirical risk minimizers or idealized predictors; this paper follows the training trajectory and identifies the parameter carrying adaptation.Source: §5 · p. 9
  • Prior gradient-based attention analyses include task allocation where head assignments are fixed by task structure, unlike changing-family contexts requiring learned routing.The related work also studies convergence for linear attention and learning shared low-dimensional structure.Source: §5 · p. 9
  • Adaptive estimation includes symmetric-location problems, and the rate gap here stems from the irregularity of the uniform family.The model is a small Set Transformer with either a softmax mixture-of-experts layer or a GLU.Source: §5 · p. 9

6 DISCUSSION

The analysis uses simplifying assumptions on training and architecture; separating training into stages slows it down but makes the problem easier to study.

6 DISCUSSION

  • The analysis relies on simplifying assumptions about both the training procedure and architecture.The authors report that removing fixed parameters makes the problem easier to learn in practice but leaves many possible optimal solutions, complicating analysis.Source: §6 · p. 9
  • Two-stage training slows training while making the problem easier to study.The authors identify extending the analysis to larger distribution families as a natural next step.Source: §6 · p. 9
  • A simple architecture has an upper bound on prior variance σ²_μ, and handling arbitrary σ²_μ remains a next-step architecture goal.The paper describes this bound as seeming unavoidable for a simple architecture.Source: §6 · p. 9

7 CONCLUSION

The conclusion presents stagewise gradient flow as learning adaptive location estimation without family labels, with distinct guarantees and routing behavior for the two architectures.

7 CONCLUSION

  • Stagewise gradient flow learns an adaptive location estimator without observing family labels, with Gaussian efficiency, nearly optimal uniform rates, and order-optimal mixture rates.The mixture guarantee applies when component variance shrinks; softmax enforces normalization and exact translation equivariance, while the GLU learns approximate normalization before selecting experts.Source: §7 · p. 10

A OUTLINE

  • The proof map assembles the main theorem from fresh-context risk, gating analysis, and expert-parameter analysis.Appendices also collect auxiliary estimates, Gaussian-mixture benchmarks, further experiments, LayerNorm discussion, and related work.Source: A OUTLINE · p. 14

B.4 ATTENTION-EXPERT DYNAMICS

  • After gates are fixed, expert-training error is quadratic in gate deviation and uniform over λ, so the bound applies throughout attention training.The attention-expert dynamics provide the accuracy guarantee for this fixed-gating stage.Source: §B.4 · p. 17

B.5 THEOREM ASSEMBLY

This proof appendix assembles theorem guarantees using estimator-accuracy arguments, population-flow calculations, concentration bounds, and staged gate-flow analyses. It establishes finite-time stability and trajectory bounds for the reduced flow, then uses loss descent and stationary-point exclusion to bound the full-flow gate.

B.5 THEOREM ASSEMBLY

  • The proof obtains the stated test-risk guarantees at finite target time tλ with joint probability at least 1−η; fixed-location bounds follow from translation identities.The proof freezes the gating parameters and applies the uniform training-task conclusion before invoking the test-risk result.Source: §B.5 · pp. 17–18

C ACCURACY OF THE FINAL ESTIMATOR

  • Centering removes location exactly, and the estimator’s deviation from the selected expert is quadratic in gate-weight error uniformly over λ.The centered features and gates are independent of the location parameter.Source: C ACCURACY OF THE FINAL ESTIMATOR · pp. 18–19

D.1 DERIVATION OF THE POPULATION GRADIENT FLOW; D.1.1 FLOW REPARAMETERIZATION

  • The population flow is reparameterized in gate coordinates by the chain rule, while layer normalization preserves the orders of the gradient and loss expansions.The expansions use Taylor expansion of the gates and population moment features.Source: §D.1 · pp. 19, 22, §D.1.1 · p. 25

D.2.1 UNIFORM CONCENTRATION BOUND FOR THE EMPIRICAL GRADIENT

  • With probability at least 1−δ, the empirical gradient bound holds simultaneously over gate parameters in balls of radius r=O(log n), both gates, and all coordinates.The concentration result conditions on a task family and assumes m and δ^-1 are polynomial in n.Source: §D.2.1 · p. 25

D.2.2 UNIFORM CONCENTRATION FOR THE LOSS

  • With probability at least 1−δ, the empirical loss concentration bound holds uniformly over the specified gate-parameter set.The result is stated for all θ1 and θ2 in the radius-r balls under the lemma’s assumptions.Source: §D.2.2 · pp. 27–28

D.3.1 PROOF OF LEMMA B.7

  • Stationarity rules out two interior gates in the same context, and convexity with strict convexity on selected coordinates yields a unique optimum.Source: §D.3.1 · p. 29

D.3.2 PROOF OF LEMMA B.8

  • A local PL inequality keeps the gates in [1/4, 3/4] and yields an O(n^-1) residual bound up to tcrit for the empirical flow.The proof also establishes exponential decay of the sum residuals and bounds the population-flow residual over this time horizon.Source: §D.3.2 · pp. 29–30

D.3.3 PROOF OF LEMMA B.9

  • The reduced selection flow tends to I2: the analysis excludes same-expert corners and shows the wrong-expert corner strictly maximizes truncated loss.The reduced flow cannot converge to that wrong-expert corner from its interior initialization.Source: §D.3.3 · p. 31
  • O(ε) closeness to the reduced slow flow and a bounded fast coordinate r prevent exit on fixed rescaled intervals.The proof extends the stability argument to population and empirical flows, with empirical gradient errors contributing o(1).Source: §D.3.3 · pp. 31–32
  • 2C0²n²σ²_mu bounds the time after tcrit to reach the S-neighborhood, while the trajectory remains in D under the stated conditions.The result applies for sufficiently small fixed σ²_mu and then sufficiently large n.Source: §D.3.3 · pp. 35–36

D.3.4 PROOF OF LEMMA B.10

  • O(n^-1) neighborhoods of stationary points with binary gates are excluded by loss descent, yielding the full-flow gate bound.The proof transfers the exclusion to the full flow using the trajectory's proximity to a stationary point and the triangle inequality.Source: §D.3.4 · pp. 37–38

D.3.5 PROOF OF LEMMA B.11

The proof bounds the learned estimator’s risk and shows that gradient flow remains near the optimal stationary point. It also carries the analysis to Gaussian and mixture estimation and to longer test contexts.

D.3.5 PROOF OF LEMMA B.11

  • The truncated, full population, and empirical gradient flows cannot exit the specified neighborhood after entering it.The truncated-flow result follows from the loss barrier around the optimal stationary point.Source: §D.3.5 · pp. 40–41

D.3.7 PROOF OF COROLLARY 3.2

  • At longer test length n′, the proof applies the corollary argument using n′ test observations.The population gates depend on context data through a quantity that does not depend on context length.Source: §D.3.7 · p. 42

E PROOF OF THE SOFTMAX MOE PART OF LEMMA 3.3

The proof analyzes softmax-router gradient flow with expert parameters fixed, then establishes finite-time gate accuracy. The same gate and parameter bounds hold for empirical flow with high probability.

E PROOF OF THE SOFTMAX MOE PART OF LEMMA 3.3

  • By time tθ = O(n^(B+1)), the population gates are within O(n^-B) of their family-specific targets, with router parameters bounded by O(log n).These bounds hold for every fixed B > 0 and sufficiently large n.Source: E PROOF OF THE SOFTMAX MOE PART OF… · p. 44
  • The reduced softmax-router flow converges to the identity gate matrix as time tends to infinity.The proof also gives finite-time convergence bounds for the empirical router flow with probability at least 1 − 1/n.Source: E PROOF OF THE SOFTMAX MOE PART OF… · p. 44
  • With probability at least 1 − 1/n, empirical flow attains the same finite-time gate and parameter bounds as population flow.The softmax weights sum to one, so the normalization-error term vanishes in the empirical-flow argument.Source: E PROOF OF THE SOFTMAX MOE PART OF… · pp. 44–45

F.1 PROOF OF LEMMA 3.4

This appendix establishes supporting results on gating dynamics, expert-output covariance, and asymptotic normality. It derives expert covariance and hierarchy conditions for Gaussian, uniform, and Gaussian-mixture inputs, and proves a concentration bound for softmax estimates.

F.1 PROOF OF LEMMA 3.4

  • At the target time, |λ1(tλ)| = o(1), because the flow drives |λ1| into every fixed neighborhood of zero before then.For each fixed δ in (0, 1), |λ1| enters [0, δ] within 2n/δ^3 time after tθ and cannot leave before tλ.Source: §F.1 · p. 45

F.2 EXPERT-PARAMETER GRADIENTS

  • With m = eΩ(n^{1+ε}) independent tasks, the empirical-gradient and loss bounds hold with probability at least 1−η; in particular, the loss is at most C/n when λ2≤1.The result is uniform over the parameter class specified in Lemma F.4.Source: §F.2 · pp. 48, 50, 51

F.3 GEOMETRY OF z

  • For d≥3, the normalized feature geometry satisfies ZGG=ZAA=1 and 0<ZGA<1; for mixtures with σn→0, λmin(Z) is eventually bounded below by a positive constant.For uniform data, these geometry quantities do not depend on n.Source: §F.3 · p. 51

F.4 EXPERT COVARIANCE

  • For sub-Gaussian Y, Lemma F.6 gives an expert-covariance approximation using the covariance obtained from the multivariate delta method.The proof concludes by combining the terms to establish the stated approximation.Source: §F.4 · p. 52

F.4.1 EXPERT HIERARCHY

  • Lemma F.7 establishes asymptotic normality for the full stacked vector of expert outputs, with covariance Σh specified in terms of the moment-generating function.The vector includes outputs indexed by λ1, …, λ4+d.Source: §F.4.1 · pp. 52–53

F.4.2 GAUSSIAN EXPERT COVARIANCE

  • For Y∼N(0,1), the moment-generating function is MY(t)=e^(t^2/2).Source: §F.4.2 · p. 53
  • For Gaussian inputs, the proof evaluates expert covariance using the Gaussian moment-generating function and the general covariance formula.The tilted mean is λ.Source: §F.4.2 · pp. 53–54

F.4.3 UNIFORM EXPERT COVARIANCE

  • For uniform inputs, the expert hierarchy holds for λ_i in [−1, 1], and the self-variance decreases on (0, 1); the hierarchy can fail at larger λ.The larger λ values reached in the second stage require a separate lemma because the bounded-range argument no longer applies.Source: §F.4.3 · pp. 55–56

F.4.4 GAUSSIAN-MIXTURE EXPERT COVARIANCE

  • For sufficiently small mixture variance σ_n, the expert hierarchy holds throughout [−L_max, L_max], with L_max growing as σ_n shrinks.The mixture observations are mean-zero, variance-one, and 1-sub-Gaussian.Source: §F.4.4 · pp. 56–57

F.5.1 SOFTMAX CONCENTRATION

  • For independent mean-zero 1-sub-Gaussian inputs, softmax concentration follows by truncating the ratios and bounding numerator and denominator deviations.The stated deviation bound has failure probability at most 4 exp(−nM^2/(32R^(2k)e^(2R|λ|))), in addition to the probability that any observation exceeds R.Source: §F.5.1 · p. 58

G ADDITIONAL RESULTS

The additional results state an exit-time bound for a Lyapunov function and give conditions for a slow system to track its reduced dynamics.

G ADDITIONAL RESULTS

  • While a trajectory remains in the region, the Lyapunov function decreases at least at rate γ, bounding its first exit time.The proof uses C < W(x(t)) ≤ W(x(0)) − γt until exit.Source: G ADDITIONAL RESULTS · pp. 58–59
  • Tikhonov’s theorem gives a tracking result when the fast subsystem has a uniformly attracting equilibrium and the stated regularity conditions hold.The passage also specifies initial proximity and boundedness conditions for the result.Source: G ADDITIONAL RESULTS · p. 59

H.1 GAUSSIAN AND UNIFORM BENCHMARKS

This block states Gaussian and uniform location-estimation benchmarks, introduces a Fisher-information lower-bound argument for Gaussian mixtures, and gives joint-training experiment details.

H.1 GAUSSIAN AND UNIFORM BENCHMARKS

  • The sample mean has risk n^-1 under both the Gaussian and uniform location families.The Gaussian benchmark passage states that this risk is attained by the sample mean.Source: §H.1 · p. 59

H.2 GAUSSIAN-MIXTURE BENCHMARK

  • For the Gaussian-mixture location model, the section derives Fisher information and applies the Cramér–Rao inequality to unbiased estimators.Source: §H.2 · p. 59

I ADDITIONAL EXPERIMENTS

  • The joint-training experiments use 8,000 fixed tasks with contexts of length 150 and run 600,000 gradient-descent steps.The setup uses both Softmax MoE and GLU gates, updating θ and λ jointly while keeping ω fixed.Source: I ADDITIONAL EXPERIMENTS · p. 60

I.1 JOINT TRAINING, GAUSSIAN-MIXTURE ALTERNATIVE

The appendix experiments examine joint and stagewise training, a less constrained end-to-end architecture, LayerNorm’s role, and related work on learned adaptation.

I.1 JOINT TRAINING, GAUSSIAN-MIXTURE ALTERNATIVE

  • Both architectures improve on the sample mean for Gaussian-mixture contexts, but remain above the component-mean error and the lower bound.The gates move toward the identity, while the second gate’s attention scale settles near 2, above the stated stopping target.Source: §I.1 · p. 60, Figure 4 · p. 61, Figure 5 · p. 61

I.3 STAGEWISE TRAINING

  • In stagewise training, the gates approach a sum of one before specializing toward different experts for Gaussian and uniform contexts.The second stage improves uniform-context expert and estimator errors toward the sample mid-range error, while Gaussian error remains approximately unchanged.Source: Figure 8 · p. 62, §I.3 · p. 62, Figure 9 · p. 63

I.4 STAGEWISE TRAINING, GAUSSIAN-MIXTURE ALTERNATIVE

  • In the stagewise Gaussian-mixture experiment, second-stage expert training reduces alternative-family error while leaving Gaussian error approximately unchanged.Mixture error remains above the component-mean error and lower bound; first-stage gates move toward the appropriate expert with residual normalization error.Source: §I.4 · p. 63, Figure 12 · p. 64, Figure 11 · p. 64

I.5 RELAXED ARCHITECTURE

  • The less constrained end-to-end architecture keeps Gaussian validation error near the sample-mean error and lowers uniform error toward the sample mid-range error.Source: §I.5 · p. 64, Figure 13 · p. 65

K EXTENDED RELATED WORK

  • Prior analyses establish adaptive predictors without following training trajectories, whereas this paper tracks an identifiable attention-scale parameter that moves estimators across a family.Source: K EXTENDED RELATED WORK · p. 66
  • Unlike prior fixed-task head allocation, this setting requires learned routing because the distribution family varies across contexts.The related-work discussion also frames the architecture comparison as enforcing expert-weight normalization versus learning it.Source: K EXTENDED RELATED WORK · p. 66

Related papers

Loading 2610.07804v1…