Nonparametric In-Context Learning under Growing Geometric Complexity: Minimax Optimality and Local Geometry-Adaptivity of Transformers¶

Conference: NeurIPS2026 (Accepted)
arXiv: 2609.31458
Code: https://github.com/seojaehee02/nonparametric_ICL
Area: Learning Theory
Keywords: in-context learning, nonparametric regression, manifold mixtures, minimax optimality, local geometry-adaptivity

TL;DR¶

The paper establishes the minimax rate of nonparametric in-context learning on heterogeneous manifold mixtures whose complexity grows with sample size, and proves that a structure-informed transformer approximates a tangent local-polynomial estimator attaining that rate; optimality of a trained predictor still requires sufficient training tasks and a sufficiently small empirical-risk optimization gap.

Background & Motivation¶

Nonparametric in-context learning (ICL) treats a prompt as regression data from a new task: a model receives labeled context examples and an unlabeled query, then estimates the new function in a forward pass without updating its parameters. Existing theory explains local-polynomial regression in Euclidean spaces and kernel-style prediction on a single manifold, but representations may contain regions with different intrinsic dimensions, smoothness levels, and probability masses. Describing error through the ambient dimension or one maximum intrinsic dimension discards both how often queries enter a region and how many context samples that region receives.

The additional challenge is not merely combining several fixed manifolds. This paper allows the number of components to increase, mixture masses to decrease, and reach and separation to shrink. It must control whether local windows mix components, whether estimated coordinates remain stable, and whether higher-order regression can still exploit function smoothness. Low-order kernel averaging is insufficient for arbitrary positive smoothness, while estimating a tangent space before regression introduces data-dependent coordinate error. Moreover, proving that a transformer represents an estimator does not prove that SGD finds it.

Core idea: under scale conditions that keep local geometry resolvable, estimate higher-order local geometry before regression in stable tangent coordinates, average component-specific statistical difficulty by query probability, and account separately for representation, cross-task generalization, and optimization error.

Method¶

Overall Architecture¶

The input consists of a task's noisy numerical covariates and responses, together with a noisy query covariate. The target is the function value at the generating latent query point, not its noisy response. Tasks share a structural environment and sampling law but independently draw new regression functions. The analysis establishes an aggregate lower bound, a statistical estimator attaining it, a transformer comparator implementing that estimator, and a conditional risk bound for selecting a predictor from finitely many pretraining tasks.

Two kinds of oracle information must be distinguished. Algorithm 1 receives the query component's intrinsic dimension, smoothness, and bandwidth, but neither its tangent space nor the task function; it fits the tangent space from context covariates. The comparator in Theorem 3 additionally uses the specified structural model, encoding a geometry-dependent component selector, structural parameters, and candidate grids in its weights. The query selects the relevant state. Parameters are fixed before prompts are sampled and shared across functions, but may change with the structural environment.

Thus, local geometry-adaptivity means estimating the current prompt's tangent space and regression coefficients within a structure-informed construction. It is not a prior-free guarantee for completely unknown environments, smoothness levels, and mixture masses. The contribution is primarily statistical and representational; the explanation below follows the proof objects rather than drawing the theorem sequence as a trained network architecture.

Key Designs¶

1. Aggregate local difficulty: retain dimension, smoothness, and rarity together

For a query in component \(k\), the expected number of context samples from that component is \(N_k=n\pi_{k,n}\), whereas only \(N_kh_k^{d_k}\) fall inside the local window. The bandwidth \(h_k=N_k^{-1/(2\alpha_k+d_k)}\) balances squared bias and response-noise variance. The aggregate risk benchmark is

\[ \mathfrak{R}_n=\sum_{k=1}^{K_n}\pi_{k,n}(n\pi_{k,n})^{-2\alpha_k/(2\alpha_k+d_k)}. \]

A small-mass component is harder conditionally, but queries also visit it less often. This weighted sum retains both effects; rarity does not automatically imply domination of overall risk. A balanced homogeneous mixture reduces to \((n/K_n)^{-2\alpha/(2\alpha+d)}\). For \(K_n\asymp n^\eta\), it becomes \(n^{-(1-\eta)2\alpha/(2\alpha+d)}\), but only for sequences satisfying the geometric conditions, with \(\eta\leq1-\kappa\).

Assumption 3 makes “growing but still tractable” precise. With \(r_n=\max_k h_k\), every component must have polynomially growing effective sample size, reach and separation above the local regression scale, and covariate perturbations below the corresponding bias scale:

\[ N_k\geq n^\kappa,\qquad \tau_{0,n}\wedge\delta_{0,n}\geq a_{\rm sc}^{-1}r_n,\qquad \sigma_{k,n}\leq a_{\rm sc}h_k^{\alpha_k\vee1}. \]

Here \(a_{\rm sc}\) is a sufficiently small constant depending on fixed model bounds, and \(\kappa\in(0,1)\). This excludes extremely rare components that are barely observed in context and implies \(K_n\leq n^{1-\kappa}\). Local chart derivative bounds and chart radii remain uniform; global reach and separation may shrink. The result does not cover arbitrary geometric degeneration.

The lower-bound proof places disjoint bumps at each component's own resolution and combines all components into one Assouad family. A submodel with zero covariate perturbation and admissible bounded smooth response noise uses a Hellinger translation bound to keep neighboring functions difficult to distinguish. One testing argument accumulates all component contributions simultaneously, rather than adding independently established lower bounds without justification.

2. Higher-order geometry fitting: prevent coordinate error from destroying smoothness gains

At query \(x\), displacements are first normalized as \(z_i=(X_i-x)/h_x\). A candidate contains a normal query offset, an orthogonal projector of the query dimension, and higher-order tensors describing the local surface. Local graph residuals are evaluated on a deterministic \(n^{-1}\) grid. Geometry fitting uses degree \(s_x=\lceil\alpha_x\rceil\), while regression uses degree \(p_x=\lceil\alpha_x\rceil-1\). Geometry needs one additional degree because, after rescaling the graph and dividing by bandwidth, its remainder must remain controlled at the regression resolution.

The geometry fit uses covariates only. Its loss is evaluated inside an ambient-space local window. Candidates less than \(\Delta_{{\rm tan},n}=n^{-3}\) above the minimum loss receive nonnegative thresholded weights, which are normalized before averaging their projectors. The average need not remain an orthogonal projector, so spectral rounding selects its top \(d_x\) directions. A fixed eigengap on the good event stabilizes this step. The grid cardinality is polynomial in sample size for fixed ambient dimension and degree, but its exponent depends on ambient dimension; this is not an inexpensive high-dimensional search guarantee.

Local-polynomial features are then formed under the estimated projector. Weights combine an ambient-displacement cutoff and a projected-displacement kernel. The former prevents distant observations from appearing nearby after projection and ensures a single-component window under separation conditions; the latter tapers weights within the window. The fitted normal-equation intercept predicts the query value, followed by clipping to the function bound. Ambient monomials are redundant on a low-dimensional tangent space, so the estimator uses a pseudoinverse rather than assuming an invertible ambient Gram matrix. A positive-spectrum bound on the identifiable polynomial subspace is sufficient to identify the intercept.

Let \(\delta_x\) be the operator-norm difference between the estimated projector and the true tangent projector at the generating latent query point. The key conditional mean-square error propagation is

\[ \mathbb{E}\bigl[(\widehat f_{\rm plug}^{\Pi}(x)-f(x^\star))^2\mid k,x^\star,\xi_{n+1}\bigr]\lesssim h_k^{2\alpha_k}+h_k^2\mathbb{E}[\delta_x^2\mid k,x^\star,\xi_{n+1}]+\sigma_{k,n}^2+\frac{1}{N_kh_k^{d_k}}. \]

Why does the logarithmic cost of tangent estimation not necessarily degrade the final statistical rate? The comparison polynomial has a bandwidth factor in its Lipschitz constant in normalized coordinates. Coordinate error therefore acquires a bandwidth factor before entering regression, producing \(h_k^2\mathbb E\delta_x^2\) after squaring. Local sample growth and perturbation conditions absorb these terms. The proof also constructs a Gram event uniform over all projectors near the true projector, handling the dependence created by using the same covariates for geometry estimation and regression. Response noise can then be analyzed conditionally on those covariates.

3. Structure-informed compilation: implement the estimator with stable charts and numerical workspace

The theoretical transformer's fixed affine input interface stores examples and the query in \(n+1\) rows, protects raw data and row-type markers, and uses additional coordinates as numerical registers. A structural selector determines dimension, degree, bandwidth, and mass normalization from the query. Local graph losses depend on fixed-dimensional weighted polynomial moments, so each candidate need not independently traverse every observation.

Communication uses uniform softmax: setting query/key matrices to zero, gating values with binary markers, and applying deterministic scaling implements exact query broadcasting and sample averaging. Localization comes from FFN-computed value cutoffs, not exactly zero softmax probabilities. Row-wise ReLU arithmetic modules perform candidate comparison, spectral decoding, and polynomial calculations. Two polynomial-width scratch banks compile these modules into residual blocks. One head communicates an entire register vector, keeping head count constant.

A tangent space does not admit a globally continuous choice of basis, so the construction uses multiple coordinate charts. A determinant threshold gives positive weight only to well-conditioned charts, each providing orthonormal coordinates for the same projected subspace. The reduced regression dictionary depends on intrinsic dimension and regression degree. On the good event, every positive-weight chart returns the same intercept as the ambient pseudoinverse, so chart averaging does not change the target estimator. Scaling both the Gram matrix and response vector by component mass maintains a fixed positive spectral lower bound, while this common factor cancels in the solution.

Spectral maps and inverse decoders are globally defined and bounded, but accurate approximation is promised only on their good spectral domains. Even a singular zero-weight chart returns a bounded value, and clipping controls small leakage from approximate weights. For every fixed \(A_0>0\), Theorem 3 gives mean-square approximation error \(O(n^{-A_0})\) relative to the statistical estimator, depth \(O(\log n)\), and polynomial width, workspace, and parameter magnitudes. Greater accuracy increases resource exponents, and geometric calculations retain ambient-dimensional costs.

This is an existence construction of a comparator. Its weights contain structural-model knowledge; it does not prove that SGD learns the selector from arbitrary training data. One parameter vector works uniformly over task functions and admissible density and perturbation laws within that environment, not over all unknown geometric environments.

4. Near empirical risk minimization: separate cross-task learning from within-task information

Pretraining draws \(\Gamma\) independent tasks from the same structural environment. Each supplies a prompt with \(n\) examples and a query response; testing draws a fresh function. The context size \(n\) determines how much of the current function is observed, while \(\Gamma\) determines the cross-task statistical cost of selecting an in-context prediction procedure. The covering analysis applies to the entire bounded-parameter softmax class, with freely varying query/key matrices, not only to the special uniform-attention weights used in the construction.

If the training output's empirical risk is within a deterministic tolerance \(\varepsilon_{{\rm opt},\Gamma}\) of the class minimum, Theorem 4 gives

\[ \mathbb{E}_{\mathcal D_\Gamma^{\rm tr}}R_{P,\rho_f}^{\star}(\widehat g_\Gamma)\lesssim\mathfrak{R}_n+n^{-A_0}+\frac{\mathfrak H_{n,\Gamma}^{\rm end}+1}{\Gamma}+\varepsilon_{{\rm opt},\Gamma}. \]

The four terms represent unavoidable within-task statistical error, implementation error, finite-task complexity, and optimization gap. The entropy proxy is at most \(C_{\rm ent}n^{c_{\rm ent}}\{1+\log(e\Gamma)\}\). Corollary 1 gives sufficient conditions \(A_0\geq2\alpha_{\max}/(2\alpha_{\max}+1)\), \(\Gamma\gtrsim n^{c_{\rm ent}}\log(en)\mathfrak R_n^{-1}\), and \(\varepsilon_{{\rm opt},\Gamma}\lesssim\mathfrak R_n\). These are not tight necessary task budgets, and they do not guarantee that AdamW or SGD reaches the tolerance.

The matching lower bound holds for every pretraining budget, but it is worst-case over task distributions. The hard prior draws an independent function index for each new task, so pretraining cannot supply the current function's missing labels. It does not imply the same error floor for every distribution: a distribution concentrated on the zero function has zero latent risk.

Loss & Training¶

The theoretical training objective is squared error against the noisy query response, whereas risk targets the latent function value. Conditionally mean-zero query noise makes the two differ by a predictor-independent noise variance:

\[ R_Y(g)=R_{P,\rho_f}^{\star}(g)+\mathbb E\epsilon_{n+1}^2. \]

This permits training on observable labels without counting irreducible response noise as latent prediction error. The latent query point is the generating point and need not equal the observed query's nearest manifold projection. The paper does not assume universal identifiability of tubular-noise models.

The empirical predictor has width 64, 4 workspace blocks, 4 heads per module, 32 learned workspace tokens, FFN width 256, and 286,401 parameters. Each block applies workspace-to-context cross-attention, workspace self-attention, and a GELU FFN, with LayerNorm/dropout disabled. An affine readout and tanh on the first workspace token produce the prediction. This is not the theoretical fixed-affine-interface, residual-ReLU compiled class.

Each budget trains from scratch with seed 0 and traverses the task pool once. Accumulating 25 single-prompt microbatches gives batch size 25. AdamW uses initial learning rate \(3\times10^{-4}\), weight decay \(10^{-4}\), cosine decay to zero, and gradient-norm clipping at 1; testing uses the final checkpoint. Budgets of 100,000, 200,000, 500,000, and 1,000,000 correspond to 4,000, 8,000, 20,000, and 40,000 updates. Increasing tasks also increases optimization, so the experiment does not isolate their contributions.

Key Experimental Results¶

Main Results¶

Appendix C fixes ambient dimension 64, context size 40,000, and 12 components with dimensions 2/4/6/8 and smoothness levels 0.75/1.5/3. Mixture masses range from 0.02 to 0.20, giving expected component context counts of 800–8,000. Ellipsoid shapes and sampling densities are heterogeneous. Covariate-noise radii follow \(h_k^{\alpha_k\vee1}/n\), response-noise variance is 0.01, and testing evaluates noiseless latent values.

The LP-CV baseline receives the query component's dimension, smoothness, mass, and true latent-query tangent space, tuning bandwidth separately for each stratum. It is not the estimated-tangent oracle in Algorithm 1. Its multiplier grid is 2/4/6/8/10/12/16/24, with 1,200 independent tuning prompts. All methods share 2,000 test prompts.

All MSEs below are in units of \(10^{-3}\); brackets are 95% percentile-bootstrap intervals. LP-CV is 0.690 [0.645, 0.739] at every budget.

Pretraining tasks	Workspace MSE	Transformer / LP-CV	Optimizer updates
100,000	14.455 [13.619, 15.342]	20.948 [19.177, 22.776]	4,000
200,000	0.919 [0.849, 0.995]	1.332 [1.222, 1.461]	8,000
500,000	0.538 [0.503, 0.575]	0.780 [0.727, 0.837]	20,000
1,000,000	0.505 [0.472, 0.540]	0.733 [0.686, 0.782]	40,000

Ablation Study¶

The paper does not report module-removal ablations. The following component analysis is retained without presenting it as mechanistic attribution. The budget is 1,000,000; MSEs remain in units of \(10^{-3}\), and ratios use unrounded values.

Intrinsic dimension	Smoothness	Test queries	LP-CV MSE	Transformer MSE	Ratio
2	0.75	41	0.606	0.448	0.739
2	1.5	60	0.376	0.310	0.825
2	3	97	0.545	0.480	0.880
4	0.75	72	0.435	0.371	0.853
4	1.5	106	0.632	0.538	0.851
4	3	217	0.845	0.418	0.495
6	0.75	117	0.692	0.474	0.684
6	1.5	175	0.882	0.404	0.458
6	3	317	0.605	0.546	0.903
8	0.75	159	0.667	0.515	0.772
8	1.5	252	0.648	0.433	0.667
8	3	387	0.780	0.679	0.871

The boundaries of theoretical and empirical conclusions can also be compared directly:

Object	Established result	What it does not establish
All measurable prompt predictors	Aggregate lower bound on a heterogeneous Assouad family	Equal difficulty for every fixed task distribution
Statistical estimator with known structural parameters	Aggregate upper bound despite estimated tangent spaces	Automatic adaptation to completely unknown dimension, smoothness, and mass
Structure-informed ReLU transformer comparator	Inverse-polynomial mean-square approximation, logarithmic depth, polynomial resources	Inexpensive geometric search with a fixed small network
Near ERM over the full bounded-parameter class	Four-term bound for statistical, implementation, finite-task, and optimization errors	SGD/AdamW convergence to a near ERM of the full class
Trained GELU workspace model	Improvement with increased joint budget on a fixed task family	Empirical verification of the asymptotic minimax rate under growing component complexity

Key Findings¶

At 500,000 tasks, the paper reports MSE reductions of 21.98% relative to LP-CV and 41.45% relative to the 200,000-budget model. The paired risk-ratio interval is below 1.
The 1,000,000-budget model reduces MSE by 26.75% relative to LP-CV and 6.10% relative to the 500,000-budget model. Percentages and ratios use the source's unrounded calculations; exact agreement should not be demanded from the three-decimal table entries.
Point estimates improve over LP-CV in all 12 components, but each has only 41–387 queries and no reported component-level confidence interval. Significant improvement in every component is therefore not established.
The 5,000 bootstrap replicates quantify test-prompt uncertainty conditional on fitted seed-0 models and selected bandwidths, not training-seed or bandwidth-selection uncertainty. Fixed context size and 12 fixed components also do not test a geometric sequence growing with sample size.

Highlights & Insights¶

Conditional and overall difficulty are distinct. The benchmark includes both reduced effective sample size and reduced query probability, making it useful for long-tail regions in mixture environments rather than describing the whole task through a worst-case dimension.
Propagation of geometric error matters more than geometric error alone. A bandwidth-squared factor suppresses stochastic tangent error and explains why estimated coordinates retain higher-order smoothness gains. This is more informative than tangent-estimator consistency alone.
Singularity does not imply unpredictability. Redundant ambient features create zero eigenvalues, yet the intercept remains identifiable. Multi-chart reduced solutions turn that statistical identifiability into a stable network implementation.
Representation and training success are accounted for separately. The comparator, class capacity, and near-ERM gap have independent roles, preventing an expressible optimal statistical procedure from being mistaken for one automatically learned by standard training.

Limitations & Future Work¶

The structural selector encodes geometry and component parameters. Genuine adaptation to unknown smoothness, mass, and changing environments remains unresolved; learning the upstream representation is also outside the analysis.
Feasible-scale conditions exclude extremely rare components, components closer than regression resolution, and excessive covariate perturbations. Results cannot be extrapolated to intersecting manifolds or arbitrary noise.
Polynomial workspace exponents depend on ambient dimension and degree. The theoretical construction provides neither a practical high-dimensional search nor sharp resource complexity guarantees.
The experiment explicitly does not certify all finite-sample constants of Assumption 3 at context size 40,000, and its empirical architecture is not the proved compiled class.
One training seed, synthetic fixed geometry, and jointly increasing task and update budgets limit empirical conclusions. Future experiments should separately vary context size, growing component count, and optimization budget, while adding repeated training runs and actual module ablations.

vs Kim et al. (2024): Earlier nonparametric ICL analysis uses learned basis representations. This paper emphasizes growing heterogeneous manifold mixtures and component-specific effective sample sizes rather than absorbing geometric differences into one global rate.
vs Ching et al. (2026): Euclidean local-polynomial ICL already exploits arbitrary positive Hölder smoothness. This paper adds usable estimated intrinsic coordinates and error-propagation control, but its resource bounds are more conservative, not a universal computational improvement.
vs Shen et al. (2025/2026): Related noisy-manifold regression and structured-manifold ICL primarily concern a single manifold and smoothness exponents at most 1. This paper addresses heterogeneous higher-order regression, but different perturbation models and structural information prevent claiming that it subsumes all earlier guarantees.
vs classical manifold regression and tangent estimation: Bickel–Li, Cheng–Wu, and Aamari–Levrard provide statistical and geometric foundations. This paper connects these tools into a prompt-dependent estimator, transformer implementation, and cross-task generalization argument.

Rating¶

Novelty: 4/5. Aggregate heterogeneous difficulty, growing-geometry conditions, and higher-order regression in estimated coordinates provide a clear theoretical advance.
Experimental Thoroughness: 4/5. Protocols, comparisons, and test uncertainty are explicit, but fixed geometry and one seed cannot verify asymptotic or optimization guarantees.
Writing Quality: 4/5. Structural information, theoretical versus empirical architectures, and quantifier boundaries are clearly stated, although the constructive appendices are intricate.
Value: 4/5. A useful statistical framework for geometry-aware ICL, with greater present value for theoretical understanding than direct guidance for large-model training.