The Depth-to-Width Interplay in Self-Attention

Introduction

The golden age of deep learning has popularized the depth-efficiency notion: From an expressiveness standpoint, increasing a neural network’s size by adding more layers (deepening) is advantageous relatively to other parameter increase alternatives, such as increasing the dimension of the internal representation (widening). Beyond overwhelming empirical signals for this notion (Simonyan and Zisserman, 2014; He et al., 2016), depth-efficiency was theoretically supported from a variety of angles (Cohen et al., 2016; Eldan and Shamir, 2016; Raghu et al., 2017; Daniely, 2017).

Diminishing returns in the case of very deep networks were mainly attributed to optimization issues, and indeed the alleviation of these issues has allowed network depths to mount from $10$s to $100$s and beyond (He et al., 2016), enabling deep convolutional networks (ConvNets) to advance the state-of-the-art in computer vision applications. However, as the field matured, a more nuanced perspective emerged. Empirical (Zagoruyko and Komodakis, 2016; Wu et al., 2019) and theoretical (Lu et al., 2017) studies suggest that the interplay between depth and width may be more subtle. Recently, a method for increasing width and depth in tandem (“EfficientNet" by Tan and Le (2019)) has lead to the state-of-the-art on ImageNet while using a ConvNet with a fraction of the parameters used by previous leaders. Our work provides principled guidelines for increasing width and depth in tandem in self-attention networks.

Since the introduction of the Transformer (Vaswani et al., 2017), along with its encoder-only variant, BERT (Devlin et al., 2019), self-attention based deep learning architectures have taken over the field of natural language processing (Liu et al., 2019; Radford et al., 2019; Yang et al., 2019; Raffel et al., 2019a; Clark et al., 2020). However, in contrast to the depth “arms race" that took place in the ConvNet case, the leading self-attention networks are not much deeper than the original BERT model. In fact, even the strongest self-attention models trained to date, which increased the $0.3$B parameter count of BERT-large by factors of $100$s to $11$B (Raffel et al., 2019a) and $175$B (Brown et al., 2020), have only increased its depth by factors of $2$ and $4$, respectively. The remaining size increase stems from an increase in layer widths, clearly countering the depth-efficiency notion.

A recent empirical ablation study by Kaplan et al. (2020) provides support for the above signal. Figure 1(a), taken from this study, leads the authors to conclude that the overall (non-embedding) network size, given by $12\cdot L\cdot d_{x}^{2}$ where $L$ is the number of self-attention layers (network depth) and $d_{x}$ is the hidden representation dimension (network width), is the main predictor of performance regardless of the depth-to-width ratio. This suggests that depth may not play as crucial a role in self-attention networks as it does in convolutional networks.

In this paper, we address the above question of the depth-to-width interplay in self-attention networks, and reveal fundamental subtleties in the above picture. We predict that self-attention will exhibit two qualitatively different depth-efficiency and depth-inefficiency behaviors, in two distinct parameter regimes, as depicted in figure 1(b). After presenting our theoretical analysis in sections 2-4, we provide a thorough empirical evaluation in section 5, which validates our predicted trends for self-attention networks of depths $6$ to $48$. Importantly, our theoretical and empirical results provide quantitative guidelines for optimal depth-to-width parameter allocation given a fixed parameter budget (see for example table 1). The current challenge of reaching beyond $1$-Trillion parameter language models renders informed guidelines of how to increase self-attention depth and width in tandem a mandatory ingredient. Our results clearly show that the optimal path towards the $1$-Trillion parameter mark includes massive widening.

We analyze self-attention networks in which all non-linear activations and normalization operations are removed. Otherwise, the analyzed class (presented in section 2) has the regular deep multi-headed Key/Query/Value structure of common self-attention. After presenting this class in detail, we point to recent studies which demonstrate that normalization and position-wise activations are much less pertinent to the ability of self-attention to correlate inputs than its core connectivity, described in full by our analyzed model. More generally, removing non-linearities for analysis of deep network connectivity traits is a common simplifying assumption: results on expressiveness and optimization of fully-connected (Saxe et al., 2013; Kawaguchi, 2016; Hardt and Ma, 2016), convolutional (Cohen et al., 2016), and recurrent (Khrulkov et al., 2018; Levine et al., 2018a) networks have been attained via this technique. Trade-offs between the depth and width of fully-connected neural networks have been recently examined from theoretical (Fan et al., 2020; Bu et al., 2020) and empirical (Nguyen et al., 2020) perspectives. To the best of our knowledge, our theoretical analysis is the first to address the question of parameter allocation between depth and width in self-attention networks.

We employ the tool of a function’s separation rank with respect to subsets of its inputs, which quantifies its ability to model input dependencies (presented in section 3). The separation rank was employed for attaining theoretical insights on the dependencies modeled by convolutional and recurrent networks (Cohen and Shashua, 2017; Levine et al., 2018a).

Rather than reinforcing the seemingly plausible hypothesis for the trend in figure 1(a), by which widening a self-attention network is as effective as deepening it, we confirm the contrary. We show that the operation of stacking self-attention layers is so effective that it quickly saturates a capacity of the network’s width. We establish in section 4 the existence of a depth threshold which depends logarithmically on the width $d_{x}$, denoted $L_{\textrm{th}}(d_{x})\sim\log(d_{x})$. Below the threshold, we prove that depth-efficiency takes place in self-attention networks: a network of depth $L\leq L_{\textrm{th}}(d_{x})$ cannot be replicated by a shallower network, unless the latter’s width grows double-exponentially with $L$. We prove the above by showing that the separation rank of functions realized by self-attention networks grows double-exponentially with depth, but only polynomially with width, shedding light on the effectiveness of the self-attention mechanism in modeling input interactions when recursively repeated. However, we show that this overwhelming advantage of depth is quickly replaced by a balanced growth. We prove that for self-attention networks with $L>L_{\textrm{th}}(d_{x})$ the ability to model input dependencies, as modeled by the separation rank, increases similarly with depth and width. We corroborate our theoretical findings empirically, as shown in the example of figure 1(b) and more extensively in section 5.

2 An overview of our empirical approach and findings

For two networks with the same parameter count but of different depths $L_{1}<L_{2}$ and widths $d_{2}<d_{1}$, our theory indicates that: (1) there is no advantage to the deeper network when its dimension $d_{2}$ is too small (width caps the benefit of the added layers of depths $L_{1}+1,...,L_{2}$), but (2) the deeper network should outperform the shallower one when its width $d_{2}$ is large enough such that the added layers are in the depth-efficiency regime.

Traces of this predicted phenomenon appear in existing literature: A closer look at depths $L\leq 6$ in the experiment of Kaplan et al. (2020) in figure 1(a) reveals depth-efficiency, though the conclusion of that study was of an overall depth inefficiency of self-attention, based on the behavior of their $L>6$ curve. In section 5 we demonstrate empirically that the more nuanced transition between depth-efficiency and inefficiency, as predicted by our theory, affects commonly used self-attention depths of $L=6,12,18,24,30,36,48$ (a representative plot from our experiments is given in figure 1(b)). The experiments reveal a third regime of “width-efficiency": a network can be too deep for a given parameter budget, and under-perform relatively to a shallower and wider network of the same size (see figure 2).

We fit the network sizes at which a transition between the different depth-efficiency regimes occur to an exponential form, predicted by our theory (see figure 3). This allows us to extrapolate the depth-efficiency behavior of larger architectures, and project practical guidelines for the architectural design of contemporary huge language models. Table 1 shows our suggested depths and widths for models of sizes used in the recent GPT3 paper (Brown et al., 2020). It seems that popular self-attention architectures at all sizes trained up to GPT3’s crossing of the 100B parameter threshold, could generally benefit from deepening, with the appropriate widening (indicated by our guidelines). With that, our results clearly indicate the importance of widening self-attention networks when aiming for the 1 Trillion parameter mark. We project the optimal architecture at that size to have depth 95 and width 30K, wider than any self-attention network trained to date.

The self-attention mechanism

Differentiable attention models in which the output attends over all LSTM-based input representations have been introduced in the context of machine translation (Bahdanau et al., 2014). Self-attention (also referred to as intra-attention), which relates different inputs to each other, was first employed for machine reading (Cheng et al., 2016), and soon thereafter shown to be useful for a variety of language applications when operating over LSTM-based representations (Parikh et al., 2016; Paulus et al., 2017; Lin et al., 2017). Vaswani et al. (2017) were the first to demonstrate that a model based solely on attention, the Transformer, can be better than LSTM based networks. The Transformer’s encoder, BERT (Devlin et al., 2019), based entirely on self-attention, has demonstrated unprecedented performance across natural language understanding tasks.

2 The analyzed architecture

where the Feed-Forward matrices can be now effectively embedded within $W^{\textrm{O},l,h}$. Our analysis below treats a deep multi-headed self-attention network that is attained by a concatenation of $L$ such layers. Importantly, the resultant “linearized" network form, where activations and normalizations are removed, is by no means a linear mapping over the network input – every layer integrates $3$ copies of its input in the above non-linear fashion.

By recursively applying eq. (3) $L$ times we attain the analyzed depth-$L$ self-attention network. We denote the function realized by a network with embedding dimension $d_{x}$ and $H$ attention heads per layer at output location $i\in[N]$ by:

3 Relaxations

Empirical evidence indicates that while the ReLU activations and softmax normalization contribute to performance (layer-norm mainly contributes to optimization), the basic mechanism in eqs. (3) and (4) above captures the defining self-attention characteristic of integrating the inputs with each other in a flexible manner:

The ReLU activation relaxation: Press et al. (2019) demonstrate that a “self-attention first" BERT variant that first performs all of the self-attention operations (eq. (1)) consecutively, and only then performs all of the position-wise feed-forward+ReLU operations, achieves comparable language modeling performance relatively to the Baseline, which takes the regular approach of interleaving these functionalities (i.e., concatenating the BERT’s layer described in eq. (2)). They report that the interleaved Baseline achieves a perplexity score of $18.63\pm 0.26$ on the WikiText-103 test (Merity et al., 2016) when averaged over $5$ random seeds, while the “self-attention first" model achieves a perplexity score of $18.82$ on this test set. The best pre-Transformer perplexity result on the WikiText-103 test, reported by an LSTM-based architecture, was $29.2$ (Rae et al., 2018). Since ReLU and feed-forward do not mix different locations, this outcome directly implies that the self-attention mechanism itself provides all of the elaborate input integration which differentiates BERT from previous architectures.

The softmax normalization relaxation: Initially, an intuitive interpretation of attention as distributing “fractions" of an overall attention budget among inputs was given to its actual operation of dynamically linking input and output locations. The intuitive interpretation, tightly linked to the need to transform the Key/Query similarity score into a distribution, has been recently challenged, as a growing body of work shows that the attention weights distribution does not directly correlate with predictions (Jain and Wallace, 2019; Pruthi et al., 2019; Brunner et al., 2020). Moreover, Richter and Wattenhofer (2020) recently point out undesirable traits of the softmax operation, demonstrating that its property of confining the outcome to the convex hull of its inputs unnecessarily limits the expressibility of the self-attention mechanism. They experiment on a suite of synthetic tasks with a BERT variant in which the softmax normalization is removed, and find it to perform on par on almost all examined tasks. When replacing the softmax with other normalizations they report improvements. Finally, completely linearized attention (softmax removed) was employed on real tasks as means of reducing costs, since the softmax operation cost scales with the input size (de Brébisson and Vincent, 2016; Wang et al., 2020).

The goal of the above points is not to advocate modifications in BERT’s non-linearity or normalization operations (we leave that to other works), but to note that while these are under examination and are susceptible for alteration, the connectivity of self-attention, manifested by eqs. (3) and (4) , is the core mechanism driving its functionality. Our results, to be presented in section 4, demonstrate how conclusions drawn by directly analyzing this mechanism accord with the operation of commonly employed self-attention networks.

A measure of capacity for modeling input dependencies

In this section, we introduce the separation rank of the function realized by a self-attention network as a measure that quantifies its ability to model dependencies between subsets of its variable set $\{{\mathbf{x}}^{j}\}_{j=1}^{N}$. We will use this measure in order to establish the two depth-efficiency/ inefficiency regimes in self-attention. The separation rank, introduced in Beylkin and Mohlenkamp (2002) for high-dimensional numerical analysis, was employed for various applications, e.g., chemistry (Harrison et al., 2003), particle engineering (Hackbusch, 2006), and machine learning (Beylkin et al., 2009). Importantly, the separation rank has been established as a measure of dependencies modeled by deep convolutional and recurrent networks w.r.t. their inputs (Cohen and Shashua, 2017; Levine et al., 2018a, b).

If the separation rank of a function w.r.t. a partition of its input is equal to $1$, the function is separable, meaning it cannot take into account consistency between the values of $\{{\mathbf{x}}^{j}\}_{j\in A}$ and those of $\{{\mathbf{x}}^{j}\}_{j\in B}$. In a statistical setting, if $y$ is a probability density function, this would mean that $\{{\mathbf{x}}^{j}\}_{j\in A}$ and $\{{\mathbf{x}}^{j}\}_{j\in B}$ are statistically independent. The higher $sep(y;A,B)$ is, the farther $y$ is from this situation, i.e. the more it models dependency between $\{{\mathbf{x}}^{j}\}_{j\in A}$ and $\{{\mathbf{x}}^{j}\}_{j\in B}$, or equivalently, the stronger the correlation it induces between the inputs indexed by $A$ and those indexed by $B$.

The fixed connectivity of ConvNets has been shown to yield high separation ranks w.r.t. partitions which separate neighboring inputs (e.g., where all odd positions are in $A$ and all even positions are in $B$), while suffering from low separation ranks w.r.t. partitions which separate distant inputs (e.g., where $A=1,...,\nicefrac{{N}}{{2}}$ and $B=\nicefrac{{N}}{{2}}+1,...,N$). Our analysis establishes a qualitatively different trait for self-attention networks, which treat all balanced partitions alike:

Accordingly, we will omit the specification of the partition in future uses, denoting $sep(y^{i,L,d_{x},H,\Theta}_{p})$ as the separation rank of $y^{i,L,d_{x},H,\Theta}_{p}$ w.r.t. any balanced partition of the inputs.

This result accords with the intuition regarding the flexibility of the attention mechanism – it does not integrate the input in a predefined pattern like convolutional networks, but dynamically learns to correlate any inter-dependent subsets of the inputs. Natural text exhibits non-smooth non-local dependency structures, as correlations between input segments can abruptly rise and decay with distance. The fact that self-attention facilitates all correlation patterns equally poses it as a more natural architecture for language modeling related tasks. Convolutional networks, with their local connectivity, may have the right inductive bias for imagery data, but partitions unfavored by them may reflect more erratic correlations that are nonetheless relevant for natural language inputs.

However, the above property of indifference to the input partition is not enough for succeeding at tasks with elaborate input dependencies, since a function with equally low separation ranks for all input partitions has limited ability to model such dependencies. In the following section, we analyze how different architectural parameters affect the ability of self-attention networks to correlate their inputs, and by bounding their separation ranks, we establish the different depth-efficiency regimes in self-attention networks.

The effect of depth in self-attention networks

In this section, we present tight bounds on the separation rank of self-attention networks, which reveal two qualitatively different regimes. In the first regime of $L<\log_{3}(d_{x})$, analyzed in subsection 4.1, we establish that deepening is clearly preferable to widening. In the second regime of $L>\log_{3}(d_{x})$, analyzed in subsection 4.2, we show that deepening and widening play a similar role in enhancing the expressiveness self-attention networks.

The recursive structure of deep self-attention hints at an exponential increase of input mixing with depth: The output of each layer is introduced $3$ times into the Key/Query/Value computation made by the subsequent layer. In this subsection, we formalize this intuition for self-attention networks of sufficient width, $d_{x}>3^{L}$. Theorem 1 below bounds the separation rank of such networks. Subsequent to its statement and brief outline of its proof, we explicitly show in corollary 1 the implied double-exponential requirement from a bounded depth network attempting to replicate a deeper one.

For $p\in[d_{x}]$, let $y^{i,L,d_{x},H,\Theta}_{p}$ be the scalar function computing the $p$th entry of an output vector at position $i\in[N]$ of the depth-$L$ self-attention network with embedding dimension $d_{x}$ and $H$ attention heads per layer, defined in eqs. (3) and (4). Let $sep(y^{i,L,d_{x},H,\Theta}_{p})$ be its separation rank (section 3). If $L,d_{x}$ obey $L<\log_{3}\left(d_{x}\right)$, then the following holds almost everywhere in the network’s learned parameter space, i.e. for all values of the weight matrices (represented by $\Theta$) but a set of Lebesgue measure zero:

with $a=-L+\left[2-\log_{3}2\right]$. (note that $\log_{3}\left(d_{x}-H\right)+a>0$ in this regime of $L<\log_{3}(d_{x})$).

We provide below a short proof sketch of the lower bound in the above theorem. The derivation of the upper bound is more straightforward, and is left for the appendix, along with a formal proof of the lower bound.

Proof sketch for the lower bound in theorem 1: We make use of grid tensor based function discretization (Hackbusch, 2012) – The function realized by a self-attention network is evaluated for a set of points on an exponentially large grid in the input space, and the outcomes are stored in a matrix ${\mathcal{M}}\left(y_{p}^{i,L,d_{x},H,\Theta}\right)$, which we prove upholds: $\textrm{rank}\left[{\mathcal{M}}\left(y_{p}^{i,L,d_{x},H,\Theta}\right)\right]\leq sep(y_{p}^{i,L,d_{x},H,\Theta})$, i.e., its rank lower bounds the separation rank. Since the entries of ${\mathcal{M}}\left(y_{p}^{i,L,d_{x},H,\Theta}\right)$ vary polynomially with the self-attention network’s weights, we show that it suffices to find a single network weights assignment $\Theta$ for which the rank of the matrix is greater than the desired lower bound, in order to prove the case for almost all of the configurations of the network’s learned weights (but a set of measure zero). Thus, we prove the lower bound in theorem 1 by choosing a simple weight assignment that still represents the self-attention connectivity, and showing that for this value of $\Theta$, $\textrm{rank}\left[{\mathcal{M}}\left(y_{p}^{i,L,d_{x},H,\Theta}\right)\right]$ achieves the lower bound, in turn lower bounding the separation rank. ∎

Theorem 1 bounds the separation rank of a deep self-attention network of sufficient width between two functions that grow double-exponentially with depth and polynomially with width, tightly describing its behavior w.r.t. depth and width. Because equivalence cannot hold between two functions of different separation ranks, the above result implies a double-exponential requirement from the width of a shallow network attempting to replicate the deep one, and clear depth efficiency holds:

With probability $1$, the function realized upon randomization of the weights of a deep self-attention network defined in eqs. (3) and (4) with depth $L^{\textrm{deep}}$ and width $d_{x}^{\textrm{deep}}>3^{L^{\textrm{deep}}}$, may only be realized by a shallower network with depth $L^{\textrm{shallow}}=\nicefrac{{L^{\textrm{deep}}}}{{d}}$ and width $d_{x}^{\textrm{shallow}}=wd_{x}^{\textrm{shallow}}$, where $d>1,w>1$ (i.e., the deep network is deeper by a factor of $d$ and the shallow network is wider by a factor of $w$), if the following holds:

2 Depth in-efficiency in self-attention

Beyond establishing depth-efficiency in early self-attention layers, the above analysis sheds light on the contribution of a self-attention network’s depth to its ability to correlate input subsets. The separation rank (w.r.t. any partition) of a single layer, given by eq. (3), is only linear in $H$ and $d_{x}$, showcasing a limitation of the class of functions realized by single self-attention layers to model elaborate input dependencies. Theorem 1 quantifies the double exponential growth of this capacity measure with the number of stacked self-attention layers. The following theorem shows that this growth is capped by the dimension of the internal representation:

For $y^{i,L,d_{x},H,\Theta}_{p}$ as defined in theorem 1, if $L>\log_{3}\left(d_{x}\right)$, then the following holds almost everywhere in the network’s learned parameter space, i.e. for all values of the weight matrices (represented by $\Theta$) but a set of Lebesgue measure zero:

with corrections on the order of $L$: $b_{1}=-L\left(\frac{H}{2}+1\right)$, $c_{1}=L$, and on the order of $d_{x}\log_{3}(d_{x})$: $b_{2}=-d_{x}\left(1+\frac{1}{2}\log_{3}\left(\frac{d_{x}-H}{2}\right)\right)$, $c_{2}=-2d_{x}\cdot\log_{3}\nicefrac{{d_{x}}}{{2\sqrt{2e}}}+\log_{3}d_{x}$.

We provide below a proof sketch of the upper bound in the above theorem. The formal proof, along with the proof of the lower bound, which is similar to the one illustrated above for the lower bound in theorem 1, are left for the appendix.

Proof sketch for the upper bound in theorem 2: By observing that $y^{i,L,d_{x},H,\Theta}_{p}$ is a polynomial of degree $2C=3^{L}-1$ ($C$ is introduced in eq. (4)), we find a kernel $\psi\left({\mathbf{x}}^{1},...,{\mathbf{x}}^{N}\right)$ that maps the input into a space where each of the output monomials is a linear functional. We find a basis for the subspace $V$ spanned by the output monomials, and bound the separation rank of each element in that basis by a constant. The dimension of $V$ is exponential in $Nd_{x}$ and polynomial in $3^{L}-1$, providing equal groundings for depth and width. A careful analysis that exploits the sums over the indices $j_{1},...,j_{C}$ in eq. (4), removes the dependence on $N$. ∎

Theorem 2 states that when the network’s depth passes a width dependent threshold, the separation rank turns from increasing polynomially with width and double-exponentially with depth to increasing-exponentially with width and depth together. Thus, while an increase in network size increases its capacity to model input dependencies, our result shows that there is no longer a clear cut advantage of depth in this respect:

Let ${\mathbf{y}}^{\textrm{deep}}$ denote the function realized by a deep self-attention network at any output location $i\in[N]$, defined in eqs. (3) and (4) with depth and width denoted $L^{\textrm{deep}},d_{x}^{\textrm{deep}}$ such that $L^{\textrm{deep}}>\log_{3}d_{x}^{\textrm{deep}}$. Denote $\beta_{1}:=\frac{\log_{3}d_{x}^{\textrm{deep}}}{L^{\textrm{deep}}}<1$. Then, there exists $\beta_{2}=O(log(H)\cdot log(d_{x}^{\textrm{deep}})\cdot log(L^{\textrm{deep}}))$ such that the function realized by a network of depth: $L^{\textrm{shallow}}=\beta_{1}\cdot L^{\textrm{deep}}+\beta_{2}$, and width: $d_{x}^{\textrm{shallow}}=3^{\beta_{2}}d_{x}^{\textrm{deep}}$, denoted ${\mathbf{y}}^{\textrm{shallow}}$, has higher separation rank, i.e.:

The above corollary, which follows from theorems 1 and 2, shows that the separation rank of a function realized by a self-attention network of arbitrary depth $L>\log_{3}(d_{x})$ can be surpassed by a shallower network of polynomial width, contrarily to the established behavior for networks of depth $L<\log_{3}(d_{x})$.

We leave it as an open conjecture that a polynomially sized shallower network can exactly replicate the operation of a deeper network in this regime. With that, we point out that a variety of results which directly bound different complexity measures of deep networks have been put forward, shedding light on their operation (Montufar et al., 2014; Bianchini and Scarselli, 2014; Raghu et al., 2017; Serra et al., 2017; Inoue, 2019). Bounds on the separation rank have been used to explain the operation of more veteran architectures, and we find them to be particularly relevant in the case of self-attention: this complexity measure quantifies the amount of input inter-dependency induced by the network, directly reflecting a widespread intuition on the success behind the self-attention mechanism.

Depth-efficiency regimes in common self-attention networks

In the previous sections, we analyzed a simplified version of self-attention networks (described in section 2). For this class, we proved the existence of the two different depth-efficiency/inefficiency regimes in self-attention networks, and further quantified the transition point between regimes to be exponential in network width (and accordingly in network size). In this section, we demonstrate that our theoretical predictions are manifested in common self-attention networks: the experiments below were conducted over common self-attention architectures which include all operations that were omitted in our theoretical analysis. We describe the training setup in section 5.1, the experiments in section 5.2, and the projection regrading optimal depth to width ratios (see table 1) in section 5.3.

We trained common self-attention architectures of depths $L=6,12,18,24,30,36,48$ and varying widths, such that the network sizes range between $10^{6}$ and $6\cdot 10^{8}$ (full details on the widths of the trained architectures are given in the appendix). We trained decoder-only (unidirectional) models, by optimizing the autoregressive log-likelihood of the training examples. We used a smaller than usual vocabulary size of $2000$ so that the vocabulary embedding parameters, given by $d_{x}\cdot V$ for a vocabulary of size $V$, would constitute a small fraction of the learned parameters for all data points. Autoregressive models were shown to work well even on character level vocabularies (e.g., (Peters et al., 2018)); due to modeling a joint distribution over the text, they are less sensitive to vocabulary size than bidirectional models (Levine et al., 2021).

Our training set was English Wikipedia, BookCorpus and OpenWebText. We report the loss on a held out test set of size $170$K sequences. Notably, we estimated the variance of the pretraining and evaluation procedure by rerunning $11$ of the trained architectures three times each, and found it to be very low – the reported test loss is stable up to its third digit. The remainder of the training details are given in the appendix.

2 Experiments

Figure 2 shows that the predicted devision into two depth-efficiency/inefficiency regimes indeed takes place in common self-attention architectures. When comparing depths $(L^{\textrm{shallow}},L^{\textrm{deep}})=\{(6,12),(12,24),(24,48)\}$, a qualitatively different depth-efficiency behavior is observed as the network size varies. For smaller network sizes, deepening is not favorable over widening. Our theoretical analysis predicts this, showing that when the width of the deeper network is not large enough it can not use its excess layers efficiently. However, when the network’s size is increased by widening, a transition into the depth-efficiency regime is clearly demonstrated: for the same parameter budget the deeper network performs better. Once the deeper network becomes wide enough, such that the depth threshold for depth-efficiency surpasses $L^{\textrm{shallow}}$, it is significantly more expressive.

2.2 Transition between regimes depends exponentially on depth

Importantly, beyond a qualitative match to the two predicted depth-efficiency/inefficiency behaviors, the experiments corroborate our prediction for an exponential dependence of the “depth-efficiency width" — the width for which a network becomes depth-efficient — on the network’s depth. By quantifying this exponential behavior (figure 4), we attain practical guidelines for depth-to-width parameter allocation in a self-attention network of a given size.

Per network depth, we examine the width in which it diverges from the subsequent trained depth, i.e., we examine the following pairs of trained adjacent depths: $(6,12),(12,18),(18,24),(24,30),(30,36),(36,48)$. For each pair, we estimate the shallower network’s transition width (marking the crossing between gray and white areas in figure 2) as the average of its width in two points: the first point in which the shallower network under-performs in a statistically significant manner (see standard deviation estimation in the appendix), and the point to its left in which the performance of the two is not distinguishable. We take the empirical error of this estimation to be the distance between the two points.

Our theoretical results in section 4 predict that the above empirically estimated transition should occur when the shallower network’s width $d_{x}$ is exponential in its depth $L$. Accordingly, we fit a linear dependence of the log of the width on the depth and receive the fit coefficients $(a,b)$: $\log\left(d_{x}\right)=a+b\cdot L$. The linear fit, shown in Figure 3(a) yields measures of $R^{2}=0.998$ and $\chi_{\textrm{red}}^{2}=0.854$ (see further details in the appendix). These measures imply a good compatibility of the theoretically predicted dependence to the measurements, and further reinforce the practical use we make of the fit parameters $a$ and $b$ hereinafter, for predicting the network size for which the regime transition occurs per depth.

Specifically, we insert $d_{x}=e^{a}\cdot e^{bL}$ into the dependence $N=12\cdot L\cdot d_{x}^{2}$ and calculate the transition size and its propagated uncertainty as:

Figure 3(b) shows the empirical transition sizes per depth on top of the projection and its error, calculated by eq. 10 with the fit parameters in eq. 11. Networks to the left of the curve are too shallow given their parameter budget, and can be improved by deepening at the expense of their width.

2.3 “Width-efficiency" in small network sizes

Our experiments reveal an empirical phenomenon that was not predicted by our theory. We established in section 4 that depth does not have an advantage when the width is too small, but our bounds do not separate wider networks from deeper ones in this depth-inefficiency regime. A surprising phenomenon is seen in figures 2(b,c): for small enough network sizes, deeper self-attention networks perform worse than shallow ones. We leave a theoretical treatment of this regime for future work.

The above “width-efficiency" empirical phenomenon leads to an important observation: for a given network size, a certain network can be too shallow, as we predicted theoretically and corroborated empirically above, but it can also be too deep. In other words, the region to the right of the fit curve in figure 3(b) includes networks that can be improved by widening at the expense of their depth. This implies that rather than representing a minimal depth per given self-attention network size, the curve in figure 3(b) represents the area of an optimal depth per network size. We provide a demonstration of this idea in figure 3(c), which clearly shows that when comparing networks of depths $L=6,12,24,48$, each one would be best to use in a different range of network sizes (the color in each range corresponds to the best performing depth in that range).

3 Projecting to larger networks

Beyond reflecting our theoretical predictions, the fit in figure 3 can be used to project beyond the scope of our experiments in order to shed light on architecture design decisions made for much larger self-attention networks, like the contemporary huge Transformer-based language models (Brown et al., 2020; Raffel et al., 2019b; Rosset, 2020). Figure 4 shows the extrapolation of the fitted function and the uncertainty up to networks of depth $100$. Notably, despite the uncertainty growing as the scope extends, $\frac{\Delta N_{\textrm{Transition}}(L=100)}{N_{\textrm{Transition}}(L=100)}=0.2$, i.e., the predictions for the optimal network size in the $L=100$ case are likely to be accurate within $20\%$ of the predicted size, yielding meaningful and unforeseen practical implications.

For example, when examining the architecture of GPT3, the deepest self-attention network trained to date with $96$ layers, we get $N_{\textrm{Transition}}(96)=1.17\pm 0.23\cdot 10^{12}$, or over a Trillion parameters. This places GPT3 with its 175B parameters significantly below our fit, suggesting that it may be too deep given its parameter budget. In fact, the optimal depth for GPT3’s size is predicted to be $L=80$, since $N_{\textrm{Transition}}(80)=1.65\pm 0.25\cdot 10^{11}$. Table 1 includes further suggestion for huge models following our fit, including a suggestion to deepen networks on the left of the curve in figure 4. With high certainty given our experimental data, the optimal model size increase towards 1 Trillion parameter models and beyond is via widening.

Discussion

An apparent “depth-inefficiency" of self-attention networks was pointed out by prior works – in contrast to other successful deep learning architectures, in the case of self-attention there does not seem to be a clear advantage to deepening vs. widening. Our theoretical analysis clearly reflects this behavior in one parameter setting, but suggests an important nuance regarding its origins, while predicting a separate “depth-efficiency" regime in another parameter setting. Rather than an obvious explanation for the observed depth inefficiency, by which the self-attention mechanism does not benefit much from the operation of compounding, our analysis strongly points at the converse: self-attention is so effective at integrating its inputs, that it very quickly reaches saturation in the amount of dependencies that can be supported by the representation dimension.

Thus, for early self-attention compounding, we prove a rapid growth in expressiveness with depth, and specifically in the ability to flexibly correlate between any input locations, which can not be accounted for by any reasonable widening. However, our analysis pinpoints a transition in which the capacity of width to support the above rapid growth exhausts. Thus, when the width of a self-attention network is not large enough, the above depth-efficiency disappears – deepening and widening become equivalent in terms of expressiveness.

We did not find a result which directly upper bounds depth-efficiency in other architecture classes. Works by Sharir and Shashua (2018); Levine et al. (2019) show an exponential growth with depth of a measure related to the separation rank in certain classes of convolutional networks. Comparing this with the double-exponential growth shown in theorem 1 for early self-attention layers, it may be conjectured that convolutional networks seemingly benefit more from depth than self-attention does because their separation rank grows less rapidly, so they do not saturate some width dependent threshold as quickly as self-attention does. We leave these investigations for future work.

The experiments presented in section 5 reveal a qualitative and quantitative match to our theoretical predictions. Beyond reinforcing the validity of our theoretical interpretation, our comprehensive experimental setup allowed us to extrapolate and project depth-to-width trade-offs in huge self-attention networks, that are currently being trained as powerful language models. For example, GPT3, the deepest self-attention network trained to date with $96$ layers, has matched this depth with an unprecedented width of $12$K. However, our projections clearly show that for this number of layers the network should be much wider. In fact, the logarithmic dependence that we establish between the optimal depth and width clearly dictates that size increase should be mainly via widening from this point ($\mathbf{\sim 100B}$ models) onwards. This is good news from an engineering perspective: width can be increased more efficiently than depth in terms of parallelization. The multi-million dollar price tag on these architectures, along with the race to push the envelope towards $1$-Trillion parameter models and beyond, make such informed guidelines an essential ingredient.

Beyond elucidating the behavior of vanilla self-attention architectures, our work theoretically motivates architectural changes that can provide the next leap in self-attention network expressiveness. By indicating the network width as the limiting factor for depth-efficiency, our analysis encourages the development of methods for increasing network width with low expenses. For example, we point at the concept of ShuffleNet (Ma et al., 2018), which has proven to be efficient for convolutional networks. They increase the representation dimension while using only a fraction of it for computation in each layer. This way, the computational costs are contained, but the width related theoretical limitations, posed by our work, are relaxed. Recently, Fedus et al. (2021) trained a $1$-Trillion parameter model via a related approach which learns to choose the subset of parameters to apply in each layer. Indeed, we view our work as part of an effort to provide timely interpretations as feedback for the tremendous empirical pull in our field.

Acknowledgments

We thank Daniel Jannai for assistance in the experiments, and Jared Kaplan for the permission to use the figure in Kaplan et al. (2020). This research was supported by the ERC (European Research Council) and the ISF (Israel Science Foundation). Experiments were performed with Cloud TPUs and supported by Google’s TensorFlow Research Cloud (TFRC). Yoav Levine was supported by the Israel Academy of Sciences Adams fellowship.

References

Appendix A Upper bounds on the separation rank

In this subsection, we prove facts on the general structure of the function realized by the analyzed self-attention architecture that will be of use to us in the upcoming proofs. For a cleaner presentation, we will rewrite eq. (3) in vectorized notation:

We begin by proving a lemma that reveals the structure of $\mathbf{g}^{L}$ presented in eq. (4):

Defining $C\left(L\right)\coloneqq\frac{3^{L}-1}{2}$, any depth $L$ composition of the self-attention layers defined in eq. (3) can be written as:

Now, substituting in the induction hypothesis on the structure of $Y^{\left(L\right)}\left(X\right)$ yields:

Finally unifying the summations over $h,h_{1}h_{2},h_{3}$ to single sum over $\left[H\right]^{\left[C\left(L\right)\cdot 3+1=C\left(L+1\right)\right]}$ gives

Note that the number of $M$ units, each with a summation on a different index $j\in[N]$, is $3C\left(L\right)+1=C\left(L+1\right)$, implying $C\left(L\right)=\frac{3^{L}-1}{2}$ as needed. ∎

Defining $C\left(L\right)\coloneqq\frac{3^{L}-1}{2}$, any depth $L$ composition of $L$ self-attention layers can be written as:

To get the required form, we will use lemma 1 above and write the matrix multiplication in eq. (13) explicitly.

In the next two subsections, we will use the above lemma 1 to prove the two competing upper bounds on the separation rank of self-attention networks.

A.2 Proof of the upper bound in theorem 1

In the following theorem, we show how an upper bound on the separation rank is implied by the form of eq. (13) in the statement of lemma 1.

Defining $C\left(L\right)\coloneqq\frac{3^{L}-1}{2}$, for any depth $L\geq 1$ input size $N>1$ partition $P\mathbin{\mathaccent 0{\cdot}\cup}Q=\left[N\right]$ and output locations $i\in\left[N\right],\>p\in\left[d_{x}\right]$, the following holds:

We begin by writing the matrix multiplication in eq. (13) explicitly.

Therefore, rewriting the summation to be over $\{P_{c}\in\{P,Q\}\}_{c=1}^{C(L)}$ that correspond to the two partition segments $P/Q$.

Now we reorder the above sum by summing over indices of swaps between $P$ and $Q$, i.e. $\beta\in[C]$ such that $P_{\beta}\neq P_{\beta+1}$, and split the multiplication $\prod_{c=1}^{C\left(L\right)}$ according to the crossing indices:

Where we assume w.l.o.g that $i\in P$ and therefore $P_{\beta_{1}},P_{\beta_{1}+1},\dots,P_{\beta_{0}-1},P_{\beta_{0}}=P$. The above reordering allows pushing the summation of non swapping $r_{c}$ indices into the $P,Q$ parentheses:

Since the separation rank of each term in the above summation is $1$, we proved the following upper bound on the separation rank:

We note that unlike the $d_{a}$ case, the same $H$ index can affect nonconsecutive $M^{\left(c_{1},h\right)},M^{\left(c_{2},h\right)}$, therefore we can’t simply push the $h$ indices as done for the $r$ indices in eq. (17). ∎

From here, the upper bound in theorem 1 follows by

A.3 Proof of the upper bound in theorem 2

Defining $C\left(L\right)\coloneqq\frac{3^{L}-1}{2}$, for any depth $L\geq 1$ input size $N>1$ partition $P\mathbin{\mathaccent 0{\cdot}\cup}Q=\left[N\right]$ and output locations $i\in\left[N\right],\>p\in\left[d_{x}\right]$, the following holds:

We begin by opening the inner products in eq. (16), explicitly writing the indices:

And separating between coefficients and ${\mathbf{x}}$’s:

Now we can group monomials by the powers $n_{1},\dots,n_{d_{x}}$ of each coordinate

Finally, we need to bound the separation rank of $\chi_{n_{1}.\dots,n_{d_{x}},\alpha_{C\left(L\right)+1}}$. W.l.o.g we choose the partition $P=\left\{1,\dots,\frac{N}{2}\right\},Q=\left\{\frac{N}{2}+1,\dots,N\right\}$ and $i\in P$ then we can divide the powers between $P,Q$ in the following way:

Thus, since each summand is of separation rank $1$, the separation rank of $\chi_{n_{1}.\dots,n_{d_{x}},\alpha_{C\left(L\right)+1}}$ is bounded by the number of summands:

where the inequality followed from lemma 3. Since we have at most $d_{x}\left(\binom{d_{x}}{2C\left(L\right)}\right)$ different $\chi$ we conclude that:

From here, theorem 2 follows by the multiset identity in lemma 4:

A.4 The effect of residual connections

Having upper-bounded the separation rank of the deep self-attention network defined in section 2.2 of the main text, we comment on the effect of adding residual connections over each layer, as is done in the regular network (described in section 2.1 of the main text). Consider a network composed of a concatenation of the building blocks shown in figure 5(a), taken from Veit et al. (2016). A building block in layer $l$ includes a module $f_{l}$, which in our case is the self-attention layer given in eq. (3) of the main text,We have embedded the Feed-Forward layer within $W^{\textrm{O}}$ due to the linearity of the analyzed model. and a skip connection which adds $f_{l}$’s input to its output (circles denote addition). Veit et al. (2016) propose an unraveled view of such a network, shown in figure 5(b), which we will employ in the proof of theorem 5 for clarity of presentation.

We begin by proving a lemma that quantifies how the separation rank of the composition of a self-attention layer over a function is related to the function’s separation rank:

The lemma follows by multiplying the number of summed terms by an upper bound on the separation rank of each summed term, $K^{3}$. ∎

We now prove a theorem which establishes that the integration of skip connections modifies the upper bound in theorem 1 of the main text by a small factor.

For $p\in[d_{x}]$, let $y^{i,L}_{p}$ be the scalar function computing the $p$th entry of an output vector at position $i\in[N]$ of the depth-$L$ residual network depicted in figure 5, where $f_{l}$ is the self-attention layer in eq. (3) of the main text. Then:

Comparing this dependence to the upper bound in the theorem 1 of the main text, given in eq. (18), this theorem implies that the effect of residual connections is insignificant to our analysis.

Observing figure 5(b) which gives the $L=3$ example, we upper bound the separation rank of the entire unraveled network by noting that its output is composed from $L+1$ additions of outputs from branches of depth $l=0,....,L$ ( being the direct link of the input to the output), such that schematically the separation rank at the output of the entire network can be upper bounded by:

where we denoted $\textrm{longest branch}(L)$ as the function at the output of $f_{L}$, before the addition with the other branches. Noting that the input to $f_{L}$ can be recursively viewed as an output of an unraveled network of depth $L-1$, we bound the separation rank of the function at the input to $f_{L}$ by $L\cdot Sep(\textrm{longest branch}(L-1))$. Since $f_{L}$ is a self-attention layer, Lemma 2 implies that $Sep(\textrm{longest branch}(L))\leq\frac{Nd_{x}^{4}}{H}\left(L\cdot Sep(\textrm{longest branch}(L-1))\right)^{3}$. Continuing recursively, and inserting the stopping condition $Sep(\textrm{longest branch}(L=1))=\frac{Nd_{x}^{4}}{H}$ (since the input to $f_{1}$ is a specific entry of the input to the entire network, of separation rank 1), we attain:

We now prove a theorem which establishes that the integration of skip connections modifies the upper bound in theorem 2 of the main text by a small factor.

Defining $C\left(L\right)\coloneqq\frac{3^{L}-1}{2}$, for $p\in[d_{x}]$, let $y_{p,\text{residual}}^{i,L,d_{x},H,\Theta}$ be the scalar function computing the $p$th entry of an output vector at position $i\in[N]$ of the depth-$L$ residual network depicted in figure 5, where $f_{l}$ is the self-attention layer in eq. (3) of the main text. Then for any partition $P\mathbin{\mathaccent 0{\cdot}\cup}Q=\left[N\right]$, the following holds:

Comparing this dependence to the upper bound in the theorem 2 of the main text, given in eq. (19), the above theorem implies that the effect of residual connections is insignificant to our analysis.

We will adapt the proof of theorem 4. All of the arguments remain unchanged, except that we obtain the network structure via lemma 5 instead of lemma 1. Following the new structure we will have two additional summations, one over $j$ and one over $\alpha$ (see lemma 5), as well as an additional input $X$ factor. We will leave the summation over $j$ during the whole proof, thus multiplying the separation rank by at most $C\left(L\right)$. Note that similarity to the $h$ summation, the summation over $\alpha$ has no influence on the separation rank, since it collapses into a single coefficient $\mathcal{T}$. Finally the unput $X$ factor contribute at most $1$ to the separation rank, therefore we can bound the separation rank by $C\left(L\right)+1$ times the bound in eq. (19). ∎

Finally, since the upper bounds undergo such minor increases in the presence of skip connections, the lower bounds can be left with no further tightening, without affecting the analysis and its conclusions.

A.5 Technical lemmas

(inequality of arithmetic and geometric multiset coefficient means)

Define $f_{t}\coloneqq\prod_{j=1}^{k}\left(r_{j}+t\right)$ and $\psi\coloneqq\prod_{t=1}^{n-1}f_{t}$ than by the inequality of arithmetic and geometric means

One can see that when $M$ divided by $k$ it hold that

$\left(\binom{n}{k}\right)\leq\left(\frac{2e\left(n+k\right)}{n}\right)^{n}$

: by using the inequality $\binom{n}{k}\leq\left(\frac{en}{k}\right)^{k}$ we have

We now prove a lemma which reveals the alternation to the network structure as expressed in eq. (13) when taking skip connections into account.

Defining $C\left(L\right)\coloneqq\frac{3^{L}-1}{2}$, any depth $L$ skip connection composition of the self-attention layers defined in eq. (3) can be written as:

Now, rewriting $Y^{\left(L\right)}$ as $X+\left(Y^{\left(L\right)}-X\right)$ yields:

Now, substituting in the induction hypothesis on the structure of $Y^{\left(L\right)}\left(X\right)$ yields:

Similarly to eq. (14) each of the $8$ terms in the outer summation is of the required form, thus we complete the proof. ∎

Appendix B Lower bounds on the separation rank

We begin by laying out basic concepts in tensor theory required for the upcoming analysis. The core concept of a tensor may be thought of as a multi-dimensional array. The order of a tensor is defined to be the number of indexing entries in the array, referred to as modes. The dimension of a tensor in a particular mode is defined as the number of values taken by the index in that mode. If ${\mathcal{A}}$ is a tensor of order $N$ and dimension $M_{i}$ in each mode $i\in[N]$, its entries are denoted ${\mathcal{A}}_{d_{1}...d_{N}}$, where the index in each mode takes values $d_{i}\in[M_{i}]$.

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

The following claim establishes a fundamental relation between a function’s separation rank (see section 3) and the rank of the matrix obtained by the corresponding grid tensor matricization. This relation, which holds for all functions, is formulated below for functions realized by self-attention networks:

where ${\mathcal{A}}(y^{i,L,d_{x},H,\Theta}_{p})$ is the grid tensor of $y^{i,L,d_{x},H,\Theta}_{p}$ with respect to the above template vectors.

where $V^{\nu}$ and $U^{\nu}$ are the tensors holding the values of $g_{\nu}^{I}$ and $g_{\nu}^{J}$, respectively, at the points defined by the template vectors. Under the matricization according to the $(I,J)$ partition, it holds that $\llbracket V^{\nu}\rrbracket_{I,J}$ and $\llbracket U^{\nu}\rrbracket_{I,J}$ are column and row vectors, respectively, which we denote by ${\mathbf{v}}_{\nu}$ and ${\mathbf{u}}_{\nu}^{T}$. It follows that the matricization of the grid tensor is given by:

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

Claim 1 assures us that the separation rank of the function realized by a self-attention network is lower bounded by the rank of the matrix obtained by the corresponding grid tensor matricization, for any choice of template vectors. Specifically:

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

Noting that the operation of raising a rank $r$ matrix to the Hadamard power of $p$ results in a matrix upper bounded by $\left(\kern-2.70003pt\left(\genfrac{}{}{0.0pt}{}{r}{p}\right)\kern-2.70003pt\right)$ (see proof in Amini et al. (2012) for example) with the notation of the multiset coefficient $\left(\kern-2.70003pt\left(\genfrac{}{}{0.0pt}{}{n}{k}\right)\kern-2.70003pt\right):=\binom{n+k-1}{k}$, and that the rank of $VV^{T}$ is upper bounded by $\nicefrac{{(d_{x}-H)}}{{2}}$, we choose the dimension $\nicefrac{{M}}{{2}}=\left(\kern-2.70003pt\left(\genfrac{}{}{0.0pt}{}{\nicefrac{{\left(d_{x}-H\right)}}{{2}}}{3^{L-2}}\right)\kern-2.70003pt\right)$ to facilitate the rank increase.

and the log of this bounds the expressions in the theorems’ lower bounds, where for each regime the tighter lower bound is used.

for $\{{\mathbf{a}}^{(k)}\}_{k=1}^{\left(\kern-1.89001pt\left(\genfrac{}{}{0.0pt}{}{d}{\lambda}\right)\kern-1.89001pt\right)}$ and $\{{\mathbf{b}}^{(k)}\}_{k=1}^{\left(\kern-1.89001pt\left(\genfrac{}{}{0.0pt}{}{d}{\lambda}\right)\kern-1.89001pt\right)}$ which are two sets of linearly independent vectors.

For $\alpha,\beta\in[\left(\kern-2.70003pt\left(\genfrac{}{}{0.0pt}{}{d}{\lambda}\right)\kern-2.70003pt\right)]$, observing an entry of $\left(VV^{\top}\right)^{\odot\lambda}$:

We show this is the case for $V$ in which the rows are each associated with one of $\left(\kern-2.70003pt\left(\genfrac{}{}{0.0pt}{}{d}{\lambda}\right)\kern-2.70003pt\right)$ configurations of distributing $d$ integer numbers that sum up to $\lambda$, i.e., in which each row is associated with specific $\left\{q_{1}^{\alpha},\ldots,q_{d}^{\alpha}\geq 0,\sum_{r=1}^{d}q_{r}^{\alpha}=\lambda\right\}$. Explicitly, we take the rows ${\mathbf{v}}_{r}^{\left(\alpha\right)}$ to be:

Given this $V$, each vector in the above defined set $\left\{{\mathbf{a}}^{(k_{1},\ldots,k_{d})}\right\}_{k_{1}+\cdots+k_{d}=\lambda}$ is equal to:

Observing that the factor attained from the normalization depends only on the rows and doesn’t vary with the different vectors labeled by $(k_{1},\ldots,k_{d})$, we note it does not affect their linear dependence (amounts to a multiplication by a diagonal matrix with non-zero entries on the diagonal - does not affect the rank).

Since the elements of $A$ are polynomial in $\Omega$, then as lemma 7 shows, it is sufficient to show that there exists a single contributor to the determinant of $A$ that has the highest degree of $\Omega$ in order to ensure that the matrix is fully ranked for all values of $\Omega$ but a finite set, so $\Omega$ should simply be chosen to be any number that is outside of this set. Observing the summands of the determinant, i.e. $\Omega^{\sum_{q_{1}+\cdots+q_{d}=\lambda}\left\langle{{\mathbf{q}}},{\sigma({\mathbf{q}})}\right\rangle}$, where $\sigma$ is a permutation on the columns of $A$, lemma 8 assures us the existence of a strictly maximal contributor, satisfying the conditions of lemma 7, thus the set $\left\{\hat{{\mathbf{a}}}^{(k_{1},\ldots,k_{d})}\right\}_{k_{1}+\cdots+k_{d}=\lambda}$ is linearly independent, and the lower bounds in the theorems follow. ∎

B.3 Technical lemmas

The following lemma details the assignment of the self-attention network weights and the choice of template vectors which help us establish theorem 1.

We present below a choice of weights and template vectors that yields the stated form for a sub-matrix of $\llbracket{\mathcal{A}}(y^{i,L,d_{x},H,\Theta}_{p})\rrbracket_{I,J}$. Subsequently we will plug these values into the self-attention operation stated in eq. (3), and prove that this form follows.

Though the proof has many technical details, it has 3 essential parts. We first choose the weights of the first layer so that the outputs in all locations are the same and equal to a summation of the input vectors. Because the weight matrices are not $d_{x}\times d_{x}$ but are decomposed through the attention dimension $d_{a}\times d_{x}$ or $d_{x}\times d_{a}$, then we divide the coordinates of the $d_{x}$-dimensional vectors into contiguous segments of length $d_{a}$, and set the weights to either project these segments to the $d_{a}$-dimensional space or invert this mapping with added zero-padding. For the second part, we set the key and query matrices to use the same “projections” we used in the first layer to compute inner-products between each segment, while setting the value and output matrices to preserve each head’s segment (with zero-padded coordinates). For the remainder of the network’s layers, we use the previous step to compute increasingly larger powers of the norm of the vector computed in the first layer, by reconstructing the squared-norm from the inner products of each segment. The template vectors (and parameters) are chosen such that the square of this norm will be equal to $VV^{T}$.

In the above, we denoted the complex root of $-1$ as $\mathbf{i}$, to differentiate it from the index $i$. The choice of template vectors:

where $\phi(j)\equiv\left\lfloor\nicefrac{{j-1}}{{d_{a}}}\right\rfloor\cdot(d_{a}-1)+(j-1\bmod d_{a})+1$.

W.l.o.g. we can assume that $I=\{1,\ldots,\nicefrac{{N}}{{2}}\},J=\{\nicefrac{{N}}{{2}}+1,\ldots,N\}$. We examine the sub-matrix defined by the following indices:

With all of the above in place, we are ready to prove that the resulting sub-matrix has the form of eq. (28). We begin with the output of the first self-attention layer:

where $(1)$ is because $W^{Q,1,h}=W^{K,1,h}$ are matrices that are zero everywhere except for entry $(1,d_{a})$, $(2)$ because when summing over the locations, only $i_{1}$ and $j_{1}$ are different from $M+1$, and $(3)$ because applying the value and output matrices on any template vector ${\mathbf{u}}$ results in:

At this point, notice that for any $i\in[N]$, ${\mathbf{y}}^{(1,i)}$ is the same, and we denote it with ${\mathbf{v}}$. Note that it is a vector composed of $H$ $d_{a}$-dimensional sub-vectors, each composed of a $\frac{d_{a}-1}{2}$-dimensional sub-vector and its complement in the next $\frac{d_{a}-1}{2}$ indices, followed by a fixed value of $1$.

Next, we will compute the result of the second layer, where we use the fact that every position is equal to ${\mathbf{v}}$ to drop the reference to a specific location $i$, i.e., ${\mathbf{y}}^{(l,i)}={\mathbf{y}}^{(l)}$:

Using the initial conditions of $\alpha_{3}=4$ and $\beta_{3}=2$, we get that $\alpha_{l}=\frac{3^{l-1}-1}{2},\beta_{l}=3^{l-2}-1$. For the $L$’th layer, the only difference is that $W^{V,L,h}$ is defined such that it returns a 1-hot vector that picks the $d_{a}$’th element of the previous step. Putting it all together we get:

Next, we show two lemmas that aid in the proof of the lower bound. We first quote an identity by which for a matrix with entries that are polynomials in $x$, if a single contributor to the determinant has the highest degree of $x$, then the matrix is fully ranked for all values of $x$ but a finite set.

The following quoted lemma, establishes a relation referred to as the vector rearrangement inequality, which helped us ensure that our matrix of interest upholds the conditions of lemma 7 and is thus fully ranked.

We rely on theorem 368 in (Hardy et al., 1952), which implies that for a set of non-negative numbers $\{a^{(1)},\ldots,a^{(N)}\}$ the following holds for all $\sigma\in S_{N}$:

with equality obtained only for $\sigma$ which upholds $\sigma(i)=j\iff a^{(i)}=a^{(j)}$. The above relation, referred to as the rearrangement inequality, holds separately for each component $j\in[\bar{R}]$ of the given vectors:

Appendix C Proof of Proposition 1 on the separation rank symmetry

Assuming that $sep(y;A,B)=R$, then there exist $g_{1},\dots,g_{R},g^{\prime}_{1},\dots,g^{\prime}_{R}$ s.t.

Appendix D Experimental details

We conducted the network training described in section 5 of the main text with Adam optimizer for $1M$ steps and a batch size of $512$ sequences of $128$ tokens. All experiments used a learning rate schedule with a $12000$ step linear warm-up followed by a cosine decay to zero. In order to increase width without changing other architectural parameters, we kept the number of heads per layer constant at $2$ (experimental evidence indicates that many heads per layer are not crucial (Michel et al., 2019; Kaplan et al., 2020), as does our theoretical analysis which shows that the number of heads per layer affects the separation rank logarithmically).

Table 2 shows the per-depth widths of the trained architecture. More experiments were conducted per adjacent depth pairs in order to identify the transition point accurately, and reduce the error bars in figure 3. Table 3 details the different standard deviation of repeating the training and evaluation experiment $3$ times per the given architectures.

The estimated experimental transition points between the two depth-efficiency regimes that were collected according to the procedure described in section 5.2.2 are given in table 4(a). For the linear fit we set $x_{i}$ to be the depth and $y_{i}$ the log of the estimated width at the measured transition point (with an empirical error $\sigma_{i}$ calculated as in table 4(a)).

The $\chi^{2}_{red}$ measure is calculated by:

where $\hat{y_{i}}$ is the predicted value of the $i$-th sample according to the fitting function given by the fit parameters $a$ and $b$ in eq. 11, and $df=n-m=3$ is the number of observations $n=5$ minus the number of fitted parameters $m=2$. The attained value of $\chi_{\textrm{red}}^{2}=0.854$ indicates a good fit for such a low $n$, though hinting at a slight overestimation of the empirical errors. This may arise due to the limitations in attaining very dense measurements around the transition points (though as can be seen in table 2, we made an effort to sample the loss densely around the transitions).

where $\bar{y}=\frac{1}{n}\sum_{i=1}^{n}{\left(\frac{1}{\sigma_{i}^{2}}y_{i}\right)}$. The attained value of $R^{2}=0.998$ indicates a good linear fit of the data.

$\Delta\left(\log d_{x}\right)=\frac{1}{d_{x}}\Delta d_{x}$