$O(n)$ Connections are Expressive Enough: Universal Approximability of Sparse Transformers

Introduction

Transformer networks and their variants have played a key role in the recent advancement of the state of the art in many natural language processing tasks, such as machine translation , language modeling , and question answering . The key component of these networks is the self-attention layer , which updates the embeddings of the input tokens based on their context. Naturally, the self-attention layer also plays the key role in the analysis of Transformers ; for example, Yun et al. show that Transformers can approximate any continuous sequence-to-sequence functions (i.e., universal approximation), by proving that self-attention layers can compute contextual mappings of the input embeddings.

On the other hand, the self-attention layer is also the main bottleneck in scaling these models. It involves computation of pairwise inner products between input tokens, which results in quadratic computational complexity $O(n^{2})$ in the length of the input sequence $n$. To mitigate this issue, researchers have developed methods to sparsify the pairwise interactions/connections in self-attention layers to reduce the computational complexity and/or improve model interpretability, and have shown successful empirical results on tasks with long sequence lengths . For example, Child et al. propose sparse Transformers for sequence generation. One of the sparsity patterns considered in is the Strided pattern, where the sparse attention layers alternate between two patterns: each token attends to only i) $w$ local neighbors, and then ii) one after every $w$ tokens in a strided manner. By choosing $w=O(\sqrt{n})$, they propose sparse attention layers with $O(n^{3/2})$ connections and show improvements on both speed and performance over the dense Transformer.

In the existing results, the rule of thumb for designing sparsity patterns (e.g., Strided) is connectivity; the intuition is that if each token can attend to the other tokens in multiple “hops,” then the resulting sparse Transformers do not lose much expressive power. However, there has been no formal justification for this intuition. How does sparsifying the interaction in the self-attention layers affect the model’s expressive power and ability to learn? What are the sparsity levels at which the model still retains its rich expressive power, and how is it affected by the sparsity pattern? Such fundamental questions about sparse attention models still remain unanswered.

In this paper, we take the first step towards a theoretical understanding of sparse Transformers.

We propose a unified framework to analyze sparse Transformers, which generalizes the existing approaches that sparsify attention layers (§ 3.1).

We propose a set of intuitive conditions on the sparsity pattern (Assumption 1) and the probability map (Assumption 2). Then, in Theorem 1, we show that Sparse Transformers, of fixed width and arbitrary depth, satisfying these conditions are universal approximators of any continuous sequence-to-sequence functions for any given fixed sequence length (§ 3.2 and § 3.3).

We next show some examples of existing sparse Transformers that satisfy these conditions, and hence have universal approximability (§ 3.4). Surprisingly, we show that there are sparse Transformers with only $O(n)$ connections per self-attention layer (instead of $n^{2}$) that have enough expressive power to approximate arbitrary continuous functions (Corollary 2).

We report experimental results on standard NLP tasks using sparse Transformers, comparing different sparsity patterns/levels (§ 5).

Preliminaries and related works

In this section, we summarize the notation we will use throughout the paper, give a brief overview of Transformers, and then discuss existing efforts to sparsify the self-attention mechanism.

2 Transformers and their universal approximation power

3 Sparse Transformers

The first category reduces computation by making ${\bm{A}}^{i}$ sparse in a pre-determined manner. Each token in the sequence only attends to a fixed smaller set of other tokens instead of the whole sequence . In some papers, auxiliary tokens are added to improve connectivity between existing tokens while maintaining sparsity . One drawback of these approaches is that the sparsity pattern is independent of input, so it cannot adapt to the data. To remedy this issue, proposes to learn local attention span from data. In a concurrent paper, Zaheer et al. propose the BigBird sparsity pattern which falls into this category. For BigBird, the authors show its theoretical properties such as universal approximation and Turing completeness, as well as its superior empirical performance. We note that our paper focuses on universal approximation for a broader class of sparse Transformers, by proposing a unifying framework to analyze them.

The second category studies making ${\bm{A}}^{i}$ sparse after the full ${\bm{A}}^{i}$ has been computed . Here, the focus is not on the computational gain via sparsity, because the full score matrix ${\bm{A}}^{i}$ has to be computed first; rather, the goal here is to make attention layers more interpretable, as well as to improve performance. This line of works modifies $\sigma_{\rm S}$ in (1a) to other probability maps, by using top-$k$ elements or adopting sparser variants such as sparselin-gen or $\alpha$-entmax . Compared to the first category, this approach has an advantage that sparsity patterns are adaptive to data.

The last category attempts to get the best of both worlds. This line of works tries to learn sparsity patterns from data using extra components predicting the connection between tokens, e.g., $k$-means clustering , LSTM , or locality-sensitive hashing . This way, one can adaptively determine the sparsity patterns before computing the score matrix. However, the drawback of this approach is that one needs extra computation to train/run these additional components, which may be expensive.

Universal approximation theorem for sparse Transformers

In this section, we derive a unifying framework to study sparse Transformers. We then propose a set of conditions on the sparse self-attention layers, and prove that the sparse Transformers satisfying theses conditions are universal approximators of any continuous sequence-to-sequence functions. Finally, we show some examples of existing sparse Transformers that satisfy these conditions.

We modify the Transformer block in (1) to the following sparse Transformer block (${\rm STB}$):

where the sets $\mathcal{A}_{k}^{l}\subseteq[n]$, for $k\in[n]$ and $l\in[p]$, define the $p$ sparsity patterns (formally defined below), which are indexed by $l\in[p]$. Moreover, the parameter dimensions stay the same as in (1).

Note that there are three main modifications from the dense Transformer.

(Sparsity patterns) Note that ${\rm SHead}^{i,l}({\bm{X}})_{k}$ denotes the $k$-th column of the $i$-th sparse attention head. Unlike dense Transformers, the inner product of the $k$-th query vector ${\bm{W}}_{Q}^{i}{\bm{X}}_{k}$ is taken only with ${\bm{W}}_{K}^{i}{\bm{X}}_{\mathcal{A}_{k}^{l}}$, the key vectors of tokens in the set $\mathcal{A}^{l}_{k}\subseteq[n]$. Hence, instead of all $n$ tokens, the $k$-th token computes attention scores with only tokens in $\mathcal{A}^{l}_{k}$. For $l\in[p]$, we refer to the collection of the index sets $\{\mathcal{A}^{l}_{k}\}_{k\in[n]}$, or simply $\{\mathcal{A}^{l}_{k}\}$, as a sparsity pattern. As a result, ${\rm SHead}^{i,l}({\bm{X}})_{k}$ is a linear combination of columns in ${\bm{W}}_{V}^{i}{\bm{X}}_{\mathcal{A}_{k}^{l}}$, rather than the whole sequence.

(Probability map) After computing the attention score matrix, the dense Transformer (1) uses the softmax operator $\sigma_{\rm S}$ to get a column stochastic matrix. In the sparse Transformers, we generalize $\sigma_{\rm S}$ to $\rho$. The probability map $\rho$ is any map that takes a matrix as input and outputs a column stochastic matrix.

As a sanity check, by choosing $p=1$, $\mathcal{A}_{k}^{1}=[n]$ for all $k\in[n]$, and $\rho=\sigma_{\rm S}$, we recover the dense Transformer (1). Note also that the sparse Transformer formulation covers the first and second categories of existing results discussed in § 2.3. The first category corresponds to choosing a predetermined sparsity pattern(s) $\{\mathcal{A}_{k}^{l}\}$, while setting $\rho=\sigma_{\rm S}$. The second category corresponds to opting for a probability map $\rho$ other than softmax $\sigma_{\rm S}$, while maintaining $\mathcal{A}_{k}^{1}=[n]$ for all $k\in[n]$.

In this paper, we assume for simplicity that all sparse attention heads ${\rm SHead}^{1,l},\dots,{\rm SHead}^{h,l}$ in a single layer have identical sparsity patterns $\{\mathcal{A}^{l}_{k}\}$. However, since our result only requires two sparse attention heads per layer (as we will see in Theorem 1), our result can be easily extended to the case that allows multiple sparsity patterns in a single layer.

Similar to $\mathcal{T}^{h,m,r}$ in § 2.2, we define the class of functions represented by sparse Transformers. We hide the dependence of this class on the sparsity patterns and probability map to simplify the notation.

2 Conditions on sparsity patterns and probability map

In this section, we define a set of conditions on the sparsity patterns $\{\mathcal{A}^{l}_{k}\}$ and the probability map $\rho$ that ensures that the sparse Transformer universally approximate the function class $\mathcal{F}$ (cf. § 2.2).

For $k\in[n]$ and the index sets $\{\mathcal{A}^{l}_{k}\}_{l\in[p]}$, we define a sequence of sets $\{\mathcal{S}^{t}_{k}\}_{t\geq 1}$ in a recursive way:

The set $\mathcal{S}^{t}_{k}$ is the set of all tokens that the $k$-th token can directly/indirectly attend to, after $t$ sparse attention layers with sparsity patterns cycling through $\{\mathcal{A}^{1}_{k}\},\{\mathcal{A}^{2}_{k}\},\dots,\{\mathcal{A}^{p}_{k}\}$. We now state our conditions on sparsity patterns.

The sparsity patterns $\{\mathcal{A}_{k}^{l}\}$ satisfy the following:

For all $k\in[n]$ and $l\in[p]$, we have $k\in\mathcal{A}_{k}^{l}$.

There exists a permutation $\gamma:[n]\to[n]$ such that, for all $i\in[n-1]$, $\gamma(i)\in\bigcup_{l=1}^{p}\mathcal{A}^{l}_{\gamma(i+1)}$.

Assumption 1.1 is equivalent to saying that every token always attends to itself. Assumption 1.2 requires that there is a chain of direct connections that covers all $n$ tokens; note that the set $\bigcup_{l=1}^{p}\mathcal{A}^{l}_{\gamma(i+1)}$ is the set of all tokens that the $\gamma(i+1)$-th token directly attends to. To elaborate more about the chain, consider a directed graph with $n$ vertices corresponding to the $n$ tokens. For any $j\in\bigcup_{l=1}^{p}\mathcal{A}^{l}_{k}$, we add a directed edge $j\to k$. Given a graph constructed this way, Assumption 1.2 requires that the graph has a Hamiltonian path $\gamma(1)\to\gamma(2)\to\cdots\to\gamma(n)$. Assumption 1.3 requires that after $s$ sparse attention layers, every token can attend to all the other tokens, either directly or indirectly.

As we discuss in § 3.4, the statements in Assumption 1 are natural enough to be satisfied by many existing sparsity patterns studied in the literature. In fact, Assumption 1.3 is necessary for universal approximation. If $p=1$, $n=2$, $\mathcal{A}^{1}_{1}=\{1\}$ and $\mathcal{A}^{1}_{2}=\{1,2\}$, then the first token never attends to the second, so this sparse Transformer cannot approximate a function whose first output token is dependent on both input tokens. The other two assumptions are required in parts of our proof, which involve “propagating information” over all the tokens in a sequential manner.

We now state the assumption on the probability map $\rho[\cdot]$. For this, we define $\sigma_{\rm H}[\cdot]$ to be the hardmax operator, which outputs the one-hot representation of the $\operatorname*{arg\,max}$ entry for each column of the input matrix. Since $\rho$ is a column-wise operator that outputs a column-stochastic matrix, we state the assumption for the operation of $\rho$ on a single column.

For any $\zeta>0$ and $\eta\in(0,1]$, $\exists~{}t>0$ such that, for any column input ${\bm{v}}$ satisfying ${v}_{j^{*}}-\max_{j\neq j^{*}}{v}_{j}\geq\zeta$ (where $j^{*}=\operatorname*{arg\,max}_{j}{v}_{j}$), we have $\rho[t{\bm{v}}]_{j^{*}}\geq 1-\eta$ and $\sum_{j\neq j^{*}}\rho[t{\bm{v}}]_{j}\leq\eta$.

Assumption 2 requires that, for inputs that have some margin between the unique maximum entry and the other entries, $\rho[\cdot]$ can closely approximate the behavior of the hardmax operator by scaling its input by a positive factor $t$. This assumption is satisfied by softmax $\sigma_{\rm S}$ and other sparse variants such as sparselin-gen and $\alpha$-entmax, as we show in § B of the supplementary material.

It is straightforward to check that the dense Transformer, which corresponds to $p=1$, $\mathcal{A}_{k}^{1}=[n]$, and $\rho[\cdot]=\sigma_{\rm S}[\cdot]$ in our framework, satisfies both Assumptions 1 and 2.

3 Sparse Transformers are universal approximators

The key justifying intuition for adopting sparse attention layers is that, if each token can attend to the other tokens in multiple hopsNote that this corresponds to our Assumption 1.3., then these models do not lose too much expressive power. However, turning this intuition into a rigorous analysis is not straightforward. Moreover, recent results show that limited width can render universal approximation impossible even with arbitrary depth , highlighting the challenges in analyzing sparse (limited “width”) Transformers.

We now state our main theorem, which shows that if the sparsity patterns $\{\mathcal{A}_{k}^{l}\}$ and the probability map $\rho$ satisfy Assumptions 1 and 2, sparse Transformers with $h=2$ attention heads of size $m=1$, and hidden layer width $r=4$ are universal approximators of continuous sequence-to-sequence functions on any compact domain (recall that $\mathcal{F}$ denotes the class of such continuous functions).

Consider any $f\in\mathcal{F}$, and the class of sparse Transformers $\mathcal{ST}^{2,1,4}$ (cf. (3.1)) with the underlying sparse attention layers satisfying Assumptions 1 and 2. Then, for any $\epsilon>0$ and $1\leq p<\infty$, there exists a function $g\in\mathcal{ST}^{2,1,4}$ such that

As discussed earlier, dense Transformers satisfy Assumptions 1 and 2, which means that Theorem 1 subsumes the existing result for dense Transformers. We note that the required $h$, $m$, and $r$ in Theorem 1 are independent of $d$, $n$, or the sparsity patterns. We provide a high-level proof sketch of Theorem 1 in § 4.1. There, we also discuss how many layers are sufficient for $\epsilon$-approximation of $f$, and show that Theorem 1 requires only $p$ times more self-attention layers than Yun et al. .

We would like to emphasize that Theorem 1 provides the first formal evidence that well-designed sparse attention layers do not limit Transformer’s universal approximation power. In § 3.4, we show a surprising fact that some existing sparse self-attention layers with only $O(n)$ connections (as opposed to $n^{2}$ in regular self-attention layers) retain enough expressive power to approximate $\mathcal{F}$. Combined with the number of layers analyzed in § 4.1, this means that our analysis reduces the connections per layer from $n^{2}$ to $O(n)$, with only $p$ times more attention layers. This advantage of sparse Transformers over their dense counterpart becomes even stronger with increasing sequence length $n$, providing a theoretical support for the adoption of sparsity for tasks with long sequence lengths.

On a final note, Theorem 1 views the sequence length $n$ as a fixed constant. Hence, our result does not contradict a recent paper by Hahn which studies the limitation of Transformers for varying $n$. Also, our analysis applies to the encoder part of the Transformer network .

4 Analysis of existing sparse Transformers

By Theorem 1, any sparse Transformer that satisfies our Assumptions 1 and 2 has universal approximation ability. In this section, we give some examples of such sparse Transformers.

Child et al. propose two kinds of $2$-step sparsity patterns (i.e., $p=2$) for sequence generation tasks, namely Strided and Fixed patterns. We consider the extension of their auto-regressive patterns (i.e., attending only to past tokens) to the whole sequence. In the Strided pattern, a token first attends to its $w$ neighbors and then attends to one token after every $w$ tokens in a strided manner. The sparsity pattern for the $k$-th token reads

In the Fixed pattern, we divide the token into segments of length $w$. A token in a segment has access to other tokens in the same segment, and then the last tokens of the other segments:

The Strided and Fixed patterns satisfy both Assumption 1 and 2 for all values of $w$. Specifically, Assumption 1.3 holds with $s=2$, because any token can directly/indirectly access all the tokens in two hops. As for Assumption 1.2, the identity permutation $\gamma(i)=i$ suffices to satisfy the assumption for both patterns. By choosing $w=O(\sqrt{n})$, sparse Transformers with the Strided and Fixed patterns achieve universal approximation power with $O(n^{3/2})$ connections per attention layer.

Guo et al. consider the Star sparsity pattern where they add an auxiliary relay token that attends to all the tokens, and the other tokens attend only to $2w$ neighboring tokens and the relay token. There is only one sparsity pattern, so $p=1$. The Star sparsity pattern can be written as

where $w\geq 1$. For any fixed $w$, this sparse Transformer has $O(n)$ connections per attention layer, and it satisfies both assumptions. Specifically, Assumption 1.2 is satisfied with the identity permutation, i.e., $\gamma(i)=(i)$ for $i\in[n]$. Since any token can access other tokens within two hops, Assumption 1.3 is satisfied with $s=2$. This demonstrates that $O(n)$ connections per layer suffice for sparse attention layers to have universal approximation power. One can similarly check that the sliding window sparsity patterns with/without global attention, proposed in Longformer , also satisfy the assumptions with $O(n)$ connections. For the BigBird sparsity pattern , it is also straightforward to check that a combination of its window attention and global attention satisfies Assumption 1 with $O(n)$ connections. We state this interesting observation as a corollary below.

There exist sparse Transformers with $O(n)$ connections per self-attention layer that are universal approximators in the sense of Theorem 1.

Recall that another line of results that replaces softmax $\sigma_{\rm S}$ with sparse variants $\rho$ also fits into our formulation, with $p=1$ and $\mathcal{A}^{1}_{k}=[n]$. As we show in § B, these alternative $\rho$’s satisfy Assumption 2. Thus, by Theorem 1, these models also have the universal approximation property.

Proof sketch and discussion

Step 2. We then approximate $\overline{f}\in\overline{\mathcal{F}}(\delta)$ with a sparse Transformer network with a slightly modified architecture. In this architecture, we replace ${\rm ReLU}$ in the feed-forward layer with any piecewise linear activation $\phi\in\Phi$, where $\Phi$ denotes the class of (possibly discontinuous) piecewise linear functions with three pieces. We also replace $\rho$ in the sparse attention layer with the hardmax $\sigma_{\rm H}$ operator. We refer to the function class represented by the modified sparse Transformer as $\overline{\mathcal{ST}}^{h,m,r}$. By a careful construction, Lemma 3 shows that any $\overline{f}\in\overline{\mathcal{F}}(\delta)$ can be exactly represented by the modified Transformer. To this end, we first carefully choose the positional embedding ${\bm{E}}$. We then quantize the inputs using feed-forward layers (Lemma 6), construct a contextual mapping using self-attention layers to map the quantized inputs to unique “ids” (Lemma 7), and then construct a value mapping with feed-forward layers to map the ids to desired output values (Lemma 8). See § D and § E in the supplementary material for details.

Step 3. The final step is to approximate the function $\overline{g}\in\overline{\mathcal{ST}}^{2,1,1}$ with a sparse Transformer $g\in\mathcal{ST}^{2,1,4}$. This is done by approximating $\phi$ and $\sigma_{\rm H}$ with ${\rm ReLU}$ and $\rho$, respectively, while carefully bounding the accumulation of errors introduced by the approximation. See § F in the supplementary material for the details.

For $\overline{g}\in\overline{\mathcal{ST}}^{2,1,1}$ in Lemma 3, there exists $g\in\mathcal{ST}^{2,1,4}$ such that $\mathsf{d}_{p}(\overline{g},g)\leq\epsilon/2$.

Combining these three steps, we establish that $\mathsf{d}_{p}(f,g)\leq\mathsf{d}_{p}(f,\overline{f})+\mathsf{d}_{p}(\overline{f},\overline{g})+\mathsf{d}_{p}(\overline{g},g)\leq\epsilon$. ∎

How many layers are sufficient? In § D, Lemmas 6–8 show that we need $\frac{dn}{\delta}$ sparse Transformer blocks (2) for quantization, $\frac{p(n-1)}{\delta^{d}}+s$ for the contextual mapping, and $\frac{n}{\delta^{dn}}$ for the value mapping. Recall that $p$ is from (2), $s$ is from Assumption 1, and $\delta$ is from Step 1 above. In comparison, § C of shows that the dense counterpart requires $\frac{dn}{\delta}$, $\frac{n}{\delta^{d}}+1$, and $\frac{n}{\delta^{dn}}$ Transformer blocks (1) for the three corresponding lemmas. Note two observations: 1) The value mapping dominates the depth, and its depth requirements are identical for the two cases; and 2) For contextual mappings (where the attention layers are used), we need roughly $p$ times more layers for sparse models. Recall from § 3.4 that $p$ is usually a small constant. These observations mean that sparse Transformers can achieve universal approximation using depth of the same order in $d$, $n$ and $\delta$ as the dense Transformers.

2 Key challenges in the proof

While the high level outline of the proof is similar to the one for dense Transformers , the proof in crucially relies on having all connections for computing attention in each layer, which we do not have in sparse Transformers. The sparsity in attention mechanism and the choice of general probability map $\rho$ pose nontrivial challenges in the proof. We highlight the key differences below.

Establishing the Step 2 of the dense result relies on constructing a contextual mapping using attention layers. A contextual mapping is a function that maps tokens in different sequences to unique values, thereby allowing Transformers to distinguish the same token appearing in different contexts. A crucial ingredient in the construction of such a mapping is a shift operation implemented with two attention heads in an attention layer. This shift operation involves each token taking the maximum and minimum over the entire sequence, which obviously cannot be done with sparse Transformers as it would require each token to attend to all the other tokens in the sequence. We circumvent this issue by carefully choosing the positional embedding ${\bm{E}}$ dependent on $\gamma$ (cf. Assumption 1.2), and ensuring that a similar shift operation is applied in a desired order even under sparsity.

As the final phase of the contextual mapping in , a single attention layer shifts the entire sequence by the maximum over the sequence. Again, this cannot be directly implemented due to sparsity. Using Assumption 1.3, we instead prove that by stacking $s$ sparse layers, one can successfully implement a similar operation that shifts the entire sequence by the maximum over the whole sequence, up to some controlled errors. This way, we overcome the difficulties posed by the sparsity and construct a new version of contextual mappings. The details can be found in § E.2 of the supplementary material.

Moreover, the proof of Step 3 in uses the simple fact that softmax can approximate hardmax arbitrarily closely. Since we do not restrict ourselves to softmax and generalize the probability map, a more careful argument is required. Since there are many layers in the network $\overline{g}$, it turns out that approximating it with an original sparse Transformer in $\mathcal{ST}^{2,1,4}$ requires carefully controlling the approximation errors accumulated over layers. The proof of Lemma 4 in § F of the supplementary material shows that this is indeed possible by utilizing Assumption 2.

Experiments

We now present our experimental study comparing different design and implementation choices, including sparsity patterns and levels, on four tasks: i) a synthetic copying task, ii) language modeling, iii) translation, and iv) GLUE tasks. Our goal is to understand the effect of such choices while employing sparse Transformers to the tasks with small sequence lengths, complementing the existing results for sparse Transformers on long sequence tasks.

We consider four sparsity patterns: Strided (4), Fixed (5), Star (6) and Random. The first three patterns are proposed in and ; we test them for different values of $w$. In case of the Random pattern, given a sparsity level, we make connections uniformly at random. Following , Strided and Fixed patterns are tested for three different head configurations: i) Sequential, where the sparse attention layers alternate between $\{\mathcal{A}_{k}^{1}\}$ and $\{\mathcal{A}_{k}^{2}\}$, as described in the previous sections; ii) Union, where all sparse attention layers use the sparsity pattern $\{\mathcal{A}_{k}^{1}\cup\mathcal{A}_{k}^{2}\}$; and iii) Multihead, where half of the attention heads in every attention layer use $\{\mathcal{A}_{k}^{1}\}$ and the other half use $\{\mathcal{A}_{k}^{2}\}$. Note that, given the same sequence length, Union is less sparse than the other two configurations. Thus, to ensure fair comparisons, we compare different configurations based on their sparsity levels.

We use maximum sequence length 256 in all our experiments, except 128 for GLUE tasks. For the copying task, we experiment with only one sparse Transformer block (cf. Eq (2)), with varying numbers of attention layers with $4$ attention heads. For language modeling and translation, we use the Tensor2Tensor framework and employ 12-block and 6-block (respectively) Transformers with $8$ attention heads per block. For GLUE tasks, we experiment with the $\text{BERT}_{\text{BASE}}$ model. For more details of the setup, see § G of the supplementary material.

2 Results

Copying task. We consider a synthetic copying task proposed in , where the input sequence has the format $0\mathbf{s}0\mathbf{s}$, where $\mathbf{s}$ is a 127 length sequence of symbols in $$. The models have to predict (copy) the second part, given the first half of the input. This task tests the ability of sparse Transformers to communicate the information. Table 1 presents the results for this task. Except for the Star and Random patterns, we can see that the networks learn to copy the sequences with four sparse attention layers. One possible explanation for the bad performance of Star is that, except for the relay token, it only attends to local neighbors while the task requires to copy distant tokens.

Language modeling. We conduct the language modeling experiments on the One Billion Word Benchmark which has almost one billion tokens and a vocabulary of more than 800K unique tokens. In Figure 1a, we plot the perplexity against the sparsity level. We observe that the Strided pattern and the Star achieve the best performance across all sparsity levels. For both the Strided and Fixed patterns, the Union configuration shows the best performance.

Translation. For the translation task, we train the model on WMT18 English-Czech (en-cs) dataset and test it on the Newstest 2015 dataset. We plot the BLEU score against the sparsity level in Figure 1b. We apply the same sparsity pattern to both the encoder and the decoder. The Strided and Fixed patterns with Union configuration show the best scores, which are similar to the dense attention. The Union configuration is also the least sensitive to the sparsity levels.

GLUE Tasks. We experiment with the $\text{BERT}_{\text{BASE}}$ model and report results on two sentence-pair classification tasks: MNLI (Figure 2a) and XNLI (Figure 2b). We plot the average accuracy of three runs on the dev set against the sparsity level. Additional results of the CoLA and MRPC tasks are reported in § H of the supplementary material.

In all tasks, the Random pattern performs worse than the deterministic patterns, demonstrating the need for a careful design of sparsity patterns. Overall, our experiments suggest that the design of the optimal sparsity patterns is heavily dependent on specific tasks. For example, the Star pattern shows the best performance on the language modeling task, while having trouble with copying, translation, and BERT experiments. Among the three head configurations tested for Strided and Fixed, the Union performs the best in language modeling and translation but suffers in BERT tasks. In translation experiments, we see an interesting trend that the performance of Multihead configuration improves as sparsity increases. We conjecture that this is due to the fact that in Strided and Fixed, we have $|\mathcal{A}_{k}^{1}|=O(w)$ and $|\mathcal{A}_{k}^{2}|=O(\nicefrac{{n}}{{w}})$ (cf. Eqs (4) and (5)), so the sparsest choice of $w=O(\sqrt{n})$ is the one with the best “balance” between $|\mathcal{A}_{k}^{1}|$ and $|\mathcal{A}_{k}^{2}|$.

Conclusion

Recently, sparse Transformers have received a lot of attention as they enable more efficient/faster attention mechanisms for the tasks with very long sequence lengths. We take an initial step to provide a theoretical understanding of these models. We provide a unifying framework that captures existing sparse attention models, and prove a universal approximation theorem for sparse Transformers which holds under intuitive conditions on sparsity patterns and probability maps. We also carry out experiments comparing different sparsity patterns and levels on standard NLP tasks. We hope that this work will shed light on the understanding of sparsity in attention layers, and provide guidance for the design of sparse attention models.

Broader Impact

This work studies theoretical aspects of a class of widely used neural network models in NLP and related areas. Since we do not propose a new method nor a new dataset, we expect that the impact of this work on ethical aspects and future societal consequences will be small, if any. Other than that, this work brings new insights into the sparsity in attention models, hence may make an impact on the study of faster and more efficient NLP models.

Acknowledgments and Disclosure of Funding

CY acknowledges partial support as a graduate Research Assistant from the NSF Grant (CAREER 1846088). CY also acknowledges Korea Foundation for Advanced Studies for their support.

References

Appendix A Outline and notation

The supplementary material is organized as follows. First, § B proves that the softmax operator as well as its sparse versions indeed satisfy Assumption 2. Next, § C provides formal statements of Step 1 in the proof sketch (§ 4.1). The outline of proof of Lemma 3 (Step 2 in the proof sketch) is presented in § D, followed by a separate section (§ E) proving the three key sublemmas in the proof. The proof of Step 3, Lemma 4, is given in § F. Lastly, § G and § H present the detailed setup of our experiments and additional experiment results, respectively.

Appendix B Sparse probability maps satisfy Assumption 2

In this section, we show that the softmax operator $\sigma_{\rm S}$ as well as the probability maps $\rho$ used to replace softmax in the existing approaches, namely softmax with only top-$k$ inputs , sparselin-gen , and $\alpha$-entmax , all satisfy Assumption 2. We restate the assumption for reader’s convenience: See 2 As in the assumption, we only consider the operation of these probability maps on a single vector, as they are applied column-wise. For each of the probability maps, we will show that for any $\zeta>0$ and $\eta\in(0,1]$, we can choose $t>0$ that satisfies the conditions of Assumption 2.

We assume without loss of generality that the entry of ${\bm{v}}$ is in decreasing order, where the first two entries satisfy ${v}_{1}-{v}_{2}\geq\zeta$. For any such $\zeta>0$ and any $0<\eta\leq 1$, our aim is to show the existence of $t>0$ such that $\sigma_{\rm S}[t{\bm{v}}]_{1}=\frac{\exp(t{v}_{1})}{\sum_{i=1}^{n}\exp(t{v}_{i})}\geq 1-\eta$. Then, $\sum_{j=2}^{n}\sigma_{\rm S}[t{\bm{v}}]_{j}\leq\eta$ follows.

Now, since ${v}_{i}\leq{v}_{1}-\zeta$ for $i\in[2:n]$, note that

Since $\frac{1}{1+(n-1)\exp(-t\zeta)}$ is an increasing function in $t>0$, one can increase $t$ sufficiently large to make it greater than $1-\eta$.

The same argument holds for the softmax with top-$k$ inputs, used in . By the assumption on ${\bm{v}}$, entries ${v}_{1},\dots,{v}_{k}$ are the top $k$ components. Thus,

can be satisfied by choosing large enough $t>0$.

B.2 Sparselin-gen

We now consider the case where $\rho$ is sparselin-gen , which was used to sparsify the attention score matrices in . Given a regularization parameter $\lambda\in[0,1)$, the sparselin-gen used in is defined as

Now, assume without loss of generality that the entry of ${\bm{v}}$ is in decreasing order, where the first two entries satisfy ${v}_{1}-{v}_{2}\geq\zeta$. For any such $\zeta>0$ and any $0<\eta\leq 1$, our aim is to show the existence of $t>0$ such that $\rho[t{\bm{v}}]_{1}\geq 1-\eta$. This is done by choosing $t=\frac{1-\eta}{\zeta}$. To see this, notice that if ${v}_{j}$’s are in decreasing order, then $\rho[{\bm{v}}]_{j}$ are also in decreasing order. Now consider

If $\rho[t{\bm{v}}]_{2}=0$, then $\rho[t{\bm{v}}]_{j}=0$ for all $j=3,\dots,n$, and $\rho[t{\bm{v}}]_{1}=1\geq 1-\eta$. If $\rho[t{\bm{v}}]_{2}>0$, then

B.3 α𝛼\alpha-entmax

Next, we consider the case where $\rho$ is $\alpha$-entmax , which was used to sparsify the attention score matrices in . Given a parameter $\alpha\geq 1$, the $\alpha$-entmax is defined as

where $\Delta^{n-1}$ is the probability simplex and $H_{\alpha}$ is the Tsallis continuous family of entropies

As shown in , the solution of $\alpha$-entmax is equal to softmax if $\alpha=1$, and otherwise ($\alpha>1$) it is given in the form

Again, assume without loss of generality that the entry of ${\bm{v}}$ is in decreasing order, where the first two entries satisfy ${v}_{1}-{v}_{2}\geq\zeta$. For any such $\zeta>0$ and any $0<\eta\leq 1$, our aim is to show the existence of $t>0$ such that $\rho[t{\bm{v}}]_{1}\geq 1-\eta$. This is done by choosing $t=\nicefrac{{1}}{{\zeta(\alpha-1)}}$.

Note that $(\alpha-1)t({v}_{1}-{v}_{2})\geq 1$ due to our choice of $t$. Then, we will show that with such a $t$, $\rho[t{\bm{v}}]_{1}=1$ must hold. For the sake of contradiction, suppose not: $\rho[t{\bm{v}}]_{1}<1$. Then, by monotonicity of $\rho[t{\bm{v}}]_{j}$, we have $\rho[t{\bm{v}}]_{2}>0$. This means

in particular, we have $(\alpha-1)t{v}_{2}-\tau(t{\bm{v}})>0$. However, recall that $(\alpha-1)t({v}_{1}-{v}_{2})\geq 1$, which implies $(\alpha-1)t{v}_{1}-\tau(t{\bm{v}})>1$. This results in

thus contradicting $\rho[t{\bm{v}}]_{1}<1$. Therefore, $\rho[t{\bm{v}}]_{1}=1$ must hold.

Appendix C Details of the Step 1 in the proof sketch (§ 4.1)

We start by formally defining the function class $\overline{\mathcal{F}}(\delta)$.

For any $f\in\mathcal{F}$ and $\epsilon>0$, there exists a small enough $\delta>0$ such that there exists $\overline{f}\in\overline{\mathcal{F}}(\delta)$ such that $\mathsf{d}_{p}(f,\overline{f})\leq\epsilon/2$.

Note that for any ${\bm{X}}\in{\bm{G}}+[0,\delta)^{d\times n}$ we have $\left\|{{\bm{X}}-{\bm{G}}}\right\|_{\infty}<\delta$, so we have

Appendix D Proof of Lemma 3 (Step 2 in § 4.1)

In this section, we describe in further details how modified sparse Transformers (the class $\overline{\mathcal{ST}}^{2,1,1}$) are able to exactly express arbitrary piecewise constant functions in $\overline{\mathcal{F}}(\delta)$. We show that we can compute a contextual mapping of the entire input sequences without relying on dense self-attention layers. The token-wise feed-forward layers then transform these contextual mappings to the desired output sequence.

Choose the positional embedding ${\bm{E}}$ according to $\gamma$ in Assumption 1.2. After addition, each column of the input ${\bm{X}}_{k}+{\bm{E}}_{k}$ are in disjoint intervals.

Given the input ${\bm{X}}+{\bm{E}}$, a series of modified feed-forward layers quantizes it so that each entry of the quantized input has a value in $\{0,\delta,\dots,n-\delta\}$ (Lemma 6).

Next, a series of modified sparse self-attention layers takes the quantized input ${\bm{H}}$ and implement a contextual mapping $q$ such that, for different quantized input sequences ${\bm{H}}$ and ${\bm{H}}^{\prime}$, all the elements in $q({\bm{H}})$ and $q({\bm{H}}^{\prime})$ are distinct (Lemma 7).

Finally, a series of modified feed-forward layers maps each element in the context id $q({\bm{H}})$ to the desired output value of $\overline{f}\in\overline{\mathcal{F}}$ at the input ${\bm{X}}$ (Lemma 8).

We defer the proofs of Lemmas 6, 7, and 8 to a separate section: see § E.

Before discussing the details of each step, we note that although a Transformer network stacks self-attention and feed-forward layers in an alternate manner, we can use a series of arbitrary number of the same layers, thanks to skip connections. The outline of the proof is similar to , but key component in their proof called selective shift operation relies on the fact that each token can attend to the entire sequence; this is not true in sparse Transformers, which poses a nontrivial challenge. We overcome this issue by a more careful construction of the positional embedding ${\bm{E}}$ and sparse self-attention layers.

Recall from Assumption 1.2 that there exists a permutation $\gamma:[n]\to[n]$ such that for all $i\in[n-1]$, $\gamma(i)$ is one of the tokens that the $\gamma(i+1)$-th token directly attends to. Using this permutation $\gamma$, we choose the columns of positional embedding ${\bm{E}}$ in the following way:

As a result, the $\gamma(1)$-th column of ${\bm{X}}+{\bm{E}}$ will be in the range $[n-1,n)^{d}$, and similarly ${\bm{X}}_{\gamma(i)}+{\bm{E}}_{\gamma(i)}\in[i-2,i-1)^{d}$ for $i\in[2:n]$. This means that the entries corresponding to different tokens lie be in disjoint intervals of the form $[j,j+1)$, where $j\in[0:n-1]$.

D.2 Quantization by feed-forward layers

Note from the previous step that each entry of ${\bm{X}}+{\bm{E}}$ must be in $[0,n)$. Next, we quantize this interval $[0,n)$ of input using to a set of $\delta$-grid points $\{0,\delta,\dots,n-\delta\}$. This allows us to deal with finite set of values, which proves useful in the later stages of the proof. The next lemma shows that the quantization can be carried out using a seried of the modified feed-forward layers.

D.3 Contextual mapping by sparse self-attention layers

This contextual mapping maps each unique sequence/context into different context ids, enabling the network to distinguish the same token appearing in different sequences.

D.4 Value mapping by feed-forward layers

D.5 Finishing the proof

Appendix E Proof of Lemmas 6, 7, and 8

The proof goes as follows. Using $\frac{n}{\delta}$ token-wise feed-forward layers, we implement the quantization function $g_{\rm q}^{\rm ent}$ that quantizes the first row of the input. Then we stack another $\frac{n}{\delta}$ layers to quantize the second row, and so on.

For the first row, we add $n/\delta$ layers of the following form, for $k\in[0:n/\delta-1]$.

E.2 Proof of Lemma 7

One can consider a sparse self-attention layer that consists of two such heads, with $b_{Q}<b^{\prime}_{Q}$:

The $(1,k)$-th entry of $\Psi^{l}({\bm{Z}};c,b_{Q},b^{\prime}_{Q})$ reads

This means that for input columns ${\bm{Z}}_{k}$ satisfying ${\bm{u}}^{T}{\bm{Z}}_{k}\in(b_{Q},b^{\prime}_{Q})$ only, $\Psi^{l}$ shifts up the first entry of ${\bm{Z}}_{k}$ by the difference of maximum and minimum values of ${\bm{u}}^{T}{\bm{Z}}_{j}$ over the sparsity pattern $j\in\mathcal{A}_{k}^{l}$, while leaving other columns intact. By choosing $b_{Q}$ and $b^{\prime}_{Q}$ properly, we can selectively modify certain columns without touching other columns; we refer to this operation $\Psi^{l}$ as the sparse selective shift operation, and we will see later that this is indeed the key ingredient of our proof.

In fact, this operation is a sparse version of the selective shift operation used in . Since $\mathcal{A}_{k}^{l}$ is usually only a small subset of $[n]$, one cannot calculate the maximum and minimum of ${\bm{u}}^{T}{\bm{Z}}_{j}$ over the whole sequence, as done in . Instead, we use Assumption 1.2 and a more careful choice of ${\bm{E}}$ to get around the restriction posed by sparsity.

The $(1,k)$-th entry of $\Omega^{l}({\bm{Z}};c)$ reads

So, for each column $k$, the all-max-shift operation shifts up the first entry of ${\bm{Z}}_{k}$ by the maximum value of ${\bm{u}}^{T}{\bm{Z}}_{j}$ over the sparsity pattern $j\in\mathcal{A}_{k}^{l}$. Unlike the selective shift operation, the all-max-shift operation is applied to all the columns.

Recall that the any input to this step is in

E.2.2 Construction of layers

Given these preliminaries, we now describe our construction of $g_{\rm c}$. Recall from Assumption 1.2 that the permutation $\gamma$ satisfies $\gamma(i-1)\in\bigcup_{l=1}^{p}\mathcal{A}^{l}_{\gamma(i)}$ for $i\in[2:n]$. From this, for $i\in[2:n]$ we let $l_{i}\in[p]$ be any index such that $\gamma(i-1)\in\mathcal{A}_{\gamma(i)}^{l_{i}}$. For simplicity of notation, let $z_{k}:={\bm{u}}^{T}{\bm{H}}_{k}$ for $k\in[n]$ and $\Delta=\sum_{i=0}^{d-1}\delta^{-i}$.

Next, starting from $i=2$, we want to sequentially stack $\delta^{-d}$ sparse selective shift operations

in increasing order of $b\in\left[(i-2)\Delta:\delta:(i-2)\Delta+\delta^{-d+1}-\delta\right]$. That is, we want to add sparse attention layers with sparsity patterns $\mathcal{A}_{\gamma(i)}^{l_{i}}$ that apply the selective shift operation to each possible value of $z_{\gamma(i)}$. Recall that the sparsity patterns have to cycle from $\mathcal{A}_{k}^{1}$ to $\mathcal{A}_{k}^{p}$, so we have to place other remaining $p-1$ sparsity patterns (whose indices are not $l_{i}$) in between the $\Psi^{l_{i}}$ layers. This can be done by setting all the other sparse attention layers to be the identity. This way, we stack a total of $p\delta^{-d}$ sparse attention layers for $i=2$, another $p\delta^{-d}$ for $i=3$, and so on, up to $i=n$.

After these layers, we further stack $s$ all-max-shift operations. For $i=1,\dots,s$, we add all-max-shift operations of the form

E.2.3 Selective shift operations

First consider the first $p\delta^{-d}$ layers. Omitting layers that are identity, they are essentially selective shift operations $\Psi^{l_{2}}(\cdot;\delta^{-d},b-\delta/2,b+\delta/2)$ for $b\in[0:\delta:\delta^{-d+1}-\delta]$. Since $[0:\delta:\delta^{-d+1}-\delta]$ is the set of possible values of $z_{\gamma(2)}$, these layers perform selective shift operation on the $\gamma(2)$-th column without changing the other columns. Each possible value of ${\bm{H}}_{\gamma(2)}$ undergoes one and only shift operation (by the corresponding layer with $b={\bm{u}}^{T}{\bm{H}}_{\gamma(2)}$), by which the $(1,\gamma(2))$-th entry of the input is updated.

Recall by Assumption 1.2 that $\gamma(1)\in\mathcal{A}_{\gamma(2)}^{l_{2}}$, and that $z_{\gamma(1)}$ and $z_{\gamma(2)}$ are the maximum and minimum over the whole sequence $z_{1},\dots,z_{n}$ (see (8)). By Assumption 1.1 we also have $\gamma(2)\in\mathcal{A}_{\gamma(2)}^{l_{2}}$. Since both $\gamma(1)$ and $\gamma(2)$ are in $\mathcal{A}_{\gamma(2)}^{l_{2}}$, the maximum and minimum value of $z_{j}:={\bm{u}}^{T}{\bm{H}}_{j}$’s over $j\in\mathcal{A}_{\gamma(2)}^{l_{2}}$ are $z_{\gamma(1)}$ and $z_{\gamma(2)}$, respectively. Therefore, the $(1,\gamma(2))$-th entry of the input matrix is shifted up as follows:

Let $\widetilde{{\bm{H}}}_{\gamma(2)}$ be the $\gamma(2)$-th column after the shift operation has shifted ${H}_{1,\gamma(2)}$ to $\widetilde{{H}}_{1,\gamma(2)}$. Then, define

Note that $\widetilde{z}_{\gamma(2)}>z_{\gamma(1)}$ because

which is true. Therefore, $\widetilde{z}_{\gamma(2)}$ becomes the new maximum among the current values $z_{\gamma(1)},\widetilde{z}_{\gamma(2)},z_{\gamma(3)},\dots,z_{\gamma(n)}$, and the new minimum element is $z_{\gamma(3)}$.

We now consider the next $p\delta^{-d}$ layers, which are essentially $\Psi^{l_{3}}(\cdot;\delta^{-d},b-\delta/2,b+\delta/2)$ for $b\in[\Delta:\delta:\Delta+\delta^{-d+1}-\delta]$. They apply the shift operation to the $\gamma(3)$-th column. Since we have $\gamma(2),\gamma(3)\in\mathcal{A}_{\gamma(3)}^{l_{3}}$, the shift operation similarly yields

We can also show $\widetilde{z}_{\gamma(3)}>\widetilde{z}_{\gamma(2)}$, because

So after this operation $\widetilde{z}_{\gamma(3)}$ and $z_{\gamma(4)}$ are the new maximum and minimum over the updated sequence $z_{\gamma(1)},\widetilde{z}_{\gamma(2)},\widetilde{z}_{\gamma(3)},z_{\gamma(4)},\dots,z_{\gamma(n)}$.

The same process continues. The next $p\delta^{-d}$ layers shifts the $\gamma(4)$-th columns and results in $\widetilde{z}_{\gamma(4)}$ which is greater than $\widetilde{z}_{\gamma(3)}$. After the first $p(n-1)\delta^{-d}$ layers, all columns except $\gamma(1)$-th column have been shifted, resulting in $z_{\gamma(1)},\widetilde{z}_{\gamma(2)},\dots,\widetilde{z}_{\gamma(n)}$ satisfying

Let us denote the output of the $p(n-1)\delta^{-d}$-th layer as $\widetilde{{\bm{H}}}$.

is one-to-one. Recall that for each column ${\bm{H}}_{k}$, the map ${\bm{H}}_{k}\mapsto{\bm{u}}^{T}{\bm{H}}_{k}=:z_{k}$ is one-to-one. Also, permutation of columns is one-to-one, which implies that it suffices to show that the map $\begin{bmatrix}z_{\gamma(1)}&\dots&z_{\gamma(n)}\end{bmatrix}\mapsto\widetilde{z}_{\gamma(n)}$ is one-to-one.

Suppose we have two sequences $\begin{bmatrix}z_{\gamma(1)}&\dots&z_{\gamma(n)}\end{bmatrix}$ and $\begin{bmatrix}z^{\prime}_{\gamma(1)}&\dots&z^{\prime}_{\gamma(n)}\end{bmatrix}$ that map to the same value of $\widetilde{z}_{\gamma(n)}=\widetilde{z}^{\prime}_{\gamma(n)}$. Then,

Suppose $z_{\gamma(n)}\neq z^{\prime}_{\gamma(n)}$. Since they both lie inside $[(n-2)\Delta:\delta:(n-2)\Delta+\delta^{-d+1}-\delta]$, we have

Note that all the terms other than $z_{\gamma(n)}-z^{\prime}_{\gamma(n)}$ are of “coarser resolution.” For example, the first term

in the summation can only take values $0,\delta^{-d+1},-\delta^{-d+1},2\delta^{-d+1},-2\delta^{-d+1},\dots$, so it can never cancel the difference $z_{\gamma(n)}-z^{\prime}_{\gamma(n)}$ and make the sum $\widetilde{z}_{\gamma(n)}-\widetilde{z}^{\prime}_{\gamma(n)}$ zero. This implies that $z_{\gamma(n)}=z^{\prime}_{\gamma(n)}$ must hold.

Next, suppose $z_{\gamma(n-1)}\neq z^{\prime}_{\gamma(n-1)}$. Since we have $z_{\gamma(n)}=z^{\prime}_{\gamma(n)}$,

E.2.4 All-max-shift operations

Next, we explain the operation of the $s$ all-max-shift layers. Recall from Assumption 1.3 that any token can attend to all the other tokens after $s$ steps, either directly or indirectly. Also recall from the last subsection that the input to the first all-max-shift layer is $\widetilde{{\bm{H}}}$, and the maximum entry of ${\bm{u}}^{T}\widetilde{{\bm{H}}}$ is $\widetilde{z}_{\gamma(n)}$, the unique id for input ${\bm{H}}$. From the statement of Lemma 7, the output after the $s$ all-max-shift operations for input ${\bm{H}}$ is denoted as $g_{\rm c}({\bm{H}})$. In this subsection, we show that through $s$ all-max-shift operations, the maximum $\widetilde{z}_{\gamma(n)}$ will propagate to all tokens and be a “dominant” term, which determines the interval that ${\bm{u}}^{T}g_{\rm c}({\bm{H}})$ lies in. As a result, we can show Properties 7.1 and 7.2 of $g_{\rm c}$ at the end.

Note that the unique id $\widetilde{z}_{\gamma(n)}$ has the following upper bound:

where we used $\Delta:=\sum_{i=0}^{d-1}\delta^{-i}=\frac{\delta^{-d}-1}{\delta^{-1}-1}\leq\delta^{-d}-1$. A similar bound

also holds from a similar derivation. Next, recall from Assumption 1.3 the definitions

and that there exists $s\geq 1$ such that, for all $k\in[n]$, $\mathcal{S}^{s}_{k}=[n]$. Finally, the following inequality will be useful throughout: for any integer $s\geq 1$,

Let us now describe the operation that the all-max-shift layers $\Omega^{(i-1)\text{ mod }p+1}(\cdot;2sn\delta^{-nd-1})$, $i=1,\dots,s$, carry out.

The input to the first all-max-shift layer is $\widetilde{{\bm{H}}}$. Let the output of the layer be ${\bm{M}}^{1}$. Recall that ${\bm{u}}^{T}\widetilde{{\bm{H}}}$ consists of values $z_{\gamma(1)},\widetilde{z}_{\gamma(2)},\dots,\widetilde{z}_{\gamma(n)}$, which are all strictly greater than 0 and strictly less than $n\delta^{-nd}$ (by (10)). So, for each column $k\in[n]$, the layer update reads

where $j_{k}^{1}:=\operatorname*{arg\,max}_{j\in\mathcal{A}_{k}^{1}}{\bm{u}}^{T}\widetilde{{\bm{H}}}_{j}$. After the update, ${\bm{u}}^{T}{\bm{M}}^{1}_{k}$ is “dominated” by $2sn\delta^{-nd-1}{\bm{u}}^{T}\widetilde{{\bm{H}}}_{j_{k}^{1}}$, meaning that for any $k,k^{\prime}\in[n]$,

This is because the minimum gap between different values of ${\bm{u}}^{T}\widetilde{{\bm{H}}}_{j_{k}^{1}}$ is at least $\delta$, and we have

so if ${\bm{u}}^{T}\widetilde{{\bm{H}}}_{j_{k}^{1}}<{\bm{u}}^{T}\widetilde{{\bm{H}}}_{j_{k^{\prime}}^{1}}$, that solely determines the order ${\bm{u}}^{T}{\bm{M}}_{k}<{\bm{u}}^{T}{\bm{M}}_{k^{\prime}}$ because ${\bm{u}}^{T}\widetilde{{\bm{H}}}_{k}$ cannot reverse it. Also, by the definition of $j_{k}^{1}$, for any index set $\mathcal{B}\in[n]$ we have

If $s\geq 2$, we move on to the second layer.

At the second all-max-shift, we have sparsity patterns $\mathcal{A}_{k}^{1\textup{ mod }p+1}$. Let us the output of this layer as ${\bm{M}}^{2}$. For each column $k\in[n]$, the layer update reads

where $j_{k}^{2}:=\operatorname*{arg\,max}_{j\in\mathcal{A}_{k}^{1\textup{ mod }p+1}}{\bm{u}}^{T}{\bm{M}}^{1}_{j}$. If we look at the update more closely, we can apply (13) and get

Again, the last term dominates the rest of the terms in ${\bm{u}}^{T}{\bm{M}}^{2}_{k}$, because the minimum gap between different values of $\max_{j\in\mathcal{S}_{k}^{2}}{\bm{u}}^{T}\widetilde{{\bm{H}}}_{j}$ is at least $\delta$, and

The last inequality holds due to inequality (12), because

If $s\geq 3$, we move on to the third layer, which outputs ${\bm{M}}^{3}$. Similarly, we can show that ${\bm{u}}^{T}{\bm{M}}^{3}_{k}$ is dominated by $(2sn\delta^{-nd-1})^{3}\max_{j\in\mathcal{S}_{k}^{3}}{\bm{u}}^{T}\widetilde{{\bm{H}}}_{j}$ because the rest of the terms in ${\bm{u}}^{T}{\bm{M}}^{3}_{k}$ is strictly upper-bounded

which can then be shown to be smaller than $(2sn\delta^{-nd-1})^{3}\cdot\delta$:

The last inequality is due to the fact that $1+6s+12s^{2}\leq 8s^{3}$ for $s\geq 3$, which can derived from (12). Repeating this process, after all $s$ layers we get ${\bm{M}}^{s}$, and ${\bm{u}}^{T}{\bm{M}}^{s}_{k}$ is dominated by

This is because the remaining terms in ${\bm{u}}^{T}{\bm{M}}^{s}_{k}$ can be strictly upper-bounded

which is then dominated by the smallest difference possible in $(2sn\delta^{-nd-1})^{s}\widetilde{z}_{\gamma(n)}$:

The last inequality used $(1+2s)^{s}-(2s)^{s}\leq(2s)^{s}$, derived from (12).

E.2.5 Verifying Properties 7.1 and 7.2

After these all-max-shift operations, we define the output ${\bm{M}}^{s}$ of the last all-max-shift layers to be the output of the function $g_{\rm c}$ for input ${\bm{H}}$, i.e., $g_{\rm c}({\bm{H}}):={\bm{M}}^{s}$.

This is because anything added by the all-max-shift operations is an integer multiple of $2sn\delta^{-nd}$, and ${\bm{u}}^{T}\widetilde{{\bm{H}}}_{k}<n\delta^{-nd}<2n\delta^{-nd}$ for all $k$. Recall that $\widetilde{{\bm{H}}}$ is the input matrix for the first max-shift operation, and that the components of ${\bm{u}}^{T}\widetilde{{\bm{H}}}$ are $z_{\gamma(1)},\widetilde{z}_{\gamma(2)},\dots,\widetilde{z}_{\gamma(n)}$, which were shown to be distinct by (9). Since ${\bm{u}}^{T}g_{\rm c}({\bm{H}})_{k}$ produce distinct outputs for a ${\rm mod}$ operation, they themselves have to distinct. This proves Property 7.1.

Also, by the “domination” argument in the previous subsection, the output $g_{\rm c}({\bm{H}})$ has the property that for any column, ${\bm{u}}^{T}g_{\rm c}({\bm{H}})_{k}$ lies inside an interval determined by $\widetilde{z}_{\gamma(n)}$, the unique id for the input ${\bm{H}}$:

E.3 Proof of Lemma 8

To prove this lemma, we implement a token-wise function that maps

This layer updates any column $j$ of its input ${\bm{Z}}$ that satisfies ${\bm{u}}^{T}g_{\rm c}({\bm{H}})_{k}-\delta/2\leq{\bm{u}}^{T}{\bm{Z}}_{j}<{\bm{u}}^{T}g_{\rm c}({\bm{H}})_{k}+\delta/2$, without modifying any other columns that are out of this range.

Appendix F Proof of Lemma 4 (Step 3 in § 4.1)

In this section, we describe how the modified sparse Transformer network $\overline{g}\in\overline{\mathcal{ST}}^{2,1,1}$ constructed in Lemma 3 can be approximated with an original sparse Transformer network $g\in\mathcal{ST}^{2,1,4}$. Recall that $\overline{g}$ is a “modified” sparse Transformer network, which employ the hardmax $\sigma_{\rm H}$ operators in place of $\rho$ operators in sparse self-attention layers and piecewise linear activations $\phi\in\Phi$ instead of ${\rm ReLU}$s in feed-forward layers. The goal of this lemma is to approximate the function $\overline{g}=g_{\rm v}\circ g_{\rm c}\circ g_{\rm q}\in\overline{\mathcal{ST}}^{2,1,1}$ with a standard sparse Transformer $g=\widetilde{g}_{\rm v}\circ\widetilde{g}_{\rm c}\circ\widetilde{g}_{\rm q}\in\mathcal{ST}^{2,1,4}$ with accuracy $\mathsf{d}_{p}(\overline{g},g)\leq\epsilon/2$. As the construction of $\overline{g}$ consists of three steps, we will approximate each of them step by step. The whole intuition behind the proof is that as long as we are considering $L_{p}$ approximation, we can approximate $\sigma_{\rm H}$ and $\phi\in\Phi$ as closely as we want with $\rho$ and ${\rm ReLU}$s, respectively. However, as the proof will show, controlling the aggregated error over layers is not a trivial job.

We first consider approximating $g_{\rm q}$ from Lemma 6 with a standard feed-forward layer counterpart, $\widetilde{g}_{\rm q}$. Recall from § E.1 that the modified feed-forward layers used in $g_{\rm q}$ are of the form

for $i\in[d]$ and $k\in[0:n/\delta-1]$. Note that the activation $\phi\in\Phi$ can be closely approximated by three ${\rm ReLU}$s:

where $0<\alpha<1$. Note that $\widetilde{\phi}_{\alpha}(t)=\phi(t)$ except for an interval $((1-\alpha)\delta,\delta)$, and by shrinking $\alpha>0$ this interval can be made arbitrarily small. Consider approximating the layers (14) with standard feed-forward layers, by replacing $\phi$ with its approximation $\widetilde{\phi}_{\alpha}$. Let the resulting function be $\widetilde{g}_{\rm q}\in\mathcal{ST}^{2,1,3}$.

Then, it is easy to check that $g_{\rm q}({\bm{X}}+{\bm{E}})=\widetilde{g}_{\rm q}({\bm{X}}+{\bm{E}})$ holds if all coordinates of ${\bm{X}}\in[0,1)^{d\times n}$ are in the intervals of the form $[k\delta,(k+1-\alpha)\delta]$ for some $k\in[0:n/\delta-1]$; i.e., the intervals in which $\widetilde{\phi}_{\alpha}$ perfectly approximates $\phi$. The Lebesgue measure of the set of such inputs ${\bm{X}}$ is

Let us now consider approximating the contextual mapping $g_{\rm c}$ in Lemma 7, constructed using the hardmax $\sigma_{\rm H}$ operators, with the standard sparse self-attention layers employing $\rho$ operator. We will call the approximation $\widetilde{g}_{\rm c}$. Recall that $\rho$ satisfies Assumption 2: See 2 This means that $\rho$ can closely approximate $\sigma_{\rm H}$ in the sense that whenever the input vector ${\bm{v}}$ to the $\rho$ operator has a maximum element ${v}_{j^{*}}$ by some margin $\zeta$, then the $j^{*}$-th component of the output $\rho[t{\bm{v}}]$ is close to $1$, while the other components of $\rho[t{\bm{v}}]$ are close to .

Recall that $g_{\rm c}$ consists of two parts. The first part is a composition of sparse selective shift operations, and the second is a composition of all-max-shift operations. We will first examine how “errors” are introduced when $\sigma_{\rm H}$ is replaced with $\rho$ in both operations, discuss how the errors accumulate, and show how to choose the right $\zeta$ and $\eta$ to control the errors in the approximation $\widetilde{g}_{\rm c}$.

Recall that the key component in both the selective shift operation and all-max-shift operation is the sparse attention head $\psi^{l}(\cdot)$, which computes its $k$-th column as the following:

Now suppose we replaced $\sigma_{\rm H}$ with $\rho$ satisfying Assumption 2. Suppose each entry in ${\bm{u}}^{T}{\bm{Z}}$ differs at least by $\delta$, which is true in the construction of $g_{\rm c}$. We choose $\zeta=\delta/2$ and some $0<\eta<1$, and corresponding $t>0$. Then, replace $\sigma_{\rm H}[\cdot]$ with $\rho[t\cdot]$ and define

If ${\bm{u}}^{T}{\bm{Z}}_{k}>b_{Q}$, it is easy to check that $\widetilde{\psi}^{l}({\bm{Z}};b_{Q})_{k}$ satisfies

Similarly, if ${\bm{u}}^{T}{\bm{Z}}_{k}<b_{Q}$, we have

Now consider the approximate sparse selective shift operator $\widetilde{\Psi}^{l}$, implemented with $\widetilde{\psi}^{l}$. For $b_{Q}<b^{\prime}_{Q}$, we define