The Impact of Depth on Compositional Generalization in Transformer Language Models (2024)

Abstract: To process novel sentences, language models (LMs) must generalize compositionally—combine familiar elements in new ways. What aspects of a model’s structure promote compositional generalization? Focusing on transformers, we test the hypothesis, motivated by theoretical and empirical work, that deeper transformers generalize more compositionally. Simply adding layers increases the total number of parameters; to address this confound between depth and size, we construct three classes of models which trade off depth for width such that the total number of parameters is kept constant (41M, 134M and 374M parameters). We pretrain all models as LMs and fine-tune them on tasks that test for compositional generalization. We report three main conclusions: (1) after fine-tuning, deeper models generalize more compositionally than shallower models do, but the benefit of additional layers diminishes rapidly; (2) within each family, deeper models show better language modeling performance, but returns are similarly diminishing; (3) the benefits of depth for compositional generalization cannot be attributed solely to better performance on language modeling. Because model latency is approximately linear in the number of layers, these results lead us to the recommendation that, with a given total parameter budget, transformers can be made shallower than is typical without sacrificing performance.

## introduction the number of possible sentences in natural language is enormous; regardless of the size of its training set, a language model (lm) will regularly encounter sentences it has never seen before. the ability to interpret such sentences relies on compositional generalization: the capacity to combine familiar words and syntactic structures in new in this paper, we test the hypothesis that increasing a transformer's depth-the number of layers it has-improves its performance on tasks that require compositional generalization. this hypothesis is motivated both by theoretical work, which has shown that adding layers increases the expressive capacity of neural networks in general @xcite and transformers in particular @xcite , and by experimental work suggesting that deeper models generalize more compositionally than shallower ones @xcite . while existing empirical work lends some credibility to this hypothesis, to directly confirm it we must address the confound between depth and size (number of parameters). as each additional layer introduces a new set of parameters, deeper models are also larger, all else being equal, and lms' performance on a wide variety of tasks is known to be correlated with their size @xcite @xcite . to address this confound, we construct classes of models with an equal total number of parameters but differing depths; we do so by reducing the model's feed-forward dimension to compensate for added depth. we pretrain all models as language models and fine-tune them on four compositional generalization tasks: the semantic parsing tasks cogs @xcite , cogs-vf @xcite and geoquery @xcite , and the english passivization portion of multilingual transformations @xcite . in all of these tasks, the model is trained on a particular data distribution and is expected to generalize to another distribution by composing together familiar elements in novel ways. in addition to any possible direct effect on compositional generalization, depth may also be correlated with other factors which may themselves predict compositional generalization, such as language modeling loss during pretraining or in-distribution performance on the fine-tuning task. this complicates the interpretation of any relationship we might find between depth and generalization performance. to address this concern, we also investigate and correct for the effect of depth on language modeling performance and in-distribution loss. we report the following findings, across three model size classes (41m, 134m, and 374m parameters): ## training and inference latency what are the practical implications of the fact that the benefit of depth saturates after a handful of layers? empirically, the compute cost of training and running our equal-parameter models exhibits a strongly linear relationship with depth. figure 8 shows (for our largest models) that the latency during training grows linearly with the depth of the model. the causes of this penalty are two-fold. first, our choice to use narrower feed-forward dimension for deeper models to maintain constant total parameter count leads to a slightly higher floating point operation (flop) count for deeper models. to see this, we start from the cost formula introduced in @xcite : the second and likely more important reason that deeper models are slower is because the computations at layer depend on the results of the computations at layer -1. because of this sequential dependency, parallelism cannot be applied across layers @xcite . combined with the diminishing utility of depth for performance, the linear latency cost of depth as total size is kept constant leads us to the following practical recommendations: these benefits are not confined to training: deeper models also incur a per-layer cost during inference. this means that the penalty a deeper model pays must be amortized over the lifetime cost of using a model for inference. here too, reducing the gpu-time spent on inference by using shallower models of comparable performance could reduce compute costs, both for cloud-served models and for on-device inference @xcite @xcite . ## analysis of feed-forward transforms at extreme feed-forward ratios, we observe that model performance degrades. we investigate the role that the feed-forward ratio plays in the observed performance of our models. when ff is smaller than model , this transform is lossy, but small values of ff could nevertheless impact performance even when ff is still larger than model . to determine if that is the case, we conduct rank analysis on the affine transforms which comprise the feed-forward block. for a given affine transform , we compute the ordered singular values { 1 , 2 , . . . , } @xmath0, ff ) is the rank of @xmath1. we then normalize each value by dividing by the ℓ 1 norm of { 1 , . . . , } to calculate how much of 's image is accounted for by the best -rank approximation of @xmath2. figure 9 shows how in deeper models (i.e., those with increasingly large model / ff ratios) the transforms become increasingly skewed away from making full-use of their available ranks. ## conclusion compositional generalization is essential for interpreting novel sentences. what aspects of the transformer lm architecture contribute to an inductive bias favoring compositional generalization? in a controlled experiment that teases apart depth from total number of parameters, we find that deeper transformers show better compositional generalization, and better language modeling performance, independent of their total number of parameters. at the same time, in most cases the usefulness of adding layers decreases rapidly as models get deeper: comparatively shallow models can achieve generalization accuracy on compositional tasks that is comparable to that of much deeper models, and language modeling perplexity within a few percentage points of the best-in-class model. because deeper transformers have higher latency, this indicates that for a given parameter budget, shallower models can be significantly faster with a minimal sacrifice in performance. alternative approaches to controlling for total size. our approach to controlling for total parameter count necessitates making depth-width tradeoffs. an alternative approach would be to construct universal transformers @xcite , where each model in a size class has a transformer layer with the same parameters repeated layers times. such a weight-sharing approach would allow for deeper models to have arbitrarily-wide feedforward networks, mitigating the impact of making models too narrow. while such weight sharing prevents models from performing different computation in different layers, such restriction may in fact be beneficial for compositional generalization where similar computations (e.g., combining two syntactic phrases to a larger phrase) may need to apply recursively at different scales. pretraining corpus effects. we consider models pretrained on natural-language data. for our particular choice of compositional generalization experiments, the presence of lexical items in both the pretraining corpus and the generalization datasets represents a potential confounder of generalization performance which could be mitigated by modifying compositional datasets @xcite . more generally, we do not study how the distribution of pretraining data affects the inductive biases conferred to lms @xcite . as a particular area of interest for future work, we point out the hypothesis that including source code in the pretraining corpus @xcite will improve compositional generalization. fine-tuning vs. in-context learning. we use fine-tuning to adapt our pretrained models to the compositional tasks. due to its computational cost and task-specificity, fine-tuning is less useful in practice than in-context learning as model size grows @xcite . because in-context learning only becomes reliable at scales far larger than we are able to train, we did not explore the effect of depth on compositional generalization accuracy in in-context learning @xcite ; we point this out as an avenue for future research.