Universal Transformers

1 INTRODUCTION . Convolutional and fully-attentional feed-forward architectures like the Transformer have recently emerged as viable alternatives to recurrent neural networks ( RNNs ) for a range of sequence modeling tasks , notably machine translation ( Gehring et al. , 2017 ; Vaswani et al. , 2017 ) . These parallel-in-time architectures address a significant shortcoming of RNNs , namely their inherently sequential computation which prevents parallelization across elements of the input sequence , whilst still addressing the vanishing gradients problem as the sequence length gets longer ( Hochreiter et al. , 2003 ) . The Transformer model in particular relies entirely on a self-attention mechanism ( Parikh et al. , 2016 ; Lin et al. , 2017 ) to compute a series of context-informed vector-space representations of the symbols in its input and output , which are then used to predict distributions over subsequent symbols as the model predicts the output sequence symbol-by-symbol . Not only is this mechanism straightforward to parallelize , but as each symbol ’ s representation is also directly informed by all other symbols ’ representations , this results in an effectively global receptive field across the whole sequence . This stands in contrast to e.g . convolutional architectures which typically only have a limited receptive field . Notably , however , the Transformer with its fixed stack of distinct layers foregoes RNNs ’ inductive bias towards learning iterative or recursive transformations . Our experiments indicate that this inductive ∗ Equal contribution , alphabetically by last name . †Work performed while at Google Brain . bias may be crucial for several algorithmic and language understanding tasks of varying complexity : in contrast to models such as the Neural Turing Machine ( Graves et al. , 2014 ) , the Neural GPU ( Kaiser & Sutskever , 2016 ) or Stack RNNs ( Joulin & Mikolov , 2015 ) , the Transformer does not generalize well to input lengths not encountered during training . In this paper , we introduce the Universal Transformer ( UT ) , a parallel-in-time recurrent self-attentive sequence model which can be cast as a generalization of the Transformer model , yielding increased theoretical capabilities and improved results on a wide range of challenging sequence-to-sequence tasks . UTs combine the parallelizability and global receptive field of feed-forward sequence models like the Transformer with the recurrent inductive bias of RNNs , which seems to be better suited to a range of algorithmic and natural language understanding sequence-to-sequence problems . As the name implies , and in contrast to the standard Transformer , under certain assumptions UTs can be shown to be Turing-complete ( or “ computationally universal ” , as shown in Section 4 ) . In each recurrent step , the Universal Transformer iteratively refines its representations for all symbols in the sequence in parallel using a self-attention mechanism ( Parikh et al. , 2016 ; Lin et al. , 2017 ) , followed by a transformation ( shared across all positions and time-steps ) consisting of a depth-wise separable convolution ( Chollet , 2016 ; Kaiser et al. , 2017 ) or a position-wise fully-connected layer ( see Fig 1 ) . We also add a dynamic per-position halting mechanism ( Graves , 2016 ) , allowing the model to choose the required number of refinement steps for each symbol dynamically , and show for the first time that such a conditional computation mechanism can in fact improve accuracy on several smaller , structured algorithmic and linguistic inference tasks ( although it marginally degraded results on MT ) . Our strong experimental results show that UTs outperform Transformers and LSTMs across a wide range of tasks . The added recurrence yields improved results in machine translation where UTs outperform the standard Transformer . In experiments on several algorithmic tasks and the bAbI language understanding task , UTs also consistently and significantly improve over LSTMs and the standard Transformer . Furthermore , on the challenging LAMBADA text understanding data set UTs with dynamic halting achieve a new state of the art . 2 MODEL DESCRIPTION . 2.1 THE UNIVERSAL TRANSFORMER . The Universal Transformer ( UT ; see Fig . 2 ) is based on the popular encoder-decoder architecture commonly used in most neural sequence-to-sequence models ( Sutskever et al. , 2014 ; Cho et al. , 2014 ; Vaswani et al. , 2017 ) . Both the encoder and decoder of the UT operate by applying a recurrent neural network to the representations of each of the positions of the input and output sequence , respectively . However , in contrast to most applications of recurrent neural networks to sequential data , the UT does not recur over positions in the sequence , but over consecutive revisions of the vector representations of each position ( i.e. , over “ depth ” ) . In other words , the UT is not computationally bound by the number of symbols in the sequence , but only by the number of revisions made to each symbol ’ s representation . In each recurrent time-step , the representation of every position is concurrently ( in parallel ) revised in two sub-steps : first , using a self-attention mechanism to exchange information across all positions in the sequence , thereby generating a vector representation for each position that is informed by the representations of all other positions at the previous time-step . Then , by applying a transition function ( shared across position and time ) to the outputs of the self-attention mechanism , independently at each position . As the recurrent transition function can be applied any number of times , this implies that UTs can have variable depth ( number of per-symbol processing steps ) . Crucially , this is in contrast to most popular neural sequence models , including the Transformer ( Vaswani et al. , 2017 ) or deep RNNs , which have constant depth as a result of applying a fixed stack of layers . We now describe the encoder and decoder in more detail . ENCODER : Given an input sequence of length m , we start with a matrix whose rows are initialized as the d-dimensional embeddings of the symbols at each position of the sequence H0∈Rm×d . The UT then iteratively computes representations Ht at step t for all m positions in parallel by applying the multi-headed dot-product self-attention mechanism from Vaswani et al . ( 2017 ) , followed by a recurrent transition function . We also add residual connections around each of these function blocks and apply dropout and layer normalization ( Srivastava et al. , 2014 ; Ba et al. , 2016 ) ( see Fig . 2 for a simplified diagram , and Fig . 4 in the Appendix A for the complete model. ) . More specifically , we use the scaled dot-product attention which combines queries Q , keys K and values V as follows ATTENTION ( Q , K , V ) = SOFTMAX ( QKT√ d ) V , ( 1 ) where d is the number of columns of Q , K and V . We use the multi-head version with k heads , as introduced in ( Vaswani et al. , 2017 ) , MULTIHEADSELFATTENTION ( Ht ) = CONCAT ( head1 , ... , headk ) WO ( 2 ) where headi = ATTENTION ( HtW Q i , H tWKi , H tWVi ) ( 3 ) and we map the state Ht to queries , keys and values with affine projections using learned parameter matrices WQ∈Rd×d/k , WK ∈Rd×d/k , WV ∈Rd×d/k and WO∈Rd×d . At step t , the UT then computes revised representationsHt∈Rm×d for allm input positions as follows Ht = LAYERNORM ( At+TRANSITION ( At ) ) ( 4 ) where At = LAYERNORM ( ( Ht−1+P t ) +MULTIHEADSELFATTENTION ( Ht−1+P t ) ) , ( 5 ) where LAYERNORM ( ) is defined in Ba et al . ( 2016 ) , and TRANSITION ( ) and P t are discussed below . Depending on the task , we use one of two different transition functions : either a separable convolution ( Chollet , 2016 ) or a fully-connected neural network that consists of a single rectified-linear activation function between two affine transformations , applied position-wise , i.e . individually to each row of At . P t ∈ Rm×d above are fixed , constant , two-dimensional ( position , time ) coordinate embeddings , obtained by computing the sinusoidal position embedding vectors as defined in ( Vaswani et al. , 2017 ) for the positions 1≤ i≤m and the time-step 1≤ t≤T separately for each vector-dimension 1≤j≤d , and summing : P ti,2j =sin ( i/10000 2j/d ) +sin ( t/100002j/d ) ( 6 ) P ti,2j+1 =cos ( i/10000 2j/d ) +cos ( t/100002j/d ) . ( 7 ) After T steps ( each updating all positions of the input sequence in parallel ) , the final output of the Universal Transformer encoder is a matrix of d-dimensional vector representations HT ∈Rm×d for the m symbols of the input sequence . DECODER : The decoder shares the same basic recurrent structure of the encoder . However , after the self-attention function , the decoder additionally also attends to the final encoder representation HT of each position in the input sequence using the same multihead dot-product attention function from Equation 2 , but with queries Q obtained from projecting the decoder representations , and keys and values ( K and V ) obtained from projecting the encoder representations ( this process is akin to standard attention ( Bahdanau et al. , 2014 ) ) . Like the Transformer model , the UT is autoregressive ( Graves , 2013 ) . Trained using teacher-forcing , at generation time it produces its output one symbol at a time , with the decoder consuming the previously produced output positions . During training , the decoder input is the target output , shifted to the right by one position . The decoder self-attention distributions are further masked so that the model can only attend to positions to the left of any predicted symbol . Finally , the per-symbol target distributions are obtained by applying an affine transformation O ∈ Rd×V from the final decoder state to the output vocabulary size V , followed by a softmax which yields an ( m×V ) -dimensional output matrix normalized over its rows : p ( ypos|y [ 1 : pos−1 ] , HT ) = SOFTMAX ( OHT ) 1 ( 8 ) To generate from the model , the encoder is run once for the conditioning input sequence . Then the decoder is run repeatedly , consuming all already-generated symbols , while generating one additional distribution over the vocabulary for the symbol at the next output position per iteration . We then typically sample or select the highest probability symbol as the next symbol . 2.2 DYNAMIC HALTING . In sequence processing systems , certain symbols ( e.g . some words or phonemes ) are usually more ambiguous than others . It is therefore reasonable to allocate more processing resources to these more ambiguous symbols . Adaptive Computation Time ( ACT ) ( Graves , 2016 ) is a mechanism for dynamically modulating the number of computational steps needed to process each input symbol 1Note that T here denotes time-step T and not the transpose operation . Previous best results : . Our Results : . ( called the “ ponder time ” ) in standard recurrent neural networks based on a scalar halting probability predicted by the model at each step . Inspired by the interpretation of Universal Transformers as applying self-attentive RNNs in parallel to all positions in the sequence , we also add a dynamic ACT halting mechanism to each position ( i.e . to each per-symbol self-attentive RNN ; see Appendix C for more details ) . Once the per-symbol recurrent block halts , its state is simply copied to the next step until all blocks halt , or we reach a maximum number of steps . The final output of the encoder is then the final layer of representations produced in this way .

Summary: This paper extends Transformer by recursively applying a multi-head self-attention block, rather than stack multiple blocks in the vanilla Transformer. An extra transition function is applied between the recursive blocks. This combines the idea from RNN and attention-based models. But the RNN structure here is not applied to the input sequence, but to the sequence of blocks inside the Transformer encoder/decoder. In addition, it also uses a dynamic adaptive computation time (ACT) halting mechanism on each position, as suggested by the previous ACT paper. In fact, it can be seen as a memory network with a dynamic number of hops at the symbol level.