Research map / Architectures and efficiency / Attention and transformers
On Understanding Attention-Based In-Context Learning for Categorical Data
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.
Extends functional gradient-descent views of in-context learning to categorical outcomes, proposing a self-attention plus cross-attention block that exactly implements multi-step GD inference for softmax likelihoods, evaluated on synthetic data, ImageNet in-context classification, and TinyStories language generation.
Problem
Existing functional-GD interpretations of Transformer in-context learning assume real-valued outputs with Gaussian likelihoods, whereas language involves categorical (discrete token) outcomes; it is unclear how attention can implement GD for categorical observations.
Approach
Shows that the categorical cross-entropy GD update requires the nonlinear expectation over class embeddings, which can be computed by a softmax cross-attention layer over learned embedding vectors; each attention block pairs a self-attention layer (kernel update of the latent function plus erasure head) with this cross-attention layer, so multi-step GD is exact.
Key results
- On synthetic data, GD and Trained TF with one self-attention layer show close agreement in accuracy and NLL once enough training contexts are used, with GD needing far fewer (L<5000 vs >25,000 for two-layer Trained TF).
- On in-context ImageNet classification with VGG features, GD attention with two to three blocks matches linear probing top-1 accuracy almost exactly.
- GPT-4o scoring on 200 story endings: a GD model with an added feedforward network performs almost the same as a Transformer, with about 8K attention-weight parameters versus 6M for the Transformer.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| add a second attention block (GD step) | GD attention model · Synthetic categorical data | Top-1 predictive accuracy | Helped |
| add attention blocks beyond one | GD attention model · Synthetic categorical data | negative log-likelihood | Helped |
| use linear attention instead of nonlinear attention | GD attention model · Synthetic categorical data | predictive accuracy | Hurt |
| use exponential attention | GD attention model · Synthetic categorical data | predictive accuracy | Hurt |
| increase the GD attention model from one to two or three blocks | GD attention model · ImageNet in-context classification | Top-1 predictive accuracy | Helped |
| use one GD self-attention layer instead of linear probing | GD attention model · ImageNet in-context classification | Top-1 predictive accuracy | Hurt |
| add a feedforward network above the GD attention output (ablation) | Softmax GD language model · Tiny Stories and Children Stories generation | GPT-4o score | Helped |
| use the GD model instead of the comparable Transformer | Single-layer language model · Tiny Stories and Children Stories generation | GPT-4o score | Hurt |
| use a single self-attention layer | GD attention model · Synthetic categorical data | Top-1 predictive accuracy | No clear effect |
Surprising findings
- Adding a feedforward layer to the GD-based language model nearly closes its gap with the Transformer, though the GD model alone scores much lower.
Limitations
- GD language model without feedforward layers underperforms the Transformer on GPT-4o scores, and generations often repeat.
- Synthetic experiments include a deliberate model mismatch between the attention kernel and the data-generating function.
- Language experiments use only single-layer models on a narrow story corpus.
Experimental setup
Datasets: Synthetic categorical GD data ImageNet TinyStories Children Stories
Benchmarks: ImageNet in-context classification (synthetic-to-ImageNet test classes) TinyStories and Children Stories story continuation
Metrics: top-1 accuracy negative log-likelihood GPT-4o quality score out of 10
Baselines: Trained TF Linear probing Single-layer Transformer
Builds on: Transformers learn in-context by gradient descent Linear attention as in-context regression Trained Transformer (TF) vs GD transformer comparison
Closest papers by meaning
- Training Dynamics of In-Context Learning in Linear Attention (ICML 2025 · similarity 0.88)
- In-Context Linear Regression Demystified: Training Dynamics and Mechanistic Interpretability of Multi-Head Softmax Attention (ICML 2025 · similarity 0.87)
- Is attention required for ICL? Exploring the Relationship Between Model Architecture and In-Context Learning Ability (ICLR 2024 · similarity 0.87)
- In-Context Learning as Conditioned Associative Memory Retrieval (ICML 2025 · similarity 0.86)
- The mechanistic basis of data dependence and abrupt learning in an in-context classification task (ICLR 2024 · similarity 0.86)
- In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness (NeurIPS 2024 · similarity 0.86)
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.