PaperlayerView exact revision on arXiv

Source-linked AI summary

One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention

Arvind Mahankali, Tatsunori B. Hashimoto, Tengyu Ma

arXiv:2307.03576v1cs.LG

TL;DR

The paper asks why restricted one-layer linear-attention transformers learn particular in-context regression algorithms, a question left open by empirical findings. It analyzes global pre-training-loss minimizers on synthetic noisy regression tasks and shows that isotropic covariates yield one-step gradient descent, non-isotropic covariates yield pre-conditioned descent, and nonlinear responses preserve the linear-descent outcome.

  • Problem

    Theory explaining which algorithms restricted one-layer linear-attention transformers learn from synthetic regression data remains limited.

  • Method

    The paper mathematically analyzes global minima of the pre-training loss for one-layer linear self-attention trained on synthetic noisy linear-regression data.

  • Results

    The minimizer implements one-step least-squares gradient descent for isotropic covariates, pre-conditioned gradient descent for non-isotropic covariates, and still one-step gradient descent when responses come from supported nonlinear target functions.

  • Takeaways & Limitations

    The learned algorithm depends strongly on covariate covariance but is comparatively insensitive to the response distribution within the studied nonlinear family.

  • Takeaways & Limitations

    The nonlinear-response result assumes a target-function family with rotation-invariance and symmetry under negation, while the covariance result uses a Gaussian covariate and matched inverse-covariance weight distribution.

Abstract

from arXiv · show

Recent works have empirically analyzed in-context learning and shown that transformers trained on synthetic linear regression tasks can learn to implement ridge regression, which is the Bayes-optimal predictor, given sufficient capacity [Akyürek et al., 2023], while one-layer transformers with linear self-attention and no MLP layer will learn to implement one step of gradient descent (GD) on a least-squares linear regression objective [von Oswald et al., 2022]. However, the theory behind these observations remains poorly understood. We theoretically study transformers with a single layer of linear self-attention, trained on synthetic noisy linear regression data. First, we mathematically show that when the covariates are drawn from a standard Gaussian distribution, the one-layer transformer which minimizes the pre-training loss will implement a single step of GD on the least-squares linear regression objective. Then, we find that changing the distribution of the covariates and weight vector to a non-isotropic Gaussian distribution has a strong impact on the learned algorithm: the global minimizer of the pre-training loss now implements a single step of $\textit{pre-conditioned}$ GD. However, if only the distribution of the responses is changed, then this does not have a large effect on the learned algorithm: even when the response comes from a more general family of $\textit{nonlinear}$ functions, the global minimizer of the pre-training loss still implements a single step of GD on a least-squares linear regression objective.

1 Introduction

The paper studies why restricted one-layer linear-attention transformers learn particular in-context regression algorithms. It proves that isotropic covariates yield one-step gradient descent, while covariance changes induce preconditioning and response-function changes have weaker effects.

  • Motivation: Prior empirical work found that transformers trained on synthetic linear-regression prompts can implement interpretable algorithms such as gradient descent or ordinary least squares.One-layer linear self-attention without an MLP was empirically associated with one gradient-descent step.
  • Motivation: The paper addresses limited theory by analyzing global minima of the pre-training loss for one-layer transformers with linear self-attention.The setup uses noisy synthetic linear-regression data and restricted model capacity.
  • Main results: The global pre-training-loss minimizer implements one gradient-descent step on a least-squares linear-regression objective under isotropic Gaussian covariates.This mathematically matches the earlier empirical finding for one-layer linear self-attention.
  • Main results: When covariates are non-isotropic, the global minimizer instead implements one pre-conditioned gradient-descent step.The learned algorithm is therefore sensitive to the covariance structure of the covariates.
  • Main results: Changing the response distribution has a weaker effect: even nonlinear target functions can still produce one-step gradient descent on a linear-regression objective.The paper contrasts this response-distribution effect with the stronger effect of changing covariate covariance.
  • Related work: Related theory independently obtains one-step gradient descent for isotropic covariates and pre-conditioned gradient descent for more general covariance structures.Other concurrent work also studies multi-layer extensions and related critical points.

2 Setup

The setup uses a single-head linear self-attention layer to predict a noisy regression response from in-context support examples and a query covariate. Training and testing draw sequences from the same synthetic distribution.

  • Architecture: The analysis excludes self-attention from a token to itself and represents the final output through the linear head applied to the last token.The token construction includes an extra zero coordinate for the query token without revealing its unknown response.
  • Model: The model is a one-layer, single-head linear self-attention transformer with key, query, and value matrices and a linear head applied to the final token.The instantiated width is d+1 for d-dimensional input covariates.
  • Data distribution: In the baseline distribution, covariates and the per-sequence weight vector are independent standard Gaussians, with y_i = w^T x_i + ε_i and Gaussian noise variance σ^2.The weight vector is freshly sampled for each sequence and the noise variables are i.i.d.
  • Data distribution: Each sequence contains support pairs (x_i, y_i) and a query covariate whose response must be predicted.The support covariates are x_1,...,x_n, while x_{n+1} is the query exemplar.

3 Main Result for Linear Models

For isotropic Gaussian linear-regression data, the pre-training-loss global minimizer implements one gradient-descent step. The proof links the transformer's effective predictor to a scaled empirical gradient update and shows this construction is globally optimal.

  • The global-loss minimizer is a one-layer linear self-attention transformer implementing one gradient-descent step on least-squares linear regression.
  • The construction matches the empirical one-step GD mechanism previously identified for this transformer architecture.
  • Rotational invariance of the Gaussian covariate and weight distributions makes the relevant expected matrices scalar multiples of the identity, enabling the global-optimality argument.
  • Starting from w0 = 0, the resulting predictor corresponds to the first iterate w1 = η Σ_i y_i x_i of gradient descent.
  • The proof compares the transformer’s effective predictor with the ridge-regression Bayes-optimal predictor and shows their expected losses differ through a simpler scaled predictor.
  • The loss depends entirely on the distance between the effective linear predictor and ηX^T y, so matching this quantity characterizes global minimizers.

4 Results for Different Data Covariance Matrices

With non-identity covariate covariance, the global pre-training-loss minimizer implements preconditioned gradient descent, showing that covariate distribution affects the learned algorithm.

  • The global minimizer implements one step of gradient descent with preconditioning when covariates have non-identity covariance.The preconditioner is Σ^-1.
  • The covariates are sampled from N(0, Σ), with Σ positive semidefinite, while weights are sampled from N(0, Σ^-1).
  • For a query token, the transformer produces the prediction associated with the preconditioned gradient-descent solution.
  • The result is obtained by reducing the non-isotropic setting to the earlier setting through a change of variables.

5 Results for Nonlinear Target Functions

Under rotationally invariant and sign-symmetric distributions of nonlinear target functions, the global loss minimizer still implements one step of least-squares gradient descent.

  • The nonlinear-target setting allows target functions such as fully connected neural networks with arbitrary depth and width.
  • The nonlinear response model is yi = f(xi) + ǫi, with Gaussian noise and f drawn from a family satisfying the stated assumptions.
  • The target-function distribution must be invariant under input rotations and symmetric under negation.
  • Under these assumptions, the global minimizer still implements one step of gradient descent on the least-squares linear regression objective.
  • The proof extends earlier arguments using symmetry properties of odd and even-degree response monomials.

6 Conclusion

The paper establishes how one-layer linear self-attention learns gradient-based algorithms under different data distributions, with covariate changes altering the algorithm more than nonlinear response changes.

  • For isotropic Gaussian covariates, the global pre-training-loss minimum corresponds to one GD step on least-squares linear regression.
  • For non-isotropic covariates, the global minimum instead corresponds to pre-conditioned GD.
  • For isotropic covariates with nonlinear responses, the global minimum still corresponds to one GD step on least-squares linear regression.
  • The study is limited to single-head linear self-attention, leaving multi-head and multi-layer settings as future directions.

A Missing Proofs from Section 3

The missing proofs establish the isotropic-case result by exploiting rotational symmetry, moment cancellations, and equality of gradients for equivalent objectives.

  • Rotational invariance makes the conditional data moments and related matrix expectations scalar multiples of the identity.
  • The proof compares the transformer objective with a least-squares objective by showing their gradients with respect to the effective parameters coincide.
  • Odd powers of the responses have zero expectation in the isotropic Gaussian setting, allowing corresponding blocks to be ignored.
  • The lower-left and upper-right block contributions match, yielding equality of gradients with respect to both parameter sets.
  • Once the effective parameter satisfies ηX^T y, the transformer output realizes the gradient-descent predictor.

B Missing Proofs for Section 4

The proof reduces the non-isotropic Gaussian setting to the standard-Gaussian case through a change of variables, showing that the transformed ridge-regression predictor is optimal and corresponds to a one-step pre-conditioned update.

  • Reduction to the standard-Gaussian case: The proof writes each covariate as x_i = Σ^1/2u_i, with isotropic Gaussian u_i, and collects the transformed covariates into U.This change of variables converts the non-isotropic data distribution into the standard-Gaussian setting.
  • Reduction to the standard-Gaussian case: Ridge regression on the transformed variables uses (U⊤U + σ^2I)^−1U⊤y, which is equivalent to expressing the predictor in the original covariates through X = UΣ^1/2.The proof explicitly uses U = XΣ^−1/2 to relate the transformed and original parameterizations.
  • Transformer reparameterization: After changing variables in the loss, the transformer parameters can be reparameterized so that the transformed problem has the same loss as the standard-Gaussian problem.The argument uses the shared distribution of (u_i, y_i) and the parameter change involving h, W_V, W_K, and W_Q.
  • Optimal predictor: The transformed ridge solution gives the conditional-mean predictor for the next response because u_{n+1} is an invertible function of x_{n+1}.The proof identifies E[y_{n+1} | u_{n+1}, eD] with the transformed ridge predictor evaluated at u_{n+1}.
  • Transformer reparameterization: The proof concludes that the corresponding transformer produces the standard-Gaussian theorem's output on the transformed input, with the learning rate determined by the transformed parameters.The output equivalence is established before the proof concludes.

C Missing Proofs from Section 5

The Section 5 proofs use conditional response moments and rotational symmetry to characterize the optimal effective linear predictor, then conclude that this predictor is necessary and sufficient for minimizing pre-training loss.

  • Characterizing the minimizer: For fixed training data, the proof studies the expected squared-loss objective as a function of the predictor u and identifies the data-dependent minimizer.The objective is decomposed into terms involving the predictor and a constant independent of u.
  • Characterizing the minimizer: The loss decomposition separates predictor-dependent terms from a constant that depends on the training data but not on u.This reduction allows the minimizer to be analyzed directly through the remaining quadratic terms.
  • Rotational symmetry: Under the response-function assumption, the conditional response covariance matrix has diagonal noise contribution σ^2 and off-diagonal terms determined by Ef[f(x_i)f(x_j)].The proof derives this structure using independent, zero-mean noise.
  • Rotational symmetry: Rotational invariance implies that the expected matrix X⊤M(X)X is a scalar multiple of the identity.The argument applies a common rotation to all covariates and uses invariance of the response moments.
  • Rotational symmetry: The proof shows that the relevant predictor map is equivariant under simultaneous rotations of the covariates.This rotational equivariance is used with the matrix identities to constrain the form of the optimizer.
  • Theorem 3: Theorem 3 follows because the effective linear predictor ηX⊤y is necessary and sufficient for minimizing the pre-training loss.This is the final optimization conclusion of the Section 5 argument.
Loading 2307.03576v1…