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

Transformers have revolutionized natural language processing, demonstrating remarkable in-context learning abilities without parameter updates. While much progress has been made in understanding their linear function learning, their capacity for nonlinear functions remains less clear. This paper addresses this gap by analyzing the fundamental limitations of linear self-attention (LSA) in nonlinear tasks, establishing that it cannot outperform simple linear predictors. To overcome this, the authors introduce a novel Transformer architecture combining LSA with GLU-like feed-forward layers, enabling the model to emulate gradient descent on polynomial kernels. Theoretical analysis proves that a single-layer model is constrained by its embedding dimension, limiting its expressivity for high-order functions. However, stacking multiple layers allows the distributed computation of richer kernel features, effectively performing block-coordinate descent in a high-dimensional feature space. Numerical experiments on random quadratic functions confirm that deep models significantly outperform shallow ones, with errors decreasing as O(1/n). These findings highlight the critical role of feed-forward layers in enabling nonlinear representations within Transformers, providing a scalable pathway for learning complex functions in context. The work bridges kernel theory and deep learning, offering new insights into the design of expressive, efficient models capable of tackling sophisticated nonlinear tasks in various domains. Future directions include extending the framework to broader function classes, improving training efficiency, and validating on real-world large-scale datasets, ultimately pushing the boundaries of what Transformer models can achieve in nonlinear in-context learning.