In-Context Nonlinear Regression

• In-context nonlinear regression is a method where transformer models predict continuous targets directly from a prompt of labeled examples without modifying their parameters.
• The approach employs explicit transformer constructions—such as polynomial, spline, and Gaussian-kernel regression—to emulate nonlinear function fitting within the model’s forward pass.
• Key insights include the importance of prompt locality, the interplay between attention and feed-forward layers, and strategies like depth replacement to overcome embedding limitations.

In-context nonlinear regression is the problem of predicting a continuous target for a query input from a prompt of labeled examples, without updating model parameters at test time. In the modern transformer literature, the prompt typically has the form $(x_1,y_1,\ldots,x_n,y_n,x)$ or $P=((x_1,y_1),\ldots,(x_N,y_N),x_{N+1})$, and the central theoretical question is how a sequence model can realize a genuinely nonlinear regressor rather than a linear least-squares surrogate. Recent work studies this question through explicit transformer constructions for polynomial, spline, and Gaussian-kernel regression; through analyses of prompt locality in imbalanced settings; and through training-dynamics results showing how attention concentrates on relevant prompt tokens during learning (Sun et al., 30 Jan 2025, Hsu et al., 6 May 2026, Yan et al., 8 May 2026, Nejjar et al., 2024, Li et al., 28 Jul 2025).

1. Formal problem setting and prompt-based regression

A common formulation draws a task function $f\sim\mathcal{D}_F$, context inputs $x_i\stackrel{i.i.d.}{\sim}\mathcal{D}_X$, and labels $y_i=f(x_i)$, then asks a transformer $G_\theta$ to predict $f(x)$ from the prompt

$\mathbf{s}^f=(x_1,y_1,\ldots,x_n,y_n,x).$

Training is by empirical risk minimization over a prompt dataset

$$\Gamma=\{(\mathbf{s}^{f_\ell},y^\ell)\}_{\ell=1}^L, \qquad G_{\hat\theta}^\Gamma=\arg\min_{G_\theta\in\mathcal{G}} \mathcal{R}^\Gamma(G_\theta),$$

with

$$\mathcal{R}^\Gamma(G_\theta)=\frac{1}{L}\sum_{\ell=1}^L |G_\theta(\mathbf{s}^{f_\ell})-y^\ell|^2, \qquad \mathcal{R}(G_\theta)=\mathbb{E}_{\mathbf{s}^f}|G_\theta(\mathbf{s}^f)-f(x)|^2.$$

This formulation makes context length $P=((x_1,y_1),\ldots,(x_N,y_N),x_{N+1})$0 and training set size $P=((x_1,y_1),\ldots,(x_N,y_N),x_{N+1})$1 the two primary sample-complexity variables (Hsu et al., 6 May 2026).

A second formulation, used for in-context Gaussian-kernel ridge regression, treats the prompt as $P=((x_1,y_1),\ldots,(x_N,y_N),x_{N+1})$2 training tokens and one query token: $P=((x_1,y_1),\ldots,(x_N,y_N),x_{N+1})$3 with kernel

$P=((x_1,y_1),\ldots,(x_N,y_N),x_{N+1})$4

kernel matrix $P=((x_1,y_1),\ldots,(x_N,y_N),x_{N+1})$5, and predictor

$P=((x_1,y_1),\ldots,(x_N,y_N),x_{N+1})$6

The in-context viewpoint is that the forward pass should approximate this predictor by emulating a numerical solver rather than by explicitly solving the linear system $P=((x_1,y_1),\ldots,(x_N,y_N),x_{N+1})$7 (Yan et al., 8 May 2026).

In imbalanced regression, the prompt is not merely a few-shot conditioning device but a distribution-selection mechanism. The setup distinguishes a source set $P=((x_1,y_1),\ldots,(x_N,y_N),x_{N+1})$8 and a target set $P=((x_1,y_1),\ldots,(x_N,y_N),x_{N+1})$9, with source and target distributions $f\sim\mathcal{D}_F$0 and $f\sim\mathcal{D}_F$1. A pre-trained transformer $f\sim\mathcal{D}_F$2 predicts

$f\sim\mathcal{D}_F$3

where $f\sim\mathcal{D}_F$4 is a retrieved local context. The associated expected prediction error is decomposed as

$f\sim\mathcal{D}_F$5

so both context size and context composition affect the bias-variance tradeoff (Nejjar et al., 2024).

2. Expressivity barriers and the role of transformer sublayers

A central impossibility result concerns linear self-attention. For a linear transformer with only LSA layers,

$f\sim\mathcal{D}_F$6

stacking more LSA layers does not break the linearity barrier. Proposition 1 / Proposition $f\sim\mathcal{D}_F$7 states

$f\sim\mathcal{D}_F$8

Accordingly, pure LSA is no better than a best linear predictor over the original input, even when the target function is quadratic or more complex (Sun et al., 30 Jan 2025).

One response is to assign nonlinear capability to feed-forward layers. The bilinear Transformer studied in "On the Role of Transformer Feed-Forward Layers in Nonlinear In-Context Learning" alternates a bilinear feed-forward layer

$f\sim\mathcal{D}_F$9

with linear self-attention: $x_i\stackrel{i.i.d.}{\sim}\mathcal{D}_X$0 In this construction, the feed-forward layer actively creates products of features, i.e. nonlinear monomials, while attention performs the optimizer-like update in the lifted space (Sun et al., 30 Jan 2025).

A different constructive line assigns the nonlinear feature map directly to attention. "Understanding In-Context Learning for Nonlinear Regression with Transformers: Attention as Featurizer" introduces an Interaction Lemma for ReLU attention heads, using

$x_i\stackrel{i.i.d.}{\sim}\mathcal{D}_X$1

to isolate token interactions and write the result into selected matrix entries. On this basis, attention is used as an arithmetic primitive to construct polynomial and spline bases in context, and the paper reports that attention-only models without FFN still perform well (Hsu et al., 6 May 2026).

The literature therefore does not support a single universal architectural assignment. One constructive account states that nonlinear in-context regression requires a GLU-like bilinear feed-forward layer because LSA alone is fundamentally limited, while another shows that ReLU attention can itself realize the nonlinear feature map under explicit parameterizations. A plausible implication is that the division of labor between attention and feed-forward layers is architecture- and proof-model-dependent rather than absolute.

3. Solver interpretations: from polynomial kernel regression to Gaussian KRR

For quadratic targets, the bilinear-then-LSA construction implements one preconditioned gradient descent step on a polynomial kernel regression objective. Writing

$x_i\stackrel{i.i.d.}{\sim}\mathcal{D}_X$2

the update is

$x_i\stackrel{i.i.d.}{\sim}\mathcal{D}_X$3

with preconditioner $x_i\stackrel{i.i.d.}{\sim}\mathcal{D}_X$4 determined by attention parameters. The bilinear layer maps each input to a quadratic feature vector

$x_i\stackrel{i.i.d.}{\sim}\mathcal{D}_X$5

and LSA then produces

$x_i\stackrel{i.i.d.}{\sim}\mathcal{D}_X$6

This is exactly the form of a kernel regression / gradient-descent estimator in the quadratic feature space (Sun et al., 30 Jan 2025).

When the embedding dimension is insufficient to represent all quadratic monomials at once, depth replaces width. The block-coordinate update is

$x_i\stackrel{i.i.d.}{\sim}\mathcal{D}_X$7

and Theorem $x_i\stackrel{i.i.d.}{\sim}\mathcal{D}_X$8 states that there exists a deep bilinear Transformer with $x_i\stackrel{i.i.d.}{\sim}\mathcal{D}_X$9 such that

$y_i=f(x_i)$0

where $y_i=f(x_i)$1 is exactly the iterate produced by block-coordinate descent. The same paper also gives a width upper bound and lower bound: if $y_i=f(x_i)$2, then one bilinear block can achieve $y_i=f(x_i)$3; if $y_i=f(x_i)$4, then any bilinear block has loss at least

$y_i=f(x_i)$5

These results formalize the claim that depth can replace width in nonlinear in-context regression (Sun et al., 30 Jan 2025).

For Gaussian kernels, the solver interpretation changes from one-step gradient descent to iterative linear-system solving. "Transformers Can Implement Preconditioned Richardson Iteration for In-Context Gaussian Kernel Regression" studies the system

$y_i=f(x_i)$6

and shows that a standard transformer can implement preconditioned Richardson iteration

$y_i=f(x_i)$7

Under bounded-data assumptions, the paper constructs a single-head transformer with $y_i=f(x_i)$8 blocks and MLP width $y_i=f(x_i)$9 that achieves $G_\theta$0-accurate prediction for prompts of length $G_\theta$1. In this decomposition, softmax attention produces a row-normalized Gaussian-kernel operator for cross-token interaction, while ReLU MLP layers approximate the intra-token scalar arithmetic needed for the update (Yan et al., 8 May 2026).

These two lines share a common mechanistic template: the transformer forward pass is interpreted as an unrolled optimizer or solver. The difference lies in the object being solved: polynomial-kernel regression via one-step or block-coordinate updates in a lifted monomial space, versus Gaussian-kernel ridge regression via a convergent iterative method on the kernel linear system.

4. Prompt locality, task distributions, and context selection

The prompt distribution is a first-order determinant of in-context nonlinear regression performance. In the quadratic ICL analysis, prompts are structured as

$G_\theta$2

with inputs sampled i.i.d. from a Gaussian,

$G_\theta$3

or in some results $G_\theta$4, and with target functions given by random quadratics

$G_\theta$5

The same work also studies a non-Gaussian induced distribution after the feed-forward layer, for example

$G_\theta$6

and shows that optimization becomes harder for nonlinear ICL under such feature distributions (Sun et al., 30 Jan 2025).

In imbalanced regression, the decisive issue is locality. The model is assumed $G_\theta$7-Lipschitz in the context examples,

$G_\theta$8

and the paper compares transformer prediction against the local averaging estimator

$G_\theta$9

Its main theoretical conclusion is that global context can be harmful in imbalanced regression, whereas using only the nearest neighbors of the query better approximates the local conditional label distribution $f(x)$0. In many-shot regions, expected error is relatively flat or stable as context size increases; in few-shot regions, adding more context eventually worsens performance, producing a U-shaped curve (Nejjar et al., 2024).

Empirically, the same work reports that in-context learning is highly competitive and often superior in minority / tail regions, that localized context retrieval is crucial, and that using the full training set as context is often worse than a carefully selected neighborhood. On AgeDB-DIR, IMDB-WIKI-DIR, STS-B-DIR, and six tabular datasets, the strongest pattern is not that ICL uniformly dominates in every regime, but that it tends to outperform conventional in-weight regression in imbalanced or low-data regions, especially in few-shot and medium-shot slices (Nejjar et al., 2024).

A related distributional assumption appears in the training-dynamics analysis of one-layer softmax attention. There the prompt contains $f(x)$1 input-response pairs and one query,

$f(x)$2

and a concentration event

$f(x)$3

ensures each feature appears with approximately the expected frequency. With high probability,

$f(x)$4

which stabilizes the analysis of attention concentration (Li et al., 28 Jul 2025).

5. Generalization guarantees and training dynamics

A constructive finite-sample theory is available for nonlinear feature-building transformers. For functions supported on

$f(x)$5

with well-conditioned feature covariance

$f(x)$6

the approximation and statistical errors are separated. The approximation term satisfies

$f(x)$7

while the statistical term is controlled by a Rademacher-complexity bound, leading to the summary rate

$f(x)$8

up to degree- and norm-dependent factors. The same framework extends to linear splines, again with a $f(x)$9 approximation term and a $\mathbf{s}^f=(x_1,y_1,\ldots,x_n,y_n,x).$0 statistical term (Hsu et al., 6 May 2026).

The training-dynamics theory addresses a different question: how a transformer acquires in-context nonlinear regression through gradient descent. "Provable In-Context Learning of Nonlinear Regression with Transformers" studies a one-layer softmax attention model

$\mathbf{s}^f=(x_1,y_1,\ldots,x_n,y_n,x).$1

trained from zero initialization on a class of non-degenerate $\mathbf{s}^f=(x_1,y_1,\ldots,x_n,y_n,x).$2-Lipschitz nonlinear functions. The query prediction can be written

$\mathbf{s}^f=(x_1,y_1,\ldots,x_n,y_n,x).$3

so correctness is equivalent to $\mathbf{s}^f=(x_1,y_1,\ldots,x_n,y_n,x).$4 for the relevant feature $\mathbf{s}^f=(x_1,y_1,\ldots,x_n,y_n,x).$5 (Li et al., 28 Jul 2025).

The paper gives two-stage convergence analyses in two regimes separated by the threshold $\mathbf{s}^f=(x_1,y_1,\ldots,x_n,y_n,x).$6. In the flat regime,

$\mathbf{s}^f=(x_1,y_1,\ldots,x_n,y_n,x).$7

and in the sharp regime,

$\mathbf{s}^f=(x_1,y_1,\ldots,x_n,y_n,x).$8

At convergence,

$\mathbf{s}^f=(x_1,y_1,\ldots,x_n,y_n,x).$9

The main message is that the Lipschitz constant $$\Gamma=\{(\mathbf{s}^{f_\ell},y^\ell)\}_{\ell=1}^L, \qquad G_{\hat\theta}^\Gamma=\arg\min_{G_\theta\in\mathcal{G}} \mathcal{R}^\Gamma(G_\theta),$$0 governs the convergence dynamics by controlling both feature separation in function values and the scale of attention-weight gradients (Li et al., 28 Jul 2025).

These generalization and dynamics results are complementary. The former gives explicit end-to-end bounds as a function of context length and training set size for constructed feature maps; the latter explains how training drives query attention toward relevant prompt tokens for unseen nonlinear tasks. Together they shift the theory of in-context regression beyond the linear least-squares setting.

6. Related notions of context in nonlinear regression

The transformer literature uses "context" to mean an explicit prompt of labeled examples, but related nonlinear regression work employs broader notions of context. One example is NOVAS, introduced in "An Algorithm for Nonlinear, Nonparametric Model Choice and Prediction" (Ferraty et al., 2014). NOVAS is a nonlinear, nonparametric model-selection procedure that ranks single variables by a leave-one-out score

$$\Gamma=\{(\mathbf{s}^{f_\ell},y^\ell)\}_{\ell=1}^L, \qquad G_{\hat\theta}^\Gamma=\arg\min_{G_\theta\in\mathcal{G}} \mathcal{R}^\Gamma(G_\theta),$$1

then performs a sequential combinatorial search using the $$\Gamma=\{(\mathbf{s}^{f_\ell},y^\ell)\}_{\ell=1}^L, \qquad G_{\hat\theta}^\Gamma=\arg\min_{G_\theta\in\mathcal{G}} \mathcal{R}^\Gamma(G_\theta),$$2-screening rule. Its purpose is to detect variables that are weak by themselves but useful in combination, while avoiding redundant variables once better combinations are formed. This is not in-context learning in the transformer sense, but it addresses nonlinear regression under combinatorial contextual interactions among predictors (Ferraty et al., 2014).

Another broader use of context appears in "Nonlinear Regression without i.i.d. Assumption" (Xu et al., 2018). There, the data are partitioned into groups interpreted as local contexts or sampling regimes, and the objective becomes

$$\Gamma=\{(\mathbf{s}^{f_\ell},y^\ell)\}_{\ell=1}^L, \qquad G_{\hat\theta}^\Gamma=\arg\min_{G_\theta\in\mathcal{G}} \mathcal{R}^\Gamma(G_\theta),$$3

The method replaces average-loss fitting by a groupwise max-of-means mini-max formulation and solves the resulting problem by a linearization, quadratic regularization, and dual QP. This suggests a robust, context-conditioned principle for nonlinear regression under distributional heterogeneity rather than prompt-based adaptation (Xu et al., 2018).

A further related direction is nonlinear function-on-function regression by RKHS. In that setting, the predictor $$\Gamma=\{(\mathbf{s}^{f_\ell},y^\ell)\}_{\ell=1}^L, \qquad G_{\hat\theta}^\Gamma=\arg\min_{G_\theta\in\mathcal{G}} \mathcal{R}^\Gamma(G_\theta),$$4 lies in a first-layer Hilbert space $$\Gamma=\{(\mathbf{s}^{f_\ell},y^\ell)\}_{\ell=1}^L, \qquad G_{\hat\theta}^\Gamma=\arg\min_{G_\theta\in\mathcal{G}} \mathcal{R}^\Gamma(G_\theta),$$5, the response $$\Gamma=\{(\mathbf{s}^{f_\ell},y^\ell)\}_{\ell=1}^L, \qquad G_{\hat\theta}^\Gamma=\arg\min_{G_\theta\in\mathcal{G}} \mathcal{R}^\Gamma(G_\theta),$$6 lies in $$\Gamma=\{(\mathbf{s}^{f_\ell},y^\ell)\}_{\ell=1}^L, \qquad G_{\hat\theta}^\Gamma=\arg\min_{G_\theta\in\mathcal{G}} \mathcal{R}^\Gamma(G_\theta),$$7, and a second-layer RKHS $$\Gamma=\{(\mathbf{s}^{f_\ell},y^\ell)\}_{\ell=1}^L, \qquad G_{\hat\theta}^\Gamma=\arg\min_{G_\theta\in\mathcal{G}} \mathcal{R}^\Gamma(G_\theta),$$8 is built from a kernel on $$\Gamma=\{(\mathbf{s}^{f_\ell},y^\ell)\}_{\ell=1}^L, \qquad G_{\hat\theta}^\Gamma=\arg\min_{G_\theta\in\mathcal{G}} \mathcal{R}^\Gamma(G_\theta),$$9, such as

$$\mathcal{R}^\Gamma(G_\theta)=\frac{1}{L}\sum_{\ell=1}^L |G_\theta(\mathbf{s}^{f_\ell})-y^\ell|^2, \qquad \mathcal{R}(G_\theta)=\mathbb{E}_{\mathbf{s}^f}|G_\theta(\mathbf{s}^f)-f(x)|^2.$$0

The conditional mean takes the form

$$\mathcal{R}^\Gamma(G_\theta)=\frac{1}{L}\sum_{\ell=1}^L |G_\theta(\mathbf{s}^{f_\ell})-y^\ell|^2, \qquad \mathcal{R}(G_\theta)=\mathbb{E}_{\mathbf{s}^f}|G_\theta(\mathbf{s}^f)-f(x)|^2.$$1

and the method accommodates irregular and sparse observations through RKHS reconstruction and regularization (Sang et al., 2022). A plausible implication is that transformer-based in-context nonlinear regression and RKHS-based nonlinear regression share a common structural theme: nonlinear prediction is reduced to linear estimation in a lifted feature space, with the main differences lying in how the feature map is constructed and how inference is performed.

Across these related areas, the unifying issue is not merely nonlinearity but the dependence of prediction on local, retrieved, grouped, or feature-lifted information. In the transformer setting, that dependence is enacted by the prompt; in broader nonlinear regression, it appears as subset selection, groupwise robustness, or kernelized functional conditioning.