Title: 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. * Equal contribution, alphabetically by last name. † Work performed while at Google Brain. An example from tasks 3: (requiring three supportive facts to solve) Story: Mary got the milk.
John moved to the bedroom. Daniel journeyed to the office. John grabbed the apple

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 The Universal Transformer repeatedly refines a series of vector representations for each position of the sequence in parallel, by combining information from different positions using self-attention (see Eqn 2) and applying a recurrent transition function (see Eqn 4) across all time steps 1 ≤ t ≤ T . We show this process over two recurrent time-steps. Arrows denote dependencies between operations. Initially, h 0 is initialized with the embedding for each symbol in the sequence. h t i represents the representation for input symbol 1 ≤ i ≤ m at recurrent time-step t. With dynamic halting, T is dynamically determined for each position (Section 2.2).

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 H 0 ∈ R m×d . The UT then iteratively computes representations H t 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

EQUATION

)

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) ,

EQUATION

where

EQUATION

and we map the state H t to queries, keys and values with affine projections using learned parameter matrices

W Q ∈ R d×d/k , W K ∈ R d×d/k , W V ∈ R d×d/k and W O ∈ R d×d .

At step t, the UT then computes revised representations H t ∈ R m×d for all m input positions as follows

EQUATION

where

EQUATION

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 A t .

P t ∈ R m×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: A complete version can be found in Appendix A. The Universal Transformer with dynamic halting determines the number of steps T for each position individually using ACT (Graves, 2016) .

P t i,2j = sin(i/10000

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 H T ∈ R m×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 H T 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 ∈ R d×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 y pos |y [1:pos-1] ,H T = SOFTMAX(OH T ) 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 (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 task2 . 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.

As originally proposed, models can either be trained on each task separately ("train single") or jointly on all tasks ("train joint"). Table 1 summarizes our results. We conducted 10 runs with different initializations and picked the best model based on performance on the validation set, similar to previous work. Both the UT and UT with dynamic halting achieve state-of-the-art results on all tasks in terms of average error and number of failed tasks3 , in both the 10K and 1K training regime (see Appendix E for breakdown by task).

To understand the working of the model better, we analyzed both the attention distributions and the average ACT ponder times for this task (see Appendix F for details). First, we observe that the attention distributions start out very uniform, but get progressively sharper in later steps around the correct supporting facts that are required to answer each question, which is indeed very similar to how humans would solve the task. Second, with dynamic halting we observe that the average ponder time (i.e. depth of the per-symbol recurrent processing chain) over all positions in all samples in the test data for tasks requiring three supporting facts is higher (3.8±2.2) than for tasks requiring only two (3.1±1.1), which is in turn higher than for tasks requiring only one supporting fact (2.3±0.8). This indicates that the model adjusts the number of processing steps with the number of supporting facts required to answer the questions. Finally, we observe that the histogram of ponder times at different positions is more uniform in tasks requiring only one supporting fact compared to two and three, and likewise for tasks requiring two compared to three. Especially for tasks requiring three supporting facts, many positions halt at step 1 or 2 already and only a few get transformed for more steps (see for example Fig 3 ). This is particularly interesting as the length of stories is indeed much higher in this setting, with more irrelevant facts which the model seems to successfully learn to ignore in this way.

Similar to dynamic memory networks (Kumar et al., 2016) , there is an iterative attention process in UTs that allows the model to condition its attention over memory on the result of previous iterations. Appendix F presents some examples illustrating that there is a notion of temporal states in UT, where the model updates its states (memory) in each step based on the output of previous steps, and this chain of updates can also be viewed as steps in a multi-hop reasoning process.

3.2. SUBJECT-VERB AGREEMENT:

Next, we consider the task of predicting number-agreement between subjects and verbs in English sentences (Linzen et al., 2016) . This task acts as a proxy for measuring the ability of a model to capture hierarchical (dependency) structure in natural language sentences. We use the dataset provided by (Linzen et al., 2016) and follow their experimental protocol of solving the task using a language modeling training setup, i.e. a next word prediction objective, followed by calculating the ranking accuracy of the target verb at test time. We evaluated our model on subsets of the test data with different task difficulty, measured in terms of agreement attractors -the number of intervening nouns with the opposite number from the subject (meant to confuse the model). For example, given the sentence The keys to the cabinet4 , the objective during training is to predict the verb are (plural). At test time, we then evaluate the ranking accuracy of the agreement attractors: i.e. the goal is to rank are higher than is in this case.

Our results are summarized in Table 2 . The best LSTM with attention from the literature achieves 99.18% on this task (Yogatama et al., 2018) , outperforming a vanilla Transformer (Tran et al., 2018) . UTs significantly outperform standard Transformers, and achieve an average result comparable to the current state of the art (99.2%). However, we see that UTs (and particularly with dynamic halting) perform progressively better than all other models as the number of attractors increases (see the last row, ∆).

3.3. LAMBADA LANGUAGE MODELING:

The LAMBADA task (Paperno et al., 2016 ) is a language modeling task consisting of predicting a missing target word given a broader context of 4-5 preceding sentences. The dataset was specifically designed so that humans are able to accurately predict the target word when shown the full context, but not when only shown the target sentence in which it appears. It therefore goes beyond language Table 3 : LAMBADA language modeling (LM) perplexity (lower better) with accuracy in parentheses (higher better), and Reading Comprehension (RC) accuracy results (higher better). '-' indicates no reported results in that setting.

modeling, and tests the ability of a model to incorporate broader discourse and longer term context when predicting the target word.

The task is evaluated in two settings: as language modeling (the standard setup) and as reading comprehension. In the former (more challenging) case, a model is simply trained for next-word prediction on the training data, and evaluated on the target words at test time (i.e. the model is trained to predict all words, not specifically challenging target words). In the latter setting, introduced by Chu et al. Chu et al. (2017) , the target sentence (minus the last word) is used as query for selecting the target word from the context sentences. Note that the target word appears in the context 81% of the time, making this setup much simpler. However the task is impossible in the remaining 19% of the cases.

The results are shown in Table 3 . Universal Transformer achieves state-of-the-art results in both the language modeling and reading comprehension setup, outperforming both LSTMs and vanilla Transformers. Note that the control set was constructed similar to the LAMBADA development and test sets, but without filtering them in any way, so achieving good results on this set shows a model's strength in standard language modeling.

Our best fixed UT results used 6 steps. However, the average number of steps that the best UT with dynamic halting took on the test data over all positions and examples was 8.2±2.1. In order to see if the dynamic model did better simply because it took more steps, we trained two fixed UT models with 8 and 9 steps respectively (see last two rows). Interestingly, these two models achieve better results compared to the model with 6 steps, but do not outperform the UT with dynamic halting. This leads us to believe that dynamic halting may act as a useful regularizer for the model via incentivizing a smaller numbers of steps for some of the input symbols, while allowing more computation for others.

3.4. ALGORITHMIC TASKS:

We trained UTs on three algorithmic tasks, namely Copy, Reverse, and (integer) Addition, all on strings composed of decimal symbols ('0'-'9'). In all the experiments, we train the models on sequences of length 40 and evaluated on sequences of length 400 (Kaiser & Sutskever, 2016) . We train UTs using positions starting with randomized offsets to further encourage the model to learn position-relative transformations. Results are shown in Table 4 . The UT outperforms both LSTM and vanilla Transformer by a wide margin on all three tasks. The Neural GPU reports perfect results on this task (Kaiser & Sutskever, 2016) , however we note that this result required a special curriculum-based training protocol which was not used for other models.

3.5. LEARNING TO EXECUTE (LTE):

As another class of sequence-to-sequence learning problems, we also evaluate UTs on tasks indicating the ability of a model to learn to execute computer programs, as proposed in (Zaremba & Sutskever, 2015) . These tasks include program evaluation tasks (program, control, and addition), and memorization tasks (copy, double, and reverse).

We use the mix-strategy discussed in (Zaremba & Sutskever, 2015) to generate the datasets. Unlike (Zaremba & Sutskever, 2015) , we do not use any curriculum learning strategy during training and we make no use of target sequences at test time. Tables 5 and 6 present the performance of an LSTM model, Transformer, and Universal Transformer on the program evaluation and memorization tasks, respectively. UT achieves perfect scores in all the memorization tasks and also outperforms both LSTMs and Transformers in all program evaluation tasks by a wide margin.

3.6. MACHINE TRANSLATION:

We trained a UT on the WMT 2014 English-German translation task using the same setup as reported in (Vaswani et al., 2017) in order to evaluate its performance on a large-scale sequence-to-sequence task.

Results are summarized in Table 7 . The UT with a fully-connected recurrent transition function (instead of separable convolution) and without ACT improves by 0.9 BLEU over a Transformer and 0.5 BLEU over a Weighted Transformer with approximately the same number of parameters (Ahmed et al., 2017) .

Model BLEU:

Universal Transformer small 26.8 Transformer base (Vaswani et al., 2017) 28.0 Weighted Transformer base (Ahmed et al., 2017) 28.4 Universal Transformer base 28.9

Table 7 : Machine translation results on the WMT14 En-De translation task trained on 8xP100 GPUs in comparable training setups. All base results have the same number of parameters.

4. DISCUSSION:

When running for a fixed number of steps, the Universal Transformer is equivalent to a multi-layer Transformer with tied parameters across all its layers. This is partly similar to the Recursive Transformer, which ties the weights of its self-attention layers across depth (Gulcehre et al., 2018) 5 . However, as the per-symbol recurrent transition functions can be applied any number of times, another and possibly more informative way of characterizing the UT is as a block of parallel RNNs (one for each symbol, with shared parameters) evolving per-symbol hidden states concurrently, generated at each step by attending to the sequence of hidden states at the previous step. In this way, it is related to architectures such as the Neural GPU (Kaiser & Sutskever, 2016) and the Neural Turing Machine (Graves et al., 2014) . UTs thereby retain the attractive computational efficiency of the original feedforward Transformer model, but with the added recurrent inductive bias of RNNs. Furthermore, using a dynamic halting mechanism, UTs can choose the number of processing steps based on the input data.

The connection between the Universal Transformer and other sequence models is apparent from the architecture: if we limited the recurrent steps to one, it would be a Transformer. But it is more interesting to consider the relationship between the Universal Transformer and RNNs and other networks where recurrence happens over the time dimension. Superficially these models may seem closely related since they are recurrent as well. But there is a crucial difference: time-recurrent models like RNNs cannot access memory in the recurrent steps. This makes them computationally more similar to automata, since the only memory available in the recurrent part is a fixed-size state vector. UTs on the other hand can attend to the whole previous layer, allowing it to access memory in the recurrent step.

Given sufficient memory the Universal Transformer is computationally universal -i.e. it belongs to the class of models that can be used to simulate any Turing machine, thereby addressing a shortcoming of the standard Transformer model6 . In addition to being theoretically appealing, our results show that this added expressivity also leads to improved accuracy on several challenging sequence modeling tasks. This closes the gap between practical sequence models competitive on large-scale tasks such as machine translation, and computationally universal models such as the Neural Turing Machine or the Neural GPU (Graves et al., 2014; Kaiser & Sutskever, 2016) , which can be trained using gradient descent to perform algorithmic tasks.

To show this, we can reduce a Neural GPU to a Universal Transformer. Ignoring the decoder and parameterizing the self-attention module, i.e. self-attention with the residual connection, to be the identity function, we assume the transition function to be a convolution. If we now set the total number of recurrent steps T to be equal to the input length, we obtain exactly a Neural GPU. Note that the last step is where the Universal Transformer crucially differs from the vanilla Transformer whose depth cannot scale dynamically with the size of the input. A similar relationship exists between the Universal Transformer and the Neural Turing Machine, whose single read/write operations per step can be expressed by the global, parallel representation revisions of the Universal Transformer. In contrast to these models, however, which only perform well on algorithmic tasks, the Universal Transformer also achieves competitive results on realistic natural language tasks such as LAMBADA and machine translation.

Another related model architecture is that of end-to-end Memory Networks (Sukhbaatar et al., 2015) . In contrast to end-to-end memory networks, however, the Universal Transformer uses memory corresponding to states aligned to individual positions of its inputs or outputs. Furthermore, the Universal Transformer follows the encoder-decoder configuration and achieves competitive performance in large-scale sequence-to-sequence tasks.

This paper introduces the Universal Transformer, a generalization of the Transformer model that extends its theoretical capabilities and produces state-of-the-art results on a wide range of challenging sequence modeling tasks, such as language understanding but also a variety of algorithmic tasks, thereby addressing a key shortcoming of the standard Transformer. The Universal Transformer combines the following key properties into one model:

Weight sharing: Following intuitions behind weight sharing found in CNNs and RNNs, we extend the Transformer with a simple form of weight sharing that strikes an effective balance between inductive bias and model expressivity, which we show extensively on both small and large-scale experiments.

Conditional computation: In our goal to build a computationally universal machine, we equipped the Universal Transformer with the ability to halt or continue computation through a recently introduced mechanism, which shows stronger results compared to the fixed-depth Universal Transformer.

We are enthusiastic about the recent developments on parallel-in-time sequence models. By adding computational capacity and recurrence in processing depth, we hope that further improvements beyond the basic Universal Transformer presented here will help us build learning algorithms that are both more powerful, data efficient, and generalize beyond the current state-of-the-art.

The code used to train and evaluate Universal Transformers is available at https: //github.com/tensorflow/tensor2tensor (Vaswani et al., 2018) . 

APPENDIX B ON THE COMPUTATIONAL POWER OF UT VS TRANSFORMER:

With respect to their computational power, the key difference between the Transformer and the Universal Transformer lies in the number of sequential steps of computation (i.e. in depth). While a standard Transformer executes a total number of operations that scales with the input size, the number of sequential operations is constant, independent of the input size and determined solely by the number of layers. Assuming finite precision, this property implies that the standard Transformer cannot be computationally universal. When choosing a number of steps as a function of the input length, however, the Universal Transformer does not suffer from this limitation. Note that this holds independently of whether or not adaptive computation time is employed but does assume a non-constant, even if possibly deterministic, number of steps. Varying the number of steps dynamically after training is enabled by sharing weights across sequential computation steps in the Universal Transformer.

An intuitive example are functions whose execution requires the sequential processing of each input element. In this case, for any given choice of depth T , one can construct an input sequence of length N > T that cannot be processed correctly by a standard Transformer. With an appropriate, input-length dependent choice of sequential steps, however, a Universal Transformer, RNNs or Neural GPUs can execute such a function.

APPENDIX C UT WITH DYNAMIC HALTING:

We implement the dynamic halting based on ACT (Graves, 2016) as follows in TensorFlow. In each step of the UT with dynamic halting, we are given the halting probabilities, remainders, number of updates up to that point, and the previous state (all initialized as zeros), as well as a scalar threshold between 0 and 1 (a hyper-parameter). We then compute the new state for each position and calculate the new per-position halting probabilities based on the state for each position. The UT then decides to halt for some positions that crossed the threshold, and updates the state of other positions until the model halts for all positions or reaches a predefined maximum number of steps: Listing 1: UT with dynamic halting.

The following shows the computations in each step:

1 d e f u t _ w i t h _ d y n a m i c _ h a l t i n g ( s t a t e , s t e p , h a l t i n g _ p r o b a b i l i t y , 2 r e m a i n d e r s , n _ u p d a t e s , p r e v i o u s _ s t a t e ) : s t i l l _ r u n n i n g = t f . c a s t ( 8 t f . l e s s ( h a l t i n g _ p r o b a b i l i t y , 1 . 0 ) , t f . f l o a t 3 2 ) 9 # Mask o f i n p u t s which h a l t e d a t t h i s s t e p 10 n e w _ h a l t e d = t f . c a s t ( 11 t f . g r e a t e r ( h a l t i n g _ p r o b a b i l i t y + p * s t i l l _ r u n n i n g , t h r e s h o l d ) , 12 t f . f l o a t 3 2 ) * s t i l l _ r u n n i n g 13 # Mask o f i n p u t s which h a v e n ' t h a l t e d , and d i d n ' t h a l t t h i s s t e p 14 s t i l l _ r u n n i n g = t f . c a s t ( 15 t f . l e s s _ e q u a l ( h a l t i n g _ p r o b a b i l i t y + p * s t i l l _ r u n n i n g , 16 t h r e s h o l d ) , t f . f l o a t 3 2 ) * s t i l l _ r u n n i n g 17 # Add t h e h a l t i n g p r o b a b i l i t y f o r t h i s s t e p t o t h e h a l t i n g Listing 2: Computations in each step of the UT with dynamic halting. 

:

Note that T here denotes time-step T and not the transpose operation.

We experimented with different hyper-parameters and different network sizes, but it always overfits.

Defined as > 5% error.

Cabinet (singular) is an agreement attractor in this case.

Note that in UT both the self-attention and transition weights are tied across layers.

Appendix B illustrates how UT is computationally more powerful than the standard Transformer.

https://research.fb.com/downloads/babi

http://clic.cimec.unitn.it/lambada/appendix_onefile.pdf

(e) Step 1 (f) Step 2

Figures and Tables:

Figure1: The Universal Transformer repeatedly refines a series of vector representations for each position of the sequence in parallel, by combining information from different positions using self-attention (see Eqn 2) and applying a recurrent transition function (see Eqn 4) across all time steps 1 ≤ t ≤ T . We show this process over two recurrent time-steps. Arrows denote dependencies between operations. Initially, h 0 is initialized with the embedding for each symbol in the sequence. h t i represents the representation for input symbol 1 ≤ i ≤ m at recurrent time-step t. With dynamic halting, T is dynamically determined for each position (Section 2.2).

Figure 2: The recurrent blocks of the Universal Transformer encoder and decoder. This diagram omits position and time-step encodings as well as dropout, residual connections and layer normalization.A complete version can be found in Appendix A. The Universal Transformer with dynamic halting determines the number of steps T for each position individually using ACT(Graves, 2016).

Figure 3: Ponder time of UT with dynamic halting for encoding facts in a story and question in a bAbI task requiring three supporting facts.

Figure 4: The Universal Transformer with position and step embeddings as well as dropout and layer normalization.

# While-l o o p s t o p s when t h i s p r e d i c a t e i s FALSE 2 # i . e . a l l ( ( p r o b a b i l i t y < t h r e s h o l d ) & ( c o u n t e r < m a x _ s t e p s ) ) a r e f a l s e 3 d e f s h o u l d _ c o n t i n u e ( u0 , u1 , h a l t i n g _ p r o b a b i l i t y , u2 , n _ u p d a t e s , u3 ) : 4 r e t u r n t f . r e d u c e _ a n y l e s s ( h a l t i n g _ p r o b a b i l i t y , t h r e s h o l d ) , 7 t f . l e s s ( n _ u p d a t e s , m a x _ s t e p s ) ) ) 8 # Do w h i l e l o o p i t e r a t i o n s u n t i l p r e d i c a t e a b o v e i s f a l s e 9 ( _ , _ , _ , r e m a i n d e r , n _ u p d a t e s , n e w _ s t a t e ) = t f . w h i l e _ l o o p ( 10 s h o u l d _ c o n t i n u e , u t _ w i t h _ d y n a m i c _ h a l t i n g , ( s t a t e , 11 s t e p , h a l t i n g _ p r o b a b i l i t y , r e m a i n d e r s , n _ u p d a t e s , p r e v i o u s _ s t a t e ) )

C a l c u l a t e t h e p r o b a b i l i t i e s b a s e d on t h e s t a t e 4 p = c o mm o n _ l a ye r s . d e n s e ( s t a t e , 1 , a c t i v a t i o n = t f . nn . s i g m o i d , 5 u s e _ b i a s = T r u e ) 6 Mask f o r i n p u t s which h a v e n o t h a l t e d y e t 7

18# p r o b a b i l i t i e s f o r t h o s e i n p u t s which h a v e n ' t h a l t e d y e t 19 h a l t i n g _ p r o b a b i l i t y += p * s t i l l _ r u n n i n g 20 # Compute r e m a i n d e r s f o r t h e i n p u t s which h a l t e d a t t h i s s t e p 21 r e m a i n d e r s += n e w _ h a l t e d * ( 1h a l t i n g _ p r o b a b i l i t y )22 # Add t h e r e m a i n d e r s t o t h o s e i n p u t s which h a l t e d a t t h i s s t e p23 h a l t i n g _ p r o b a b i l i t y += n e w _ h a l t e d * r e m a i n d e r s 24 # I n c r e m e n t n _ u p d a t e s f o r a l l i n p u t s which a r e s t i l l r u n n i n g 25 n _ u p d a t e s += s t i l l _ r u n n i n g + n e w _ h a l t

Figure 7: Visualization of the attention distributions, when encoding the question: "Where was the apple before the 23

Average error and number of failed tasks (> 5% error) out of 20 (in parentheses; lower is better in both cases) on the bAbI dataset under the different training/evaluation setups. We indicate state-of-the-art where available for each, or '-' otherwise.

Table content: <table/>

Accuracy on the subject-verb agreement number prediction task (higher is better).

Table content: <table><tr><td>Model</td><td colspan="3">LM Perplexity &amp; (Accuracy)</td><td/><td colspan="2">RC Accuracy</td></tr><tr><td/><td>control</td><td>dev</td><td>test</td><td colspan="2">control dev</td><td>test</td></tr><tr><td>Neural Cache (Grave et al., 2016)</td><td>129</td><td>139</td><td>-</td><td>-</td><td>-</td><td>-</td></tr><tr><td colspan="2">Dhingra et al. Dhingra et al. (2018) -</td><td>-</td><td>-</td><td>-</td><td>-</td><td>0.5569</td></tr><tr><td>Transformer</td><td colspan="6">142 (0.19) 5122 (0.0) 7321 (0.0) 0.4102 0.4401 0.3988</td></tr><tr><td>LSTM</td><td colspan="6">138 (0.23) 4966 (0.0) 5174 (0.0) 0.1103 0.2316 0.2007</td></tr><tr><td>UT base, 6 steps (fixed)</td><td colspan="6">131 (0.32) 279 (0.18) 319 (0.17) 0.4801 0.5422 0.5216</td></tr><tr><td>UT w/ dynamic halting</td><td colspan="6">130 (0.32) 134 (0.22) 142 (0.19) 0.4603 0.5831 0.5625</td></tr><tr><td>UT base, 8 steps (fixed)</td><td>129(0.32)</td><td colspan="3">192 (0.21) 202 (0.18) -</td><td>-</td><td>-</td></tr><tr><td>UT base, 9 steps (fixed)</td><td>129(0.33)</td><td colspan="3">214 (0.21) 239 (0.17) -</td><td>-</td><td>-</td></tr></table>

Accuracy (higher better) on the algorithmic tasks. * Note that the Neural GPU was trained with a special curriculum to obtain the perfect result, while other models are trained without any curriculum.

Table content: <table><tr><td>Model</td><td/><td>Copy</td><td/><td>Reverse</td><td/><td>Addition</td></tr><tr><td/><td colspan="6">char-acc seq-acc char-acc seq-acc char-acc seq-acc</td></tr><tr><td>LSTM</td><td>0.45</td><td>0.09</td><td>0.66</td><td>0.11</td><td>0.08</td><td>0.0</td></tr><tr><td>Transformer</td><td>0.53</td><td>0.03</td><td>0.13</td><td>0.06</td><td>0.07</td><td>0.0</td></tr><tr><td colspan="2">Universal Transformer 0.91</td><td>0.35</td><td>0.96</td><td>0.46</td><td>0.34</td><td>0.02</td></tr><tr><td>Neural GPU  *</td><td>1.0</td><td>1.0</td><td>1.0</td><td>1.0</td><td>1.0</td><td>1.0</td></tr><tr><td/><td/><td>Copy</td><td/><td>Double</td><td/><td>Reverse</td></tr><tr><td>Model</td><td colspan="6">char-acc seq-acc char-acc seq-acc char-acc seq-acc</td></tr><tr><td>LSTM</td><td>0.78</td><td>0.11</td><td>0.51</td><td>0.047</td><td>0.91</td><td>0.32</td></tr><tr><td>Transformer</td><td>0.98</td><td>0.63</td><td>0.94</td><td>0.55</td><td>0.81</td><td>0.26</td></tr><tr><td colspan="2">Universal Transformer 1.0</td><td>1.0</td><td>1.0</td><td>1.0</td><td>1.0</td><td>1.0</td></tr></table>

Character-level (char-acc)  and sequence-level accuracy (seq-acc) results on the Memorization LTE tasks, with maximum length of 55.

Table content: <table><tr><td/><td/><td>Program</td><td/><td>Control</td><td/><td>Addition</td></tr><tr><td>Model</td><td colspan="6">char-acc seq-acc char-acc seq-acc char-acc seq-acc</td></tr><tr><td>LSTM</td><td>0.53</td><td>0.12</td><td>0.68</td><td>0.21</td><td>0.83</td><td>0.11</td></tr><tr><td>Transformer</td><td>0.71</td><td>0.29</td><td>0.93</td><td>0.66</td><td>1.0</td><td>1.0</td></tr><tr><td colspan="2">Universal Transformer 0.89</td><td>0.63</td><td>1.0</td><td>1.0</td><td>1.0</td><td>1.0</td></tr></table>



Table content: <table/>

e d 26 # Compute t h e w e i g h t t o be a p p l i e d t o t h e new s t a t e and o u t p u t :

Table content: <table><tr><td>27</td><td>#</td><td>0 when t h e i n p u t h a s a l r e a d y h a l t e d ,</td></tr><tr><td>28</td><td>#</td><td>p when t h e i n p u t h a s n ' t h a l t e d y e t ,</td></tr><tr><td>29</td><td>#</td><td>t h e r e m a i n d e r s when i t h a l t e d t h i s s t e p .</td></tr><tr><td>30</td><td colspan="2">u p d a t e _ w e i g h t s = t f . e x p a n d _ d i m s ( p  *  s t i l l _ r u n n i n g +</td></tr><tr><td>31</td><td/><td>n e w _ h a l t e d  *  r e m a i n d e r s , -1)</td></tr><tr><td>32</td><td colspan="2"># Apply t r a n s f o r m a t i o n t o t h e s t a t e</td></tr><tr><td>33</td><td colspan="2">t r a n s f o r m e d _ s t a t e = t r a n s i t i o n _ f u n c t i o n ( s e l f _ a t t e n t i o n ( s t a t e ) )</td></tr><tr><td>34</td><td colspan="2"># I n t e r p o l a t e t r a n s f o r m e d and p r e v i o u s s t a t e s f o r non-h a l t e d i n p u t s</td></tr><tr><td>35</td><td colspan="2">n e w _ s t a t e = ( ( t r a n s f o r m e d _ s t a t e  *  u p d a t e _ w e i g h t s ) +</td></tr></table>

Peer review (iclr2019):
 This paper describes a transformer with recurrent structure to take advantage of self-attention mechanism. The number of recurrences can be dynamically determined through ACT-like halting depending on the difficulty of the input. A series of experiments on language modeling tasks have been demonstrated to show promising performances. The overall concerns about this paper is that while the performances are quite promising, the theoretical claims and comparisons in the discussion section are of question. The authors attempt to provide connections to other networks (i.e., Natural GPU, RNN) since UT is an amalgamation of both transformers and RNN, they sound a little “hand-wavy” (i.e., comments about UT effectively interpolating between the feed-forward, fixed-depth Transformer and a gated recurrent architecture). In short, while empirically completely acceptable, intuitively or theoretically it is hard to grasp why UT is superior other than the dynamic/sharing layers across t (not time). I believe that improving this aspect could make this paper even better. Based on the comments below and the responses with the authors, I am willing to improve my score. Pros: 1. The best of both worlds from parallelizable transformer and recurrent structure for repeated self-attention mechanism. Essentially, the “depth” of the transformer can vary if we “unroll” the recurrent stacks. 2. Extensive experiments showing the performance of UT. 3. Analysis of the effect of the recurrent aspect of UT and how it can vary depending on the task difficulty. Comments/cons: 1. I am having trouble understanding the “universal” aspect of the transformer. Is this because the variability of the depth of UT (since “given sufficient memory” was mentioned)? If so, such characteristic of “computational universality” does not seem much unique to UT compared to infinite memory for a transformer or a simple RNN across stack (i.e., input is the while sequence and recurrent step is through the stack analogous to UT stack). Please comment on this. 2. It is nice to see many experiments, but without preexisting knowledge about the datasets and their tasks, I can only make relative judgements based on the provided comparisons against other methods. It would be nice to see slightly more detailed descriptions of each task (particularly LAMBADA LM), not necessarily in the main paper (due to space) but in the appendix if possible for improved self-containedness. 3. In the discussion, the crucial difference between UT and RNN is that RNN is stated to be that RNN cannot access memory in the recurrent steps while UT can. This seems to be the case for not just UT but any Transformer-type model by construction. 4. The authors stated that the “recurrent step” for RNN is through time (as the authors stated) while the “recurrent step” in UT is not through time. While this claim is completely correct itself, the RNN’s inability to access memory in its “recurrent steps” was compared with how UT could still access memory throughout its “recurrent steps”. In this sense, we may argue that the UT cannot access memory across its own t (stacking across t). I am not sure if it is fair to make such implications by putting both “recurrent steps” to be of same nature and pointing out one’s weakness. Perhaps the authors could comment on this. Minor: 1. Table 2.: Best Stack-RNN for 1 attractor is the highest but not bold-faced.