Scaling Laws and Representation Learning in Simple Hierarchical Languages: Transformers vs. Convolutional Architectures

TL;DR

The paper tackles how neural language models learn hierarchical structure and scaling laws when trained for next-token prediction on data generated by the Random Hierarchy Model (RHM), a tractable ensemble of probabilistic context-free grammars. It develops a correlation-based theory of representation learning and contrasts convolutional networks with locality/weight sharing against transformers with global self-attention, predicting faster, structure-aligned scaling for CNNs. The authors derive and validate power-law scaling relations for the excess test loss under online training, linking them to the ability to reconstruct latent hierarchies via token correlations and to translation-invariant architectural priors. This work elucidates how architectural biases interact with data statistics to shape neural scaling laws, with implications for designing models that efficiently learn compositional structure in hierarchical data and for understanding when attention-based models offer advantages beyond generic sequence modeling.

Abstract

How do neural language models acquire a language's structure when trained for next-token prediction? We address this question by deriving theoretical scaling laws for neural network performance on synthetic datasets generated by the Random Hierarchy Model (RHM) -- an ensemble of probabilistic context-free grammars designed to capture the hierarchical structure of natural language while remaining analytically tractable. Previously, we developed a theory of representation learning based on data correlations that explains how deep learning models capture the hierarchical structure of the data sequentially, one layer at a time. Here, we extend our theoretical framework to account for architectural differences. In particular, we predict and empirically validate that convolutional networks, whose structure aligns with that of the generative process through locality and weight sharing, enjoy a faster scaling of performance compared to transformer models, which rely on global self-attention mechanisms. This finding clarifies the architectural biases underlying neural scaling laws and highlights how representation learning is shaped by the interaction between model architecture and the statistical properties of data.

Figures (8)

• Figure 1: Comparison of structural hierarchies in (left) natural language and (right) an instance of the Random Hierarchy Model (RHM).
• Figure 2: Training dynamics of a depth-$4$ Transformer trained on RHM data with $L=4$, $s=2$, $v=24$, and $m=6$. At each step of training, the test cross-entropy loss is evaluated on a test set of size $2^{15}$ and averaged over $8$ independent realisations of the RHM. This procedure applies to all figures in the paper unless otherwise stated. Colored dashed lines show the averaged $s^{\ell}$-gram losses defined in \ref{['eq:slgram-losses']}: as the number of training steps increases, the Transformer’s performance transitions between successive approximations.
• Figure 3: Scaling of the excess test loss, $\mathcal{L} - \mathcal{L}_{\infty}$, as a function of the number of training steps for depth-$4$ Transformers trained on RHMs with shared tree structure ($L=4$, $s=2$) but varying values of $m$ and $v$ (indicated in the legend). Each parameter set is associated with a distinct color. Solid lines show empirical results, while dashed lines represent scaling predictions from \ref{['eq:scaling-law-general']}.
• Figure 4: Training dynamics of depth-$4$ models on RHM data with $L\,{=}\,4$, $s\,{=}\,2$, $v\,{=}\,16$, and $m\,{=}\,4$. Left: Test cross-entropy loss as a function of the number of training steps for depth-$4$ Transformers trained with Adam and vanilla SGD, and for depth-$4$ CNNs trained with SGD. Black dashed lines indicate the $s^{\ell}$-gram losses. CNNs exhibit a much faster loss decay compared to Transformers, and a constant delay is observed between the Adam- and SGD-trained Transformers, with Adam reaching lower losses earlier. Right: Scaling of the excess test loss, $\mathcal{L} - \mathcal{L}_{\infty}$, as a function of the number of training steps. The Adam curve is rescaled to overlap with the Transformer SGD curve, highlighting their similar scaling behaviour. By contrast, CNNs display a different scaling behaviour. Both CNN and Transformer scalings are compared with the corresponding theoretical predictions reported in the legend.
• Figure 5: Scaling of the excess test loss, $\mathcal{L} - \mathcal{L}_{\infty}$, as a function of the number of training steps for depth-$4$ CNNs trained on RHMs with shared tree structure ($L=4$, $s=2$) but varying values of $m$ and $v$ (indicated in the legend). Each parameter set is associated with a distinct colour. Solid lines show empirical results, while dashed lines represent scaling predictions from \ref{['eq:scaling-law-improved']}.