# On the Role of Transformer Feed-Forward Layers in Nonlinear In-Context Learning

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

## Key facts

- Cited by 7 papers in the map (data through 14 Oct 2026).
- Builds on 40 of its 47 references that are in the map.
- First posted on arXiv 30 Jan 2025.

## Builds on (most cited first, 30 of 40)

- [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)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [GPT-4 Technical Report](https://mlsanity.com/paper/2303.08774/) (2023)
- [Neural Machine Translation by Jointly Learning to Align and Translate](https://mlsanity.com/paper/1409.0473/) (2014)
- [Robust Speech Recognition via Large-Scale Weak Supervision](https://mlsanity.com/paper/2212.04356/) (2022)
- [Emergent Abilities of Large Language Models](https://mlsanity.com/paper/2206.07682/) (2022)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality](https://mlsanity.com/paper/2405.21060/) (2024)
- [In-context Learning and Induction Heads](https://mlsanity.com/paper/2209.11895/) (2022)
- [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)
- [Language Modeling with Gated Convolutional Networks](https://mlsanity.com/paper/1612.08083/) (2016)
- [What learning algorithm is in-context learning? Investigations with linear models](https://mlsanity.com/paper/2211.15661/) (2022)
- [The Dawn of LMMs: Preliminary Explorations with GPT-4V(ision)](https://mlsanity.com/paper/2309.17421/) (2023)
- [Gated Delta Networks: Improving Mamba2 with Delta Rule](https://mlsanity.com/paper/2412.06464/) (2024)
- [Gated Linear Attention Transformers with Hardware-Efficient Training](https://mlsanity.com/paper/2312.06635/) (2023)
- [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)
- [Looped Transformers as Programmable Computers](https://mlsanity.com/paper/2301.13196/) (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)
- [MiniMax-01: Scaling Foundation Models with Lightning Attention](https://mlsanity.com/paper/2501.08313/) (2025)
- [Inductive Biases and Variable Creation in Self-Attention Mechanisms](https://mlsanity.com/paper/2110.10090/) (2021)
- [The mechanistic basis of data dependence and abrupt learning in an in-context classification task](https://mlsanity.com/paper/2312.03002/) (2023)
- [How Transformers Learn Causal Structure with Gradient Descent](https://mlsanity.com/paper/2402.14735/) (2024)
- [In-Context Convergence of Transformers](https://mlsanity.com/paper/2310.05249/) (2023)
- [Statistically Meaningful Approximation: a Case Study on Approximating Turing Machines with Transformers](https://mlsanity.com/paper/2107.13163/) (2021)

## Built on by (most cited first, 7 of 7)

- [Gating is Weighting: Understanding Gated Linear Attention through In-context Learning](https://mlsanity.com/paper/2504.04308/) (2025)
- [Understanding In-Context Learning for Nonlinear Regression with Transformers: Attention as Featurizer](https://mlsanity.com/paper/2605.05176/) (2026)
- [Mamba Can Learn Low-Dimensional Targets In-Context via Test-Time Feature Learning](https://mlsanity.com/paper/2510.12026/) (2025)
- [Bayesian Optimality of In-Context Learning with Selective State Spaces](https://mlsanity.com/paper/2602.17744/) (2026)
- [Transformers Can Implement Preconditioned Richardson Iteration for In-Context Gaussian Kernel Regression](https://mlsanity.com/paper/2605.08475/) (2026)
- [Why LLMs Fail at Causal Discovery and How Interventional Agents Escape](https://mlsanity.com/paper/2605.27567/) (2026)
- [When Does Depth Matter For In-Context Learning? Adaptive Inference in Deep Transformers](https://mlsanity.com/paper/2606.16694/) (2026)

## Abstract (arXiv)

Transformer-based models demonstrate a remarkable ability for in-context learning (ICL), where they can adapt to unseen tasks from a few prompt examples without parameter updates. Recent research has illuminated how Transformers perform ICL, showing that the optimal linear self-attention (LSA) mechanism can implement one step of gradient descent for linear least-squares objectives when trained on random linear regression tasks. Building on this, we investigate ICL for nonlinear function classes. We first prove that LSA is inherently incapable of outperforming linear predictors on nonlinear tasks, underscoring why prior solutions cannot readily extend to these problems. To overcome this limitation, we analyze a Transformer block consisting of LSA and feed-forward layers inspired by the gated linear units (GLU), which is a standard component of modern Transformers. We show that this block achieves nonlinear ICL by implementing one step of gradient descent on a polynomial kernel regression loss. Furthermore, our analysis reveals that the expressivity of a single block is inherently limited by its dimensions. We then show that a deep Transformer can overcome this bottleneck by distributing the computation of richer kernel functions across multiple blocks, performing block-coordinate descent in a high-dimensional feature space that a single block cannot represent. Our findings highlight that the feed-forward layers provide a crucial and scalable mechanism by which Transformers can express nonlinear representations for ICL.