Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Probing In-Context Learning: Impact of Task Complexity and Model Architecture on Generalization and Efficiency

Binwen Liu, Peiyu Xu, Quan Yuan, Yihong Chen

TL;DR

The paper investigates in-context learning (ICL) by varying task complexity (linear, Gaussian kernel, nonlinear dynamics) and model architectures (GPT-2 Transformer, FlashAttention, Hyena, Mamba). Using synthetic data and a curriculum that scales context and input dimensionality, it demonstrates that architecture exerts a strong influence on ICL performance, with Mamba excelling in temporally structured tasks and Transformers offering broad generalization, while Hyena and FlashAttention present trade-offs in efficiency and low-data sensitivity. Key mechanisms identified include locality-induced shortcuts in kernel tasks, improved nonlinear separability via input scaling, and curriculum-driven alignment that enhances convergence. These results offer practical guidance for selecting and combining architectural biases to improve robust ICL across diverse function classes.

Abstract

We investigate in-context learning (ICL) through a meticulous experimental framework that systematically varies task complexity and model architecture. Extending beyond the linear regression baseline, we introduce Gaussian kernel regression and nonlinear dynamical system tasks, which emphasize temporal and recursive reasoning. We evaluate four distinct models: a GPT2-style Transformer, a Transformer with FlashAttention mechanism, a convolutional Hyena-based model, and the Mamba state-space model. Each model is trained from scratch on synthetic datasets and assessed for generalization during testing. Our findings highlight that model architecture significantly shapes ICL performance. The standard Transformer demonstrates robust performance across diverse tasks, while Mamba excels in temporally structured dynamics. Hyena effectively captures long-range dependencies but shows higher variance early in training, and FlashAttention offers computational efficiency but is more sensitive in low-data regimes. Further analysis uncovers locality-induced shortcuts in Gaussian kernel tasks, enhanced nonlinear separability through input range scaling, and the critical role of curriculum learning in mastering high-dimensional tasks.

Probing In-Context Learning: Impact of Task Complexity and Model Architecture on Generalization and Efficiency

TL;DR

The paper investigates in-context learning (ICL) by varying task complexity (linear, Gaussian kernel, nonlinear dynamics) and model architectures (GPT-2 Transformer, FlashAttention, Hyena, Mamba). Using synthetic data and a curriculum that scales context and input dimensionality, it demonstrates that architecture exerts a strong influence on ICL performance, with Mamba excelling in temporally structured tasks and Transformers offering broad generalization, while Hyena and FlashAttention present trade-offs in efficiency and low-data sensitivity. Key mechanisms identified include locality-induced shortcuts in kernel tasks, improved nonlinear separability via input scaling, and curriculum-driven alignment that enhances convergence. These results offer practical guidance for selecting and combining architectural biases to improve robust ICL across diverse function classes.

Abstract

We investigate in-context learning (ICL) through a meticulous experimental framework that systematically varies task complexity and model architecture. Extending beyond the linear regression baseline, we introduce Gaussian kernel regression and nonlinear dynamical system tasks, which emphasize temporal and recursive reasoning. We evaluate four distinct models: a GPT2-style Transformer, a Transformer with FlashAttention mechanism, a convolutional Hyena-based model, and the Mamba state-space model. Each model is trained from scratch on synthetic datasets and assessed for generalization during testing. Our findings highlight that model architecture significantly shapes ICL performance. The standard Transformer demonstrates robust performance across diverse tasks, while Mamba excels in temporally structured dynamics. Hyena effectively captures long-range dependencies but shows higher variance early in training, and FlashAttention offers computational efficiency but is more sensitive in low-data regimes. Further analysis uncovers locality-induced shortcuts in Gaussian kernel tasks, enhanced nonlinear separability through input range scaling, and the critical role of curriculum learning in mastering high-dimensional tasks.
Paper Structure (48 sections, 4 equations, 18 figures)

This paper contains 48 sections, 4 equations, 18 figures.

Table of Contents

  1. Introduction
  2. Related Work
  3. Approach
  4. Problem Setup
  5. Function Families
  6. Linear Regression.
  7. Gaussian Kernel Regression.
  8. Nonlinear Dynamical Systems.
  9. Model Architectures
  10. Training Procedure
  11. Curriculum Learning
  12. Evaluation Criteria
  13. Baseline Estimators.
  14. Experiments
  15. Task and Dataset
  16. ...and 33 more sections

Figures (18)

  • Figure 1: FlashAttention mechanism dao2022flashattention. The design tiles attention computation to avoid memory bottlenecks, achieving high throughput on modern hardware.
  • Figure 2: Hyena recurrence poli2023hyena. Combines implicit long convolutions with multiplicative gating, allowing attention-like behavior without quadratic cost.
  • Figure 3: Mamba architecture gu2023mamba. Uses state-space sequence modeling (SSM) with gating and convolution to replace self-attention.
  • Figure 4: Summary of GPT-2 results on linear and Gaussian kernel regression tasks.
  • Figure 5: Results of Nonlinear Dynamics Trained with GPT-2 Architecture
  • ...and 13 more figures