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

• URL: http://arxiv.org/abs/2501.18187v3
• Date: Sat, 11 Oct 2025 04:44:57 GMT
• Title: On the Role of Transformer Feed-Forward Layers in Nonlinear In-Context Learning
• Abstract summary: 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.<n>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.
• Score: 31.908611830179108
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Plug-and-Play Spiking Operators: Breaking the Nonlinearity Bottleneck in Spiking Transformers [32.743518029177146]
  We propose a plug-and-play framework that implements spike-friendly approximations for Transformer nonlinearities and integrates into existing ANN-to-SNN pipelines.<n>Our method decomposes these nonlinear computations into three recurring primitives -- division, exponentiation, and $ell$ norms -- and realizes them via population groups.<n>Experiments show that selectively replacing the targeted nonlinear operators incurs less than a $1% accuracy drop across all evaluated tasks.
  arXiv  Detail & Related papers  (2026-05-19T06:59:46Z)
• Optimality and NP-Hardness of Transformers in Learning Markovian Dynamical Functions [32.71332125930795]
  Transformer architectures can solve unseen tasks based on input-output pairs in a given prompt due to in-context learning (ICL)<n>We investigate Markovian function learning through a structured ICL setup to reveal underlying optimization behaviors.
  arXiv  Detail & Related papers  (2025-10-21T13:42:48Z)
• Provable In-Context Learning of Nonlinear Regression with Transformers [58.018629320233174]
  In-context learning (ICL) is the ability to perform unseen tasks using task-specific prompts without updating parameters.<n>Recent research has actively explored the training dynamics behind ICL.<n>This paper investigates more complex nonlinear regression tasks, aiming to uncover how transformers acquire in-context learning capabilities.
  arXiv  Detail & Related papers  (2025-07-28T00:09:28Z)
• Exact Learning Dynamics of In-Context Learning in Linear Transformers and Its Application to Non-Linear Transformers [1.7034813545878589]
  Transformer models exhibit remarkable in-context learning (ICL) Our work offers an exact dynamical model for ICL and theoretically grounded tools for analyzing complex transformer training.
  arXiv  Detail & Related papers  (2025-04-17T13:05:33Z)
• Re-examining learning linear functions in context [1.8843687952462742]
  In-context learning (ICL) has emerged as a powerful paradigm for easily adapting Large Language Models (LLMs) to various tasks.<n>We explore a simple model of ICL in a controlled setup with synthetic training data.<n>Our findings challenge the prevailing narrative that transformers adopt algorithmic approaches to learn a linear function in-context.
  arXiv  Detail & Related papers  (2024-11-18T10:58:46Z)
• Pretrained transformer efficiently learns low-dimensional target functions in-context [40.77319247558742]
  We show that a nonlinear transformer optimized by gradient descent learns $f_*$ in-context with a prompt length that only depends on the dimension of the distribution of target functions $r$. Our result highlights the adaptivity of the pretrained transformer to low-dimensional structures of the function class, which enables sample-efficient ICL.
  arXiv  Detail & Related papers  (2024-11-04T19:24:39Z)
• Can Looped Transformers Learn to Implement Multi-step Gradient Descent for In-context Learning? [69.4145579827826]
  We show a fast flow on the regression loss despite the gradient non-ity algorithms for our convergence landscape. This is the first theoretical analysis for multi-layer Transformer in this setting.
  arXiv  Detail & Related papers  (2024-10-10T18:29:05Z)
• Asymptotic theory of in-context learning by linear attention [33.53106537972063]
  In-context learning is a cornerstone of Transformers' success. Questions about the necessary sample complexity, pretraining task diversity, and context length for successful ICL remain unresolved.
  arXiv  Detail & Related papers  (2024-05-20T03:24:24Z)
• Your Transformer is Secretly Linear [7.935853865895353]
  We analyze embedding transformations between sequential layers, uncovering a near-perfect linear relationship. We show that removing or linearly approximating some of the most linear blocks of transformers does not affect significantly the loss or model performance. In our pretraining experiments on smaller models we introduce a cosine-similarity-based regularization, aimed at reducing layer linearity.
  arXiv  Detail & Related papers  (2024-05-19T22:44:00Z)
• How Do Nonlinear Transformers Learn and Generalize in In-Context Learning? [82.51626700527837]
  Transformer-based large language models displayed impressive in-context learning capabilities, where a pre-trained model can handle new tasks without fine-tuning. We analyze how the mechanics of how Transformer to achieve ICL contribute to the technical challenges of the training problems in Transformers.
  arXiv  Detail & Related papers  (2024-02-23T21:07:20Z)
• How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations [98.7450564309923]
  This paper takes initial steps on understanding in-context learning (ICL) in more complex scenarios, by studying learning with representations. We construct synthetic in-context learning problems with a compositional structure, where the label depends on the input through a possibly complex but fixed representation function. We show theoretically the existence of transformers that approximately implement such algorithms with mild depth and size.
  arXiv  Detail & Related papers  (2023-10-16T17:40:49Z)
• How Many Pretraining Tasks Are Needed for In-Context Learning of Linear Regression? [92.90857135952231]
  Transformers pretrained on diverse tasks exhibit remarkable in-context learning (ICL) capabilities. We study ICL in one of its simplest setups: pretraining a linearly parameterized single-layer linear attention model for linear regression.
  arXiv  Detail & Related papers  (2023-10-12T15:01:43Z)
• Pessimistic Nonlinear Least-Squares Value Iteration for Offline Reinforcement Learning [53.97335841137496]
  We propose an oracle-efficient algorithm, dubbed Pessimistic Least-Square Value Iteration (PNLSVI) for offline RL with non-linear function approximation. Our algorithm enjoys a regret bound that has a tight dependency on the function class complexity and achieves minimax optimal instance-dependent regret when specialized to linear function approximation.
  arXiv  Detail & Related papers  (2023-10-02T17:42:01Z)
• Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection [88.23337313766353]
  This work first provides a comprehensive statistical theory for transformers to perform ICL. We show that transformers can implement a broad class of standard machine learning algorithms in context. A emphsingle transformer can adaptively select different base ICL algorithms.
  arXiv  Detail & Related papers  (2023-06-07T17:59:31Z)
• Theoretical Characterization of the Generalization Performance of Overfitted Meta-Learning [70.52689048213398]
  This paper studies the performance of overfitted meta-learning under a linear regression model with Gaussian features. We find new and interesting properties that do not exist in single-task linear regression. Our analysis suggests that benign overfitting is more significant and easier to observe when the noise and the diversity/fluctuation of the ground truth of each training task are large.
  arXiv  Detail & Related papers  (2023-04-09T20:36:13Z)
• Sample-Efficient Reinforcement Learning Is Feasible for Linearly Realizable MDPs with Limited Revisiting [60.98700344526674]
  Low-complexity models such as linear function representation play a pivotal role in enabling sample-efficient reinforcement learning. In this paper, we investigate a new sampling protocol, which draws samples in an online/exploratory fashion but allows one to backtrack and revisit previous states in a controlled and infrequent manner. We develop an algorithm tailored to this setting, achieving a sample complexity that scales practicallyly with the feature dimension, the horizon, and the inverse sub-optimality gap, but not the size of the state/action space.
  arXiv  Detail & Related papers  (2021-05-17T17:22:07Z)
• Sparse Quantized Spectral Clustering [85.77233010209368]
  We exploit tools from random matrix theory to make precise statements about how the eigenspectrum of a matrix changes under such nonlinear transformations. We show that very little change occurs in the informative eigenstructure even under drastic sparsification/quantization.
  arXiv  Detail & Related papers  (2020-10-03T15:58:07Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.