Understanding the Staged Dynamics of Transformers in Learning Latent Structure

1. Problem Statement

While large language models (transformers) have demonstrated the ability to discover latent structures (e.g., syntax, causal graphs) from context, the dynamics of how they acquire these capabilities remain poorly understood. Existing theories often suggest models merely "remix" training data rather than truly generalizing. Furthermore, previous studies on "staged learning" (where performance jumps in discrete steps rather than continuously) occur in complex settings where it is difficult to attribute specific learning stages to specific structural components.

The authors aim to:

1. Investigate the precise progression of latent structure learning in a controlled environment.
2. Determine if learning occurs in discrete stages corresponding to specific sub-skills.
3. Analyze the asymmetry between composition (combining simple rules) and decomposition (inferring atomic rules from complex examples).
4. Identify the plasticity windows (specific timeframes and layers) required for these stages to occur using causal interventions.

2. Methodology

Dataset: The Alchemy Benchmark

The authors utilize the Alchemy benchmark (Wang et al., 2021), a meta-reinforcement learning task with a well-defined latent structure.

• Structure: A "chemistry" is a directed graph of 8 stone states arranged in a cubic structure.
• Features: Stones have 4 features: Color, Size, Roundness, and Reward.
• Transitions: Applying "potions" (edges) transitions the stone between states. Crucially, a single potion may change multiple perceptual features simultaneously, but follows a consistent latent rule across all chemistries.
• Control: The authors use complete chemistries (no missing edges in the ground truth) to ensure the model must learn the underlying transition rules rather than memorizing specific paths.

Model Architecture

• A small decoder-only transformer (~2M parameters).
• 4 hidden layers, embedding dimension 256, feedforward dimension 512.
• Trained for 1000 epochs with a 90/10 train/validation split.

Experimental Tasks

The authors factorize the task into three specific formulations to isolate different learning dynamics:

1. Withheld Potion Pair (Latent Discovery): The model sees transitions for all potion pairs except one. It must infer the effect of the missing pair. This tests the ability to align two disconnected halves of the graph.
2. Composition: The model is given all 1-hop (single-step) transitions and must solve multi-hop queries (2 to 5 hops). This tests the ability to compose known rules.
3. Decomposition: The model is given multi-hop support examples (2 to 5 hops) and must infer the underlying 1-hop transition rules to solve a query. This tests the ability to decompose complex sequences.

Analytical Framework: Event Factorization

To understand what is being learned at each stage, the authors decompose the overall accuracy ($P[C]$) into conditional probabilities of interpretable events:

• Event A (In-Support): The prediction is within the set of valid stones for the current chemistry.
• Event B/R (Correct Half/Reachable): The prediction is in the correct "half" of the graph (for withheld tasks) or within the set of stones reachable by the query sequence.
• Event C (Exact Match): The prediction is the correct target stone.
• Formula: $P[C] = P[A] \times P[B|A] \times P[C|A \cap B]$.

By tracking the rise of these components, the authors map performance plateaus and jumps to specific sub-skills.

Causal Interventions

To identify plasticity windows, the authors freeze specific layers (embedding or transformer layers) at different epochs during training. If freezing a layer at a specific time significantly delays or prevents the completion of a learning stage, that layer and time window are deemed critical for that stage.

---

3. Key Contributions & Results

A. Discrete Staged Learning

The model does not learn continuously; it exhibits staged dynamics characterized by plateaus followed by sudden jumps. These stages correspond to the sequential acquisition of sub-skills:

1. Stage 1: Learning to restrict predictions to the valid support set ($P[A]$).
2. Stage 2: Learning the coarse structure (e.g., identifying the correct "half" of the graph or the set of reachable stones).
3. Stage 3: Refining to the exact target ($P[C]$).

B. Asymmetry: Composition vs. Decomposition

• Composition (Robust): When composing 1-hop rules to solve multi-hop queries, the model shows invariance to task complexity. Learning speed is nearly identical for 2-hop to 5-hop queries. The model successfully composes rules once the latent structure is understood.
• Decomposition (Fragile): When decomposing multi-hop sequences to infer 1-hop rules, convergence is significantly delayed as the support hop length increases. The model struggles to extract atomic rules from complex sequences, leading to prolonged stagnation in the "reachable" stage ($P[B|A]$).

C. The "Adjacency Bias" Phenomenon

In the Withheld Potion task, the authors observed a counterintuitive dip in performance during Stage 2.

• The model initially learns to predict stones based on reward adjacency (e.g., if a stone has reward +15, neighbors have +1).
• Because the reward structure creates a bias toward the "incorrect half" of the graph (where 2 out of 3 neighbors reside), the model's accuracy on the "correct half" temporarily drops before it learns the full latent graph structure and corrects this bias.

D. Layer-Specific Plasticity Windows

Using freezing interventions, the authors identified that:

• Plasticity is temporally bounded: Each layer has a specific window during which it must remain trainable to support the transition to the next stage.
• Layer Sensitivity: Freezing transformer layers early in training causes massive delays or non-convergence. Freezing them after the plasticity window closes has negligible impact.
• Embedding vs. Transformer: The embedding layer has a much shorter plasticity window compared to transformer layers.

---

4. Significance and Implications

1. Mechanistic Insight into Generalization: The paper provides evidence that transformers do not just memorize data; they learn latent structures in a structured, coarse-to-fine manner. The "stages" are not random noise but correspond to the acquisition of specific, interpretable sub-skills.
2. Understanding "Grokking": The findings align with the phenomenon of "grokking" (sudden generalization after a plateau), offering a granular view of what changes during that transition (e.g., shifting from reward-bias to structural alignment).
3. Compositionality Limits: The results highlight a fundamental asymmetry in current transformer architectures: they are robust at composing known rules but struggle significantly with decomposing complex examples to infer those rules. This suggests a bottleneck in reverse-engineering latent structures from complex contexts.
4. Training Dynamics & Optimization: The identification of plasticity windows suggests that training strategies could be optimized by dynamically adjusting learning rates or layer freezing schedules based on the specific stage of learning, rather than using static schedules.
5. Benchmarking: The work establishes a rigorous methodology for analyzing learning dynamics by factorizing accuracy into structural components, a method applicable to other latent structure learning tasks.

In summary, this paper moves beyond observing that transformers learn, to explaining how they learn, revealing a complex interplay of staged acquisition, structural biases, and layer-specific plasticity requirements.