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

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. Notably, recent research has provided insight into how the Transformer architecture can 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 upon this understanding of linear ICL, we investigate ICL for nonlinear function classes. We first show that LSA is inherently incapable of solving problems that go beyond linear least-squares objectives, underscoring why prior solutions cannot readily extend to nonlinear ICL tasks. To overcome this limitation, we investigate a mechanism combining LSA with feed-forward layers that are inspired by the gated linear units (GLU) commonly found in modern Transformer architectures. We show that this combination empowers the Transformer to perform nonlinear ICL, specifically by implementing one step of gradient descent on a polynomial kernel regression loss. Furthermore, we show that multiple blocks of our GLU-LSA model implement block coordinate descent in this polynomial kernel space. Our findings highlight the distinct roles of attention and feed-forward layers, demonstrating that the feed-forward components provide a mechanism by which Transformers gain nonlinear capabilities for ICL.