Out-of-distribution generalization via composition: a lens through induction heads in Transformers

TL;DR

This paper investigates how large language models generalize to out-of-distribution tasks by composing self-attention layers, revealing a shared latent space that facilitates rule inference and OOD generalization.

Contribution

It uncovers the role of induction heads and a shared latent subspace in enabling OOD generalization through compositional mechanisms in Transformers.

Findings

Models learn rules by composing self-attention layers.

A shared latent subspace aligns early and later layers.

Composition via induction heads improves OOD generalization.

Abstract

Large language models (LLMs) such as GPT-4 sometimes appear to be creative, solving novel tasks often with a few demonstrations in the prompt. These tasks require the models to generalize on distributions different from those from training data -- which is known as out-of-distribution (OOD) generalization. Despite the tremendous success of LLMs, how they approach OOD generalization remains an open and underexplored question. We examine OOD generalization in settings where instances are generated according to hidden rules, including in-context learning with symbolic reasoning. Models are required to infer the hidden rules behind input prompts without any fine-tuning. We empirically examined the training dynamics of Transformers on a synthetic example and conducted extensive experiments on a variety of pretrained LLMs, focusing on a type of components known as induction heads. We found that OOD generalization and composition are tied together -- models can learn rules by composing two self-attention layers, thereby achieving OOD generalization. Furthermore, a shared latent subspace in the embedding (or feature) space acts as a bridge for composition by aligning early layers and later layers, which we refer to as the common bridge representation hypothesis.