mlsanityData to 7 Oct 2026

In-Context Learning of a Linear Transformer Block: Benefits of the MLP Component and One-Step GD Initialization

arXiv 2402.14951 ↗First posted 22 Feb 2024stat.ML cs.CL cs.LGcited by 30

Not on the map. Why?

Abstract (arXiv)

We study the in-context learning (ICL) ability of a Linear Transformer Block (LTB) that combines a linear attention component and a linear multi-layer perceptron (MLP) component. For ICL of linear regression with a Gaussian prior and a non-zero mean, we show that LTB can achieve nearly Bayes optimal ICL risk. In contrast, using only linear attention must incur an irreducible additive approximation error. Furthermore, we establish a correspondence between LTB and one-step gradient descent estimators with learnable initialization (mathsfGD-β), in the sense that every mathsfGD-β estimator can be implemented by an LTB estimator and every optimal LTB estimator that minimizes the in-class ICL risk is effectively a mathsfGD-β estimator. Finally, we show that mathsfGD-β estimators can be efficiently optimized with gradient flow, despite a non-convex training objective. Our results reveal that LTB achieves ICL by implementing mathsfGD-β, and they highlight the role of MLP layers in reducing approximation error.

Key facts. Cited by 30 papers in the map (data through 7 Oct 2026). Builds on 28 of its 37 references that are in the map. First posted on arXiv 22 Feb 2024.

Builds on

28 of 37 references.

References · 28

  1. 2017Attention Is All You Needcited by 59,932
  2. 2020Language Models are Few-Shot Learnerscited by 32,199
  3. 2023LLaMA: Open and Efficient Foundation Language Modelscited by 13,958
  4. 2022PaLM: Scaling Language Modeling with Pathwayscited by 5,304
  5. 2017Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networkscited by 4,861
  6. 2022Locating and Editing Factual Associations in GPTcited by 2,609
  7. 2021Pre-train, Prompt, and Predict: A Systematic Survey of Prompting Methods in Natural Language Processingcited by 2,222
  8. 2019HuggingFace's Transformers: State-of-the-art Natural Language Processingcited by 2,201
  9. 2020Transformer Feed-Forward Layers Are Key-Value Memoriescited by 1,307
  10. 2022Transformers learn in-context by gradient descentcited by 714
  11. 2022What Can Transformers Learn In-Context? A Case Study of Simple Function Classescited by 700
  12. 2022What learning algorithm is in-context learning? Investigations with linear modelscited by 671
16 more
  1. 2021Knowledge Neurons in Pretrained Transformerscited by 590
  2. 2021MetaICL: Learning to Learn In Contextcited by 450
  3. 2023Trained Transformers Learn Linear Models In-Contextcited by 331
  4. 2023Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selectioncited by 291
  5. 2023Transformers learn to implement preconditioned gradient descent for in-context learningcited by 274
  6. 2023Transformers as Algorithms: Generalization and Stability in In-context Learningcited by 220
  7. 2023One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attentioncited by 168
  8. 2021Meta-learning via Language Model In-context Tuningcited by 163
  9. 2023Pretraining task diversity and the emergence of non-Bayesian in-context learning for regressioncited by 138
  10. 2023Supervised Pretraining Can Learn In-Context Reinforcement Learningcited by 137
  11. 2023How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?cited by 93
  12. 2023Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretrainingcited by 82
  13. 2023In-Context Learning through the Bayesian Prismcited by 78
  14. 2023How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representationscited by 72
  15. 2023Understanding In-Context Learning in Transformers and LLMs by Learning to Learn Discrete Functionscited by 62
  16. 2023Transformers can optimally learn regression mixture modelscited by 15

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2024: 11; 2025: 11; 2026: 8.

30 papers.

Sort:

By area:

  1. 2024Asymptotic theory of in-context learning by linear attentioncited by 54
  2. 2024Active-Dormant Attention Heads: Mechanistically Demystifying Extreme-Token Phenomena in LLMscited by 54
  3. 2025Training Dynamics of In-Context Learning in Linear Attentioncited by 44
  4. 2024Transformers are Minimax Optimal Nonparametric In-Context Learnerscited by 36
  5. 2024Pretrained transformer efficiently learns low-dimensional target functions in-contextcited by 29
  6. 2024On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixturescited by 15
  7. 2024How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regressioncited by 13
  8. 2024Transformers Handle Endogeneity in In-Context Linear Regressioncited by 8
  9. 2025On the Role of Transformer Feed-Forward Layers in Nonlinear In-Context Learningcited by 7
  10. 2025Pretrain-Test Task Alignment Governs Generalization in In-Context Learningcited by 6
  11. 2024Borrowing Treasures from Neighbors: In-Context Learning for Multimodal Learning with Missing Modalities and Data Scarcitycited by 5
  12. 2024On Mesa-Optimization in Autoregressively Trained Transformers: Emergence and Capabilitycited by 5
18 more
  1. 2024In-Context Learning of Linear Systems: Generalization Theory and Applications to Operator Learningcited by 5
  2. 2025Out-of-Distribution Generalization of In-Context Learning: A Low-Dimensional Subspace Perspectivecited by 5
  3. 2025On the Robustness of Transformers against Context Hijacking for Linear Classificationcited by 4
  4. 2025In-Context Occam's Razor: How Transformers Prefer Simpler Hypotheses on the Flycited by 4
  5. 2025Optimization-Inspired Few-Shot Adaptation for Large Language Modelscited by 3
  6. 2025Test time training enhances in-context learning of nonlinear functionscited by 3
  7. 2025The Initialization Determines Whether In-Context Learning Is Gradient Descentcited by 3
  8. 2026Rethinking Zero-Shot Time Series Classification: From Task-specific Classifiers to In-Context Inferencecited by 3
  9. 2025Softmax ≥ Linear: Transformers may learn to classify in-context by kernel gradient descentcited by 2
  10. 2026Unlabeled Data Can Provably Enhance In-Context Learning of Transformerscited by 2
  11. 2026Multi-layer Cross-attention is Provably Optimal for Multi-modal In-context Learningcited by 2
  12. 2026Understanding Contextual Recall in Transformers: How Finetuning Enables In-Context Reasoning over Pretraining Knowledgecited by 2
  13. 2024In-context Learning for Mixture of Linear Regressions: Existence, Generalization and Training Dynamicscited by 1
  14. 2025In-Context Learning of Linear Dynamical Systems with Transformers: Approximation Bounds and Depth-Separationcited by 1
  15. 2026Data Distribution as a Lever for Guiding Optimizers Toward Superior Generalization in LLMscited by 0
  16. 2026Transformers Can Implement Preconditioned Richardson Iteration for In-Context Gaussian Kernel Regressioncited by 0
  17. 2026Unbiased Gradients, Moving Stability Boundaries: Exact Mini-Batch Geometry in Linear Self-Attentioncited by 0
  18. 2026Understanding Generalization and Forgetting in In-Context Continual Learningcited by 0