PaperlayerView exact revision on arXiv

Source-linked AI summary

The Depth-to-Width Interplay in Self-Attention

Yoav Levine, Noam Wies, Or Sharir, Hofit Bata, Amnon Shashua

arXiv:2006.12467v3cs.LGcs.CLstat.ML

TL;DR

Leading self-attention models have grown mostly through width, challenging expectations that greater depth is advantageous. The paper analyzes depth and width theoretically and tests the predicted regimes empirically. It finds that width conditions the depth advantage and projects practical allocations for very large models.

  • Problem

    Despite depth-efficiency expectations, leading self-attention models have expanded mostly by widening rather than deepening, leaving their depth-to-width interplay in question.

    Source: §1 · p. 1

  • Method

    The authors analyze separation-rank growth with depth and width, then empirically test the predicted regimes in self-attention networks.

    Source: §1 · p. 2, §1.1 · p. 3

  • Results

    Below Lth(dx) ∼ log(dx), separation rank grows double-exponentially with depth but polynomially with width; above it, dependency capacity grows similarly with both.

    Source: §1.1 · p. 3

  • Takeaways & Limitations

    For fixed-size self-attention models, allocation should account for both depth and width; the paper projects that a 1-Trillion-parameter model requires massive widening.

    Source: §1 · p. 2

    Whether a polynomially sized shallower network can exactly replicate a deeper network in this regime remains an open conjecture.

    Source: §4.2 · p. 8

Abstract

from arXiv · show

Self-attention architectures, which are rapidly pushing the frontier in natural language processing, demonstrate a surprising depth-inefficient behavior: previous works indicate that increasing the internal representation (network width) is just as useful as increasing the number of self-attention layers (network depth). We theoretically predict a width-dependent transition between depth-efficiency and depth-inefficiency in self-attention. We conduct systematic empirical ablations on networks of depths 6 to 48 that clearly reveal the theoretically predicted behaviors, and provide explicit quantitative suggestions regarding the optimal depth-to-width allocation for a given self-attention network size. The race towards beyond 1-Trillion parameter language models renders informed guidelines for increasing self-attention depth and width in tandem an essential ingredient. Our guidelines elucidate the depth-to-width trade-off in self-attention networks of sizes up to the scale of GPT3 (which we project to be too deep for its size), and beyond, marking an unprecedented width of 30K as optimal for a 1-Trillion parameter network.

1 Introduction

The introduction challenges the view that depth and width are interchangeable in self-attention, predicting that width determines when added depth is advantageous. Theory and experiments motivate quantitative depth-to-width allocations, including projections for very large models.

1 Introduction

  • Although deepening has been effective in other neural-network settings, leading self-attention models have grown mostly through width, raising questions about depth-efficiency in this architecture.The paper positions its work as principled guidance for increasing self-attention width and depth in tandem.Source: §1 · p. 1

1.1 An overview of our theoretical approach and findings

  • Below Lth(dx) ∼ log(dx), depth-efficiency holds; above it, depth and width similarly increase modeled dependencies, so added depth helps only when width is sufficient.The theory attributes the early depth advantage to separation rank growing double-exponentially with depth but only polynomially with width.Source: §1.1 · p. 3, §1.2 · p. 3

1.2 An overview of our empirical approach and findings

  • Experiments across depths 6 to 48 show the predicted transition and a width-efficiency regime, where a deeper network can underperform a shallower, wider network of equal size.The transition between regimes was fitted to an exponential form predicted by the theory.Source: §1.2 · pp. 3–4
  • For models trained up to GPT3’s 100B-parameter threshold, appropriate widening alongside deepening could generally help; at 1 Trillion parameters, projected optimal depth is 95 and width 30K.Source: §1.2 · p. 4

2 The self-attention mechanism

The section describes the Transformer’s multi-head self-attention, then defines a simplified architecture for analyzing how its connectivity combines inputs. The analysis removes normalization and activations while retaining the core input-integration mechanism.

2.1 The Transformer encoder architecture

  • Each Transformer layer combines multi-head self-attention with a residual connection, layer normalization, and a position-wise feed-forward sublayer.Per-head Key, Query, and Value projections map representations into attention space, and the output projection returns them to the representation dimension.Source: §2.1 · p. 4

2.2 The analyzed architecture

  • The analyzed network removes layer normalization, softmax normalization, and ReLU activation, and composes L self-attention layers with learned Key, Query, Value, and output matrices.Its resulting form is not a linear mapping: each layer integrates three copies of its input, and depth increases the number of inputs mixed and summed terms.Source: §2.2 · pp. 4–5

2.3 Relaxations

  • The relaxations target normalization and position-wise activations, not self-attention’s input connectivity, which the authors treat as its core functional mechanism.Prior work cited in the section reports comparable language-modeling perplexity for a self-attention-first BERT variant and the interleaved baseline.Source: §2.3 · pp. 5–6

3 A measure of capacity for modeling input dependencies

The section uses separation rank to measure how a function models dependencies across input subsets. For self-attention, this measure is invariant across balanced partitions, unlike the partition-sensitive behavior described for ConvNets.

3 A measure of capacity for modeling input dependencies

  • Separation rank is the minimum number of products of subset-specific functions needed to express a function across a partition of its inputs.A rank of one means the function is separable across the partition; higher rank indicates stronger modeled dependency between the subsets.Source: §3 · p. 6
  • Self-attention’s separation rank is invariant to the identity of a balanced input partition, unlike ConvNets’ differing ranks for nearby versus distant input splits.The result applies to balanced partitions that divide the inputs into two equal-sized subsets.Source: §3 · p. 6
  • Treating all balanced input partitions alike does not by itself ensure capacity for elaborate dependencies; the subsequent analysis examines how architecture affects this capacity.The section motivates bounding separation rank to establish depth-efficiency regimes.Source: §3 · pp. 6–7

4 The effect of depth in self-attention networks

Self-attention has two depth regimes: below a width-dependent threshold, deepening is preferable to widening; above it, both contribute similarly to expressiveness. The analysis relates these regimes to separation rank, a measure of input dependencies.

4 The effect of depth in self-attention networks

  • For L < log3(dx), deepening is clearly preferable to widening; for L > log3(dx), both play a similar role in enhancing expressiveness.These bounds reveal two qualitatively different regimes for self-attention networks.Source: §4 · p. 7

4.1 Depth-efficiency in self-attention

  • When dx > 3L, each layer’s Key/Query/Value computation reintroduces the preceding layer’s output three times, motivating exponential growth of input mixing with depth.Theorem 1 bounds separation rank in this regime; its corollary implies a double-exponential width requirement for a shallow network attempting to replicate a deeper one.Source: §4.1 · p. 7

4.2 Depth in-efficiency in self-attention

  • A single self-attention layer has separation rank only linear in H and dx, while stacking layers makes this input-dependency capacity grow double-exponentially with depth.The analysis then shows that this growth is capped by the internal representation dimension.Source: §4.2 · p. 8
  • For L > log3(dx), a polynomial-width shallower network can surpass the separation rank of a deeper network, unlike the depth-efficient regime.Theorem 2 describes the transition as separation rank shifts from polynomial growth in width and double-exponential growth in depth to exponential growth in width and depth together.Source: §4.2 · p. 8
  • Exact replication of a deeper network by a polynomially sized shallower network in the depth-inefficient regime remains an open conjecture.The result establishes a separation-rank comparison, not exact functional replication.Source: §4.2 · p. 8
  • Figure 2 validates both regimes experimentally; the transition occurs at exponentially larger network sizes as networks get deeper.The figure compares common self-attention architectures and reports agreement with the theory.Source: Figure 2 · p. 9

5 Depth-efficiency regimes in common self-attention networks

Experiments on common self-attention architectures reveal both depth-inefficient and depth-efficient regimes, with the transition depending exponentially on depth. The fitted transition also supports projections of suitable depth and width for larger models.

5 Depth-efficiency regimes in common self-attention networks

  • The empirical study tests common self-attention architectures that include operations omitted from the simplified networks analyzed theoretically.It evaluates depths 6, 12, 18, 24, 30, 36, and 48.Source: §5 · p. 9, §5.1 · p. 9

5.1 The training setup

  • The experiments train decoder-only models at depths 6–48 on autoregressive log-likelihood, using a 2,000-token vocabulary and three text corpora.Test loss is measured on 170K held-out sequences, and reruns indicate stability to the third digit.Source: §5.1 · p. 9

5.2.1 Distinct depth-efficiency regimes in self-attention

  • For smaller parameter budgets, deepening does not outperform widening, but sufficiently wide deeper models perform better at the same budget.The experiments compare depth pairs (6, 12), (12, 24), and (24, 48).Source: §5.2.1 · p. 9

5.2.2 Transition between regimes depends exponentially on depth

  • R²=0.998 and χ²_red=0.854 indicate good compatibility between measured transition widths and the predicted exponential dependence on depth.The fitted relationship is log(dx) = a + b · L.Source: §5.2.2 · pp. 9–10

5.2.3 “Width-efficiency" in small network sizes

  • At a given network size, models may be too shallow or too deep; the fitted curve marks an optimal depth range, with the best-performing depth varying by size.Models left of the curve can improve by deepening, while models to its right can improve by widening.Source: §5.2.2 · p. 10, Figure 3 · p. 10, §5.2.3 · p. 11

5.3 Projecting to larger networks

  • For 96 layers, the projected transition size is 1.17 ± 0.23 · 10^12 parameters, while GPT3 has 175B; the paper predicts optimal depth L=80 for GPT3’s size.The authors therefore suggest widening as the optimal way to scale toward 1-Trillion-parameter models and beyond.Source: §5.3 · p. 11

6 Discussion

The discussion attributes depth efficiency to rapid early growth in self-attention expressiveness, followed by width-limited saturation. It connects this account to architectural approaches that increase representation width while limiting per-layer computation.

6 Discussion

  • Early self-attention layers rapidly increase the ability to correlate arbitrary input locations, but this depth advantage disappears when width cannot support further growth.The analysis attributes depth inefficiency to saturation in the dependencies supported by representation dimension.Source: §6 · pp. 11–12
  • The discussion notes that it found no direct result upper-bounding depth efficiency in other architecture classes and leaves those investigations for future work.It contrasts a separation-rank growth result for certain convolutional networks with the double-exponential growth shown for early self-attention layers.Source: §6 · p. 12
  • The analysis motivates methods that enlarge representation width while using only a subset of it for computation in each layer, including ShuffleNet [Ma et al., 2018] and a related 1-Trillion-parameter model [Fedus et al., 2021].The authors present these approaches as containing computational costs while relaxing width-related theoretical limitations.Source: §6 · p. 12

A.1 The function realized by a deep multi-headed self-attention network

  • The appendix represents a depth-L composition as a sum over heads of products of input-dependent matrices M(c,h), with C(L)=3L−1 units.Each M(c,h) is defined as A(c,h)XXᵀB(c,h)ᵀ.Source: §A.1 · pp. 16–17

A.2 Proof of the upper bound in theorem 1

The proof derives a separation-rank upper bound by expanding the network expression and grouping terms according to the input partition. Each summand has separation rank one.

A.2 Proof of the upper bound in theorem 1

  • The proof uses the form of the network expression in Lemma 1 to establish an upper bound on separation rank.The theorem is stated for depth L≥1, input size N>1, a partition P∪Q=[N], and output locations i and p.Source: §A.2 · p. 18
  • Each summand in the resulting expansion has separation rank one, so counting summands yields an upper bound on the separation rank.The proof reorders the sum around indices where the partition segment switches between P and Q.Source: §A.2 · pp. 18–19

A.3 Proof of the upper bound in theorem 2

The proof bounds separation rank by bounding the number of terms in a polynomial expansion of the output. It then applies a multiset identity to obtain the stated theorem bound.

A.3 Proof of the upper bound in theorem 2

  • The argument bounds separation rank using the polynomial degree of the output coordinate function.The resulting estimate counts the summands of the expansion, each of separation rank one.Source: §A.3 · pp. 19–20
  • The proof concludes theorem 2 by applying the multiset identity after bounding the number of summands.The displayed bound includes 2 log3(dx + H).Source: §A.3 · pp. 19, 21

A.4 The effect of residual connections

The appendix shows that residual connections change the separation-rank upper bound only slightly, then develops tensor and grid-tensor tools for bounding separation rank and establishing generic rank results.

A.5 Technical lemmas

  • With skip connections, the network can be represented as a sum of branches of depths 0 through L, yielding an alternating structure for the proof.Each block adds a self-attention layer's input to its output.Source: §A.4 · pp. 21–22, §A.5 · p. 23

A.4 The effect of residual connections

  • Residual connections modify the separation-rank upper bound by only a small factor, leaving the analysis and its conclusions unaffected.The comparison is made against the main-text bounds; the authors therefore leave the lower bounds unchanged.Source: §A.4 · p. 22

B.1.1 Tensors and their matricization

  • A tensor is a multidimensional array, and its matricization arranges entries into a matrix whose rows and columns correspond to the two parts of a balanced index partition.Tensor order counts modes; the dimension in a mode counts the values available to that index.Source: §B.1.1 · p. 24

B.1.2 Grid tensors provide lower bounds for the separation rank

  • The rank of a grid tensor's matricization lower-bounds the corresponding self-attention function's separation rank for any choice of template vectors.The grid tensor stores function values at points formed from a fixed set of template vectors.Source: §B.1.2 · p. 24

B.1.3 Method for bounding the grid tensor’s rank

  • For a polynomial-parameter matrix, one parameter assignment attaining rank r implies rank at least r for all real parameter settings except a set of Lebesgue measure zero.The argument reduces to showing that a nonzero polynomial cannot vanish except on a measure-zero set.Source: §B.1.3 · p. 25

B.2 Proof of the lower bounds in theorems 1 and 2

The lower-bound proofs construct a high-rank submatrix of a grid-tensor matricization using chosen weights and template vectors, then extend the result to almost all weight settings.

B.2 Proof of the lower bounds in theorems 1 and 2

  • The constructed rank surpasses the theorem's lower bounds in the appropriate depth-to-width regimes, establishing those bounds for all but a measure-zero set of network weights.The proof uses the fact that one qualifying parameter assignment suffices for the almost-everywhere result.Source: §B.2 · pp. 25–26
  • The construction chooses weights and template vectors so a submatrix of the grid-tensor matricization takes a form whose rank can be bounded below.The proof builds layer outputs from segment projections and inner products, then uses later layers to produce increasing powers of a norm whose square is V V^T.Source: §B.2 · p. 26, §B.3 · p. 27

B.3 Technical lemmas

  • A unique highest-degree determinant term guarantees full rank except at finitely many parameter values, while a rearrangement inequality helps establish the needed strict inequality.The technical lemmas use this route to ensure the constructed matrix meets the full-rank condition.Source: §B.3 · p. 30

C Proof of Proposition 1 on the separation rank symmetry

For the network function considered, separation rank is invariant to which balanced partition of input locations is used.

C Proof of Proposition 1 on the separation rank symmetry

  • The separation rank is the same for every balanced partition of the input locations.The claim applies to partitions whose two parts each contain N/2 locations.Source: C Proof of Proposition 1 on the… · p. 31
  • The symmetry claim is restricted to balanced partitions, not arbitrary partitions of the input locations.Both parts of each compared partition have size N/2.Source: C Proof of Proposition 1 on the… · p. 31
  • The proof compares ranks across partitions by reindexing the inputs and then switching the roles of the two partitions to obtain the reverse inequality.The reverse inequality completes the equality argument.Source: C Proof of Proposition 1 on the… · pp. 31–32

D Experimental details

The experimental details specify a fixed training schedule and describe how architecture comparisons and repeat runs were organized.

D Experimental details

  • Training used Adam for 1M steps with batches of 512 sequences of 128 tokens, a 12000-step linear warm-up, and cosine decay to zero.The number of heads per layer was held constant at 2 while width increased.Source: D Experimental details · p. 32
  • The study ran extra experiments for adjacent depth pairs to locate the transition point more accurately and reduce figure 3 error bars.Widths for each trained depth are listed in Table 2.Source: D Experimental details · p. 32
  • Each architecture's training and evaluation experiment was repeated three times to estimate standard deviations.Those standard deviations are reported in Table 3.Source: D Experimental details · p. 32

D.1 Fit details

The fit models log transition width as a linear function of depth, using empirical errors from measured transition points. The reported R^2 and reduced chi-squared indicate a good fit, with the latter suggesting slight overestimation of the empirical errors.

D.1 Fit details

  • The fit uses depth xi and log transition width yi, with empirical error σi, and predicts yi from fit parameters a and b.The collected transition widths per depth are reported in Table 4.Source: §D.1 · p. 32, Table 4 · p. 34
  • χ2_red = 0.854 indicates a good fit for the low sample count, while hinting at slight overestimation of empirical errors.The fit has 3 degrees of freedom, from 5 observations and 2 fitted parameters.Source: §D.1 · p. 32
  • R^2 = 0.998 indicates a good linear fit of the transition data.The fit uses depth as the independent variable and log estimated transition width as the dependent variable.Source: §D.1 · p. 32
  • Widths were measured densely around potential transition points to improve data-point and empirical-error estimates.The trained networks' widths are listed in Table 2.Source: Table 2 · p. 33
  • Test-loss standard deviations summarize variability across three repeated training and evaluation runs for each width-depth point.The standard deviations are reported in Table 3.Source: Table 3 · p. 34
Loading 2006.12467v3…