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

Fuente: arXiv
Guardado en:
Detalles Bibliográficos
Autores principales: Song, Jiajun, Xu, Zhuoyan, Zhong, Yiqiao
Formato: Preprint
Publicado: 2024
Materias:
Acceso en línea:
Etiquetas: Agregar Etiqueta
Sin Etiquetas, Sea el primero en etiquetar este registro!
_version_ 1866909442913599488
author Song, Jiajun
Xu, Zhuoyan
Zhong, Yiqiao
author_facet Song, Jiajun
Xu, Zhuoyan
Zhong, Yiqiao
contents 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.
format Preprint
id arxiv_https___arxiv_org_abs_2408_09503
institution arXiv
publishDate 2024
record_format arxiv
spellingShingle Out-of-distribution generalization via composition: a lens through induction heads in Transformers
Song, Jiajun
Xu, Zhuoyan
Zhong, Yiqiao
Computation and Language
Artificial Intelligence
Machine Learning
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.
title Out-of-distribution generalization via composition: a lens through induction heads in Transformers
topic Computation and Language
Artificial Intelligence
Machine Learning
url https://arxiv.org/abs/2408.09503