The Impact of Depth and Width on Transformer Language Model Generalization

We thank the reviewer for their constructive feedback.

W. 1: Paper mostly confirms existing findings.

In our view, prior empirical findings are suggestive but incomplete as the experiments were not controlled: depth was confounded with parameter count. The primary thing which distinguishes our work here from the findings of, e.g., Mueller et al (2022) or Tay et al. (2021) is our goal of controlling for parameter count; while these previous works note that depth can be helpful, this conclusion is fundamentally confounded by the fact that deeper models in their setups are also larger models. Our core contribution is to disentangle that confound and show what the impact of depth is, independent from size. Because, as we show, the usefulness of depth saturates extremely quickly, our results suggest that adding depth does not confer the strong advantage that conventional understanding might expect; rather, adding depth does very little for all but the shallowest models, and in fact, past a certain point increasing model depth can actually be harmful. To our knowledge, these takeaways are novel. Finally, we also show that this effect is present in models’ compositional generalization abilities, and demonstrate that this effect scales with model size.

We have also since conducted an analysis on the compute efficiency of our models by depth. We find that compute efficiency–as measured by latency–is very linear in depth, indicating that models pay a constant price for each additional layer added. Combining this with our main results showing the diminishing returns on depth in terms of pure performance, our results additionally suggest that practitioners bound by resource constraints should prefer to make models much wider and shallower than convention would otherwise suggest. A revised version of the manuscript will include the specifics, and we will include a supplementary figure in our reviewer responses in the next couple of days.

W. 2: Adding other tasks.

We agree, and are in the process of evaluating our models on the BigBench-Lite suite of tasks to get a broader picture of how depth impacts model performance. A revised version of the manuscript will include a breakdown of performance-by-depth on these tasks, similar to our section on compositional generalization.

Question: Generalization beyond compositionality

In this particular study, we focused on compositional generalization; as we say in the second paragraph of the introduction, “we test the hypothesis that increasing a transformer’s depth—the number of layers it has—improves its out-of-distribution performance on tasks that require compositional generalization”. We cite theoretical and empirical studies that motivate this hypothesis. We do not have specific hypotheses as to the effect of depth on other types of generalization, though we agree that it would be interesting to study any such effects empirically in future work. We do agree with you that our scope can be made clearer in the paper; in fact, to make it clearer that we’re only making claims about compositional generalization, as opposed to other types of generalization, we have decided to add the word “compositional” to the title of the paper. We will also go through the manuscript and ensure that the scope of our claims is appropriate.

We thank the reviewer for their constructive feedback!

W. 1 The title of the paper should focus on depth, not width.

Thank you for this comment! On reflection, we tend to agree, and have decided to change the title as you propose.

W. 2: The effect of width is not extensively explored.

We agree, and this is intentional: the goal of the paper is to test the hypothesis that deeper models have stronger compositional generalization abilities. To the extent that we do vary width, we do so only to enable fair comparisons across models of different depths (by keeping the total number of parameters matched). While there may be interesting effects of width (and other hyperparameters) on generalization, we do not have particular hypotheses as to what those would be, and as such such experiments are outside the scope of the paper. We hope that by changing the title we have made the scope of the paper clearer.

W. 3: Compute performance analysis.

We fully agree that studying the compute impact of depth in addition to the performance impact is important. Since submitting the original manuscript, we have conducted an analysis of the compute cost of increasing depth as measured by latency; we plan to include this analysis in a future version of the manuscript very soon. We find that latency is roughly linear in depth, indicating that models incur a consistent compute-efficiency penalty as depth increases. This, combined with our main results which show that the performance improvement depth affords diminishes rapidly and quite early suggests that compute-constrained teams should likely prefer shallower models which are just deep enough to capture a substantial portion of the benefits depth affords, for two reasons. (1) In training, shallower models will either train in a shorter amount of time (for a fixed data budget) or can train on more data (for a fixed time budget) than deeper models of the same size; (2) In inference, shallower models have lower per-use latencies than deeper models of the same size do.

Q. 1: Standardize citation format.

Thanks for pointing this out, we have gone through the bibliography and standardized our references.

Q. 2: The tradeoff between depth and width.

As we mention in our response to Weakness 2, the main focus of our paper is on the effect of depth on generalization. To study it while keeping the comparison fair, however, we do need to trade depth off against width: for all models within a size class, we have to vary depth and width inversely with one another, which allows us to quantify how different depth-width combinations (for a fixed number of parameters) fare on our studied tasks. One particularly interesting finding we report concerning this tradeoff is the U-shaped relationship between depth and perplexity: when models become very deep their performance begins to degrade; as we mention in the paper, we hypothesize that this is due to the fact that they become too narrow.

Q. 3: Performing additional runs.

Thanks for pointing this out - we fully agree, and have been working to replicate our analyses on multiple runs. We have updated our paper to replicate each result five times (that is, we now have 5 pretraining runs per condition for the language modeling results, and we have 5 fine-tuning runs based on those 5 pretraining checkpoints for the compositional generalization results). We have also added confidence intervals to both of these sets of results to quantify measurement uncertainty. In all cases, we find that our initial results hold up after repeated trials, strengthening our confidence in the conclusions.

Q. 4: Is there an optimal depth for all sizes & domains?

Thanks for this very insightful question. Our analysis in Section 3.1, particularly the data shown in Fig. 3(b), suggests that the optimal performance–as measured by absolute score for a fixed data budget–consistently appears to be when the ratio of feed-forward dimension to embedding dimension is between 0.5 and 1.0. In our setup, this would mean that the optimal depth for a model should in fact increase with size, assuming that the embedding dimension is also scaled with parameter count, although this point may not line up with the depth of conventional existing models. For instance, among the model families we examine, the optimal 41M and 374M models are deeper than their conventional reference models. However, taking into account the reviewer’s earlier point about compute cost, the “compute-optimal” depth may actually be much shallower than convention suggests.

We thank the reviewer for their helpful comments.

W1. The paper focuses on empirical experiments

We agree that this paper is an empirical experimental study, though it is motivated by a theory (Merrill, Sabharwal & Smith, TACL, 2022), but we disagree that this is a reason to reject the paper; in our experience, ICLR regularly publishes papers that are primarily experimental, and in fact such papers may be more common than theory-heavy papers. In general, we believe that there is a place for both theory papers and empirical papers, and especially for empirical papers such as ours that test predictions made by a theory.

W2. The focus on decoder-only transformers

We agree with the reviewers that in future work it would be worthwhile to investigate additional architectures, including encoder-decoder transformers or even LSTMs, but given the vast amount of computational resources required to pretrain language models, we had to limit our scope to a single architecture. We chose to focus on decoder-only transformers as at the time of writing they were the state-of-the-art and the de-facto standard for both proprietary (PaLM, GPT-3,4) and open-source (GPT-2, LlaMA) large language models.

We thank the reviewer for their constructive feedback.

W1: The study would mostly serve as a reference for what is usually considered common knowledge

The reviewer’s summary of our results as a recommendation to, for a given parameter budget, pick depth over width, has some truth in it, but it does not fully convey the richness of results we report. First, we find that the marginal utility of added depth reduces significantly as models get deeper, such that a more accurate summary of our recommendation would need to be amended by the statement “for a given parameter budget, picking a very deep model will not help much over picking a ‘just-deep-enough’ model”. But we also find that when the corresponding feed-forward block becomes too small, it can actually be harmful to increase depth. Finally, motivated by prior theoretical claims, we show empirically that depth dramatically affects models’ compositional generalization abilities, and demonstrate that this effect scales with model size; we are aware of very little work that has investigated the sensitivity of compositional generalization to depth. Overall, then, we believe that our results significantly expand upon what anyone might take to be “common knowledge” in the field. Even if that were not the case, however, in our view there is significant value in publishing controlled scientific experiments that confirm a commonly but informally belief.

W2: It remains extremely clear from this paper that beyond very small depth (as soon as we get 3~ layers), performance doesn’t really go up with depth alone

The diminishing return to increasing depth is indeed one of the interesting findings of our paper (though the specific number of layers where we begin to observe diminishing returns varies by task and is usually greater than 3). But we are not sure why the reviewer mentions this finding as a weakness of the paper - this is an empirical scientific discovery, one we didn’t expect to find! This discovery also has applied implications: standard pretrained transformers may be deeper, and therefore more computationally expensive at inference time, than they need to be to generalize well (see also the next point).

W3: I would be interested by the following question: what about if my budget is not really in terms of number of paremeters, but rather in compute power or memory?

Great question, thank you for pointing this out. We had already begun investigating this after the submission deadline and we are able to provide an answer. A revised version of the manuscript will include an analysis of the compute performance by model depth. For the accelerators we trained on, we find a strongly-linear relationship between depth and latency when controlling for parameter count. Because latency increases linearly with depth while performance saturates quite early, compute-constrained practitioners can benefit from choosing architectures which are “just deep enough” for two reasons: (1) since shallower models are computationally cheaper, they can be trained in less time for a fixed data budget, or trained on more data for a fixed time budget; (2) for the same reason, shallower models incur less per-inference compute cost. This reduces point-of-use latency.

Taken together, these points actually point against what commonly-understood knowledge might otherwise suggest: picking depth over width is not necessarily the best choice. Rather, if one can ensure that a model is deep enough to capture most or all of the value of added depth (which happens quite early on), practitioners concerned with things like compute or latency tradeoffs should actually then choose to make models as wide as possible.

W4: p8: "when studying the the impact"

Thank you for catching the typo. It has been fixed.