Research map / Optimization / Implicit bias and generalization
Training Dynamics of In-Context Learning in Linear Attention
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Studies gradient-flow training dynamics of multi-head linear self-attention on in-context linear regression, comparing merged vs separate key-query parametrizations, deriving a single abrupt loss drop for merged and saddle-to-saddle (principal component regression) dynamics for separate.
Problem
How linear self-attention acquires in-context learning through gradient descent training, and how the key/query parametrization (merged vs separate, and rank) shapes the loss landscape and training dynamics.
Approach
Shows merged-key-query linear attention equals a two-layer linear network on cubic features, giving an analytical sigmoidal time course with one abrupt drop; separate key/query linear attention equals sums of three-layer convolutional linear networks, with exponentially many fixed points and saddle-to-saddle dynamics that reduce to scalar ODEs, progressively learning principal components.
Key results
- Merged-parametrization linear attention shows a single abrupt loss drop with an analytical sigmoidal time-course solution for white covariance and small initialization.
- Separate key/query linear attention has 2^D fixed points and D abrupt loss drops; loss at each plateau matches theory, with the m-th drop approximately the eigenvalue lambda_m.
- Early-stopped separate-parametrization models implement principal component regression with the number of components increasing over training.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Use merged rather than separate key and query weights | Linear attention trained from small initialization · In-context linear regression | loss dynamics | Mixed |
| Use separate rank-one key and query weights | Linear attention · In-context linear regression | loss dynamics | Helped |
| Use rank 2 for separate key and query weights (ablation) | ATTN_S · In-context linear regression, D=8 | plateau pattern | Mixed |
| Use rank 4 for separate key and query weights (ablation) | ATTN_S · In-context linear regression, D=8 | plateau pattern | Mixed |
| Increase the rank of separate key and query weights (ablation) | ATTN_S · In-context linear regression, D=8 | plateau duration | Helped |
| Increase initialization scale (ablation) | ATTN_M · In-context linear regression | plateau duration | Helped |
| Increase initialization scale (ablation) | ATTN_S · In-context linear regression | plateau duration | Helped |
| Add softmax attention to ATTN_M | Linear ATTN_M · In-context linear regression | loss dynamics | No clear effect |
| Add softmax attention to ATTN_S | Linear ATTN_S · In-context linear regression | loss dynamics | No clear effect |
Surprising findings
- Separate key/query parametrization yields progressive, saddle-to-saddle ICL acquisition rather than a single abrupt drop.
- Increasing initialization scale shortens plateaus, and very large initialization yields exponential-shaped (lazy) loss decay.
Limitations
- Analysis of the full time course is closed-form only for white covariance and infinitesimal initialization; general covariance relies on approximations and ansatz.
- Focuses on linear attention; softmax results are only empirical.
- Task vector distribution is zero-mean Gaussian, purely in-context task without in-weight learning.
Experimental setup
Datasets: synthetic in-context linear regression data
Metrics: squared loss loss plateau values
Baselines: softmax attention
Builds on: linear attention deep linear networks in-context linear regression
Closest papers by meaning
- In-Context Linear Regression Demystified: Training Dynamics and Mechanistic Interpretability of Multi-Head Softmax Attention (ICML 2025 · similarity 0.93)
- Fine-grained Analysis of In-context Linear Estimation: Data, Architecture, and Beyond (NeurIPS 2024 · similarity 0.89)
- How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression? (ICLR 2024 · similarity 0.89)
- In-Context Learning as Conditioned Associative Memory Retrieval (ICML 2025 · similarity 0.89)
- How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression (NeurIPS 2024 · similarity 0.89)
- In-Context Learning Strategies Emerge Rationally (NeurIPS 2025 · similarity 0.89)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.