Distinct Computations Emerge From Compositional Curricula in In-Context Learning

TL;DR

This paper explores how presenting a compositional subtask curriculum in in-context learning influences the computations learned by transformers, enabling zero-shot generalization and increased robustness.

Contribution

It demonstrates that a subtask curriculum in in-context learning leads to distinct computations and better generalization compared to direct training on complex tasks.

Findings

Models trained with a subtask curriculum can perform zero-shot inference on unseen tasks.

Curriculum-trained models are more robust with the same context length.

Different training regimes lead to diverse task representations.

Abstract

In-context learning (ICL) research often considers learning a function in-context through a uniform sample of input-output pairs. Here, we investigate how presenting a compositional subtask curriculum in context may alter the computations a transformer learns. We design a compositional algorithmic task based on the modular exponential-a double exponential task composed of two single exponential subtasks and train transformer models to learn the task in-context. We compare (a) models trained using an in-context curriculum consisting of single exponential subtasks and, (b) models trained directly on the double exponential task without such a curriculum. We show that models trained with a subtask curriculum can perform zero-shot inference on unseen compositional tasks and are more robust given the same context length. We study how the task and subtasks are represented across the two training regimes. We find that the models employ diverse strategies modulated by the specific curriculum design.