The Impact of Depth on Compositional Generalization in Transformer Language Models

• With equal parameter budgets, the study varies feed-forward width to test depth directly across three model size classes: 41M, 134M, and 374M parameters.Source: §1 · pp. 1–2
• Deeper models generally achieve lower perplexity, but the marginal benefit of additional layers diminishes rapidly and performance can degrade at large feed-forward dimensions.Source: §1 · p. 2
• Deeper transformers generally show better compositional generalization, although most gains accrue from the first few layers and often saturate quickly.Source: §1 · p. 2
• Depth continues to predict more compositional generalization after correcting for lower language-modeling perplexity and better in-distribution fine-tuning performance.Source: §1 · p. 2
• Because transformer latency is approximately linear in depth, making models wider rather than deeper can be more efficient at a fixed parameter budget.Source: §1 · p. 2

2.1 Constructing Families of Models with Equal Numbers of Parameters

• Models vary depth while holding total parameters constant by reducing feed-forward width and keeping contextual and attention dimensions fixed.This preserves the self-attention mechanism across depths and uses three model size classes.Source: §2.1 · p. 2
• The models are decoder-only GPT-style transformers implemented in t5x, with hyperparameters derived from T5 configurations and a smaller Kim–Linzen model.Source: §2.1 · p. 2
• The parameter-count equations relate layer count and feed-forward dimension, allowing added layers to be offset by narrower feed-forward blocks.Source: §2.1 · p. 2

2.2.1 Language Modeling

• For language modeling, models are pretrained on C4 for 1M steps using 1024-token contexts and batches of approximately 131k tokens.Source: §2.2.1 · p. 3

2.2.2 Compositional Generalization

• For compositional generalization, pretrained models are fine-tuned for 10,000 steps on COGS, COGS-vf, GeoQuery, and English passivization.These tasks test novel combinations of familiar elements, including lexical and structural generalization, database-query generation, and resistance to positional heuristics.Source: §2.2.2 · p. 3
• COGS-vf simplifies semantic representations by removing numbered variables, reducing output complexity but limiting representation of phenomena requiring variable-identity coordination.Source: §2.2.2 · p. 3

3.1 Language Modeling

• Deeper models achieve lower perplexity, with single-layer-to-optimal-model ratios of 1.59, 1.86, and 1.99 for the 41M, 134M, and 374M classes.The optimal depths are 5, 12, and 24 layers respectively, but diminishing returns persist even for the largest models.Source: §3.1 · pp. 3–4
• Language-modeling gains from added layers diminish rapidly: the largest improvement is from 1 to 2 layers, while 2 to 4 layers yields a smaller improvement.Performance can degrade when the feed-forward dimension becomes smaller than the model dimension, where transforms become increasingly rank-deficient.Source: §3.1 · p. 4
• The near-optimal feed-forward-ratio range expands for larger models, allowing them to trade depth for width with smaller perplexity penalties.For smaller models, performance is more sensitive to the particular ratio, and the standard ratio may not be optimal.Source: §3.1 · p. 4

3.2 Compositional Generalization

• Deeper models generalize better on compositional tasks across datasets and size classes, even after matching pretraining perplexity and in-distribution loss.The depth advantage is therefore not explained solely by stronger upstream language modeling or downstream in-distribution performance.Source: §3.2 · p. 5, §3.3 · p. 6, Figure 4 · p. 6
• Compositional-generalization gains saturate after 4–6 layers for COGS and 2–4 layers for COGS-vf and GeoQuery.English Passivization shows more variable depth effects and no clear size-independent saturation threshold.Source: §3.2 · p. 5, Figure 4 · p. 6
• Depth primarily benefits easier lexical generalization, while even the largest deep models systematically fail on structural generalization.Thus, increasing depth alone does not substantially improve the inductive bias needed for structural tasks.Source: §3.2 · pp. 5–6

3.3 Depth Effects are Independent between Upstream and Downstream Tasks

• To isolate depth, the authors compare fine-tuned models at equal validation perplexity and at equal in-distribution loss, finding deeper models still achieve higher out-of-distribution accuracy.These controls address confounds from both pretraining and fine-tuning performance.Source: §3.3 · p. 6

4 Training and Inference Latency

• Depth’s performance benefits saturate after a handful of layers, motivating comparisons between quality gains and compute costs.The practical question is whether deeper models’ diminishing utility justifies their higher latency.Source: §4 · pp. 6–7
• Training latency grows strongly linearly with depth even at fixed parameter count.The FLOP count is linear in depth, although the 2N term dominates unless depth, attention size, or context length become very large.Source: §4 · pp. 6–7
• Sequential dependencies between layers prevent parallelism across layers and are a major reason deeper models train more slowly.Narrower feed-forward dimensions in deeper equal-parameter models also slightly increase FLOP count.Source: §4 · p. 7
• For fixed data volume or GPU-hour budgets, shallower models can train faster, process more data, and retain acceptable performance.The same latency trade-off applies during inference, where reducing layers can lower cloud and on-device compute costs.Source: §4 · p. 7

5 Analysis of Feed-Forward Transforms

• Deeper models generalize better after matching in-distribution fine-tuning loss, separating the observed generalization advantage from that loss difference.Figure 7 compares models at equal in-distribution loss of 0.0002.Source: Figure 7 · p. 8
• As depth increases and d_model/d_ff grows, feed-forward input and output projections become increasingly close to rank-deficient transforms.This indicates that deeper models increasingly fail to use their available transformation ranks fully.Source: §5 · p. 8, Figure 7 · p. 8
• Rank analysis orders singular values, normalizes them by their ℓ1 norm, and measures how much of a transform’s image its best i-rank approximation captures.The analysis applies to affine transforms in the feed-forward block, with k = min(d_model, d_ff).Source: §5 · p. 8

6 Related Work

• This work complements data, scaling, and architectural approaches to compositionality by increasing depth without changing total model size.The comparison targets a specific architectural factor rather than proposing a general compositionality method.Source: §6 · p. 8
• Earlier studies associate depth with expressive capacity, downstream performance, and compositional representations, but commonly fail to control total parameter count.Other results suggest wide shallow transformers can match deeper models and that residual networks may use shallow subnetworks.Source: §6 · pp. 8–9
• Early-exit methods reduce inference computation dynamically, whereas this study evaluates shallower equal-parameter models across all inputs during training and inference.The authors do not test early-exit training or inference with their models, so additional savings remain possible.Source: §6 · p. 9

7 Conclusion

• Deeper transformers show better compositional generalization and language modeling independent of total parameter count.The conclusion is based on a controlled experiment separating depth from model size.Source: §7 · p. 9
• In most cases, adding layers becomes rapidly less useful, allowing comparatively shallow models to approach much deeper models’ accuracy and perplexity.The paper reports comparable compositional-task accuracy and perplexity within a few percentage points of the best model in a size class.Source: §7 · p. 9
• Because deeper transformers have higher latency, shallower models can be significantly faster with only a minimal performance sacrifice at a fixed parameter budget.This is the paper’s practical recommendation for model design.Source: §7 · p. 9

Attention heads

• The study does not examine how attention-head number or function changes with model depth, leaving compositional computations by individual heads unresolved.It specifically raises whether head-count importance changes in much shallower and wider models and whether heads perform compositional computations.Source: §8 · pp. 9–10

Alternative approaches to controlling for total size

• Weight sharing in Universal Transformers could permit arbitrarily wide feed-forward networks, but it also prevents different computations in different layers.The paper notes that this restriction might benefit recursive compositional operations, making the alternative design substantively different from the tested models.Source: §8 · p. 10

Pretraining corpus effects

• The experiments leave pretraining-data distribution effects unstudied, while lexical overlap between pretraining and generalization data may confound measured generalization.The paper proposes source-code-inclusive pretraining as a future direction.Source: §8 · p. 10

Fine-tuning vs. in-context learning

• The study evaluates fine-tuning rather than in-context learning because reliable in-context learning requires scales larger than the models trained here.The effect of depth on compositional generalization in in-context learning remains future work.Source: §8 · p. 10

• The researchers use open-source datasets that third parties can inspect for bias and toxicity, and release no public checkpoints.They derive their implemented models from existing publicly available T5 models.Source: §9 · p. 10
• Because the models are trained on English-only natural-language data, broader language coverage is needed to assess how language affects the results.Source: §9 · p. 10

• Table 2 reports exact hyperparameters for the trained model classes, while Table 3 displays pretraining and compositional generalization accuracy across all model sizes and tasks.Source: A Design and Result Tables · p. 13

• Information flows through each transformer component from left to right, combines with residual connections, and proceeds to subsequent layers.The diagram highlights vector dimensionality, especially the model embedding size and feed-forward internal dimension.Source: B Annotated Transformer Layer · p. 13, Figure 10 · p. 14
• Table 2 organizes models by depth across three parameter-size classes, with bold variants serving as baseline configurations.Source: Table 2 · p. 14
• Table 3 pairs validation perplexity on C4 with compositional-dataset accuracy after 10 k fine-tuning steps, using lower-is-better and higher-is-better directions respectively.Bold values mark the best performance within each size class, and each condition comes from a single run.Source: Table 3 · p. 14