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

TL;DR

This paper investigates how Transformer feed-forward layers enable nonlinear in-context learning by implementing gradient descent on kernel regression losses, revealing the importance of depth and expressivity in such models.

Contribution

It demonstrates that feed-forward layers are essential for nonlinear ICL, showing how deep Transformers distribute kernel computations to overcome single-block limitations.

Findings

Feed-forward layers enable nonlinear ICL via gradient descent on kernel losses.

Single Transformer blocks have limited expressivity due to their dimensions.

Deep Transformers distribute computations across blocks to achieve richer nonlinear representations.

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.