Research map / Architectures and efficiency / Attention and transformers
The Closeness of In-Context Learning and Weight Shifting for Softmax Regression
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.
Theoretical analysis of softmax regression showing that the data transformation induced by a single softmax self-attention layer and by one gradient-descent step on an l2 loss are both norm-bounded, with small synthetic experiments showing similar behavior.
Problem
Prior theory explains in-context learning only for linear self-attention without softmax; it is unclear how softmax attention relates to gradient descent in in-context learning.
Approach
Formulates softmax regression min_x ||<exp(Ax),1_n>^{-1} exp(Ax) - b||_2 and proves Lipschitz-style bounds showing that one GD step on x and a self-attention-style update to the document A each shift the target b by at most M times the parameter change, with M = n^1.5 exp(10R^2).
Key results
- Trained single softmax self-attention layer and one-step GD reach nearly identical losses on synthetic softmax regression with n=200, d=20 (Figure 1a).
- Model cosine similarity between GD and SA sensitivities rises to roughly 0.9-1.0 over training for n=200 and d=20 (Figure 1b), with prediction and model differences decreasing.
- Theorems give ||b_tilde - b||_2 <= M*||x_{t+1}-x_t||_2 and <= M*||A_{t+1}-A_t|| with M = n^1.5 exp(10R^2).
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| train a single self-attention layer with softmax instead of one-step gradient descent | Synthetic softmax regression tasks · Synthetic softmax regression tasks | loss | No clear effect |
| increase self-attention training steps | Single self-attention layer with softmax · Synthetic softmax regression tasks | prediction difference | Helped |
| increase self-attention training steps | Single self-attention layer with softmax · Synthetic softmax regression tasks | model cosine similarity | Helped |
| increase document length | Single self-attention layer with softmax compared with one-step gradient descent · Synthetic softmax regression tasks | training steps required for similarity | Hurt |
Surprising findings
- Larger document lengths require more training steps before the Transformer and GD models align.
Limitations
- Findings restricted to small Transformers and simple synthetic regression problems (stated by authors).
- Bounds are loose; the softmax nonlinearity means exact matching is not expected.
- Only self-attention-only Transformers are studied, with no real LLM experiments.
Experimental setup
Datasets: synthetic softmax regression tasks (randomly sampled length-n documents A and targets b)
Benchmarks: synthetic softmax regression tasks
Metrics: loss prediction difference (l2) model sensitivity cosine similarity model sensitivity difference (l2)
Baselines: one-step gradient descent on softmax regression
Builds on: Transformers learn in-context by gradient descent (von Oswald-style linear atten Data transformation view of in-context learning (Olsson/Ong-style, ONR+22) Softmax regression analysis of DLS23 Exponential regression GMS23
Closest papers by meaning
- In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness (NeurIPS 2024 · similarity 0.92)
- Transformers are Minimax Optimal Nonparametric In-Context Learners (NeurIPS 2024 · similarity 0.90)
- Understanding In-Context Learning in Transformers and LLMs by Learning to Learn Discrete Functions (ICLR 2024 · similarity 0.90)
- Unveiling Induction Heads: Provable Training Dynamics and Feature Learning in Transformers (NeurIPS 2024 · similarity 0.90)
- How Do Nonlinear Transformers Learn and Generalize in In-Context Learning? (ICML 2024 · similarity 0.90)
- In-Context Learning with Representations: Contextual Generalization of Trained Transformers (NeurIPS 2024 · 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.