Are Transformers universal approximators of sequence-to-sequence functions?

1 INTRODUCTION . 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 modified 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 fixed word embeddings as in word2vec ( Mikolov et al. , 2013 ) . Is it possible to formalize the notion of ∗Based on work performed at Google Research New York contextual mappings ? If yes , can Transformers actually compute such mappings ? Such questions still remain elusive . In this paper , we provide a mathematical definition of contextual mappings and 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 efficient architecture to compute contextual mappings , consequently , preserving the ability to universally approximate sequence-to-sequence functions ? Towards this , we explore other 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.1 SUMMARY 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.2 RELATED 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 findings 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 one-hidden-layer 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 . ( 2017 ) ; 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 , Pérez et al . ( 2019 ) have shown that Transformers with infinite precision are Turing complete , which is not the case in finite 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 Ai , j , Ai , : , and A : ,j denote its ( i , j ) -th entry , i-th row , and j-th column , respectively . We use ‖A‖p to denote the entrywise ` 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 each column summing to 1 . We similarly define σ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 1n 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 . 2 TRANSFORMER NETWORKS . A Transformer block is a sequence-to-sequence function mapping Rd×n to Rd×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 ∈ Rd×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 layers1 : Attn ( X ) = X + ∑h i=1 W iOW i VX · σ [ ( W iKX ) TW iQX ] , ( 1 ) FF ( X ) = Attn ( X ) + W2 · ReLU ( W1 ·Attn ( X ) + b11Tn ) + b21Tn , ( 2 ) where W iO ∈ Rd×m , W iV , W iK , W iQ ∈ Rm×d , W2 ∈ Rd×r , W1 ∈ Rr×d , b2 ∈ Rd , b1 ∈ Rr , 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 definition of the self-attention layer ( 1 ) is an equivalent reformulation of ( Vaswani et al. , 2017 ) , where they concatenate attention heads and multiply a matrix WO ∈ Rd×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 define the Transformer networks as the composition of Transformer blocks . The family of the sequence-to-sequence functions corresponding to the Transformers can be defined as : T h , m , r : = { g : Rd×n → Rd×n | g is a composition of Transformer blocks th , m , r ’ s } . where th , m , r : Rd×n → Rd×n denotes a Transformer block defined 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 : Rd×n → Rd×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 th , m , r defines a permutation equivariant map from Rd×n to Rd×n . As seen in above , both layers ( cf . ( 1 ) and ( 2 ) ) of a Transformer block employ parameter reuse/sharing , because each token/column undergoes the same transformations ( e.g. , W iQ , W i K , or W1 ) regardless of its position . Moreover , interactions between tokens can only be captured through pairwise dot-products in the softmax operator σ [ · ] ( cf . ( 1 ) ) . Given such limitations in a single Transformer block ’ s representation power , it is not obvious what kinds of sequence-to-sequence functions T h , m , r can approximate ; we provide the answer to this question in the next section . 3 TRANSFORMERS ARE UNIVERSAL APPROXIMATORS OF SEQUENCE-TO-SEQUENCE FUNCTIONS . In this section , we present our theorems showing that the Transformer networks are universal approximators of sequence-to-sequence functions . Let us start by defining the target function class FPE , which consists of all continuous permutation equivariant functions with compact support that 1In our proof we use bias vectors biQ for query projections in attention layers . We omit them here for brevity . map Rd×n to Rd×n . Here , continuity is defined with respect to any entry-wise ` p norm , 1 ≤ p < ∞ . Given two functions f1 , f2 : Rd×n → Rd×n , for 1 ≤ p < ∞ , we define a distance between them as dp ( f1 , f2 ) : = ( ∫ ‖f1 ( X ) − f2 ( X ) ‖pp dX ) 1/p . 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 FPE . Theorem 2 . Let 1 ≤ p < ∞ and > 0 , then for any given f ∈ FPE , there exists a Transformer network g ∈ T 2,1,4 , such that dp ( f , g ) ≤ . 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 ∈ Rd×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 : T h , m , rP : = { gP ( X ) = g ( X + E ) | g ∈ T h , m , r and E ∈ Rd×n } . Here we show that if E is trainable , these positional encodings are sufficient to remove the permutation equivariance restriction of the Transformers . Towards this , we define FCD to be the set of all continuous functions that map a compact domain in Rd×n to Rd×n . Note that FCD does not have the restriction of permutation equivariance as in FPE , but any f ∈ FCD is defined on a compact domain instead of the whole Rd×n . The following result states that , equipped with the trainable positional encodings , Transformers can approximate any sequence-to-sequence function in FCD . Theorem 3 . Let 1 ≤ p < ∞ and > 0 , then for any given f ∈ FCD , there exists a Transformer network g ∈ T 2,1,4P such that we have dp ( f , g ) ≤ . Theorems 2 and 3 provide an interesting characterization of the representation power of fixed-width Transformer networks . Since the function classes T h , m , r and T h , m , rP 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 again point out that Theorems 2 and 3 appear quite surprising at a first 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 briefly describe one of our key steps in overcoming the aforementioned restrictions and proving universal approximation power of Transformers .

Summary: C1. Transformers (without positional encodings and without layer normalization), with 2 attention heads of dimension 1 and feed-forward layers (FFN) with 4 hidden nodes, are universal approximators of continuous permutation-equivariant functions f of compact support, relative to any Lp metric (1 <= p < \infty). (Thm. 2, p. 3). (Without positional encodings, a function f computed by a Transformer is permutation equivariant: f(P(X)) = P(f(X)) for any permutation P of the columns of X, which are the vector encodings of the input tokens.)