[Paper Review] Universal Transformers

Mostafa Dehghani, Stephan Gouws+3 more|arXiv (Cornell University)|Jul 10, 2018

Topic ModelingComputer Science15 references396 citations

TL;DR

The Universal Transformer generalizes the Transformer by adding parallel-in-time recurrence and dynamic per-position halting, achieving state-of-the-art results on several algorithmic and language tasks and showing enhanced generalization and expressivity.

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.

Motivation & Objective

• Motivate the need for models that combine parallel processing with recurrent inductive bias for sequence modeling.
• Introduce the Universal Transformer (UT) as a generalization of Transformer that refines representations over depth in parallel across positions.
• Demonstrate that UTs can be Turing-complete under certain assumptions and assess their empirical performance across diverse tasks.
• Show that dynamic per-position halting improves accuracy on several tasks and analyze its effects on performance and computation.

Proposed method

• UT uses encoder and decoder with self-attention and a recurrent transition function shared across positions and time steps.
• In each recurrent step, UT applies multi-head self-attention to revise representations in parallel, followed by a transition function (depth-wise separable convolution or position-wise feedforward) and residual connections.
• Per-symbol depth is unbounded in principle, enabling dynamic computation depth via ACT-style halting.
• Positional and time-step encodings are added to inform processing at each depth step.
• The model is trained in a Transformer-like encoder–decoder setup with teacher forcing for the decoder.
• UT can be viewed as weight-tied, depth-unfolded Transformer blocks, enabling recurrence over depth rather than over sequence length.

Experimental results

Research questions

• RQ1 Can parallel-in-time recurrence with self-attention and a shared transition function improve generalization and expressivity beyond the standard Transformer?
• RQ2 Does dynamic per-position halting (adaptive computation time) improve performance on algorithmic and language tasks?
• RQ3 Under what conditions is the UT computationally universal (Turing-complete) relative to the Transformer?
• RQ4 How does UT perform on language understanding and large-scale tasks such as machine translation compared to standard Transformers and LSTMs?
• RQ5 What is the impact of recurrent depth on tasks requiring long-range reasoning and compositional generalization?

Key findings

• UT outperforms standard Transformers and LSTMs on multiple algorithmic and language tasks.
• On LAMBADA language modeling, UT achieves state-of-the-art results.
• On WMT14 English-German translation, UT with fully-connected transitions and no ACT achieves a BLEU improvement over a Transformer of similar size.
• Dynamic halting (ACT) improves accuracy on several smaller tasks and shows deeper per-symbol processing where needed, acting as a regularizer.
• UT can learn under adaptive depth, performing more steps for harder inputs and fewer steps for easier ones, while maintaining parallel computation across sequence positions.
• The UT framework is theoretically more powerful than a fixed-depth Transformer and can emulate neural GPUs and Neural Turing Machines under certain parameterizations.