A RE T RANSFORMERS UNIVERSAL APPROXIMATORS OF SEQUENCE - TO - SEQUENCE FUNCTIONS ?
bsrinadh@google.com Ankit Singh Rawat Google NY ankitsrawat@google com
Reddi Google Research NY sashank@google.com NY sanjivk@google.com
A BSTRACT
Despite the widespread adoption of Transformer models for NLP tasks, the expressive power of these models is not well-understood.
In this paper, we establish that Transformer models are universal approximators of continuous permutation equivariant sequence-to-sequence functions with compact support, which is quite surprising given the amount of shared parameters in these models.
Furthermore, using positional encodings, we circumvent the restriction of permutation equivariance, and show that Transformer models can universally approximate arbitrary continuous sequence-to-sequence functions on a compact domain.
Interestingly, our proof techniques clearly highlight the different roles of the self-attention and the feed-forward layers in Transformers.
In particular, we prove that ﬁxed width self-attention layers can compute contextual mappings of the input sequences, playing a key role in the universal approximation property of Transformers.
Based on this insight from our analysis, we consider other simpler alternatives to selfattention layers and empirically evaluate them.
I NTRODUCTION
Self-attention based Transformer networks (Vaswani et al., 2017) have been at the center of the recent progress on various natural language processing (NLP) tasks, including machine translation (Vaswani et al., 2017), language modeling (Radford et al., 2018; 2019), and question answering (Devlin et al., 2018;
Yang et al., 2019; Liu et al., 2019).
All these tasks involve learning models that map an input sequence of tokens to an output sequence of tokens.
Transformers make it feasible to train large models to approximate these sequence-to-sequence functions due to their ability to process the input tokens in a parallel way, as opposed to the sequential nature of RNNs and LSTMs.
A Transformer block consists of two kinds of layers: a self-attention layer and a token-wise feedforward layer, with skip connections present in both layers.
The self-attention layer transforms each input token embedding using a weighted combination of the embeddings of all tokens in the input sequence, where weights are generated by pairwise dot-products among the input token embeddings.
The token-wise feed-forward layer then independently processes each of these modiﬁed input token embeddings without any interaction among them.
Notably, Transformers employ parameter reuse across tokens, as both layers use the same parameters to process each token.
Moreover, Transformers have to rely solely on the pairwise dot-products to capture interaction between the input tokens.
Given the parameter sharing and limited interactions between tokens, it is natural to wonder: what class of sequence-to-sequence functions can the Transformer networks represent?
Also, what is the role of the two different kinds of layers?
Are both layers needed to obtain the representation power of Transformers?
In the existing literature, the advantage of Transformers has often been attributed to their capability of computing contextual embeddings/mappings of the input, as opposed to ﬁxed word embeddings as in word2vec (Mikolov et al., 2013).
Is it possible to formalize the notion of contextual mappings?
If yes, can Transformers actually compute such mappings?
Such questions still remain elusive.
In this paper, we provide a mathematical deﬁnition of contextual mappings and surprisingly show that multi-head self-attention layers can indeed compute contextual mappings of the input sequences.
We further show that this ability to compute contextual mappings coupled with the value mapping ability of the feed-forward layers makes Transformers universal approximators of any permutation equivariant sequence-to-sequence function.
We also improve this result using positional encodings, and show that Transformers can represent any sequence-to-sequence function; i.e., the restriction of permutation equivariance can be removed by positional encodings.
These results on universal approximation of sequence-to-sequence functions raise a natural question: is it possible to have a more efﬁcient architecture to compute contextual mappings, consequently, preserving the ability to universally approximate sequence-to-sequence functions?
Towards this, we discuss other simpler architectures that can implement contextual mappings (to some extent), and experimentally evaluate their performance.
In our experiments, we notice that the models that combine these simpler architectures with Transformers have better performance, compared to the standalone Transformers.
We conclude the paper by presenting more discussion and interesting future research directions along these lines.
1. S UMMARY OF OUR CONTRIBUTIONS
• We prove that Transformers are universal approximators of continuous and permutation equivariant sequence-to-sequence functions with compact support (Theorem 2).
We also show that, if Transformers have trainable positional encodings added to the input, then they are universal approximators of continuous sequence-to-sequence functions on a compact domain (Theorem 3).
• We formalize the notion of contextual mappings and show that the attention layers can compute contextual mappings, where each unique context is mapped to a unique vector (Lemma 6).
• We experimentally evaluate other simpler layers that can compute contextual mappings to some extent, such as bi-linear projections and separable convolutions, and show that substituting some of the self-attention layers with these layers can result in better performance (Section 5).
1. R ELATED WORKS & NOTATION
Analysis of attention-based models.
Given the popularity of Transformers, there have been numerous works trying to understand the role of attention layers in natural language processing models.
One such line of work focuses on probing the output of attention layers to understand the attention mechanism and internal language representation (Hewitt & Manning, 2019;
Clark et al., 2019; Coenen et al., 2019; Vig & Belinkov, 2019).
Although these results give valuable insights, a consistent theoretical analysis corroborating these ﬁndings is missing.
Universal approximation theorems.
Universal approximation theorems are classical results in neural network theory, dating back many decades (Cybenko, 1989; Hornik, 1991).
These results show that given unbounded width, a shallow neural network can approximate arbitrary continuous function with compact support, up to any accuracy.
Other results focusing on depth appeared more recently (Lu et al., 2017; Hanin & Sellke, 2017; Lin & Jegelka, 2018).
In particular, Lu et al.
Hanin & Sellke (2017) consider fully-connected ReLU networks whose input dimension is d , and show that networks with width d + 1 and unbounded depth are universal approximators of scalar-valued continuous functions.
Lin & Jegelka (2018) show that a residual network with one hidden neuron per residual block is a universal approximator of scalar-valued functions, given unbounded depth.
Although Transformer networks do have residual connections, due to their heavy parameter sharing, the existing analyses for residual networks do not extend to Transformers.
Sannai et al.
(2019) consider universally approximating permutation invariant/equivariant functions using fully-connected ReLU networks.
Turing completeness results on Transformers.
Recently, P´erez et al.
(2019) have shown that
Transformers with inﬁnite precision are Turing complete, which is not the case in ﬁnite precision setting (Dehghani et al., 2018).
We note that Turing completeness deals with computation on formal languages (thus discrete objects), while universal approximation focuses on functions on a continuum.
In other words, these are two different concepts; and one does not imply another.
Notation.
We use the following notation in the paper.
Given a matrix A , let A i,j , A i, : , and A : ,j denote its ( i, j ) -th entry, i -th row, and j -th column, respectively.
We use (cid:107) A (cid:107) p to denote the entrywise (cid:96) p norm of A .
Let σ [ · ] be the softmax operator, which takes a matrix as input and applies softmax operation to each column of the matrix, which results in a column stochastic matrix, i.e., a matrix that has non-negative entries with all columns summing to 1.
We similarly deﬁne σ H [ · ] to be the hardmax operator, which outputs the one-hot representation of the arg max entry for each column of the input matrix.
If there are k arg max entries, then the output is 1 /k for such entries.
We use 1 n to denote a vector of length n whose entries are all 1 .
We denote the 0-1 indicator function by 1 {·} .
We use d and n to denote the embedding dimension and the sequence length, respectively.
We assume throughout that n ≥ 2 , as the Transformers reduce to residual networks when n = 1 .
T RANSFORMER NETWORKS
A Transformer block is a sequence-to-sequence function mapping R d × n to R d × n .
It consists of two layers: a self-attention layer and a token-wise feed-forward layer, with both layers having a skip connection.
More concretely, for an input X ∈ R d × n consisting of d -dimensional embeddings of n tokens, a Transformer block with multiplicative or dot-product attention (Luong et al., 2015) consists of the following two layers 1 :
[Equation]
where W iO ∈ R d × m , W iV , W iK ,
W iQ ∈ R m × d , W 2 ∈ R d × r , W 1 ∈ R r × d , b 2 ∈ R d , b 1 ∈ R r , and
FF( X ) is the output of the Transformer block.
The number of heads h and the head size m are two main parameters of the attention layer; and r denotes the hidden layer size of the feed-forward layer.
Here, we would like to point out that our deﬁnition of the self-attention layer (1) is an equivalent reformulation of (Vaswani et al., 2017), where they concatenate attention heads and multiply a matrix
W O ∈ R d × mh to the concatenation.
One difference in our setup is the absence of layer normalization, which simplies our analysis while preserving the basic architecture of the Transformer.
We deﬁne the Transformer networks as the composition of Transformer blocks.
The family of the sequence-to-sequence functions corresponding to the Transformers can be deﬁned as:
[Equation]
where t h,m,r : R d × n → R d × n denotes a Transformer block deﬁned by an attention layer with h heads of size m each, and a feed-forward layer with r hidden nodes.
We say that a function f : R d × n → R d × n is permutation equivariant if for any permutation matrix
P , we have f ( XP ) =
f ( X ) P ; i.e., if we permute the columns of X , then the columns of f ( X ) are permuted in the same way.
A Transformer block is permutation equivariant, which we formally prove in Section A.
This consequently establishes the permutation equivariance of the class T h,m,r .
Claim 1.
A Transformer block t h,m,r deﬁnes a permutation equivariant map from R d × n to R d × n .
T RANSFORMERS ARE UNIVERSAL APPROXIMATORS OF SEQUENCE - TO - SEQUENCE FUNCTIONS
In this section, we present our results showing that the Transformer networks are universal approximators of sequence-to-sequence functions.
Let us start by deﬁning the target function class F PE , which consists of all continuous permutation equivariant functions with compact support that map R d × n to R d × n .
Here, continuity is deﬁned with respect to any entry-wise (cid:96) p norm, 1 ≤ p < ∞ .
Given two functions f 1 , f 2 : R d × n → R d × n , for 1 ≤ p < ∞ , we deﬁne a distance between them as
[Equation]
The following result shows that a Transformer network with a constant number of heads h , head size m , and hidden layer of size r can approximate any function in F PE .
Theorem 2.
Let 1 ≤ p < ∞ and (cid:15) > 0 , then for any given f ∈ F PE , there exists a Transformer network g ∈ T 2 , 1 , 4 , such that d p ( f, g ) ≤ (cid:15) .
Next, we present our theorem on Transformers with positional encodings .
In order to endow the
Transformer networks with the ability to capture the information about the position of tokens in the input sequence, it is a common practice to add positional encodings E ∈ R d × n to the input sequence before feeding it to the Transformer network (Vaswani et al., 2017; Devlin et al., 2018).
Consider the functions represented by Transformers with positional encodings:
[Equation]
Here we show that if E is trainable, these positional encodings are sufﬁcient to remove the permutation equivariance restriction of the Transformers.
Towards this, we deﬁne F CD to be the set of all continuous functions that map a compact domain in R d × n to R d × n .
Note that F CD does not have the restriction of permutation equivariance as in F PE , but any f ∈ F CD is deﬁned on a compact domain instead of the whole R d × n .
The following result states that, equipped with the trainable positional encodings, Transformers can approximate any sequence-to-sequence function in F CD .
Theorem 3.
Let 1 ≤ p < ∞ and (cid:15) > 0 , then for any given f ∈ F CD , there exists a Transformer network g ∈ T 2 , 1 , 4 P such that we have d p ( f, g ) ≤ (cid:15) .
Theorems 2 and 3 provide an interesting characterization of the representation power of ﬁxed-width
Transformer networks.
Since the function classes T h,m,r and T h,m,r P become richer as we increase the values of ( h, m, r ) , our results establish that general Transformer networks are also universal approximators of sequence-to-sequence functions.
Remarkably, none of the parameters ( h, m, r ) depend on the input sequence length n or embedding dimension d .
Here, we would like to point out that Theorems 2 and 3 appear quite surprising at a ﬁrst glance, given the parameter sharing across all the tokens in a sequence, e.g., feed-forward layers are applied token-wise and the projection matrices in the self-attention layers are the same across different tokens.
Furthermore, attention layers can only capture pairwise interaction between different tokens in the sequence.
In the next subsection, we brieﬂy describe one of our key steps in overcoming the aforementioned restrictions and proving universal approximation power of Transformers.
3. A KEY STEP : SELF - ATTENTION LAYERS CAN IMPLEMENT CONTEXTUAL MAPPINGS
Let us consider a setting where we are interested in embedding two sentences: 1) I am happy; and 2) I am Bob.
These sentences are fed to a sequence-to-sequence model as
[Equation]
where v I , v am , v happy , and v Bob denote d -dimensional embedding for the tokens ‘I’, ‘am’, ‘happy’, and ‘Bob’, respectively.
Since the word ‘I’ occurs in different contexts in these sentences, in order to implement arbitrary sequence-to-sequence functions, the sequence-to-sequence model should map the two occurrences of ‘I’ to different values.
We formally deﬁne this requirement below.
Deﬁnition 3.1 (Contextual mapping) .
Consider a ﬁnite set L ⊂ R d × n .
A map q : L → R 1 × n deﬁnes a contextual mapping if the map satisﬁes the following:
For any L ∈ L , the n entries in q ( L ) are all distinct.
For any L , L (cid:48) ∈ L , with L (cid:54)
=
L (cid:48) , all entries of q ( L ) and q ( L (cid:48) ) are distinct.
In other words, a contextual mapping maps each token (column) of L ∈ L to a unique value which depends on the entire L ; as a result, capturing the precise context of L .
This allows the subsequent token-wise function (e.g., deﬁned by the feed-forward layers in case of Transformer networks) to realize the outputs of any arbitrary sequence-to-sequence functions.
At ﬁrst thought, we can consider getting a contextual mapping by simply averaging all the tokens, because this can capture the one-word difference (e.g., “happy” vs. “Bob”) in two different contexts.
However, if there are multiple words that are different, it is not guaranteed that the average will be different.
Indeed, requiring unique mappings for all the tokens for any change in any number of tokens, is a steep requirement.
While the self-attention layer does consider pair-wise interactions among different input tokens, it is not clear if this weak form of pair-wise interaction with shared projection weights is sufﬁcient to extract the underlying context.
The following result, which we sketch here, shows that self-attention layers can implement a permutation equivariant contextual mapping over almost all elements of a grid in [0 , 1] d × n .
We defer the full statement to Section 4.2.
Lemma 6 (informal) .
Consider the grid G δ := { 0 , δ, . .
, 1 − δ } d × n .
Then, there exist a function g c : R d × n → R d × n composed of δ − d + 1 self-attention layers ( h = 2 , m = 1 ) and a vector u ∈ R d such that q ( L ) :
=
u T g c ( L ) satisﬁes the following properties, for a subset (cid:101)
G δ ⊂ G δ that contains almost all elements of G δ :
For any L ∈ (cid:101) G δ , the entries of q ( L ) are all distinct.
For any L , L (cid:48) ∈ (cid:101)
G δ such that L is not a permutation of L (cid:48) , all entries of q ( L ) , q ( L (cid:48) ) are distinct.
Lemma 6 shows that a series of self-attention layers can implement contextual mappings, despite the apparent restriction that each of them can only capture pair-wise interaction.
However, the restriction of permutation equivarance still exists because attention layers are inherently permutation equivariant.
Coupled with the ability of token-wise feed-forward layers to map different values in q ( L ) to arbitrary output values, we can prove universal approximation capability of Transformers.
3. P ROOF OF THE UNIVERSAL APPROXIMATION THEOREM (T HEOREM 2)
Next, we outline the proof of Theorem 2 in greater detail.
We refer the reader to Section C for the proof of Theorem 3, since it is a modiﬁcation of Theorem 2.
Even though Theorems 2 and 3 do not speciﬁcally mention the required depth for approximation, our proof techniques do characterize it, and we show that our construction is tight in the number of parameters.
We defer the discussion of depth to Section 4.4.
Recall that we want to show that given a function f ∈ F PE , we can ﬁnd a Transformer network g ∈ T 2 , 1 , 4 such that d p ( f, g ) ≤ (cid:15) .
Without loss of generality, we can assume that the compact support of f is contained in [0 , 1] d × n .
We achieve our desired objective in three key steps:
Step 1.
Approximate F PE with piece-wise constant functions.
We ﬁrst use (a variant of) the classical result that any continuous function can be approximated up to arbitrary accuracy by piecewise constant functions.
For δ > 0 , we deﬁne the following class of piece-wise constant functions.
[Equation]
where G δ := { 0 , δ, . . .
, 1 − δ } d × n and, for a grid point L ∈ G δ , S L := (cid:81) dj
=1 (cid:81) nk =
[ L j,k , L j,k + δ ) ⊂ [0 , 1] d × n denotes the associated cube of width δ .
Let f ∈ F PE ( δ ) be such that d p ( f, f ) ≤ (cid:15)/ 3 .
Step 2.
Approximate F PE ( δ ) with modiﬁed Transformers.
We then consider a slightly modiﬁed architecture for Transformer networks, where the softmax operator σ [ · ] and ReLU( · ) are replaced by the hardmax operator σ H [ · ] and an activation function φ ∈ Φ , respectively.
Here, the set of allowed activations Φ consists of all piece-wise linear functions with at most three pieces, where at least one piece is constant.
Let T h,m,r denote the function class corresponding to the sequence-to-sequence functions deﬁned by the modiﬁed Transformer networks.
The following result establishes that the modiﬁed Transformer networks in T 2 , 1 , 1 can closely approximate functions in F PE ( δ ) .
Proposition 4.
For each f ∈ F PE ( δ ) and 1 ≤ p < ∞ , ∃ g ∈ T 2 , 1 , 1 such that d p ( f, g ) =
O ( δ d/p ) .
Step 3.
Approximate modiﬁed Transformers with (original) Transformers.
Finally, we show that g ∈ T 2 , 1 , 1 can be approximated by T 2 , 1 , 4 .
Let g ∈ T 2 , 1 , 4 be such that d p ( g, g ) ≤ (cid:15)/ 3 .
Theorem 2 now follows from these three steps, because we have
[Equation]
Choosing δ small enough ensures that d p ( f, g ) ≤ (cid:15) .
We refer the reader to Sections B.1 and B.2 in the supplementary material for the formal statements and proofs of Steps 1 and 3 , respectively.
As for Step 2 , which is the most critical step in establishing the universal approximation property of Transformers, we provide a sketch of the proof of
Proposition 4 in the next section, and refer the reader to Section B.3 for the complete proof.
P ROOF SKETCH OF P ROPOSITION 4: DIFFERENT ROLES OF TWO LAYERS
As mentioned earlier, the heavy parameter sharing in Transformers makes the goal of universally approximating sequence-to-sequence functions seemingly difﬁcult.
Both the self-attention and the feed-forward layer weights inside a Transformer block are ﬁxed across n tokens.
In this section, we show that Transformers are able to overcome this architectural constraint, and compute contextual mappings of the entire input sequence just based on the pair-wise interactions.
The token-wise feedforward layers then transform these contextual mappings to the desired output sequence.
We highlight these inner workings of Transformers en route to proving Proposition 4.
We want to show that given a piece-wise constant function f ∈ F PE ( δ ) , there exists a modiﬁed Transformer network g ∈ T 2 , 1 , 1 that closely approximates f .
We achieve this goal by establishing the following three claims, which correspond to Lemmas 5, 6, and 7.
Given an input X ∈ R d × n , a series of feed-forward layers in the modiﬁed Transformer network can quantize X to an element L on the extended grid G + δ :=
{− δ − nd , 0 , δ, . . .
, 1 − δ } d × n . 2.
Next, a series of self-attention layers in the modiﬁed Transformer network can take the input L and implement a contextual mapping q such that, for L and L (cid:48) that are not permutation of each other, all the elements in q ( L ) and q ( L (cid:48) ) are distinct.
Finally, a series of feed-forward layers in the modiﬁed Transformer network can map elements of the contextual embedding q ( L ) to the desired output value of f ∈ F PE at the input X .
Before discussing these three claims in detail, we note that even though a Transformer network stacks self-attention and feed-forward layers in an alternate manner, the skip connections enable these networks to employ a composition of multiple self-attention or feed-forward layers.
Furthermore, as alluded earlier, these three steps clearly highlight the different roles that self-attention and feed-forward layers play in realizing the ability to universally approximate sequence-to-sequence functions: 1) self-attention layers compute precise contextual maps; and 2) feed-forward layers then assign the results of these contextual maps to the desired output values.
4. Q UANTIZATION BY FEED - FORWARD LAYERS
Since our objective in Proposition 4 is to approximate the function f ∈ F PE ( δ ) , which takes a constant value on the cubes S L ’s, the (modiﬁed) Transformer network approximating f ﬁrst quantizes the input X
according to these cubes.
In particular, we want each input X ∈ S L to be mapped to the point L .
The following result shows that a modiﬁed Transformer network can indeed implement this quantization map with a composition of multiple feed-forward layers.
Lemma 5.
Consider a scalar quantization map g ent q :
[Equation]
[Equation]
[Equation]
− δ
[Equation]
otherwise. There exists a function g q : R d × n (cid:55)→ G + δ composed of dδ + d token-wise feed-forward layers with r = 1 and activations in Φ , which employs the scalar quantization g ent q to each entry of its input.
As desired, the function g q maps any X ∈ S L to L .
Furthermore, if any element of X is not in [0 , 1] , the element is mapped to − δ − nd , indicating that X is outside the compact support of f ∈ F PE ( δ ) .
4. C ONTEXTUAL MAPPING BY SELF - ATTENTION LAYERS
In this subsection, we show that the (modiﬁed) Transformer network can compute contextual mappings (cf.
Deﬁnition 3.1) from the output L ∈ G + δ of the map g q (cf.
Section 4.1) by using a composition of self-attention layers.
The following lemma, sketched earlier in Section 3.1, shows that the (modiﬁed) Transformer networks can implement a permutation equivariant contextual mapping over almost all elements of G δ , while mapping the rest of elements in G + δ to a disjoint set.
Lemma 6.
Consider the following subset of G δ = { 0 , δ, . .
, 1 − δ } d × n :
[Equation]
Assume that n ≥ 2 and δ − 1 ≥ 2 .
Then, there exist a function g c : R d × n → R d × n composed of δ − d + 1 self-attention layers ( h = 2 , m = 1 ) that employ the σ H operator, a vector u ∈ R d , constants t l , t r ∈ R ( 0 < t l < t r ) , such that q ( L ) :=
u T g c ( L ) satisﬁes the following properties:
For any L ∈ (cid:101) G δ , the entries of q ( L ) are all distinct.
For any L , L (cid:48) ∈ (cid:101)
G δ such that L is not a permutation of L (cid:48) , all entries of q ( L ) , q ( L (cid:48) ) are distinct.
For any L ∈ (cid:101) G δ , all the entries of q ( L ) are in [ t l , t r ] .
For any L ∈ G + δ \ (cid:101) G δ , all the entries of q ( L ) are outside [ t l , t r ] .
At this point, a few remarks about the result in Lemma 6 are in order.
First, since the Transformer networks are bound to implement permutation invariant maps, we require the Property 6.2 to hold for the pair of sequences that cannot be mapped to each other via permutation of columns.
Furthermore, the self-attention layers implement the desirable contextual map for only (cid:101) G δ ⊆
G δ , where all columns of L are distinct.
Note that for small δ , G δ \ (cid:101) G δ constitutes a negligible fraction of G δ because | G δ \ (cid:101) G δ | =
O ( δ d | G δ | ) .
The function q in Lemma 6 maps the elements of G + δ \ (cid:101)
G δ outside [ t l , t r ] —the interval where the outputs of the contextual mapping for (cid:101) G δ reside.
4.2. P ROOF SKETCH OF L EMMA
Since Lemma 6 is one of the major technical contributions of this paper, we provide a short sketch of its proof.
The complete proof is presented in Section B.5.
For simplicity, we consider the case d
=
1 , so the input L ∈ G + δ is a row vector of length n .
The key idea of the proof is that, using two attention heads of size 1 , one can implement a selfattention layer that shifts up input entries that are in a speciﬁc interval , while leaving all other entries intact.
We call this the selective shift operation .
Since the entries in L are quantized, we apply the selective shift operation to 0 , δ, . .
, 1 − δ using 1 /δ attention layers.
Interestingly, the value of the largest output entry after these operations is unique for each L ∈ (cid:101) G δ up to permutations.
Using the largest entry, one can add one last layer that shifts up the entire matrix and outputs q ( L ) that satisﬁes
Properties 6.1 and 6.2 of the lemma.
More concretely, the following function Ψ : R 1 × n → R 1 × n , parametrized by b, b (cid:48)
∈ R satisfying b < b (cid:48) , can be implemented with two attention heads of size 1 with the hardmax ( σ H ) operator:
[Equation]
If we deﬁne an attention layer of the form Z (cid:55)→ Z + Ψ( Z ; b, b (cid:48) ) , then any entry Z 1 ,j in ( b, b (cid:48) ) is shifted up by max k Z 1 ,k − min k Z 1 ,k , while all the other entries stay untouched .
We can choose b and b (cid:48) to selectively shift certain entries, hence the name selective shift operation.
We stack 1 /δ self-attention layers, with attention parts δ − 1 Ψ( · ; l − δ/ 2 , l + δ/ 2) for each l ∈ { 0 , δ, . . .
, 1 − δ } , in increasing order of l .
With these layers, we can apply the selective shift operations to input entries of values 0 , δ, . .
, 1 − δ .
To see how the shift operations modify the input, now consider n = 2 for simplicity, and let L = [ ll 2 ] ∈ (cid:101) G δ .
Without loss of generality, we can assume l 1 < l 2 .
The selective shift operation is applied to l 1 ﬁrst, shifting it by δ − 1 (max L − min L ) =
δ − 1 ( l 2 − l 1 ) , resulting in (cid:101) l 1 = l 1 + δ − 1 ( l 2 − l 1 ) > l 2 .
After that, the operation on l 2 shifts it up by δ − 1 ( (cid:101) l 1 − l 2 ) .
Thus, the ﬁrst 1 /δ layers map L = [ ll 2 ] ( l 1 < l 2 ) to
[Equation]
We can show that the map from [ ll 2 ] ∈ { L ∈ (cid:101)
G δ | l 1 < l 2 } to (cid:101) l 2 is one-to-one, and that 0 < (cid:101) l 1 < (cid:101) l 2 < δ − 2 .
We then add one last layer that shifts all positive entries of (cid:101) L by δ − 3 max (cid:101)
L = δ − 3 (cid:101) l 2 , whose output we denote by q ( L ) = (cid:2) δ − 3 (cid:101)
l 2 + (cid:101) lδ − 3 (cid:101) l 2 + (cid:101) l 2 (cid:3) .
All entries of q ( L ) are in [ δ − 3 (cid:101) l 2 , δ − 3 (cid:101) l 2 + δ − 2 ) , and this interval is disjoint for different L ’s because L (cid:55)→ (cid:101) l 2 is one-to-one.
Thus, q ( L ) satisﬁes Properties 6.1 and 6.2 of the lemma.
The remaining details are in Section B.5.
4. F UNCTION VALUE MAPPING BY FEED - FORWARD LAYERS
This brings us to the ﬁnal step, which demonstrates the key utility of the feed-forward layers.
After the contextual mapping by self-attention layers, each token captures the entire context available in the input sequence.
The following result shows that token-wise application of a composition of feed-forward layers can map these tokens to the desired output values required by the function f .
Lemma 7.
Let g c : R d × n → R d × n be the function from Lemma 6.
Then, there exists a function g v : R d × n → R d × n composed of O ( n ( 1 δ ) dn /n !)
token-wise feed-forward layers ( r = 1 ) with activations in Φ such that g v is deﬁned by a token-wise function g tknv : R d → R d on each column,
[Equation]
where for all j ∈ { 1 , . . . , n } ,
[Equation]
4. T IGHTNESS OF CONSTRUCTIONS
We showed in this section that Theorem 2 requires O ( n (1 /δ )
dn /n !)
Transformer blocks for approximation, where δ is the width of the cubes.
Each transformer block is of constant width, so it has O ( d ) parameters; this means that the total number of parameters is O ( dn (1 /δ )
dn /n !) .
We note that this exponential dependence cannot be avoided in the worse case.
If we assume continuity without any additional smoothness, quantizing the domain to cubes and approximating the function with constants require memorizing ( output dim ) ×
( num cubes ) /n !
real numbers, where the factor of 1 /n ! is due to permutation equivariance.
Thus, Theorem 2 is optimal in the order of parameters.
If we compare with the residual network result (Lin & Jegelka, 2018), we can consider “ﬂattening”
X into a dn -dimensional vector and ﬁtting the function.
The proof technique in (Lin & Jegelka, 2018) requires O ((1 /δ ) dn ) layers, where each layer has O ( dn ) parameters: the total parameter requirement is O ( dn (1 /δ ) dn ) .
This shows that Transformers can approximate permutation equivariant functions in a more efﬁcient way than residual networks.
In Section C, our proof of Theorem 3 shows that we require O ( n (1 /δ ) dn ) layers to approximate continuous (not permutation equivariant) sequence-to-sequence functions.
As seen from the argument above, this construction is also optimal in the order of parameters.
D ISCUSSION AND E XPERIMENTS
As detailed in Section 4, the ability of the self-attention layers to compute contextual mappings plays a crucial role in the universal approximation property.
Interestingly, our analysis shows that replacing the dot-product attention in Transformers with any other component capable of computing contextual mappings should preserve this universal approximation property.
This leads naturally to questions about the alternative architectures that realize certain kinds of contextual mappings at different computational and memory costs.
We explore and discuss some examples of such alternatives in this section.
Our preliminary empirical study demonstrates their practical utility.
5. B I - LINEAR PROJECTION
Given token embeddings X as input, the bi-linear projection layer computes the following update.
[Equation]
The bi-linear projection layer (Gong et al., 2013) is motivated from the ability of random (Gaussian) matrices to map sparse differences to dense vectors (Ailon & Chazelle, 2009).
If there are two input contexts X 1 and X 2 that differ in one token, their difference X 1 − X 2 is sparse; however, after random projection, the difference ( X 1 − X 2 ) W P will be dense, and the numbers are distinct with high probability, implementing a form “pair-wise contextual mapping,” 2 although different from the contextual mapping in Deﬁnition 3.1.
This layer advantageously incurs smaller number of matrix multiplications as compared to the dotproduct attention.
That said, the number of parameters in this layer depend on the sequence length, [Figure] making it harder to reuse the model across tasks with different input sequence lengths.
Moreover, the weights used to compute the contextual embeddings ( W P ) are independent of the inputs ( X ), whereas in self-attention the weights ( σ [( W iK X ) T W iQ X ]) depend on X .
The ﬁrst drawback can be addressed by replacing the linear projection with a depth-wise separable convolution layer, which is discussed in the next subsection.
5. D EPTH - WISE SEPARABLE CONVOLUTIONS
A depth-wise convolution layer (Sifre & Mallat, 2014; Chollet, 2017; Kaiser et al., 2017) involves convolving each dimension of X with a corresponding convolution ﬁlter of size k :
[Equation]
where W C ∈ R
d
× k
and ( X ∗ W C )
i, : :=
X i, : ∗ ( W C )
i, : .
Unlike bi-linear projection, this layer can be used across tasks with different input sequence lengths as the number of parameters are independent of the sequence length.
While a single layer is unable to compute contextual mappings when the ﬁlter size is small, stacking multiple such layers can potentially provide a cheaper way to compute contextual mappings.
In fact, based on depth-wise separable convolutions, Wu et al.
(2019) proposed a light-weight dynamic convolution architecture that performs competitively with
Transformers on machine translation.
5. E XPERIMENTS
We now present our experiments with these other architectures, with the goal of understanding the extent to which computing contextual mappings can capture the performance of Transformers.
As discussed earlier, BProj and SepConv do not implement contextual mappings (cf.
Deﬁnition 3.1), so we do not expect that either BProj or SepConv based models to have the same performance as the expensive Transformers.
These models do not use input dependent weights to compute attention, and hence have weaker representation power.
Instead, our goal is to see if we can use these cheaper layers to replace (some of) the expensive self-attention layers.
We follow the experimental setting from Devlin et al.
(2018) to train the Transformers, with the masked language model pre-training followed by a task speciﬁc ﬁne-tuning, and work with a 12 layer architecture based on BERT BASE .
We present our results on a question answering task (SQuAD) (Rajpurkar et al., 2016) and a sentence entailment task (MNLI)
(Williams et al., 2018).
In our ﬁrst set of experiments we train models that employ BProj and SepConv layers, instead of the selfattention layer in eq.(1).
We notice that, as expected, these simpler models have weaker performance than the self-attention layer.
See Table 1 in Section D for a comparison of these models on MNLI.
Next, we swap a varying number of the ﬁrst few self-attention layers in BERT BASE with SepConv , implemented with ﬁlter reuse across dimensions (Wu et al., 2019) 3 .
Fig.
1 illustrates the performance of these hybrid models.
Interestingly, models with 1 or 2 convolution layers and rest the self-attention layers, perform better than models with only the self-attention layers.
Note that, replacing self-attention layer with SepConv also reduces the computational cost and the number of parameters.
One explanation we have is that the ﬁrst few attention layers tend to attend broadly to the whole sequence (as empirically observed in (Clark et al., 2019)), and the cheaper convolution layers can perform this job more efﬁciently.
A detailed evaluation of such hybrid architectures will be interesting future research.
Our experiments also call for a deeper understanding of the exact nature of the embeddings computed by practical attention models.
Since Transformers in practice have ﬁxed depth, we believe that they might not be able to exactly implement contextual mappings as we deﬁned in Deﬁnition 3.1.
However, there is some preliminary empirical evidence that Transformers do implement some sort of “contextual mappings.”
For example, Fig. 4 of Coenen et al.
(2019) presents visualizations of embeddings of a single word in different contexts (sentences).
They experimentally notice that Transformers, in addition to computing contextual mappings, also map a word into semantic clusters.
Formalizing and evaluating this property of Transformers is an interesting direction for future work.
We again note that Wu et al.
(2019) have proposed an alternative way to compute such embeddings based on dynamic convolution layers.
Evaluating the mappings computed by these models should shed more light on the workings of attention models and inspire efﬁcient and better performing architectures.
A P ROOF OF C LAIM
Suppose XP was given as input, where P is a permutation matrix.
First note that
[Equation]
After the softmax operation, we get
[Equation]
Then,
[Equation]
where we used P P T
= I .
Permutation equivariance of the token-wise feed-forward layer can be shown similarly:
[Equation]
where ReLU( XP ) =
ReLU( X ) P was used.
This analysis shows that the function class T h,m,r ( · ) is restricted to permutation equivariant functions.
B P ROOF DETAILS OF T HEOREM
We ﬁrst deﬁne some additional notation.
For a, b ∈ N where a ≤ b , let [ a ] = { 1 , . . .
, a } and [ a : b ] = { a, a + 1 , . .
, b − 1 , b } .
For a, b, c ∈ R where b − a > 0 is an integer multiple of c > 0 , we write [ a : c : b ] := { a, a + c, a + 2 c, . . .
, b − c, b } .
B.
A PPROXIMATING F PE WITH F PE ( δ )
Lemma 8.
For any given f ∈ F PE and 1 ≤ p < ∞ , one can ﬁnd a δ ∗ > 0 such that ∃ f ∈ F PE ( δ ∗ ) which satisﬁes d p ( f, f ) ≤ (cid:15)/ 3 .
Proof
Since f : R d × n → R d × n is a continuous function with compact support, the function is uniformly continuous.
Since continuity is deﬁned using entry-wise (cid:96) p norm, and entry-wise (cid:96) p norm is equivalent to entry-wise (cid:96) ∞ norm when the number of entries are ﬁnite, uniform continuity implies that
[Equation]
This means that given any (cid:15)/ 3 > 0 , we have such a δ > 0 .
Using this δ , we can create a grid G δ and corresponding cubes S L , as described in the main text.
For any L ∈ G δ , we deﬁne C L ∈ S L to be the center point of the cube S L .
Then, we can deﬁne a piece-wise constant approximation f ( X ) =
(cid:80) L ∈ G δ f ( C L ) 1 { X ∈ S L } .
Note that, for any X ∈ S L , we have (cid:107) X − C L (cid:107) ∞ <
δ , so by uniform continuity, we have (cid:13)(cid:13) f
( X ) − f ( X ) (cid:13)(cid:13) p
= (cid:107) f
( X ) − f ( C L ) (cid:107) p < (cid:15)/ 3 .
This proves that d p ( f, f ) < (cid:15)/ 3 .
As for permutation equivariance, since f is permutation equivariant, we have f ( C L P ) =
f ( C L ) P for any permutation matrix P .
For any X ∈ S L , we have XP ∈ S LP , so
[Equation]
Thus, the approximation f is also permutation equivariant.
This proves the lemma.
B. A PPROXIMATING T 2 , 1 , 1 WITH T 2 , 1 ,
Lemma 9.
For each g ∈ T 2 , 1 , 1 and 1 ≤ p < ∞ , ∃ g ∈ T 2 , 1 , 4 such that d p ( g, g ) ≤ (cid:15)/ 3 .
Proof
Recall that T h,m,r refers to the class of functions representable with composition of Transformer blocks with h heads of size m in self-attention layers and r hidden nodes in feed-forward layers.
The same notation holds for the modiﬁed Transformers T h,m,r .
Note that the softmax operator on a matrix A can be made arbitrarily close to hardmax by scaling up A .
That is,
[Equation]
This means that by scaling up parameters inside σ , we can approximate σ H arbitrarily closely.
Thus, the modiﬁed self-attention layers can be approximated with the original self-attention layers of the same number of heads h and head size m .
Also, any arbitrary (possibly discontinuous) piecewise linear function φ ∈ Φ can be approximated arbitrarily closely by four ReLU ’s. Note that φ ∈ Φ as at most three pieces, and at least one of the pieces is constant.
For example, consider the following function φ ∈ Φ :
[Equation]
This function can be approximated by four ReLU ’s, as claimed by the lemma:
[Equation]
 
[Equation]
Also, as we make (cid:15) → 0 , we can approximate φ as closely as possible using (cid:101) φ .
The cases where the second or third piece is constant can be shown similarly.
This means that the modiﬁed feed-forward layers (whose activation is φ ∈ Φ ) with single hidden node can be approximated with the original feed-forward layers ( ReLU ) with four hidden nodes.
Thus, given any g ∈ T 2 , 1 , 1 , there exists a function g ∈ T 2 , 1 , 4 arbitrarily close to g , by appropriately choosing the parameters to be large enough.
This ﬁnishes the proof.
B. F INISHING PROOF OF P ROPOSITION
As we have already discussed in Section 4, we establish Proposition 4 in three steps:
Given an input X , a group of feed-forward layers in the modiﬁed Transformer network can quantize X to an element L on the extended grid G + δ :=
{− δ − nd , 0 , δ, . . .
, 1 − δ } d × n . 2.
Next, a group of self-attention layers
in the modiﬁed Transformer network can take the input L and produce desirable contextual mappings q ( L ) such that, for L and ˜ L , that are not permutation of each other, all the elements in q ( L ) and q ( ˜ L ) are distinct.
Finally, a group of feed-forward layers in the modiﬁed Transformer network can map elements of the contextual embedding q ( L ) to the desirable values, i.e., the output of f ∈ F PE on the input X .
These steps are formally stated in Lemmas 5, 6, and 7 in the main text.
We present the proofs of these lemmas in the subsequent sections.
With the results established in these lemmas, we are now equipped with all the tools necessary to complete the proof of Proposition 4.
Let us recall the functions g q , g c , and g v from Lemma 5, 6, and 7, respectively.
We now show that the (modiﬁed) Transformer network g = g v ◦ g c ◦ g q approximates the underlying peicewise constant function f ∈ F PE over all points in its support except for a set of of measure
O ( δ d ) .
Consider a point X ∈ S L ⊂ [0 , 1] d × n , where L ∈ (cid:101) G δ .
By Lemma 5, we have that g q ( X ) =
L .
Thus, it follows from Lemmas 6 and 7 that
[Equation]
On the other hand, any point X ∈ (cid:83) L ∈ G δ \ (cid:101) G δ S L ∪ ( R d × n \ [0 , 1] d × n ) is mapped by g q to L ∈
G + δ \ (cid:101) G δ ; as a result, we get g v ◦ g c ◦
g q ( X ) =
g v ◦ g c
( L ) = 0 .
Therefore, we have g ( X ) =
g v ◦ g c ◦ g q ( X ) =
A L = f ( X ) for X ∈ (cid:83) L ∈ (cid:101) G δ S L , and 0 everywhere else.
Recall that f has its compact support in [0 , 1] d , thus bounded; i.e., there exists B ≥
such that (cid:107) f
( X ) (cid:107) p ≤ B .
The modiﬁed Transformer network g takes the same value as f on all points in [0 , 1] d except for a set (cid:83) L ∈ G δ \ (cid:101)
G δ S L that has measure O ( δ d ) .
This implies that d p ( f, g ) ≤ ( B p δ d ) 1 /p =
O ( δ d/p ) .
B. P ROOF OF L EMMA
The proof strategy is simple; using 1 δ + 1 token-wise feed-forward layers, we implement the quantization function g entq that works on the ﬁrst row of the input.
Then stack another 1 δ + 1 layers that quantizes the second row, and so on.
Given input X , we ﬁrst start by clipping X 1 , : in the set ( −∞ , 0) ∪ [1 , + ∞ ) and mapping the intervals to − δ − nd .
This can be done by the following layer:
[Equation]
Next, add 1 /δ layers of the following form, for k = 0 , δ, . .
, 1 − δ .
[Equation]
Each layer quantizes X 1 , : in [ kδ, kδ + δ ) to kδ , without modifying other intervals.
Note that both φ ’s used in this construction are piecewise linear functions with three pieces, and at least one of them are constant.
Thus, both φ ’s are in Φ .
We can repeat the same thing for the other rows, and at the end we will get a map from R d × n to G + δ .
B. P ROOF OF L EMMA
Selective shift operation.
Before starting the proof, we ﬁrst describe the key component of our proof, which we refer to the selective shift operation .
Consider the following function, which can be expressed with a multiplicative attention head, with head size m = 1 and hardmax σ H :
[Equation]
where u ∈ R d is a vector that we will choose later, and e (1)
=
. , 0) ∈
R
d is the standard basis vector.
To see what this function computes, ﬁrst consider the j -th column of the attention score matrix: ( u T Z ) T ( u T Z : ,j − b Q ) .
Note that, if u T Z : ,j > b Q , σ H will calculate arg max of u T Z , whereas if u T Z : ,j < b Q , it will calculate arg min .
Therefore, the (1 , j ) -th entry of ψ ( Z ; b Q ) ∈ R
d × n can be written as
[Equation]
for j ∈ [ n ] .
Note that due to e (1) , all rows of ψ ( Z ; b Q ) except the ﬁrst row are zero.
From this observation, one can deﬁne a function parametrized by b Q and b (cid:48) Q , where b Q < b (cid:48) Q , which consists of two attention heads:
[Equation]
What this means is that, if we deﬁne an attention layer of the form Z (cid:55)→ Z + Ψ( Z ; b Q , b (cid:48) Q ) , then any column Z : ,j satisfying u T Z : ,j ∈ ( b Q , b (cid:48) Q ) is shifted up in its ﬁrst coordinate Z 1 ,j by max k u T Z : ,k − min k u T Z : ,k , while all the other coordinates stay untouched .
We call this the selective shift operation, because we can choose b Q and b (
cid:48) Q to selectively shift certain entries of the input.
Bijective column id mapping.
Recall that the input to this step is from the range of g q (Lemma 5), which is G + δ = {− δ − nd , 0 , δ, . .
, 1 − δ } d × n .
Now consider L ∈ G + δ and u = (1 , δ − 1 , δ − 2 , . .
, δ − d +1 ) .
For any j ∈ [ n ] , it is easy to check two following facts:
If L i,j (cid:54) = − δ − nd for all i ∈ [ d ] , i.e., L : ,j ∈ { 0 , δ, . . .
, 1 − δ } d , then u T L : ,j ∈ [0 : δ : δ − d
+1 − δ ] , and the map L : ,j (cid:55)→ u T L : ,j from { 0 , δ, . .
, 1 − δ } d to [0 : δ : δ − d +1 − δ ] is a bijection.
If there exists i ∈ [ d ] such that L i,j = − δ − nd , then u T L : ,j ≤ − δ − nd + δ − d +1 − 1 < 0 .
Therefore, one can say that u T L : ,j gives the “column id” for each possible value of L : ,j ∈ {
0 , δ, . .
, 1 − δ } d .
The rough idea of the construction is to apply the selective shift operation to each column id, by setting u in the deﬁnition of Ψ( · ) to be (1 , δ − 1 , δ − 2 , . . .
, δ − d +1 ) and choosing b Q = l − δ/ 2 and b (cid:48)
Q = l + δ/ 2 for each l ∈ [0 : δ : δ − d +1 − δ ] .
More concretely, we stack (1 /δ )
d attention layers, with attention parts δ − d Ψ( · ; l − δ/ 2 , l + δ/ 2) for each l ∈ [0 : δ : δ − d +1 − δ ] , in increasing order of l .
After that, we add an extra single-head attention layer with attention part δ − ( n +1)
d ψ ( · ; 0) .
We now divide possible input values L ∈ G
+ δ into three disjoint categories,
and show how these layers change the input values at the end of all the layers.
Recall the hierarchy (cid:101)
G δ ⊂ G δ ⊂ G + δ .
The categories are deﬁned as follows:
L ∈ (cid:101) G δ .
All entries are between 0 and 1 − δ , and all columns are unique.
L ∈ G δ \ (cid:101) G δ .
All entries are between 0 and 1 − δ , but there are duplicate columns.
L ∈ G + δ \ G δ .
The point has at least one entry that equals to − δ − nd .
B.5. C ATEGORY
In Category 1, we have L ∈ (cid:101) G δ .
Let l j :=
u T L : ,j .
Due to permutation equivariance, we can assume without loss of generality that l j ’s are in increasing order: l 1 < l 2 < · · · < l n .
The ﬁrst (1 /δ )
d layers sweep the set [0 : δ : δ − d +1 − δ ] and apply selective shift operation on each element in the set.
This means that selective shift operation will be applied to l 1 ﬁrst, then l 2 , and then l 3 , and so on, regardless of the speciﬁc values of l j ’s.
First shift operation.
In the ﬁrst selective shift operation, the (1 , 1) -th entry of L ( L 1 , 1 ) is shifted by the operation, while the other entries are left untouched.
The updated value (cid:101) L 1 , 1 is
[Equation]
Therefore, after the operation, the output of the layer is (cid:2)(cid:101) L : ,
L
[Equation] [Equation]
L : ,n (cid:3) , and the new value of the ﬁrst column (cid:101) L : , 1 results in u
[Equation]
Let us denote the updated “column id” u T (cid:101) L : , 1 as (cid:101) l 1 .
We can show that l n < (cid:101) l 1 , because
[Equation]
Therefore, after updating,
[Equation]
and the new minimum is l 2 .
Second shift operation.
The second selective shift operation is applied to l 2 , by which only one entry L 1 , 2 will be shifted.
The updated value (cid:101) L 1 , 2 is
[Equation]
After updating, the new inner product of u and (cid:101) L : , 2 results in
[Equation]
We can show that (cid:101) l 1 < (cid:101) l 2 , because
[Equation]
⇔ [Equation] ) , and the last inequality is true because δ − d > 1 and l n > l 2 .
Since we have (cid:101) l 1 < (cid:101) l 2 , and the new maximum in u T (cid:2)(cid:101) L : ,(cid:101) L : ,
L : ,. . .
L : ,n (cid:3) is now (cid:101) l 2 , and the new minimum is l 3 .
Repeating the process.
More generally, we can repeat this process, and show that the j -th shift operation shifts L 1 ,j by δ − d ( (cid:101) l
j − 1 − l j ) , and results in the new column id (cid:101) l j := u T (cid:101) L : ,j = l j + j − 1 (cid:88) k =
[Equation]
k +1 )
+ δ − jd ( l n − l 1 ) .
In the general case, (cid:101) l j − 1 < (cid:101) l
j holds j = [2 : n ] , because (cid:101) l
j − 1 = l j − 1 + j − 1 (cid:88) k
=
[Equation]
and the last inequality holds because
[Equation]
Therefore, after the j -th selective shift operation, (cid:101) l j is the new maximum among { (cid:101) l 1 , . . .
, (cid:101) l j , l j +1 , . . .
, l n } and l j +1 is the new minimum, which makes us possible to continue the process until the n -th operation.
After n shift operations.
As a result, after the whole sweep from 0 to δ − d +1 − δ by the ﬁrst (1 /δ ) d layers, a total of n shift operations are applied, and the input L is mapped to a new point (cid:101) L ,
where u T (cid:101) L
=
(cid:2)(cid:101) l(cid:101) l. . .
(cid:101) l n (cid:3) and (cid:101) l 1 < (cid:101) l 2 < · · · < (cid:101) l n .
We can now prove the following technical lemma, whose proof is deferred to Appendix B.5.4:
Lemma 10.
After n shift operations, (cid:101) l n = u T (cid:101) L : ,n satisﬁes the following bounds:
[Equation]
Also, the map from [ ll· · · l n ] ∈ [0 : δ : δ − d +1 − δ ] (where l 1 < l 2 < · · · < l n ) to (cid:101) l n is one-to-one.
Global shifting by the last layer.
As mentioned earlier, after this sweep, there is another attention layer with attention part δ − ( n +1)
d ψ ( · ; 0) .
Since 0 < (cid:101) l 1 < · · · < (cid:101) l n , what it does to (cid:101) L is that it adds δ − ( n +1)
d max k u T (cid:101) L : ,k = δ − ( n +1)
d (cid:101) l n to each entry in the ﬁrst row of (cid:101) L .
The output of this layer is deﬁned to be the function g c ( L ) .
Now, in summary, for any L ∈ (cid:101) G δ , i ∈ [ d ] , and j ∈ [ n ] , we have g c ( L ) i,j =
[Equation] [Equation]
if i = 1 ,
L i,
[Equation]
if i ∈ [2 , d ] , and for any L ∈ (cid:101) G δ and j ∈ [ n ] ,
[Equation]
[Equation]
Checking Properties 6.1 and 6.2.
Given this result so far, it is now left to check if the constructed network is really a permutation equivariant contextual mapping, i.e., if it satisﬁes Properties 6.1 and 6.2 in Lemma 6.
First, for any L ∈ (cid:101) G δ , Property 6.1 holds because we already know (cid:101) l 1 < (cid:101) l 2 < · · · < (cid:101) l n , so they are all distinct.
As for Property 6.2, note that the upper bound on (cid:101) l n from Lemma 10 also holds for other (cid:101) l j ’s, so [Equation] [Equation]
[Equation]
, for all j ∈ [ n ] .
Now, from Lemma 10, two L , L (cid:48) ∈ (cid:101) G δ
(that are not permutations of each other) map to different (cid:101) l n and (cid:101) l (cid:48) n , and they differ at least by δ .
This means that two intervals [ δ − ( n +1) d (cid:101) l n , δ − ( n +1) d (cid:101) l n + δ − ( n +1)
d +1 )
and [ δ − ( n +1) d (cid:101) l (cid:48) n , δ − ( n +1)
d (cid:101) l (cid:48) n + δ − ( n +1)
d +1 ) are guaranteed to be disjoint, so the entries of u T g c ( L ) and u T g c ( L (cid:48) ) are all distinct.
This proves Property 6.2.
Therefore, we ﬁnished showing that the map g c ( · ) we constructed using (1 /δ )
d + 1 attention layers implements a permutation equivariant contextual mapping on (cid:101) G δ .
Checking Property 6.3.
It is now left to check if the map g c satisﬁes the other properties.
At this point, we can check Property 6.3.
From u T g c ( L ) : ,j ∈ [ δ − ( n +1) d (cid:101) l n , δ − ( n +1)
d (cid:101) l n + δ − ( n +1)
d +1 ) and
Lemma 10, we can show that for any L ∈ (cid:101) G δ , we have δ {+− 2 nd +
+}
( δ − d − 1) ≤ u T g c ( L ) : ,j < δ − ( n +1)
d ( {+δ − nd +
+}
( δ − d − 1) − δ ( δ − d − 1) 2 ) + δ − ( n +1) d + ≤ δ − (2 n +1)
d +1 ( δ − d − 1) , where we used δ − 1 ≥ 2 .
This proves that all u T g c ( L ) : ,j are between t l = δ − 2 nd +
( δ − d − 1) and t r = δ − (2 n +1)
d +1 ( δ − d − 1) .
For the remaining input points L ∈ G + δ \ (cid:101) G δ , we will check that u T g c ( L ) : ,j is outside the interval [ t l , t r ] (Property 6.4).
B.5. C ATEGORY
In Category 2, we have L ∈ G δ \ (cid:101) G δ .
Here, entries δ
For the input L in Category 2, there exist some j, j (cid:48) ∈ [ n ] , j (cid:54) =
j (cid:48) , such that l j = l j (cid:48) .
This means that when the input passes through the attention layer δ − d Ψ( · ; l j − δ/ 2 , l j + δ/ 2) , the selective shift operation for l j is applied to both j -th and j (cid:48) -th columns; the two columns are coupled together.
More generally, suppose we have n (cid:48) < n distinct columns.
If n (cid:48) = 1 .
In the extreme case of n (cid:48)
= 1 , we have max j l j = min
j l j , so the selective shift operation applied at l j does not shift the entry at all; therefore, at the end of the ﬁrst (1 /δ )
d attention layers, (cid:101)
L = L .
If 1 < n (cid:48) ≤ n − 1 .
When 1 < n (cid:48) ≤ n − 1 , let the n (cid:48) distinct values of l j ’s be l (cid:48) 1 , . .
, l (cid:48) n (cid:48) .
The shift operation is applied n (cid:48) times, to l (cid:48) 1 , . . .
, l (cid:48) n (cid:48) , and shifts one or more entries at a time.
After the ﬁrst (1 /δ )
d layers, the output (cid:101) L has n (cid:48) distinct (cid:101) l j = u T (cid:101) L : ,j , 0 ≤ (cid:101) l 1 ≤ (cid:101) l 2 ≤ · · · ≤ (cid:101) l n , whose distinct values are the same as the numbers we get when we apply shift operations to a length- n (cid:48) sequence [ l (cid:48). . .
l (cid:48) n (cid:48) ] .
Then, applying the same calculations from Category 1 shows that
[Equation]
and it follows from the upper bound in Lemma 10 that
[Equation]
Note that the RHS matches the lower bound in Lemma 10.
This implies that the value of (cid:101) l n calculated from the input L ∈ G δ \ (cid:101) G δ (Category 2) is always strictly less (by at least δ ) than that calculated from L ∈ (cid:101) G δ
(Category 1).
Checking Property 6.4.
After the global shifting by the last layer with attention part δ − ( n +1)
d ψ ( · ; 0) , we get the output g c ( L ) which satisﬁes
[Equation]
where the RHS is a lower bound on possible values of u T g c ( L ) : ,j for L ∈ (cid:101) G δ (Category 1).
This means that the entries of u T g c ( L ) for Category 2 are outside [ t l , t r ] , which satisﬁes Property 6.4.
B.5. C ATEGORY
In Category 3, we have L ∈ G + δ \ G δ ; the point L has at least one entry that equals to − δ − nd .
Let l j :=
u T L : ,j , and recall that whenever a column L : ,j has an entry that equals to − δ − nd , we have l j = u T L : ,j ≤ − δ − nd + δ − d +1 − 1 < 0 .
Assume without loss of generality that l 1 ≤ l 2 ≤ · · · ≤ l n .
Recall that the selective shift operation is applied to each element of [0 : δ : δ − d +1 − δ ] , not to negative values.
In case of Category 3, we have min k u T L : ,k = l 1 < 0 , and l 1 never gets shifted upwards, so it remains as the minimum for the whole time.
If all l j ’s are negative.
In case where all l j ’s are negative, selective shift operation never changes the input L , so we get (cid:101)
L = L .
Since we have u T (cid:101) L <
Tn (entry-wise), the last layer with attention part δ − ( n +1)
d ψ ( · ; 0) adds δ − ( n +1)
d min k u T (cid:101) L : ,k < 0 to each entry in the ﬁrst row of (cid:101) L , further pushing it to the negative side.
Therefore, the ﬁnal output g c ( L ) satisﬁes u T g c ( L ) <
Tn <
t l 1 Tn .
If not all l j ’s are negative.
Now consider the case where at least one l j is positive.
Let i be the index that satisﬁes l
i − 1 < 0 ≤ l i .
Then, selective shift operation does not affect l 1 , . . .
, l i − 1 , and then it shifts l i by
[Equation]
where we used δ − 1 ≥ 2 at the last inequality.
The next shift operations shift l i +1 , . . .
, l n by even larger amount, so at the end of the ﬁrst (1 /δ )
d layers, we have δ − ( n +1)
d +1 ≤ (cid:101) l i ≤ · · · ≤ (cid:101) l n , while (cid:101)
Shifts by the last layer.
Here, the last layer with attention part δ − ( n +1)
d ψ ( · ; 0) acts differently for negative and positive (cid:101) l j ’s. For negative (cid:101) l j ’s, it adds δ − ( n +1)
d min k (cid:101) l k = δ − ( n +1)
d l 1 <to (cid:101) l 1 , . .
, (cid:101) l i − 1 , pushing them further to the negative side.
For positive (cid:101) l j ’s, the layer adds δ − ( n +1)
d max k (cid:101) l k = δ − ( n +1) d (cid:101) l n ≥ δ − (2 n +2)
d +1 to (cid:101) l i , . . .
, (cid:101) l n , so that they are all greater than or equal to δ − (2 n +2)
d +1 .
Note that δ − (2 n +2)
d +1 > t r .
Checking Property 6.4.
Therefore, in both cases, we can see that the ﬁnal output g c ( L ) satisﬁes u T g c ( L ) : ,j / ∈ [ t l , t r ] , for all j ∈ [ n ] .
This completes the veriﬁcation of Property 6.4.
B.5. P ROOF OF L EMMA
Proof of lower and upper bounds on (cid:101) l n are straightforward:
[Equation]
For one-to-one property of the map, consider [ ll· · · l n ] and [ l (cid:48)l (cid:48)· ·
[Equation]
l (cid:48) n ] with increasing entries, which are mapped to (cid:101) l n and (cid:101) l (cid:48) n , respectively.
Suppose (cid:101) l n = (cid:101) l (cid:48) n .
By deﬁnition,
[Equation]
[Equation]
Now assume for contradiction that l n (cid:54)
=
l (cid:48) n .
Then, we have − δ −
d
+1 + δ ≤ l n − l (cid:48) n ≤ δ − d
+1 − δ .
However, the remaining terms have “coarse resolution”, and they can never cancel l n − l (cid:48) n and make the sum zero, because for example, δ − d ( l n − 1 − l n − l (cid:48) n − 1 + l (cid:48) n ) can only have values 0 , δ − d
+1 , − δ − d
+1 , 2 δ − d +1 , − 2 δ − d +1 , . . . .
Thus, l n = l (cid:48) n must hold and the ﬁrst term must be zero.
Similarly, assume that l n − 1 (cid:54) =
l (cid:48)
n − 1 .
Then, the second term is in the interval [ − δ − 2 d +1 +
δ − d +1 , δ − 2 d
+1 − δ − d +1 ] .
Again, the remaining terms cannot cancel the second term, hence l n − 1 = l (cid:48) n − 1 must hold.
We can proceed this way, and show that l j = l (cid:48)
j must hold for all j ∈ [ n ] , hence proving that the map is one-to-one.
B. P ROOF OF L EMMA
Note that | G + δ | =
( 1 δ + 1) dn , so the image of g c ( G + δ ) (from Lemma 6) has ﬁnite number of distinct real numbers.
Let M be the maximum over all these numbers.
By construction of g c , we know that
M > 0 .
To construct a function g tknv that satisﬁes the statement of the lemma, we ﬁrst implement the second part: g tknv ( g c ( L ) : ,j ) =
d if L ∈ G + δ \ (cid:101) G δ .
Note from Lemma 6 that, for any L ∈ (cid:101) G δ , we have u T g c ( L ) : ,j ∈ [ t l , t r ] for all j , and for any L ∈ G + δ \ (cid:101) G δ , u T g c ( L ) : ,j / ∈ [ t l , t r ] for all j .
Using this, we add the following feed-forward layer:
[Equation]
, t r ] .
Input to this layer is g c ( L ) .
If L ∈ (cid:101) G δ , then φ ( u T g c ( L )) =
Tn , so the output stays the same as the input.
If L ∈ G + δ \ (cid:101) G δ , then φ ( u T g c ( L ))
=
1 Tn , so all the entries of the input are shifted by − M − 1 , and become strictly negative.
Recall that by deﬁnition of (cid:101) G δ , all the entries of g c ( L ) for L ∈ (cid:101)
G δ are nonnegative.
So the next thing to do is mapping all strictly negative entries to zero.
This can be done in a similar way as
Lemma 5.
For i ∈ [ d ] , add the following layer:
[Equation]
After these d layers, the output for L ∈ G + δ \ (cid:101) G δ is a zero matrix, while the output for L ∈ (cid:101)
G δ is g c
( L ) .
Now, it is left to map g c ( L ) to A L , for L ∈ (cid:101)
G δ .
Up to permutation equivariance, each different context L maps to n unique numbers u T g c ( L ) , which are at least δ apart from each other.
The idea of value mapping is to map each unique number to the corresponding output column.
More precisely, choose any L ∈ (cid:101) G δ .
For each value of u T g c ( L ) : ,j , j ∈ [ n ] , we add one feedforward layer
[Equation]
If the input Z is a zero matrix, which is the case for L ∈ G + δ \ (cid:101) G δ , u T Z =
Tn .
Since t l is much larger than 0 , activation is all zero.
Thus, zero input matrix remains the same at the output.
If the input Z is g c ( L ) , where L ∈ (cid:101) G δ is not a permutation of L , then
[Equation]
) − u T g c ( L ) : ,j 1 Tn ) =
Tn , so g c ( L ) is left untouched.
If some other L is a permutation of L , and L : ,i =
L : ,j , then φ ( u T g c ( L ) − u T g c
( L ) : ,j 1 Tn ) =
( e ( i ) )
T , so i -th column of g c ( L ) will turn to g c ( L ) : ,i (cid:55)→ g c ( L ) : ,i + (( A L ) : ,j − g c ( L ) : ,j ) =
g c ( L ) : ,i + (( A L ) : ,i − g c
( L ) : ,i ) =
( A L ) : ,i , which is the desired output.
In conclusion, this layer maps the column g c ( L ) : ,j to ( A L ) : ,j , without affecting any other columns.
As seen above, we need one layer per each unique value of u T g c ( L ) : ,j for each L ∈ (cid:101) G δ .
Note that there are O ( n (1 /δ )
dn /n !)
such numbers, so we can use O ( n (1 /δ )
dn /n !)
layers to ﬁnish our construction.
C P ROOF OF T HEOREM
Proof of Theorem 3 can be done in a similar way as Theorem 2.
As in the proof of Theorem 2, there are three parts: Lemma 8, Proposition 4, and Lemma 9.
The statement and proof of Lemmas 8 and 9 can be done in almost the same way, this time without permutation equivariance.
For the proof of the second part, which corresponds to Proposition 4, we construct the network in a similar way.
Recall that we can assume without loss of generality that X ∈ [0 , 1] d × n .
Choose
[Equation]
[Equation]
n
[Equation] [Equation]
n
[Equation]
   .
Then, the ﬁrst column of X + E is in [0 , 1] d , second is in [1 , 2] d , and so on; this means that for all rows, the coordinates are monotonically increasing.
So we can use the same technique as the proof of Proposition 4 to divide the input values into cubes, quantize them to L , apply contextual mapping, and then value mapping.
We describe each step in the following.
C. Q UANTIZATION BY FEED - FORWARD LAYERS
In a similar way as Lemma 5, the goal of this step is to quantize the input in [0 , 1] d × [1 , 2] d × · · · × [ n − 1 , n ] d to its discrete version:
[Equation]
This can be done by dn/δ feed-forward layers.
We add dn/δ layers of the following form, for k = 0 , δ, . . . , n − δ
and i = 1 , . . .
, d :
[Equation]
i [Equation] After dn/δ layers, any input entry of X + E in [ kδ, kδ + δ ) is quantized to kδ .
C. C ONTEXTUAL MAPPING BY ATTENTION LAYERS
By Step 1, we quantized any input X + E to its quantized version.
We call this quantized version
L
[Equation]
As done in Lemma 6, we deﬁne u :=
(1 , δ − 1 , . . .
, δ − d +1 )
and l j :=
u T L : ,j , for all j ∈ [ n ] .
Note that, because L : ,j ∈ [ j − 1 : δ : j − δ ] d , we have
[Equation]
and l 1 < l 2 < · · · < l n .
Notice that this corresponds to the Category 1 in the proof of Lemma 6.
For simplicity of notation, let s j =
( j − 1) (cid:80) d − 1 k =
δ − k .
We stack n (1 /δ )
d attention layers, with attention parts δ − d Ψ( · ; l − δ/ 2 , l + δ/ 2) for each l ∈ (cid:83) nj =
[ s j : δ : s j + δ − d
+1 − δ ] , in increasing order of l .
These n (1 /δ )
d attention layers perform selective shift operations on l j ’s, in increasing order of j .
As seen in Appendix B.5.1, shift operations result in (cid:101) l 1 < (cid:101) l 2 < · · · < (cid:101) l n .
Also, the map from L to (cid:101) l n is one-to-one, which can be shown in the same way as Appendix B.5.4.
Since the range of l j ’s are a bit different, we have a different upper bound on (cid:101) l n : (cid:101) l n := l n +
[Equation]
≤
[Equation]
Finally, we add an extra single-head attention layer with attention part nδ − ( n +1)
d − 1 ψ ( · ; 0) .
We deﬁne the output of this layer as g c ( L ) .
In a similar way as Appendix B.5.1, this layer shifts all the layers by nδ − ( n +1)
d − 1 (cid:101) l n , thus making the intervals corresponding to different values of (cid:101) l n disjoint from each other.
This ensures that different contexts L are mapped to distinct numbers in u T g c ( L ) , thus implementing a contextual mapping.
C. F UNCTION VALUE MAPPING BY FEED - FORWARD LAYERS
Now, it is left to map g c ( L ) to the desired output.
As seen in the last step, each different context L maps to n unique numbers u T g c ( L ) , which are at least δ apart from each other.
The value mapping step can be done in a similar way as Lemma 7.
The construction now requires O ( n (1 /δ ) dn ) layers because there is no permutation equivariance.
[Table]
D E XPERIMENTAL SETUP
For our experiments we follow the same setting as in BERT (Devlin et al., 2018).
We ﬁrst pre-train the models on the masked language modeling task and the next sentence prediction task.
We use
English Wikipedia corpus and BooksCorpus dataset (Zhu et al., 2015) for this pre-training.
We use
BERT BASE , a 12 layer Transformer model as the baseline.
This model uses an embedding size ofand has 12 head self-attention layers and 3072 wide feed forward layers.
We train it with the Adam optimizer, with .
01 dropout and weight decay.
We do pre-training for 250k steps with a batch size of 1024 and a max sequence length of 512.
Pre-training takes around 2 days on 16 TPUv3 chips.
We take the pre-train models and ﬁnetune them on the MNLI and SQuAD datasets separately using the same hyper-parameters as in Devlin et al.
MNLI is a sentence entailment task in which, given a premise sentence, requires us to classify a hypothesis sentence into neutral, contradiction or entailment classes.
We report the classiﬁcation accuracy on this task.
SQuAD is a question answering task, in which given a paragraph and a question, requires us to identify the answer as a span of the words in the paragraph.
For this task we report both the F1 score and the Exact Match (EM) percentage.
The metrics are reported on the dev sets of these datasets.
For our experiments with the depth-wise separable convolution layers, we follow the implementation in (Wu et al., 2019).
We ﬁrst use a GLU layer followed by the convolution layer.
We use 16 separable convolution ﬁlters, of ﬁlter length 128, and reuse them, with each ﬁlter operating on 48 of thedimensions of the input.
This layer also has a skip connection and the output is normalized using layer normalization, similar to the self-attention layer.
In our experiments, we replace the selfattention layers of the Transformers, in the lower layers, with this convolution layer.
We keep the feed forward layer of the Transformer block the same.
For the experiments performed in this paper, one might consider an alternate explanation that the tasks considered maybe are easy, and do not require any advanced architecture to solve them, and even a simple architecture (bi-linear projection or separable convolution) might solve these tasks.
To rule out this case we consider an even simpler architecture, namely average attention, as a baseline for our experiments.
Average attention.
An average attention layer replaces the self-attention layer, and just computes the average of projections of all the other tokens.
That is, we replace σ [( W iK X ) T W iQ X ] in (1) with a matrix full of 1 /n .
The model still has the skip connections and the feed-forward layers like
Transformer.