Transformer-Based ICL Optimizer

• Transformer-Based ICL Optimizer is a transformer that simulates iterative optimization steps through its forward pass, using the prompt to update latent parameters without altering pretrained weights.
• It integrates advanced techniques like gradient descent, Newton iterations, and specialized layers (e.g., EWML) to perform inner optimization by encoding both demonstration data and latent state.
• Constructive results and synthetic experiments validate the approach, though challenges remain regarding prompt distribution shifts, architecture constraints, and real-world robustness.

Searching arXiv for papers on transformer-based in-context optimization and related ICL mechanisms. Searching arXiv for "transformer in-context optimizer gradient descent Newton softmax feed-forward nonlinear ICL". A Transformer-Based ICL Optimizer is a transformer whose forward pass implements an optimization procedure over a latent task-specific model or decision rule, using the prompt as data and internal activations as optimization state, without changing the transformer’s own pretrained weights. In the strongest formalizations, the context does not merely specify an input-output mapping; it induces an inner objective, and transformer depth realizes successive optimizer updates on latent parameters stored in the residual stream. This perspective has been made explicit for projected gradient descent on deep ReLU networks, for Newton-like solvers in linear regression, for polynomial-kernel regression induced by feed-forward feature lifts, and for sparse recovery via LISTA-type iterations (Wu et al., 2024).

1. Formal notion and representational structure

A canonical formalization represents a dataset $\mathcal D_n=\{(x_i,y_i)\}_{i=1}^n$ together with a query $x_{n+1}$ as a sequence matrix

$$H \coloneqq \begin{bmatrix} x_1 & x_2 & \cdots & x_n & x_{n+1}\ y_1 & y_2 & \cdots & y_n & 0\ q_1 & q_2 & \cdots & q_n & q_{n+1} \end{bmatrix}\in\mathbb R^{D\times (n+1)},$$

where the final token is the query token with missing label. In the in-context gradient-descent formulation, the inner model is $f(w,\cdot)$, the empirical loss is

$$\mathcal L_n(w)\coloneqq \frac{1}{2n}\sum_{i=1}^n \ell(f(w,x_i),y_i),$$

and the target optimization dynamics are

$w_{t+1}=w_t-\eta \nabla \mathcal L_n(w_t),$

or, with projection,

$$w_{t+1}={\rm Proj}_{\mathcal W}\bigl(w_t-\eta \nabla \mathcal L_n(w_t)\bigr).$$

The transformer is then said to be an in-context optimizer if, after each block of layers, every token contains the same updated parameter vector,

$h_i^{(Tl)}=[x_i;y_i;\bar w^{(l)};\mathbf 0;1;t_i],$

with

$$\bar w^{(l)}={\rm Proj}_{\mathcal W}\!\left(\bar w^{(l-1)}-\eta\left(\nabla \mathcal L_n(\bar w^{(l-1)})+\epsilon^{(l-1)}\right)\right), \qquad \bar w^{(0)}=\mathbf 0.$$

Under this view, the prompt encodes not only demonstrations but also a latent optimization state that evolves through depth (Wu et al., 2024).

A related but training-dynamics-oriented formulation appears in autoregressive next-token prediction. For a one-layer linear causal self-attention model trained by gradient flow on sequences generated by $x_{t+1}=Wx_t$, the learned forward pass can converge to an estimator $x_{n+1}$0 obtained by one gradient descent step on the in-context OLS objective

$x_{n+1}$1

starting from $x_{n+1}$2. In that setting, the mesa-optimization hypothesis is verified in a precise toy regime: outer autoregressive training yields an inner optimization procedure executed during inference (Zheng et al., 2024).

2. Constructive first-order optimizer circuits

The most explicit constructive result shows that a transformer can simulate $x_{n+1}$3 projected gradient-descent steps for an $x_{n+1}$4-layer ReLU network via a $x_{n+1}$5-layer construction. The inner network has parameter vector $x_{n+1}$6 with $x_{n+1}$7, and the transformer block structure is

$x_{n+1}$8

One optimization step therefore costs $x_{n+1}$9 layers, organized so that the first $$H \coloneqq \begin{bmatrix} x_1 & x_2 & \cdots & x_n & x_{n+1}\ y_1 & y_2 & \cdots & y_n & 0\ q_1 & q_2 & \cdots & q_n & q_{n+1} \end{bmatrix}\in\mathbb R^{D\times (n+1)},$$0 attention layers compute forward activations $$H \coloneqq \begin{bmatrix} x_1 & x_2 & \cdots & x_n & x_{n+1}\ y_1 & y_2 & \cdots & y_n & 0\ q_1 & q_2 & \cdots & q_n & q_{n+1} \end{bmatrix}\in\mathbb R^{D\times (n+1)},$$1, one attention layer computes activation derivatives, one MLP computes the loss derivative $$H \coloneqq \begin{bmatrix} x_1 & x_2 & \cdots & x_n & x_{n+1}\ y_1 & y_2 & \cdots & y_n & 0\ q_1 & q_2 & \cdots & q_n & q_{n+1} \end{bmatrix}\in\mathbb R^{D\times (n+1)},$$2, $$H \coloneqq \begin{bmatrix} x_1 & x_2 & \cdots & x_n & x_{n+1}\ y_1 & y_2 & \cdots & y_n & 0\ q_1 & q_2 & \cdots & q_n & q_{n+1} \end{bmatrix}\in\mathbb R^{D\times (n+1)},$$3 element-wise multiplication layers compute backprop states $$H \coloneqq \begin{bmatrix} x_1 & x_2 & \cdots & x_n & x_{n+1}\ y_1 & y_2 & \cdots & y_n & 0\ q_1 & q_2 & \cdots & q_n & q_{n+1} \end{bmatrix}\in\mathbb R^{D\times (n+1)},$$4, one attention layer aggregates the context into $$H \coloneqq \begin{bmatrix} x_1 & x_2 & \cdots & x_n & x_{n+1}\ y_1 & y_2 & \cdots & y_n & 0\ q_1 & q_2 & \cdots & q_n & q_{n+1} \end{bmatrix}\in\mathbb R^{D\times (n+1)},$$5, and one final MLP applies the projection map (Wu et al., 2024).

The theorem is constructive and quantitative. Under boundedness of data and parameters, Lipschitzness and $$H \coloneqq \begin{bmatrix} x_1 & x_2 & \cdots & x_n & x_{n+1}\ y_1 & y_2 & \cdots & y_n & 0\ q_1 & q_2 & \cdots & q_n & q_{n+1} \end{bmatrix}\in\mathbb R^{D\times (n+1)},$$6-smoothness of $$H \coloneqq \begin{bmatrix} x_1 & x_2 & \cdots & x_n & x_{n+1}\ y_1 & y_2 & \cdots & y_n & 0\ q_1 & q_2 & \cdots & q_n & q_{n+1} \end{bmatrix}\in\mathbb R^{D\times (n+1)},$$7, and MLP-realizability of $$H \coloneqq \begin{bmatrix} x_1 & x_2 & \cdots & x_n & x_{n+1}\ y_1 & y_2 & \cdots & y_n & 0\ q_1 & q_2 & \cdots & q_n & q_{n+1} \end{bmatrix}\in\mathbb R^{D\times (n+1)},$$8, there exists a transformer with $$H \coloneqq \begin{bmatrix} x_1 & x_2 & \cdots & x_n & x_{n+1}\ y_1 & y_2 & \cdots & y_n & 0\ q_1 & q_2 & \cdots & q_n & q_{n+1} \end{bmatrix}\in\mathbb R^{D\times (n+1)},$$9 hidden layers such that at layer $f(w,\cdot)$0,

$f(w,\cdot)$1

with per-step perturbation bounded by $f(w,\cdot)$2. The architecture size satisfies

$f(w,\cdot)$3

so the construction is explicit but wide, and the number of heads scales as $f(w,\cdot)$4 (Wu et al., 2024).

An important auxiliary point is architectural. The special layer

$f(w,\cdot)$5

is introduced because standard attention and MLPs cannot directly multiply two coordinates from the same token. This layer realizes token-wise self-products and enables Hadamard products required by the backprop recursion. The result is therefore a proof of possibility rather than a claim that standard pretrained transformers literally instantiate the same circuit (Wu et al., 2024).

A complementary linear result sharpens the role of the feed-forward sublayer. For a Linear Transformer Block, the optimal in-class ICL solution for linear regression with nonzero Gaussian prior mean is effectively one-step gradient descent with a learnable initialization,

$f(w,\cdot)$6

Every such $f(w,\cdot)$7 estimator can be implemented by a Linear Transformer Block, and every optimal Linear Transformer Block is effectively a $f(w,\cdot)$8 estimator. In that setting, the linear MLP is what permits encoding the nonzero prior mean; linear attention alone incurs an irreducible additive approximation error (Zhang et al., 2024).

3. Higher-order, nonlinear, and structured in-context optimization

The first-order account is not exhaustive. In in-context linear regression, trained transformers have been shown to align more closely with Iterative Newton’s Method than with gradient descent. The empirical layerwise progression matches Newton iterations linearly, with each middle layer roughly computing 3 iterations, and the paper proves that transformers can implement $f(w,\cdot)$9 Newton steps with $$\mathcal L_n(w)\coloneqq \frac{1}{2n}\sum_{i=1}^n \ell(f(w,x_i),y_i),$$0 layers. In this regime, the transformer is better described as a curvature-aware in-context optimizer that approximates the pseudoinverse of the Gram matrix $$\mathcal L_n(w)\coloneqq \frac{1}{2n}\sum_{i=1}^n \ell(f(w,x_i),y_i),$$1 via the Newton–Schulz iteration

$$\mathcal L_n(w)\coloneqq \frac{1}{2n}\sum_{i=1}^n \ell(f(w,x_i),y_i),$$2

rather than as a shallow executor of first-order GD (Fu et al., 2023).

For nonlinear ICL, linear self-attention alone is provably insufficient. Deep linear self-attention cannot beat the best linear predictor in the original input coordinates, even with arbitrarily many layers. To obtain nonlinear in-context optimization, one must combine linear self-attention with a GLU-like bilinear feed-forward layer

$$\mathcal L_n(w)\coloneqq \frac{1}{2n}\sum_{i=1}^n \ell(f(w,x_i),y_i),$$3

In the quadratic case, this layer constructs explicit degree-$$\mathcal L_n(w)\coloneqq \frac{1}{2n}\sum_{i=1}^n \ell(f(w,x_i),y_i),$$4 monomials $$\mathcal L_n(w)\coloneqq \frac{1}{2n}\sum_{i=1}^n \ell(f(w,x_i),y_i),$$5, after which attention implements one step of preconditioned gradient descent for polynomial-kernel regression,

$$\mathcal L_n(w)\coloneqq \frac{1}{2n}\sum_{i=1}^n \ell(f(w,x_i),y_i),$$6

With limited width, stacked GLU-LSA blocks implement block-coordinate descent in polynomial kernel space (Sun et al., 30 Jan 2025).

Sparse recovery yields another structured optimizer interpretation. For the LASSO objective

$$\mathcal L_n(w)\coloneqq \frac{1}{2n}\sum_{i=1}^n \ell(f(w,x_i),y_i),$$7

a $$\mathcal L_n(w)\coloneqq \frac{1}{2n}\sum_{i=1}^n \ell(f(w,x_i),y_i),$$8-layer transformer can implement a LISTA-type learning-to-optimize algorithm,

$$\mathcal L_n(w)\coloneqq \frac{1}{2n}\sum_{i=1}^n \ell(f(w,x_i),y_i),$$9

with a provable linear convergence rate in $w_{t+1}=w_t-\eta \nabla \mathcal L_n(w_t),$0: $w_{t+1}=w_t-\eta \nabla \mathcal L_n(w_t),$1 Here attention computes prompt-dependent linear update terms, the MLP implements soft-thresholding, and depth functions as optimization time (Liu et al., 2024).

4. Training dynamics, learned priors, and meta-level selection

A distinct question is whether outer training actually converges to optimizer-like behavior. In Gaussian-mixture in-context classification, a single-layer transformer trained by gradient descent converges linearly to a unique global minimizer. In the binary case, the minimizer satisfies

$w_{t+1}=w_t-\eta \nabla \mathcal L_n(w_t),$2

and in the multiclass case

$w_{t+1}=w_t-\eta \nabla \mathcal L_n(w_t),$3

The resulting inference error decomposes into training-prompt bias $w_{t+1}=w_t-\eta \nabla \mathcal L_n(w_t),$4 and test-prompt estimation error $w_{t+1}=w_t-\eta \nabla \mathcal L_n(w_t),$5, so sufficiently large training and testing prompts drive the predictor toward the ground-truth label distribution (Shen et al., 2024).

In inverse linear regression, the transformer is used not only to predict outputs but to infer the latent task parameter $w_{t+1}=w_t-\eta \nabla \mathcal L_n(w_t),$6 itself from context. In the rank-deficient regime $w_{t+1}=w_t-\eta \nabla \mathcal L_n(w_t),$7 with low-rank task covariance $w_{t+1}=w_t-\eta \nabla \mathcal L_n(w_t),$8, the learned map $w_{t+1}=w_t-\eta \nabla \mathcal L_n(w_t),$9 behaves as an implicit regularized inverse solver. The paper’s central empirical interpretation is that the transformer learns both a cross-task prior and an effective regularization strategy, and its error scales linearly with the noise level, the ratio of task dimension to context length, and the condition number of the input data (Lu et al., 17 May 2025).

A related positive result for nonlinear low-dimensional task families shows how pretraining can compress the effective inner optimization problem. For single-index targets $$w_{t+1}={\rm Proj}_{\mathcal W}\bigl(w_t-\eta \nabla \mathcal L_n(w_t)\bigr).$$0 with $$w_{t+1}={\rm Proj}_{\mathcal W}\bigl(w_t-\eta \nabla \mathcal L_n(w_t)\bigr).$$1 drawn from an $$w_{t+1}={\rm Proj}_{\mathcal W}\bigl(w_t-\eta \nabla \mathcal L_n(w_t)\bigr).$$2-dimensional subspace, a nonlinear transformer pretrained across tasks learns a prompt-time estimator whose in-context sample complexity depends on $$w_{t+1}={\rm Proj}_{\mathcal W}\bigl(w_t-\eta \nabla \mathcal L_n(w_t)\bigr).$$3, not the ambient dimension $$w_{t+1}={\rm Proj}_{\mathcal W}\bigl(w_t-\eta \nabla \mathcal L_n(w_t)\bigr).$$4. The attention layer estimates coefficients in a learned low-dimensional basis, while the nonlinear embedding layer learns the shared subspace structure during pretraining (Oko et al., 2024).

At a higher meta-level, transformers can also implement in-context algorithm selection. One constructive framework uses pre-ICL testing to route between different base learners; another uses post-ICL validation to evaluate multiple candidate predictors inside the prompt and select a near-best one. This supports the interpretation of a transformer not merely as a fixed inner optimizer but as an amortized selector over multiple in-context optimization procedures (Bai et al., 2023).

5. Empirical regimes and application domains

The explicit deep-network optimizer construction has been validated on synthetic tasks generated by random 3-layer, 4-layer, and 6-layer neural networks $$w_{t+1}={\rm Proj}_{\mathcal W}\bigl(w_t-\eta \nabla \mathcal L_n(w_t)\bigr).$$5. ReLU-transformers and Softmax-transformers, both using a GPT-2-style backbone with 12 transformer blocks, 8 heads, hidden size 256, and MLP size 256, were pretrained for 500k steps on prompts

$$w_{t+1}={\rm Proj}_{\mathcal W}\bigl(w_t-\eta \nabla \mathcal L_n(w_t)\bigr).$$6

and evaluated up to 75 context examples plus a test query. The reported result is that in-context performance closely matches direct training of the corresponding target networks, while deeper transformers improve ICL performance, supporting the depth-as-optimization-steps interpretation (Wu et al., 2024).

Wireless estimation provides a different empirical manifestation. In in-context estimation for wireless communications, a single-layer softmax-attention transformer is provably capable of computing the Bayes-optimal estimator in the large-prompt limit for a structured subclass, and the corresponding optimal attention matrix is $$w_{t+1}={\rm Proj}_{\mathcal W}\bigl(w_t-\eta \nabla \mathcal L_n(w_t)\bigr).$$7. In broader finite-context experiments, multi-layer GPT-2-style transformers efficiently solve in-context estimation problems and, with a few context examples, achieve the same performance as an estimator with perfect knowledge of the latent context (Kunde et al., 2023).

A practically oriented tabular counterpart is TabICLv2. It remains an amortized learner over train-set plus test-query inputs, but its authors explicitly distinguish it from a learned test-time optimizer in the meta-learning sense: at inference time it does not perform gradient-based finetuning or test-time adaptation. Its contribution lies instead in the optimization recipe for pretraining and in attention scaling for long contexts, including the three-stage curriculum, Muon, and QASSMax. This distinction is important because it separates inner-loop optimization in activations from system-level optimization of an ICL model (Qu et al., 11 Feb 2026).

6. Limitations, counterevidence, and open questions

The strongest positive results are highly scoped. The deep ReLU construction is highly structured, uses the nonstandard EWML layer, requires width on the order of all target-network weights plus auxiliary state, assumes bounded domains and smooth approximability, and exhibits worst-case exponential error accumulation in the number of inner GD steps. Its softmax extension is broader but less constructive, relying on a universality theorem rather than an explicit simulation of forward and backward mechanics (Wu et al., 2024).

Several papers also challenge strong algorithmic interpretations. In synthetic linear regression, transformers trained for ICL can fail under prompt-distribution shifts in ways inconsistent with exact OLS or gradient descent. One study reports that even in-distribution, once the prompt enters the overdetermined regime, OLS achieves error several orders of magnitude lower than the transformer, and out-of-distribution performance is closely tied to a spectral signature in the residual stream rather than to basis-invariant least-squares computation (Hill et al., 13 Jul 2025). Another study on univariate linear functions similarly finds strong on-distribution ICL but sharp failures under shifts in function scale, boundary-value saturation, and strong dependence on prompt ordering and prompt length, arguing against the claim that these models learn linear regression as an in-context algorithm (Naim et al., 2024).

This suggests a central fault line in the literature. One line proves that transformer architectures are expressive enough to realize optimizer circuits, and in several controlled settings training does converge to such behavior. Another line shows that practical transformers can achieve regression-like or classification-like ICL while remaining tightly coupled to the pretraining prompt distribution. The unresolved question is therefore not whether transformers can be Transformer-Based ICL Optimizers, but when standard training actually induces optimizer-like circuits rather than distribution-specific heuristics. Open problems stated across the literature include more precise weight-matrix characterizations for practical softmax transformers, reductions in hidden dimension, broader nonlinear kernel and neural-network targets, robustness beyond synthetic settings, and sharper criteria for distinguishing genuine in-context optimization from prompt-conditioned representation matching (Wu et al., 2024).