Universal Transformers

This post will discuss the Universal Transformer, which combines the original Transformer model with a technique called Adaptive Computation Time. The main innovation of Universal Transformers is to apply the Transformer components a different number of times for each symbol.

Paper Reference

Dehghani M, Gouws S, Vinyals O, Uszkoreit J, Kaiser Ł. Universal transformers.  ICLR 2019.

Background and Transformer Review

If you are not already familiar with the Transformer model, you should read through “The Transformer: Attention is All You Need.” The Universal Transformer is a simple modification of the Transformer, so it is important to understand the Transformer model first.

If you’re already familiar with the Transformer model and would like a quick review, here goes:

Figure modified from Transformer paper

The base Transformer consists of an encoder and a decoder:

• Encoder:
  • 6 encoder layers
  • Each encoder has 2 sub-layers: (1) multi-head self-attention; (2) feed-forward
• Decoder:
  • 6 decoder layers
  • Each decoder layer has 3 sub-layers: (1) masked multi-head self-attention; (2) encoder-decoder multi-head attention; (3) feed forward

Here’s a one-figure review of multi-head attention, one of the Transformer model’s key innovations:

Multi-head attention is used for encoder self-attention (which takes as input the previous encoder layer output), decoder self-attention (which takes as input the previous decoder layer output), and encoder-decoder attention (which uses the final encoder output for the keys and values and the previous decoder output as the queries.) In the figure above the parts of the model where multi-head attention is used are boxed in red on the left. On the right, the dimensions of the Tensors at each part of the multi-head attention calculation are shown.

Finally, here’s a quick review of the position-wise fully connected feed forward network that’s used in the encoder sub-layers and in the decoder sub-layers:

Motivation for the Universal Transformer

The original Transformer is a natural language processing model that processes all words in the input sequence in parallel, while making use of attention mechanisms to incorporate context. It’s faster to train than an RNN, which has to process input tokens one-by-one. It achieves good performance on language translation. However, it has poorer performance on algorithmic tasks like string copying (e.g. given ‘abc’ as input, output ‘abcabcabc.’)

The Neural GPU and Neural Turing Machine (different kinds of models) have poorer performance on language translation but good performance on algorithmic tasks.

The goal of the Universal Transformer is to achieve good performance on both language translation and algorithmic tasks with only one model. The authors of the Universal Transformer also note that it is a Turing-complete model. (“Turing-complete” means it can simulate any Turing machine, which is a formal definition of a computer.)

Overview

In the Universal Transformers paper, the authors provide a new figure to describe their model:

Figure 4 from Universal Transformer paper

However, I think that the use of a different figure style relative to the original Transformer paper obscures the key differences between the models. Thus, I have modified the figure from the original Transformers paper to more clearly emphasize the similarities and differences of the Transformer and Universal Transformer models. Key differences are emphasized in red:

They key differences between the Transformer and the Universal Transformer are as follows:

1. The Universal Transformer applies the encoder for a variable number of steps for each input token (T steps), while the base Transformer applies exactly 6 encoder layers.
2. The Universal Transformer applies the decoder for a variable number of steps for each output token (T steps), while the base Transformer applies exactly 6 decoder layers.
3. The Universal Transformer uses a slightly different input representation: it includes a “timestep embedding” in addition to the “positional encoding.”

Differences (1) and (2), the variable numbers of steps, are achieved through use of “Adaptive Computation Time” which will be described more later. In brief, Adaptive Computation Time is a dynamic per-position halting mechanism that allows for different amounts of computation on each symbol.

The Universal Transformer is a “parallel-in-time self-attentive recurrent sequence model” which is parallelizable over the input sequence. Like the base Transformer, it has a “global receptive field” (meaning it looks at a lot of words at once.) The main new idea is that in each recurrent step, the Universal Transformer iteratively refines its representations for all symbols in the sequence using self-attention followed by a “transition function” shared across all positions and time-steps.

Here is a cool animation from Oriol Vinyals (@OriolVinyalsML on Twitter) that illustrates the Universal Transformer:

The parameters of the Universal Transformer, including the self-attention and the transition weights, are tied across all positions and time-steps. If the Universal Transformer is run for a fixed number of steps (rather than a variable number of steps T), then the Universal Transformer is equivalent to a multi-layer Transformer with tied parameters across all layers.