Research map / Architectures and efficiency / Attention and transformers
Does learning the right latent variables necessarily improve in-context learning?
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.
Tests whether forcing a Transformer to infer task latents through a bottleneck (explicit model) improves ICL, finding it does not improve ID or OOD performance even though latents are recovered, because the downstream predictor fails to use them.
Problem
Whether Transformers performing in-context learning infer underlying task latents or rely on non-parametric shortcuts, and whether explicitly inferring the right latents improves ID and OOD generalization.
Approach
Minimally modifies the Transformer with a fixed-size bottleneck that separates context aggregation (infers task latents from the context) from query prediction, so the query cannot directly attend to demonstrations; compares this explicit model against a standard implicit Transformer across synthetic regression/classification and compositional reasoning tasks, with oracle variants that use the tru
Key results
- Explicit (bottleneck) models match implicit Transformers in-distribution and out-of-distribution; implicit models are slightly better ID and implicit models consistently outperform explicit ones on the reusable modular MoE compositional task.
- Explicit models accurately recover true task latents linearly (latents decodable from bottleneck), but adding an oracle prediction function is needed to obtain substantially better OOD generalization.
- Using DAS, counterfactual interventions succeed in the explicit model's bottleneck but not in the implicit model, offering better interpretability.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Add a context bottleneck to explicitly infer task latents | Transformer ICL · In-distribution tasks | task performance | No clear effect |
| Add a context bottleneck to explicitly infer task latents | Transformer ICL · Synthetic regression OOD | task performance | Hurt |
| Add a context bottleneck to explicitly infer task latents | Transformer ICL · Classification and compositional tasks OOD | task performance | No clear effect |
| Use the implicit model instead of the explicit model | Reusable modular MoE · Compositional generalization | task performance | Helped |
| Train queries near context points, encouraging nearest-neighbor shortcuts (ablation) | Implicit model · Sinusoid regression OOD with queries far from context | task performance | Hurt |
| Train queries near context points, encouraging nearest-neighbor shortcuts (ablation) | Explicit model · Sinusoid regression OOD with queries far from context | task performance | Helped |
| Add an auxiliary loss to predict the ground-truth task latent up to a linear transformation (ablation) | Explicit model · ICL tasks | task performance | No clear effect |
| Apply the auxiliary ground-truth latent loss directly to the bottleneck instead of a linear decoding (ablation) | Explicit model · Linear regression OOD | task performance | No clear effect |
| Use the ground-truth prediction function instead of learning the predictor (ablation) | Explicit model · OOD tasks except nonlinear regression with MLP latents | task performance | Helped |
| Use the ground-truth prediction function instead of learning the predictor (ablation) | Explicit model · Nonlinear regression with MLP latents | task performance | Hurt |
| Increase input dimensionality, context length, or model size (ablation) | Implicit and explicit models · Linear regression OOD | task performance | No clear effect |
Surprising findings
- Learning the correct task latents with an auxiliary loss gives no OOD improvement, showing latent inference alone is insufficient for generalization.
- Implicit models outperform explicit ones on the compositional MoE task across different proportions of seen latent combinations.
Limitations
- Explicit model architecture is minimal and not optimized; authors state the goal was not to build the best explicit model.
- Nonlinear regression with MLP latents: oracle prediction function performs poorly, suggesting latents are hard to infer there.
- Results limited to chosen synthetic and reasoning tasks at small scale (single RTX8000 GPUs).
Experimental setup
Datasets: Alchemy Raven's Progressive Matrices (symbolic version) Perturb-seq gene targeting dataset synthetic regression and classification tasks (linear, MLP, sinusoid, GP, Hodgkin-Huxley ODE) reusable modular mixture of experts (synthetic)
Benchmarks: Alchemy Raven's Progressive Matrices Perturb-seq gene targeting synthetic linear/nonlinear/sinusoid regression synthetic linear/nonlinear classification reusable modular MoE
Metrics: mean squared error accuracy relative accuracy interchange intervention accuracy latent decoding accuracy
Baselines: implicit Transformer explicit model with MLP prediction oracle with known prediction function oracle with auxiliary latent loss
Builds on: Conditional Neural Processes Distributed Alignment Search Transformer
Closest papers by meaning
- Technical Debt in In-Context Learning: Diminishing Efficiency in Long Context (NeurIPS 2025 · similarity 0.92)
- Emergence and Effectiveness of Task Vectors in In-Context Learning: An Encoder Decoder Perspective (ICML 2025 · similarity 0.91)
- When can in-context learning generalize out of task distribution? (ICML 2025 · similarity 0.89)
- In-Context Learning with Transformers: Softmax Attention Adapts to Function Lipschitzness (NeurIPS 2024 · similarity 0.89)
- In-Context Learning through the Bayesian Prism (ICLR 2024 · similarity 0.89)
- Transformers are Minimax Optimal Nonparametric In-Context Learners (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.