Universal Transformers — part 1/3

Cleaned-text code points [0, 15711). Bibliography omitted here; source.md includes it.

# Universal Transformers

###### Abstract

Recurrent neural networks (RNNs) sequentially process data by updating their state with each new data point, and have long been the de facto choice for sequence modeling tasks. However, their inherently sequential computation makes them slow to train. Feed-forward and convolutional architectures have recently been shown to achieve superior results on some sequence modeling tasks such as machine translation, with the added advantage that they concurrently process all inputs in the sequence, leading to easy parallelization and faster training times. Despite these successes, however, popular feed-forward sequence models like the Transformer fail to generalize in many simple tasks that recurrent models handle with ease, e.g. copying strings or even simple logical inference when the string or formula lengths exceed those observed at training time. We propose the Universal Transformer (UT), a parallel-in-time self-attentive recurrent sequence model which can be cast as a generalization of the Transformer model and which addresses these issues. UTs combine the parallelizability and global receptive field of feed-forward sequence models like the Transformer with the recurrent inductive bias of RNNs. We also add a dynamic per-position halting mechanism and find that it improves accuracy on several tasks. In contrast to the standard Transformer, under certain assumptions UTs can be shown to be Turing-complete. Our experiments show that UTs outperform standard Transformers on a wide range of algorithmic and language understanding tasks, including the challenging LAMBADA language modeling task where UTs achieve a new state of the art, and machine translation where UTs achieve a 0.9 BLEU improvement over Transformers on the WMT14 En-De dataset.

## 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 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 m , we start with a matrix whose rows are initialized as the d d -dimensional embeddings of the symbols at each position of the sequence H 0 ∈ ℝ m × d H^{0}\in\mathbb{R}^{m\times d} . The UT then iteratively computes representations H t H^{t} at step t t for all m 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 Q , keys K K and values V V as follows

Attention ​ ( Q , K , V ) = softmax ​ ( Q ​ K T d ) ​ V , \textsc{Attention}(Q,K,V)=\textsc{softmax}\left(\frac{QK^{T}}{\sqrt{d}}\right)V, (1)

where d d is the number of columns of Q Q , K K and V V . We use the multi-head version with k k heads, as introduced in ( Vaswani et al., 2017 ) ,

MultiHeadSelfAttention ​ ( H t ) \displaystyle\textsc{MultiHeadSelfAttention}(H^{t}) = Concat ​ ( head 1 , … , head k ) ​ W O \displaystyle=\textsc{Concat}(\mathrm{head_{1}},...,\mathrm{head_{k}})W^{O} (2) where ​ head i \displaystyle\text{where}~\mathrm{head_{i}} = Attention ​ ( H t ​ W i Q , H t ​ W i K , H t ​ W i V ) \displaystyle=\textsc{Attention}(H^{t}W^{Q}_{i},H^{t}W^{K}_{i},H^{t}W^{V}_{i}) (3)

and we map the state H t H^{t} to queries, keys and values with affine projections using learned parameter matrices W Q ∈ ℝ d × d / k W^{Q}\in\mathbb{R}^{d\times d/k} , W K ∈ ℝ d × d / k W^{K}\in\mathbb{R}^{d\times d/k} , W V ∈ ℝ d × d / k W^{V}\in\mathbb{R}^{d\times d/k} and W O ∈ ℝ d × d W^{O}\in\mathbb{R}^{d\times d} .

At step t t , the UT then computes revised representations H t ∈ ℝ m × d H^{t}\in\mathbb{R}^{m\times d} for all m m input positions as follows

H t \displaystyle H^{t} = LayerNorm ​ ( A t + Transition ​ ( A t ) ) \displaystyle=\textsc{LayerNorm}(A^{t}+\textsc{Transition}(A^{t})) (4) where ​ A t \displaystyle\mathrm{where}~A^{t} = LayerNorm ​ ( ( H t − 1 + P t ) + MultiHeadSelfAttention ​ ( H t − 1 + P t ) ) , \displaystyle=\textsc{LayerNorm}((H^{t-1}+P^{t})+\textsc{MultiHeadSelfAttention}(H^{t-1}+P^{t})), (5) where LayerNorm() is defined in Ba et al. (2016) , and Transition() and P t 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 A t A^{t} .

P t ∈ ℝ m × d P^{t}\in\mathbb{R}^{m\times 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 1\leq i\leq m and the time-step 1 ≤ t ≤ T 1\leq t\leq T separately for each vector-dimension 1 ≤ j ≤ d 1\leq j\leq d , and summing: P i , 2 ​ j t \displaystyle P^{t}_{i,2j} = sin ⁡ ( i / 10000 2 ​ j / d ) + sin ⁡ ( t / 10000 2 ​ j / d ) \displaystyle=\sin(i/10000^{2j/d})+\sin(t/10000^{2j/d}) (6) P i , 2 ​ j + 1 t \displaystyle P^{t}_{i,2j+1} = cos ⁡ ( i / 10000 2 ​ j / d ) + cos ⁡ ( t / 10000 2 ​ j / d ) . \displaystyle=\cos(i/10000^{2j/d})+\cos(t/10000^{2j/d}). (7)

After T 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 d -dimensional vector representations H T ∈ ℝ m × d H^{T}\in\mathbb{R}^{m\times d} for the m 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 H T H^{T} of each position in the input sequence using the same multihead dot-product attention function from Equation 2 , but with queries Q Q obtained from projecting the decoder representations, and keys and values ( K K and V 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 ∈ ℝ d × V O\in\mathbb{R}^{d\times V} from the final decoder state to the output vocabulary size V V , followed by a softmax which yields an ( m × V ) (m\times V) -dimensional output matrix normalized over its rows:

p ( y p ​ o ​ s | y [ 1 : p o s − 1 ] , H T ) = softmax ( O H T ) p\left(y_{pos}|y_{[1:pos-1]},H^{T}\right)=\textsc{softmax}(OH^{T}) (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 (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.

## 3 Experiments and Analysis

We evaluated the Universal Transformer on a range of algorithmic and language understanding tasks, as well as on machine translation. We describe these tasks and datasets in more detail in Appendix D .

### 3.1 bAbI Question-Answering

The bAbi question answering dataset ( Weston et al., 2015 ) consists of 20 different tasks, where the goal is to answer a question given a number of English sentences that encode potentially multiple supporting facts. The goal is to measure various forms of language understanding by requiring a certain type of reasoning over the linguistic facts presented in each story. A standard Transformer does not achieve good results on this task 2 2 2 We experimented with different hyper-parameters and different network sizes, but it always overfits. . However, we have designed a model based on the Universal Transformer which achieves state-of-the-art results on this task.

To encode the input, similar to Henaff et al. (2016) , we first encode each fact in the story by applying a learned multiplicative positional mask to each word’s embedding, and summing up all embeddings. We embed the question in the same way, and then feed the (Universal) Transformer with these embeddings of the facts and questions.