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: This paper tries to analyse the Transformer, widely applied building block of a neural network component, to improve understanding of the internals of the model. The analysis starts showing that the transformer blocks generate permutation equivalent maps and then shows that the transformer can approximate any permutation equivalent map in a compact domain with arbitrary precision. Three key steps are developed and used to prove the universal approximation of arbitrary permutation equivalent map: 1) quantization of input via feed-forward layers, 2) contextual mapping via attention layers, and 3) value mapping via feed-forward layers. By introducing positional embeddings, the paper relaxes the restriction on permutation equivalence and proves that the Transformer is a universal approximator of any sequence to sequence function.