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

arXiv: https://arxiv.org/abs/2402.14951 · HTML: https://mlsanity.com/paper/2402.14951/

## Key facts

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

## Builds on (most cited first, 28 of 28)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Locating and Editing Factual Associations in GPT](https://mlsanity.com/paper/2202.05262/) (2022)
- [Pre-train, Prompt, and Predict: A Systematic Survey of Prompting Methods in Natural Language Processing](https://mlsanity.com/paper/2107.13586/) (2021)
- [HuggingFace's Transformers: State-of-the-art Natural Language Processing](https://mlsanity.com/paper/1910.03771/) (2019)
- [Transformer Feed-Forward Layers Are Key-Value Memories](https://mlsanity.com/paper/2012.14913/) (2020)
- [Transformers learn in-context by gradient descent](https://mlsanity.com/paper/2212.07677/) (2022)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [Knowledge Neurons in Pretrained Transformers](https://mlsanity.com/paper/2104.08696/) (2021)
- [MetaICL: Learning to Learn In Context](https://mlsanity.com/paper/2110.15943/) (2021)
- [Trained Transformers Learn Linear Models In-Context](https://mlsanity.com/paper/2306.09927/) (2023)
- [Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection](https://mlsanity.com/paper/2306.04637/) (2023)
- [Transformers learn to implement preconditioned gradient descent for in-context learning](https://mlsanity.com/paper/2306.00297/) (2023)
- [Transformers as Algorithms: Generalization and Stability in In-context Learning](https://mlsanity.com/paper/2301.07067/) (2023)
- [One Step of Gradient Descent is Provably the Optimal In-Context Learner with One Layer of Linear Self-Attention](https://mlsanity.com/paper/2307.03576/) (2023)
- [Meta-learning via Language Model In-context Tuning](https://mlsanity.com/paper/2110.07814/) (2021)
- [Pretraining task diversity and the emergence of non-Bayesian in-context learning for regression](https://mlsanity.com/paper/2306.15063/) (2023)
- [Supervised Pretraining Can Learn In-Context Reinforcement Learning](https://mlsanity.com/paper/2306.14892/) (2023)
- [How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression?](https://mlsanity.com/paper/2310.08391/) (2023)
- [In-Context Learning through the Bayesian Prism](https://mlsanity.com/paper/2306.04891/) (2023)
- [Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining](https://mlsanity.com/paper/2310.08566/) (2023)
- [How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations](https://mlsanity.com/paper/2310.10616/) (2023)
- [Understanding In-Context Learning in Transformers and LLMs by Learning to Learn Discrete Functions](https://mlsanity.com/paper/2310.03016/) (2023)
- [Transformers can optimally learn regression mixture models](https://mlsanity.com/paper/2311.08362/) (2023)

## Built on by (most cited first, 30 of 33)

- [Active-Dormant Attention Heads: Mechanistically Demystifying Extreme-Token Phenomena in LLMs](https://mlsanity.com/paper/2410.13835/) (2024)
- [Asymptotic theory of in-context learning by linear attention](https://mlsanity.com/paper/2405.11751/) (2024)
- [Training Dynamics of In-Context Learning in Linear Attention](https://mlsanity.com/paper/2501.16265/) (2025)
- [Transformers are Minimax Optimal Nonparametric In-Context Learners](https://mlsanity.com/paper/2408.12186/) (2024)
- [Pretrained transformer efficiently learns low-dimensional target functions in-context](https://mlsanity.com/paper/2411.02544/) (2024)
- [On the Training Convergence of Transformers for In-Context Classification of Gaussian Mixtures](https://mlsanity.com/paper/2410.11778/) (2024)
- [How Transformers Utilize Multi-Head Attention in In-Context Learning? A Case Study on Sparse Linear Regression](https://mlsanity.com/paper/2408.04532/) (2024)
- [Transformers Handle Endogeneity in In-Context Linear Regression](https://mlsanity.com/paper/2410.01265/) (2024)
- [Pretrain-Test Task Alignment Governs Generalization in In-Context Learning](https://mlsanity.com/paper/2509.26551/) (2025)
- [On the Role of Transformer Feed-Forward Layers in Nonlinear In-Context Learning](https://mlsanity.com/paper/2501.18187/) (2025)
- [In-Context Learning of Linear Systems: Generalization Theory and Applications to Operator Learning](https://mlsanity.com/paper/2409.12293/) (2024)
- [Out-of-Distribution Generalization of In-Context Learning: A Low-Dimensional Subspace Perspective](https://mlsanity.com/paper/2505.14808/) (2025)
- [Borrowing Treasures from Neighbors: In-Context Learning for Multimodal Learning with Missing Modalities and Data Scarcity](https://mlsanity.com/paper/2403.09428/) (2024)
- [On Mesa-Optimization in Autoregressively Trained Transformers: Emergence and Capability](https://mlsanity.com/paper/2405.16845/) (2024)
- [In-Context Occam's Razor: How Transformers Prefer Simpler Hypotheses on the Fly](https://mlsanity.com/paper/2506.19351/) (2025)
- [In-Context Learning of Linear Dynamical Systems with Transformers: Approximation Bounds and Depth-Separation](https://mlsanity.com/paper/2502.08136/) (2025)
- [On the Robustness of Transformers against Context Hijacking for Linear Classification](https://mlsanity.com/paper/2502.15609/) (2025)
- [Softmax ≥ Linear: Transformers may learn to classify in-context by kernel gradient descent](https://mlsanity.com/paper/2510.10425/) (2025)
- [Optimization-Inspired Few-Shot Adaptation for Large Language Models](https://mlsanity.com/paper/2505.19107/) (2025)
- [Test time training enhances in-context learning of nonlinear functions](https://mlsanity.com/paper/2509.25741/) (2025)
- [The Initialization Determines Whether In-Context Learning Is Gradient Descent](https://mlsanity.com/paper/2512.04268/) (2025)
- [Rethinking Zero-Shot Time Series Classification: From Task-specific Classifiers to In-Context Inference](https://mlsanity.com/paper/2602.00620/) (2026)
- [Understanding Contextual Recall in Transformers: How Finetuning Enables In-Context Reasoning over Pretraining Knowledge](https://mlsanity.com/paper/2603.20969/) (2026)
- [Unlabeled Data Can Provably Enhance In-Context Learning of Transformers](https://mlsanity.com/paper/2601.10058/) (2026)
- [Multi-layer Cross-attention is Provably Optimal for Multi-modal In-context Learning](https://mlsanity.com/paper/2602.04872/) (2026)
- [Looped Transformers with Layer Normalization Provably Learn the Power Method](https://mlsanity.com/paper/2606.00605/) (2026)
- [In-context Learning for Mixture of Linear Regressions: Existence, Generalization and Training Dynamics](https://mlsanity.com/paper/2410.14183/) (2024)
- [Transformers Efficiently Perform In-Context Logistic Regression via Normalized Gradient Descent](https://mlsanity.com/paper/2605.06609/) (2026)
- [Transformers Can Implement Preconditioned Richardson Iteration for In-Context Gaussian Kernel Regression](https://mlsanity.com/paper/2605.08475/) (2026)
- [Unbiased Gradients, Moving Stability Boundaries: Exact Mini-Batch Geometry in Linear Self-Attention](https://mlsanity.com/paper/2605.21292/) (2026)

## 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.