Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Layer Specialization Underlying Compositional Reasoning in Transformers

Jing Liu

TL;DR

The paper investigates how transformers acquire hierarchical compositional reasoning by evaluating causal and masked transformers on the Random Hierarchy Model (RHM), a controlled hierarchical sequence generator. It decouples memorization, in-distribution generalization, out-of-distribution generalization, and cross-layer transfer to reveal how internal representations evolve. A key finding is the progressive emergence of layer specialization that correlates with generalization: causal models concentrate compositional processing in early layers, while masked models do so in late layers, yet achieve comparable transfer performance. This work links observable generalization capabilities to interpretable, modular internal organization and highlights architecture-dependent strategies for in-context learning and skill composition.

Abstract

Transformers exhibit compositional reasoning on sequences not observed during training, a capability often attributed to in-context learning (ICL) and skill composition. We investigate this phenomenon using the Random Hierarchy Model (RHM), a probabilistic context-free grammar that generates sequences through recursive rule application. Models are trained on subsets of sequences and evaluated across four generalization conditions: memorization, in-distribution generalization, out-of-distribution generalization with the same rules, and cross-layer transfer. Behaviorally, performance improves systematically with task complexity and the number of in-context examples, with out-of-distribution tasks requiring substantially more examples than in-distribution scenarios. Mechanistically, we identify a progressive emergence of layer specialization during training that correlates with generalization performance. Principal component analysis and attention pattern clustering reveal that transformers develop structured, hierarchically organized representations in specialized layers. These results demonstrate that transformers develop modular, interpretable mechanisms supporting compositional reasoning, linking internal algorithmic structure to observed behavioral capabilities.

Layer Specialization Underlying Compositional Reasoning in Transformers

TL;DR

The paper investigates how transformers acquire hierarchical compositional reasoning by evaluating causal and masked transformers on the Random Hierarchy Model (RHM), a controlled hierarchical sequence generator. It decouples memorization, in-distribution generalization, out-of-distribution generalization, and cross-layer transfer to reveal how internal representations evolve. A key finding is the progressive emergence of layer specialization that correlates with generalization: causal models concentrate compositional processing in early layers, while masked models do so in late layers, yet achieve comparable transfer performance. This work links observable generalization capabilities to interpretable, modular internal organization and highlights architecture-dependent strategies for in-context learning and skill composition.

Abstract

Transformers exhibit compositional reasoning on sequences not observed during training, a capability often attributed to in-context learning (ICL) and skill composition. We investigate this phenomenon using the Random Hierarchy Model (RHM), a probabilistic context-free grammar that generates sequences through recursive rule application. Models are trained on subsets of sequences and evaluated across four generalization conditions: memorization, in-distribution generalization, out-of-distribution generalization with the same rules, and cross-layer transfer. Behaviorally, performance improves systematically with task complexity and the number of in-context examples, with out-of-distribution tasks requiring substantially more examples than in-distribution scenarios. Mechanistically, we identify a progressive emergence of layer specialization during training that correlates with generalization performance. Principal component analysis and attention pattern clustering reveal that transformers develop structured, hierarchically organized representations in specialized layers. These results demonstrate that transformers develop modular, interpretable mechanisms supporting compositional reasoning, linking internal algorithmic structure to observed behavioral capabilities.
Paper Structure (12 sections, 3 figures)

This paper contains 12 sections, 3 figures.

Table of Contents

  1. Introduction
  2. Related Work
  3. Method
  4. Results
  5. Behavioral Analysis: Compositional Generalization
  6. Mechanistic Analysis: Layer Specialization
  7. Discussion
  8. Appendix
  9. Architecture
  10. Initialization
  11. Optimization and Schedule
  12. Hyperparameter Choice

Figures (3)

  • Figure 1: Example RHM derivation with depth $L=3$ and branching factor $s=2$. The next-token prediction task asks the model to predict the final observable token (red) from the preceding $d-1$ tokens. Hidden symbols are inferred via correlations between 2-tuples and the target token $x_d$.
  • Figure 2: Layer-wise specialization scores for (Left) Causal Language Models and (Right) Masked Language Models across four generalization conditions: memorization (mem), in-distribution (ind), out-of-distribution with same rules (gen same), and out-of-distribution layer transfer (transfer).
  • Figure 3: Evolution of normalized specialization scores during training for CLM and MLM architectures. The score increases rapidly during early training, plateaus during the transition from memorization to in-distribution generalization, and refines during later stages as out-of-distribution generalization emerges.