How Do Transformers Learn In-Context Beyond Simple Functions? A Case Study on Learning with Representations

TL;DR

This paper investigates how transformer models perform in in-context learning with complex, compositional representations, combining theoretical analysis and empirical experiments to reveal underlying mechanisms.

Contribution

It introduces synthetic tasks with compositional structure, demonstrates transformers can implement optimal algorithms, and uncovers mechanisms like copying and representation selection.

Findings

Transformers achieve near-optimal ICL performance on complex tasks.

Lower layers transform data while upper layers perform linear ICL.

Mechanisms such as copying and representation selection are observed in trained models.

Abstract

While large language models based on the transformer architecture have demonstrated remarkable in-context learning (ICL) capabilities, understandings of such capabilities are still in an early stage, where existing theory and mechanistic understanding focus mostly on simple scenarios such as learning simple function classes. This paper takes initial steps on understanding ICL in more complex scenarios, by studying learning with representations. Concretely, 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, composed with a linear function that differs in each instance. By construction, the optimal ICL algorithm first transforms the inputs by the representation function, and then performs linear ICL on top of the transformed dataset. We show theoretically the existence of transformers that approximately implement such algorithms with mild depth and size. Empirically, we find trained transformers consistently achieve near-optimal ICL performance in this setting, and exhibit the desired dissection where lower layers transforms the dataset and upper layers perform linear ICL. Through extensive probing and a new pasting experiment, we further reveal several mechanisms within the trained transformers, such as concrete copying behaviors on both the inputs and the representations, linear ICL capability of the upper layers alone, and a post-ICL representation selection mechanism in a harder mixture setting. These observed mechanisms align well with our theory and may shed light on how transformers perform ICL in more realistic scenarios.