Universal Transformers — part 2/3

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

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 tasks 3 3 3 Defined as > 5 % >5\% error. , 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 3.8\raisebox{0.86108pt}{$\scriptstyle\pm$}2.2 ) than for tasks requiring only two ( 3.1 ± 1.1 3.1\raisebox{0.86108pt}{$\scriptstyle\pm$}1.1 ), which is in turn higher than for tasks requiring only one supporting fact ( 2.3 ± 0.8 2.3\raisebox{0.86108pt}{$\scriptstyle\pm$}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 cabinet 4 4 4 Cabinet (singular) is an agreement attractor in this case. , 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, Δ \Delta ).

### 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 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.

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

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 8.2\raisebox{0.86108pt}{$\scriptstyle\pm$}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 ) .

## 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 5 5 Note that in UT both the self-attention and transition weights are tied across layers. . 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 feed-forward 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 model 6 6 6 Appendix B illustrates how UT is computationally more powerful than the standard Transformer. . 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 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.

## 5 Conclusion

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

#### Acknowledgements

We are grateful to Ashish Vaswani, Douglas Eck, and David Dohan for their fruitful comments and inspiration.

## Appendix A Detailed Schema of the Universal Transformer

## Appendix B On the Computational Power of UT vs Transformer