---
title: In-Context Nonlinear Regression
url: https://www.emergentmind.com/topics/in-context-nonlinear-regression
type: topic
---

# In-Context Nonlinear Regression

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 [2501.18187] [2605.05176] [2605.08475] [2405.18202] [2507.20443].

## 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 \(n\) and training set size \(L\) the two primary sample-complexity variables [2605.05176].

A second formulation, used for in-context Gaussian-kernel ridge regression, treats the prompt as \(N\) training tokens and one query token:
\[
P=\big((x_1,y_1),\ldots,(x_N,y_N),x_{N+1}\big),
\]
with kernel
\[
k(x,x')=\exp\!\left(-\frac{\|x-x'\|^2}{2\sigma^2}\right),
\]
kernel matrix \(K\), and predictor
\[
\widehat y_{N+1}=k(x_{N+1},X)^\top (K+\lambda I)^{-1}y.
\]
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 \((K+\lambda I)\alpha=y\) [2605.08475].

In imbalanced regression, the prompt is not merely a few-shot conditioning device but a distribution-selection mechanism. The setup distinguishes a source set \(D_s=\{(\mathbf{x}_i,y_i)\}_{i=1}^{n_s}\) and a target set \(D_t=\{(\mathbf{x}_i,y_i)\}_{i=1}^{n_t}\), with source and target distributions \(p_s(\mathbf{x},y)\) and \(p_t(\mathbf{x},y)\). A pre-trained transformer \(f_\theta\) predicts
\[
\hat{y}_{query}=f_\theta(\mathbf{x}_{query}\mid D_k),
\]
where \(D_k\) is a retrieved local context. The associated expected prediction error is decomposed as
\[
\operatorname{EPE}_{f_{\theta}}(\mathbf{x})
=
\operatorname{Var}_{D_n}[f_\theta(\mathbf{x}\mid D_n)]
+
\operatorname{Bias}_{D_n}^2[f_\theta(\mathbf{x}\mid D_n)]
+\sigma^2,
\]
so both context size and context composition affect the bias-variance tradeoff [2405.18202].

## 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,
\[
Z^{(\ell+1)}=\mathsf{attn}(Z^{(\ell)}),
\qquad
\mathsf{attn}(Z)=Z+\frac{1}{n}\mathbf{P}ZM\,(Z^\top \mathbf{Q} Z),
\]
stacking more LSA layers does not break the linearity barrier. Proposition 1 / Proposition \(\ref{thm:lsa-lower-bound}\) states
\[
\mathbb{E}_x\!\left[(\hat y_{\mathsf{TF}_{\mathsf{lin}}}+y_{\textsf{query}})^2\right]
\ge
\min_{\alpha,\beta}\;
\mathbb{E}_x\!\left[\left(x_{\textsf{query}}^\top \beta + \alpha + y_{\textsf{query}}\right)^2\right].
\]
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 [2501.18187].

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
\[
\mathsf{bilin}(Z)=Z+\left(
\begin{bmatrix}\mathbf{W}_1 & 0 \\ 0 & 0\end{bmatrix} Z
\right)\odot
\left(
\begin{bmatrix}\mathbf{W}_2 & 0 \\ 0 & 0\end{bmatrix} Z
\right)
\]
with linear self-attention:
\[
Z^{(2\ell+1)}=\mathsf{bilin}(Z^{(2\ell)}),
\qquad
Z^{(2\ell+2)}=\mathsf{attn}(Z^{(2\ell+1)}).
\]
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 [2501.18187].

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
\[
\mathrm{A}_{Q,K,V}(H)=VH\,\sigma\!\left((KH)^TQH\right),
\]
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 [2605.05176].

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
\[
f(x;w)=
\begin{bmatrix}1 & x^\top\end{bmatrix}
W
\begin{bmatrix}1\\x\end{bmatrix},
\qquad
L(w)=\frac{1}{2n}\sum_{i=1}^n \left(f(x_i;w)+y_i\right)^2,
\]
the update is
\[
w^{+}=\Gamma \nabla_w L(w)\big|_{w=0},
\]
with preconditioner \(\Gamma\) determined by attention parameters. The bilinear layer maps each input to a quadratic feature vector
\[
\bar x_i=(1,\ x_i[1],\dots,x_i[d],\ x_i[1]^2,\ x_i[1]x_i[2],\dots,\ x_i[d]^2),
\]
and LSA then produces
\[
\hat y = \frac{1}{n}\sum_{i=1}^n \bar x_{\textsf{query}}^\top \Gamma \bar x_i\, y_i.
\]
This is exactly the form of a kernel regression / gradient-descent estimator in the quadratic feature space [2501.18187].

When the embedding dimension is insufficient to represent all quadratic monomials at once, depth replaces width. The block-coordinate update is
\[
w^{(\ell)}_j=
\begin{cases}
w^{(\ell-1)}_j-\eta^{(\ell)} \partial_{w_j}L(w)\big|_{w=w^{(\ell-1)}}, & j\in b_\ell,\\[4pt]
w^{(\ell-1)}_j, & j\notin b_\ell,
\end{cases}
\]
and Theorem \(\ref{thm:block-construction}\) states that there exists a deep bilinear Transformer with \(\bar d=2d+1\) such that
\[
\hat y^{(\ell)}=
\begin{bmatrix}1 & x_{\textsf{query}}^\top\end{bmatrix}
W^{(\ell)}
\begin{bmatrix}1\\x_{\textsf{query}}\end{bmatrix},
\]
where \(W^{(\ell)}\) is exactly the iterate produced by block-coordinate descent. The same paper also gives a width upper bound and lower bound: if \(\bar d\ge \binom{d+2}{2}\), then one bilinear block can achieve \(\mathcal L(\mathsf{BTFB})=O(1/n)\); if \(\bar d<\binom{d+1}{2}\), then any bilinear block has loss at least
\[
\mathcal{L}(\mathsf{BTFB}) \ge \binom{d+1}{2}-\bar d.
\]
These results formalize the claim that depth can replace width in nonlinear in-context regression [2501.18187].

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
\[
A\alpha=y,
\qquad
A:=K+\lambda I,
\]
and shows that a standard transformer can implement preconditioned Richardson iteration
\[
\alpha^{(t+1)}=\alpha^{(t)}+P\big(y-A\alpha^{(t)}\big),
\qquad
r^{(t+1)}=(I-AP)r^{(t)}.
\]
Under bounded-data assumptions, the paper constructs a single-head transformer with \(O(\log(1/\varepsilon))\) blocks and MLP width \(O(\sqrt{N/\varepsilon})\) that achieves \(\varepsilon\)-accurate prediction for prompts of length \(N\). 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 [2605.08475].

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
\[
Z=
\begin{bmatrix}
1 & \cdots & 1 & 1\\
x_1 & \cdots & x_n & x_{\textsf{query}}\\
\mathbf{0} & \cdots & \mathbf{0} & \mathbf{0}\\
y_1 & \cdots & y_n & 0
\end{bmatrix},
\]
with inputs sampled i.i.d. from a Gaussian,
\[
x_1,\dots,x_{n+1}\sim \mathcal{N}(0,\Sigma),
\]
or in some results \(\mathcal{N}(0,I)\), and with target functions given by random quadratics
\[
f(x)=w_0+\sum_{i=1}^d w_i x_i+\sum_{1\le i\le j\le d}w_{ij}x_ix_j,
\qquad
w_i,w_{ij}\overset{i.i.d.}{\sim}\mathcal{N}(0,1).
\]
The same work also studies a non-Gaussian induced distribution after the feed-forward layer, for example
\[
x_i[j]=v_i[1]v_i[j], \qquad v_i[k]\stackrel{i.i.d.}{\sim}\mathcal N(0,1),
\]
and shows that optimization becomes harder for nonlinear ICL under such feature distributions [2501.18187].

In imbalanced regression, the decisive issue is locality. The model is assumed \(c\)-Lipschitz in the context examples,
\[
\left| f_{\theta}(\mathbf{x} \mid D_n) - f_{\theta}(\mathbf{x }\mid D_n') \right|
\leq
\sum_{i=1}^n c_i \mathbf{1}_{\{\mathbf{x}_i \neq \mathbf{x}_i'\}},
\qquad
c_i=\delta i^{-\alpha},\ \alpha>0.5,
\]
and the paper compares transformer prediction against the local averaging estimator
\[
\tilde{f}(\mathbf{x}\mid D_k)=\frac{1}{k}\sum_{i=1}^k y_i.
\]
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 \(p(y\mid \mathbf{x})\). 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 [2405.18202].

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 [2405.18202].

A related distributional assumption appears in the training-dynamics analysis of one-layer softmax attention. There the prompt contains \(N\) input-response pairs and one query,
\[
P=
\begin{pmatrix}
x_1 & x_2 & \cdots & x_N & x_{\mathrm{query}}\\
y_1 & y_2 & \cdots & y_N & 0
\end{pmatrix},
\qquad
y_i=f(x_i),
\]
and a concentration event
\[
\mathcal{E}^*:=\Big\{P_{1:N}:|\mathcal V_k|\in[(p_k-\delta)N,(p_k+\delta)N]\ \text{for }k\in[K]\Big\}
\]
ensures each feature appears with approximately the expected frequency. With high probability,
\[
\mathbb P(P_{1:N}\in\mathcal E^*)\ge 1-3\exp\!\left(-\frac{\delta^2N}{25}\right),
\]
which stabilizes the analysis of attention concentration [2507.20443].

## 5. Generalization guarantees and training dynamics

A constructive finite-sample theory is available for nonlinear feature-building transformers. For functions supported on
\[
F(\delta,\alpha,d,R_F)=
\left\{
f\;\middle|\;
\min_{p\in \Pi_d^\alpha}\|f-p\|_{L^\infty(X)}\le \delta,\;
\|f\|_{L^\infty(X)}\le R_F
\right\},
\]
with well-conditioned feature covariance
\[
\Sigma=\mathbb{E}_{x\sim\mathcal{D}_X}[\mathbf{v}(x)\mathbf{v}(x)^T],
\qquad
\|\Sigma^{-1}\|\le \tau,
\]
the approximation and statistical errors are separated. The approximation term satisfies
\[
\mathcal{R}(G_{\tilde\theta})
\le
C_1\left(
\frac{\alpha^2\tau^2 R_{\mathbf v}^6\log d}{n}
+\delta^2R_{\mathbf v}^4\tau^2
+\delta^2
\right),
\]
while the statistical term is controlled by a Rademacher-complexity bound, leading to the summary rate
\[
\mathbb{E}_\Gamma[\mathcal{R}(G_{\hat\theta}^\Gamma)]
=
O\!\left(
\frac{1}{n}+\frac{\sqrt{n\log(n^2L)}}{\sqrt{L}}
\right),
\]
up to degree- and norm-dependent factors. The same framework extends to linear splines, again with a \(1/n\) approximation term and a \(1/\sqrt{L}\) statistical term [2605.05176].

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
\[
F(P;Q)= y\cdot \mathrm{softmax}\big(X^\top Q\bar X\big),
\qquad
\hat y_{\mathrm{query}}=F(P;Q)_{N+1},
\]
trained from zero initialization on a class of non-degenerate \(L\)-Lipschitz nonlinear functions. The query prediction can be written
\[
\hat y_{\mathrm{query}}
=
\sum_{i=1}^N \mathrm{attn}_i^{(t)}y_i
=
\sum_{k=1}^K \mathrm{Attn}_k^{(t)}f(v_k),
\]
so correctness is equivalent to \(\mathrm{Attn}_k^{(t)}\approx 1\) for the relevant feature \(v_k\) [2507.20443].

The paper gives two-stage convergence analyses in two regimes separated by the threshold \(\Theta(1/(\Delta\delta))\). In the flat regime,
\[
T_f^*=\Theta\!\left(\frac{K\log(K\epsilon^{-1})}{\eta\delta^2L^2\Delta^2}\right),
\]
and in the sharp regime,
\[
T_s^*=\Theta\!\left(\frac{K\log(KL\Delta\epsilon^{-1})}{\eta\epsilon\delta^2L^2\Delta^2}\right).
\]
At convergence,
\[
1-\mathrm{Attn}_k^{(t)}=\mathcal O(\epsilon),
\qquad
(\hat y_{\mathrm{query}}-f(x_{\mathrm{query}}))^2=\mathcal O(\epsilon^2).
\]
The main message is that the Lipschitz constant \(L\) governs the convergence dynamics by controlling both feature separation in function values and the scale of attention-weight gradients [2507.20443].

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" [1401.8097]. NOVAS is a nonlinear, nonparametric model-selection procedure that ranks single variables by a leave-one-out score
\[
S(J)=\sum_{i\in}\{Y_i-\hat{}_{J}^{-i}(X_i^{J})\}^2\, w_{|J|}(X_i^{J}),
\]
then performs a sequential combinatorial search using the \(\sqrt{q}\)-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 [1401.8097].

Another broader use of context appears in "Nonlinear Regression without i.i.d. Assumption" [1811.09623]. There, the data are partitioned into groups interpreted as local contexts or sampling regimes, and the objective becomes
\[
\min_{\theta}\max_{1\leq j\leq N}\frac{1}{n_j}\sum_{l=1}^{n_j}(g^\theta(x_{jl})-y_{jl})^2.
\]
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 [1811.09623].

A further related direction is nonlinear function-on-function regression by RKHS. In that setting, the predictor \(X\) lies in a first-layer Hilbert space \(\mathcal H_X\), the response \(Y\) lies in \(\mathcal H_Y\), and a second-layer RKHS \(\mathcal M_X\) is built from a kernel on \(\mathcal H_X\), such as
\[
k(f,g)=\exp\{-\gamma \|f-g\|_{\mathcal H_X}^2\}.
\]
The conditional mean takes the form
\[
\mathbb E(Y\mid X=x)
=
\Sigma_{YX}\Sigma_{XX}^{\dagger}\big[k(\cdot,x)-\mu_X\big]+\mu_Y,
\]
and the method accommodates irregular and sparse observations through RKHS reconstruction and regularization [2207.08211]. 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.

Source: https://www.emergentmind.com/topics/in-context-nonlinear-regression