Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

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

Haoyuan Sun, Ali Jadbabaie, Navid Azizan

TL;DR

The paper addresses the gap in understanding nonlinear in-context learning (ICL) in Transformer models. It demonstrates that linear self-attention (LSA) alone cannot capture nonlinear target functions, and shows that a bilinear feed-forward (GLU-like) layer enables nonlinear ICL by realizing one-step gradient descent in a quadratic kernel feature space; however, a single block is fundamentally limited by its embedding dimension. By stacking multiple bilinear blocks, deep Transformers perform block-coordinate descent across a richer feature space, enabling learning of higher-degree polynomials in-context (quadratic and cubic examples), with theoretical bounds and approximate optimality analyses. The key finding is that feed-forward layers provide a scalable, essential mechanism for nonlinear ICL in Transformers, and depth distributes nonlinear kernel computation to overcome single-block bottlenecks, offering insights for designing models with stronger in-context learning capabilities. These results have implications for understanding and engineering Transformers to perform task adaptations from prompts without weight updates, particularly for complex nonlinear tasks.

Abstract

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.

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

TL;DR

The paper addresses the gap in understanding nonlinear in-context learning (ICL) in Transformer models. It demonstrates that linear self-attention (LSA) alone cannot capture nonlinear target functions, and shows that a bilinear feed-forward (GLU-like) layer enables nonlinear ICL by realizing one-step gradient descent in a quadratic kernel feature space; however, a single block is fundamentally limited by its embedding dimension. By stacking multiple bilinear blocks, deep Transformers perform block-coordinate descent across a richer feature space, enabling learning of higher-degree polynomials in-context (quadratic and cubic examples), with theoretical bounds and approximate optimality analyses. The key finding is that feed-forward layers provide a scalable, essential mechanism for nonlinear ICL in Transformers, and depth distributes nonlinear kernel computation to overcome single-block bottlenecks, offering insights for designing models with stronger in-context learning capabilities. These results have implications for understanding and engineering Transformers to perform task adaptations from prompts without weight updates, particularly for complex nonlinear tasks.

Abstract

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.

Paper Structure

This paper contains 31 sections, 13 theorems, 127 equations, 10 figures, 1 algorithm.

Table of Contents

  1. Introduction
  2. Our contributions
  3. Related works
  4. Problem Setting
  5. In-context learning
  6. Transformer architecture
  7. In-context learning objective
  8. Solving Quadratic In-context Learning
  9. Limitations of linear self-attention in nonlinear ICL
  10. Quadratic ICL with bilinear feed-forward layers
  11. Analysis of the optimal bilinear Transformer block
  12. Deep Bilinear Transformers
  13. Deep bilinear Transformers as block-coordinate descent solvers
  14. Learning higher-order polynomials in-context
  15. Conclusion and Future Work
  16. ...and 16 more sections

Key Result

Proposition 1

Consider any fixed target function $f$ and inputs $x_1, \dots, x_{n+1} \overset{i.i.d.}{\sim} \mathcal{N}(0, \Sigma)$. The in-context learning prompts are sampled according to the form as described in equ:prompt-form. Then, for any linear Transformer $\mathsf{TF_{lin}}$ (regardless of the number of

Figures (10)

  • Figure 1: Quadratic ICL test loss (over 5 trials) of both linear and bilinear Transformer models given different test prompt lengths. The linear Transformer has 6 layers, and the bilinear Transformers have $2L$ layers, which implement $L$ steps of block-coordinate descent. In Figure (a), the quadratic function has $d=3$ variables, and Figures (b-c) have $d=4$. Deep bilinear models significantly outperform the linear baseline, show improved performance with increased depth ($L$), and can learn the quadratic function even with an embedding dimension ($\bar{d}=12$) insufficient to fit the number of quadratic features (15 for $d=4$).
  • Figure 2: Cubic ICL test loss (over 5 trials) for bilinear Transformers with both dense \ref{['equ:dense-bilinear']} and sparse \ref{['equ:sparse-bilinear']} bilinear layer weights. The differences between the two models confirm our proposed mechanism through which bilinear Transformers can learn higher-order features in context.
  • Figure 3: Quadratic ICL loss of the bilinear Transformer block construction associated with Proposition \ref{['thm:bilinear-construction']}.
  • Figure 4: Quadratic ICL test loss for 8-layer bilinear Transformer models trained on training prompts of length $n_{\rm train} \in \{100, 150, 200, 250\}$, respectively.
  • Figure 5: Loss curve (over 5 trials) for the quadratic ICL task, $d = 4, L = 2$.
  • ...and 5 more figures

Theorems & Definitions (17)

  • Proposition 1
  • Theorem 2
  • Lemma 3
  • Proposition 4
  • Remark 1
  • Remark 2
  • Theorem 5
  • Corollary 6
  • Theorem 7
  • Corollary 8
  • ...and 7 more