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

TL;DR

This paper systematically examines how task complexity and model architecture influence in-context learning, revealing architecture-specific strengths and the importance of curriculum learning for complex tasks.

Contribution

It introduces new tasks and evaluates multiple architectures, demonstrating how model design impacts ICL performance and generalization.

Findings

Transformer performs robustly across tasks

Mamba excels in temporal dynamics

Hyena captures long-range dependencies

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.