Understanding the Staged Dynamics of Transformers in Learning Latent Structure

TL;DR

This paper investigates how small transformer models learn latent structures in language tasks, revealing staged learning processes, asymmetries, and layer-specific plasticity windows through controlled experiments.

Contribution

It introduces a detailed analysis of the stages and mechanisms by which transformers acquire latent structure, using the Alchemy benchmark in a controlled setting.

Findings

Transformers learn latent structure components in discrete stages.

The model robustly composes fundamental transitions but struggles with complex decomposition.

Layer-specific plasticity windows influence stage completion and learning dynamics.

Abstract

Language modeling has shown us that transformers can discover latent structure from context, but the dynamics of how they acquire different components of that structure remain poorly understood, leading to assertions that models just remix training data. In this work, we use the Alchemy benchmark in a controlled setting (Wang et al.,2021) to investigate latent structure learning. We train a small decoder-only transformer on three task variants: 1) inferring missing transitions from partial contextual information, 2) composing simple rules to solve multi-transition sequences, and 3) decomposing complex multi-step examples to infer intermediate transitions. By factorizing each task into interpretable components, we show that the model learns the different latent structure components in discrete stages. We also observe an asymmetry: the model composes fundamental transitions robustly, but struggles to decompose complex examples to discover the atomic transitions. Finally, using causal interventions, we identify layer-specific plasticity windows during which freezing substantially delays or prevents stage completion. These findings provide insight into how a transformer model acquires latent structure, offering a detailed view of how capabilities evolve during training.