PaperlayerView exact revision on arXiv

Source-linked AI summary

In-Context Convergence of Transformers

Yu Huang, Yuan Cheng, Yingbin Liang

arXiv:2310.05249v1cs.LGcs.AImath.OCstat.ML

TL;DR

Prior theoretical studies of in-context learning focused on linear transformers, leaving softmax attention dynamics under gradient descent open. This paper analyzes a one-layer softmax transformer on linear tasks with structured features, proving near-zero prediction-error convergence for balanced and imbalanced settings. The imbalanced case converges stage-wise, with dominant-feature queries preceding under-represented-feature queries.

  • Problem

    Prior theoretical studies of transformer in-context learning focused on linear transformers, leaving the learning dynamics of softmax-based transformers trained by gradient descent open.

    Source: §1 · p. 2

  • Method

    The paper analyzes a one-layer softmax transformer trained by gradient descent on linear tasks with inputs sampled from orthonormal feature vectors.

    Source: §1 · p. 2, §2.1 · p. 3

  • Results

    Balanced features yield near-minimum prediction error after two attention-dynamics phases, while imbalanced features converge stage-wise for dominant then under-represented query features.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The analysis characterizes training phases through the changing dominance of target-feature and off-target-feature attention weights.

    Source: §1 · p. 2, §4 · p. 8

    The analysis is scoped to linear tasks and inputs sampled from a set of orthonormal feature vectors with specified sampling probabilities.

    Source: §2.1 · p. 3

Abstract

from arXiv · show

Transformers have recently revolutionized many domains in modern machine learning and one salient discovery is their remarkable in-context learning capability, where models can solve an unseen task by utilizing task-specific prompts without further parameters fine-tuning. This also inspired recent theoretical studies aiming to understand the in-context learning mechanism of transformers, which however focused only on linear transformers. In this work, we take the first step toward studying the learning dynamics of a one-layer transformer with softmax attention trained via gradient descent in order to in-context learn linear function classes. We consider a structured data model, where each token is randomly sampled from a set of feature vectors in either balanced or imbalanced fashion. For data with balanced features, we establish the finite-time convergence guarantee with near-zero prediction error by navigating our analysis over two phases of the training dynamics of the attention map. More notably, for data with imbalanced features, we show that the learning dynamics take a stage-wise convergence process, where the transformer first converges to a near-zero prediction error for the query tokens of dominant features, and then converges later to a near-zero prediction error for the query tokens of under-represented features, respectively via one and four training phases. Our proof features new techniques for analyzing the competing strengths of two types of attention weights, the change of which determines different training phases.

1 Introduction

Prior theoretical work on transformer in-context learning focused on linear transformers, leaving softmax-based learning dynamics open. This work analyzes a one-layer softmax transformer trained by gradient descent and characterizes distinct convergence behavior for balanced and imbalanced features.

1 Introduction

  • Earlier theoretical studies of transformer in-context learning focused on linear transformers, leaving open how softmax-based transformers learn in context.The paper frames this as a question about softmax transformers trained via gradient descent.Source: §1 · p. 2
  • The analysis studies a single-layer softmax transformer trained by gradient descent on in-context linear regression prompts with structured feature inputs.Tokens are randomly selected from a set of feature vectors; training uses squared prediction error.Source: §1 · p. 2
  • For balanced features, attention parameters first grow rapidly while aligning queries with their target feature, then prediction error converges near its minimum.The contribution is stated for each feature k.Source: §1 · p. 2
  • With imbalanced features, query tokens for dominant features reach near-zero prediction error first, followed by under-represented features through one and four phases, respectively.The setting has one dominant feature with p1 = Θ(1) and other features under-represented.Source: §1 · p. 2
  • The proof tracks competing bilinear attention weights for target and off-target features, whose changing dominance determines different training phases.The authors suggest these analysis tools may also be useful for other transformer problems.Source: §1 · p. 2

2 Problem Setup

The setup defines in-context prompts for linear regression tasks and a one-layer softmax self-attention model. It specifies masking and parameterization choices, then trains the model by gradient descent on squared prediction error.

2.1 In-Context Learning Framework

  • Each prompt samples a task and input-query tokens, pairs inputs with their function values, and asks the in-context learner to predict the query output.The framework is adopted from [15].Source: §2.1 · p. 3
  • Tasks are linear functions with zero-mean, identity-covariance weight vectors, while inputs are sampled from orthonormal features with probabilities p_k.The feature probabilities sum to one; their balance or imbalance is studied later.Source: §2.1 · p. 3

2.2 One-Layer Transformer Architecture

  • The model masks the query column before key and value projections so the query token cannot attend to itself.The masking removes the last column of the embedding matrix.Source: §2.2 · p. 4
  • The model consolidates query and key matrices into W_KQ and fixes the value parameter ν at 1 during training.The authors state that the resulting parameterization incurs at most exponentially vanishing error relative to the general parameter space.Source: §2.2 · p. 4

2.3 Training Settings

  • Training minimizes expected squared prediction error over prompts and task weights using gradient descent with learning rate η, initializing Q at zero.The expectation includes input and query tokens and the weight vector w.Source: §2.3 · pp. 4–5

3 Main Results

The main results establish finite-time convergence for one-layer softmax transformers trained by gradient descent, with distinct dynamics for balanced and imbalanced feature distributions. Balanced features yield near-zero prediction error through two phases, while imbalanced features yield stage-wise convergence from dominant to under-represented features.

3 Main Results

  • The results characterize gradient-descent convergence for in-context learning under both balanced and imbalanced feature distributions.The imbalanced setting considers a dominant feature and under-represented features.Source: §3 · p. 5, §3.2 · p. 6

3.1 In-Context Learning with Balanced Features

  • For balanced features, gradient descent reaches near-minimum loss in polynomial time through two phases: target-feature attention grows first, then prediction error converges.The theorem also establishes attention concentration on input tokens matching the query feature.Source: §3.1 · pp. 5–6
  • With balanced features, the trained transformer can approximate test prompts for linear tasks whose task model was absent from the training distribution.This result is stated for a test prompt whose data follow the input distribution, even when its task model was unseen during training.Source: §3.1 · p. 6

3.2 In-Context Learning with Imbalanced Features

  • For imbalanced features, prediction error converges first for the dominant feature and later for under-represented features, with the latter requiring a longer, four-phase process.The resulting prediction error is near-zero for both groups despite the longer convergence time for under-represented features.Source: §3.2 · pp. 6–7

4 Overview of Training Phases

The analysis tracks bilinear attention weights because they shape attention scores and predictions; shifts in which weight leads divide learning into phases. Under-represented features pass through four phases, while other settings use fewer phases.

4 Overview of Training Phases

  • The analysis tracks bilinear attention weights because their changing dominant roles determine the attention dynamics and training phases.The weights are associated with attention scores, which in turn determine the query prediction.Source: §4 · pp. 7–8

4.1 Learning Process for Under-represented Features

  • For under-represented query features, training proceeds through four phases, ending with attention to the target feature approaching 1 and prediction error converging to a global minimum.The phases are illustrated as suppression of dominant-feature attention, a shift in leading influence, growth of target-feature attention, and convergence.Source: §4 · p. 8, §4.1 · p. 9

4.1.1 Phase I: Decrease of Dominant Feature.

  • Phase I suppresses attention to the dominant feature as its associated weight decreases faster than the other off-target weights.By the phase's end, the relevant off-target weights remain close to zero.Source: §4.1.1 · p. 9

4.1.2 Phase II: Switching of Leading Influence

  • In Phase II, the dominant-feature weight's decline slows, allowing the target-feature weight's growth to become the leading influence.This transition marks the switch in which bilinear weight drives the attention dynamics.Source: §4.1.2 · pp. 9–10

4.1.3 Phase III: Growth of Target Feature

  • In Phase III, the target-feature weight continues to grow rapidly while other off-target weights remain relatively unchanged.This growth drives the target feature's attention upward.Source: §4.1.3 · p. 10

4.1.4 Phase IV: Convergence

  • Phase IV begins after the target feature is learned to a sufficient level; prediction error then converges as the remaining attention mass diminishes.The analysis connects the target-feature weight to prediction error through the change in attention mass.Source: §4.1.4 · p. 10

4.2 Training Dynamics of Other Settings

  • For the dominant feature in imbalanced data, training enters convergence directly because its initial attention to the target is already substantial.The passage describes this setting as having target-feature attention of at least constant order from the abundance of its tokens.Source: §4.2 · p. 11
  • Balanced data begins with nearly equal feature attention, and its dynamics resemble the target-feature growth phase for under-represented features.The section says other settings use phases similar to those in the four-phase analysis.Source: §4.2 · p. 11

5 Additional Related Work

Prior work studied in-context learning's expressive capacity, generalization, and training dynamics. The closest training-dynamics studies used linear self-attention, whereas this paper analyzes softmax attention.

In-Context Learning

  • Prior work on in-context learning examined expressive power, generalization, algorithmic mechanisms, and Bayesian perspectives.The cited studies include linear-regression tasks, transformer implementations of gradient descent, and sample complexity.Source: §5 · p. 11
  • Related training-dynamics studies analyzed linear transformers, including one-layer models trained by gradient flow or pretraining loss.Those works considered linear regression tasks or random instances of linear regression.Source: §5 · pp. 11–12

Training Dynamics of Transformers

  • Other transformer-training analyses covered vision-transformer learning, shallow-model sample complexity, SGD dynamics, and self-attention optimization geometry.The passage notes that one analysis of shallow ViT relied on initialization near a target pattern, while another did not provide an SGD convergence guarantee.Source: §5 · p. 12

6 Conclusions

The paper analyzes gradient-descent training of a one-layer softmax-attention transformer for in-context learning with balanced and imbalanced features. It proves vanishing prediction error and characterizes four phases for under-represented features.

6 Conclusions

  • The paper proves convergence to vanishing in-context prediction error for balanced and imbalanced feature settings, and characterizes four phases for under-represented features.The analysis centers on changes in which of two types of bilinear attention weights plays the dominant role.Source: §6 · p. 12

A.1 Gradient Computations

  • The appendix computes gradients with respect to Q and projects them onto feature directions to analyze training updates.The parameter ν is not updated during training.Source: §A.1 · pp. 17–18

A.2 Useful Probabilistic Lemmas for Prompt

  • The appendix uses multinomial concentration to identify high-probability prompt-count events for balanced and imbalanced data.These events provide the data properties used in the subsequent analysis.Source: §A.2 · pp. 20–21

A.3 Properties of Loss Function and Prediction Error

The section characterizes population loss using a decomposition and compares attainable loss with unavoidable error. For balanced and imbalanced features, it describes the relevant error scales and assumptions used in the analysis.

A.3 Properties of Loss Function and Prediction Error

  • The population loss is decomposed into a form used to relate loss to attention scores.The passage introduces this as Lemma A.6.Source: §A.3 · p. 21

A.3.1 Loss Characterization for the Balanced Case

  • For balanced data, Llow is a lower bound on L*, and both are Θ(e^-poly(K)).Llow is the sum of unavoidable errors when the query feature was absent from the first N training samples; convergence is analyzed relative to Llow.Source: §A.3.1 · p. 22

A.3.2 Loss Characterization for the Imbalanced Case

  • For each feature in imbalanced data, the conditional prediction error's minimum and unavoidable values are both Θ(e^-poly(K)).The section defines a conditional prediction error for each query feature before characterizing these values.Source: §A.3.2 · pp. 24–25

A.4 Notations and Parameters

  • The proofs assume N = poly(K) ≫ K^3 and sufficiently large K.Source: §A.4 · p. 25

B.1 Roadmap of the Proof

The proof roadmap tracks induction hypotheses through the target-feature growth and convergence phases, combining phase analyses to establish convergence while bounding conditional loss in Stage II of Phase II and beginning the induction argument for its final step.

B.2 Phase I: Growth of Target Feature

  • In Phase I, the target-feature parameter grows faster than the competing attention parameters, dominating the learning dynamics.The phase begins at initialization and is analyzed using an induction hypothesis intended to hold throughout.Source: §B.1 · p. 27, §B.2 · p. 28

B.1 Roadmap of the Proof

  • Combining the two phases’ analyses proves the training-process convergence result stated in Theorem 3.1.The roadmap divides the balanced-case dynamics into target-feature growth and convergence phases.Source: §B.1 · p. 27

B.2.1 Technical Lemmas

  • Under Induction Hypothesis B.1, the Phase I technical lemmas give a lower bound of Ω(1) on 1 − Attn_k(t)/2 for a target-feature query.The condition is stated for balanced input sequences.Source: §B.2.1 · p. 28

B.2.2 Controlling Gradient Updates in Phase I

  • For balanced data with N ≫ K^3, the Phase I analysis controls gradient updates under Induction Hypothesis B.1.The update analysis distinguishes the target feature’s parameter from those for n ≠ k.Source: §B.2.2 · p. 30

B.2.3 End of Phase I

  • Induction Hypothesis B.1 holds through Phase I, which ends by T1,k = O(log(K)K^2/η) iterations.The phase-end argument establishes the hypothesis at each iteration through its endpoint.Source: §B.2.3 · p. 31

B.3 Phase II: Convergence: Stage I

  • After Phase I aligns a query with input tokens of its feature, Stage I of Phase II increases A_k(t) and decreases B_k,n(t), supporting loss convergence.The convergence phase is divided into two sub-stages based on parameter variation rates.Source: §B.3 · p. 32

B.3.1–B.3.3 Stage I Analysis

  • Stage I technical lemmas and gradient-update analysis proceed under Induction Hypothesis B.2 for balanced inputs, establishing the conditions used to end the stage.The supplied passages identify the technical lemmas, update analysis, and stage-end induction argument but do not provide their full bounds.Source: §B.3.1 · pp. 32–33, §B.3.2 · p. 34, §B.3.3 · pp. 35–36

B.4 Phase II: Convergence: Stage II

  • In Stage II, A_k(t) increases but is bounded by O(log(K/ϵ)), while B_k,m(t) decreases for m ≠ k.Under Induction Hypothesis B.3, the Stage II update analysis also establishes α_k(t) ≥ 0 and conditions on β_k,n(t).Source: §B.4 · p. 37, §B.4.1 · p. 37, §B.4.2 · p. 39

B.4.3 Controlling Loss in Stage II of Phase II

  • Under Induction Hypothesis B.3, for t in (eT_{ϵ^2,k}, T_{ϵ^2,k}], the loss eL_k(θ(t)) is less than p_kϵ^2 when polylog(K) ≫ log(1/ϵ).The lemma assumes k ∈ [K] and 0 < ϵ < 1.Source: §B.4.3 · p. 40

B.4.4 End of Stage II of Phase II

  • The final-step proof checks that Induction Hypothesis B.3 holds at t = eT_{ϵ^2,k} + 1 before proceeding by induction.Source: §B.4.4 · p. 41

B.5 Proof of Theorem 3.1 for Balanced Case

For balanced features, the theorem bounds the loss after a finite training time and shows attention concentrates on tokens matching the query feature. The proof also records the scale of the global minimum.

B.5 Proof of Theorem 3.1 for Balanced Case

  • At T*, the population loss is within ε of its global minimum L*.Source: §B.5 · p. 41
  • The global minimum L* is Θ(e^-poly(K)).Source: §B.5 · p. 41
  • With probability at least 1 − e^-Ω(poly(K)), attention nearly all goes to input tokens featuring the query feature; the probability is over the first N test-prompt input tokens.Source: §B.5 · p. 41

C.1 Roadmap of the Proof

This section organizes the proof of convergence for under-represented features into four phases, combining them to establish Theorem 3.2. It develops Phase III through induction conditions and technical attention and update bounds, and records an induction step at the end of Phase II.

C.1 Roadmap of the Proof

  • The four-phase roadmap moves from dominant-weight reduction, through a shift toward A(t)_k growth, to the desired attention structure and prediction-error convergence.Phase III maintains a magnitude gap between α(t)_k and other β(t)_k,n; phase IV analyzes the diminishing 1 − Attn(t)_k.Source: §C.1 · pp. 43–44

C.2 Phase I: Decrease of Dominant Feature

  • In phase I, B(t)_k,1 decreases while A(t)_k increases for under-represented feature v_k, with an induction hypothesis tracking their dynamics.The phase is framed as mitigating the dominant feature v_1’s high-occurrence bias.Source: §C.2 · p. 44

C.2.1 Technical Lemmas

  • Under the phase-I induction hypothesis, technical lemmas bound attention quantities for prompts whose query is the under-represented feature v_k.These lemmas are used to prove the phase-I induction hypothesis.Source: §C.2.1 · p. 44

C.2.2 Controlling Gradient Updates in Phase I

  • Phase-I gradient-update lemmas control the changes in α(t)_k, β(t)_k,1, and β(t)_n for n ≠ 1,k under the induction hypothesis.Source: §C.2.2 · pp. 46–48

C.2.3 End of Phase I

  • At the end of phase I, the induction hypothesis holds through T1,k, and the analysis carries its remaining claim into phase II.Source: §C.2.3 · p. 49

C.3 Phase II: Switching of Leading Influence

  • In phase II, the decline of B(t)_k,1 slows, shifting the leading influence toward growth of A(t)_k.The roadmap describes the decline of B(t)_k,1 as approaching a comparable order to the growth influence.Source: §C.1 · p. 43, §C.3 · p. 50

C.3.1 Technical Lemmas

  • Phase-II technical lemmas and gradient-update bounds analyze attention and parameter changes for under-represented query features under the phase-II induction hypothesis.Source: §C.3.1 · pp. 50–51, §C.3.2 · pp. 52, 54

C.3.3 End of Phase II

  • The phase-II induction hypothesis is established through T2,k, completing the phase’s transition analysis.Source: §C.3.3 · p. 55
  • At the end of Phase II, the proof verifies induction hypothesis C.2 at T1,k + 1 and proceeds by induction from earlier iterations.Source: §C.3.3 · p. 55

C.4 Phase III: Growth of Target Feature

  • Phase III is defined for T2,k < t ≤ T3,k, with the target feature’s A(t) growing and its effect becoming dominant.The induction hypothesis states that A(t) increases while B_k,1(t) decreases during this phase.Source: §C.4 · p. 56

C.4.1 Technical Lemmas

  • Under induction hypothesis C.3, Lemma C.13 gives the bound 1 − Attn(t)_k ≥ Ω(1) for query vk and inputs in E*_imbal.Source: §C.4.1 · p. 56

C.4.2 Controlling Gradient Updates in Phase III

  • Phase III’s update analysis bounds α(t) and β(t), including β(t) for the target feature and for features n ≠ 1, k.The stated conditions invoke induction hypothesis C.3; the supplied passages do not show the bounds’ values.Source: §C.4.2 · pp. 58–59

C.4.3 End of Phase III

  • For k > 1, induction hypothesis C.3 holds throughout T2,k < t ≤ T3,k.Lemma C.18 states that T3,k is T2,k plus an order term, but the supplied expression is incomplete.Source: §C.4.3 · p. 61

C.5 Phase IV: Convergence

Phase IV establishes that the attention structures achieved for query tokens of under-represented features are solutions toward which the algorithm converges. Its analysis tracks gradient updates under induction hypotheses to prove this final convergence phase.

C.5 Phase IV: Convergence

  • Phase IV establishes that the desired attention structures, including those for under-represented features, are solutions toward which the algorithm converges.The desired structure for query feature vk has already been achieved at t = T3,k + 1.Source: §C.5 · p. 61

C.5.1 Technical Lemmas

  • During Phase IV, Lemmas C.19 and C.20 analyze query tokens for feature vk on imbalanced inputs under Induction Hypothesis C.4.Both lemmas apply for T3,k < t ≤ T4,k.Source: §C.5.1 · p. 62

C.5.2 Controlling Gradient Updates in Phase IV

  • Under Induction Hypothesis C.4, the Phase IV update analysis establishes that α_k(t) is nonnegative.Source: §C.5.2 · p. 63

C.6 Proof of Theorem 3.2 for Under-represented Features

For under-represented query features, the theorem reports prediction-error convergence and concentration of attention on matching input tokens. The adjoining dominant-feature analysis describes a simpler, single-phase route that enters convergence immediately.

C.6 Proof of Theorem 3.2 for Under-represented Features

  • For each under-represented feature vk, the theorem states that its prediction error converges.Source: §C.6 · p. 67
  • After Tk iterations, a query token vk nearly concentrates attention on input tokens featuring vk, with probability at least 1 − e^−Ω(poly(K)).Source: §C.6 · p. 67

D Analysis for Imbalanced Case: Dominant Feature

D.4 Proof of Theorem 3.2 for Dominant Feature

For the dominant feature, the theorem gives a finite-iteration prediction-error guarantee and a matching-feature attention-concentration result. These statements apply under the stated assumptions on feature probabilities, sample size, and ε.

D.4 Proof of Theorem 3.2 for Dominant Feature

  • The theorem assumes p1 = Θ(1), pk = Θ(·) for 2 ≤ k ≤ K, N ≥ poly(K), and polylog(K) ≫ log(1/ϵ).The supplied theorem statement gives the pk scaling only in a truncated form.Source: §D.4 · p. 73
  • The dominant-feature loss satisfies L1(θ(T1)) ≤ L1* + ϵ after at most T1 gradient-descent iterations.The theorem identifies L1* = Θ(e^−poly(K)) as the global minimum of eq. (6).Source: §D.4 · p. 73
  • After Tk iterations, a query token vk nearly concentrates attention on input tokens featuring vk, with probability at least 1 − e^−Ω(poly(K)).Source: §D.4 · p. 73

Related papers

Loading 2310.05249v1…