PaperlayerView exact revision on arXiv

Source-linked AI summary

Universal Transformers

Mostafa Dehghani, Stephan Gouws, Oriol Vinyals, Jakob Uszkoreit, Łukasz Kaiser

arXiv:1807.03819v3cs.CLcs.LGstat.ML

TL;DR

The paper addresses the Transformer’s weak length generalization and missing recurrent inductive bias by introducing the Universal Transformer, a parallel-in-time recurrent self-attentive model with optional dynamic halting. UTs outperform Transformers and LSTMs across a wide range of tasks and are computationally universal under sufficient memory, while the authors identify further improvements as future work.

  • Problem

    The standard Transformer does not generalize well to input lengths beyond those seen during training and lacks RNNs’ inductive bias toward iterative or recursive transformations.

  • Method

    The Universal Transformer combines parallel self-attention, recurrent representation refinement with shared weights, and optional per-position dynamic halting.

  • Results

    UTs outperform Transformers and LSTMs across a wide range of tasks, achieve a new state of the art on LAMBADA, and are computationally universal under sufficient memory.

  • Takeaways & Limitations

    The added recurrence and conditional computation connect practical sequence models with computationally universal models while improving accuracy on challenging sequence modeling tasks.

  • Takeaways & Limitations

    The authors present further improvements in power, data efficiency, and generalization beyond the current state of the art as future work.

Abstract

from arXiv · show

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

The Transformer enables parallel sequence processing but lacks the recurrent inductive bias needed for strong length generalization. The Universal Transformer combines parallel self-attention with recurrent refinement and improves performance across challenging tasks.

  • 1 INTRODUCTION: Parallel-in-time feed-forward architectures avoid the sequential computation bottleneck of RNNs while retaining global access to sequence elements.These architectures also address vanishing gradients as sequence length increases.
  • 1 INTRODUCTION: The standard Transformer’s fixed stack of distinct layers lacks RNNs’ inductive bias toward iterative or recursive transformations.This limitation is especially relevant to algorithmic and language understanding tasks.
  • 1 INTRODUCTION: The Transformer does not generalize well to input lengths not encountered during training.The passage contrasts this behavior with Neural Turing Machines, Neural GPUs, and Stack RNNs.
  • 1 INTRODUCTION: The Universal Transformer is a parallel-in-time recurrent self-attentive generalization of the Transformer that combines global receptive fields with recurrent inductive bias.Under certain assumptions, UTs can be shown to be Turing-complete.
  • 1 INTRODUCTION: Each recurrent step refines all sequence positions in parallel through self-attention followed by a shared transition function.Dynamic per-position halting lets the model choose the required number of refinement steps for each symbol.
  • 1 INTRODUCTION: UTs outperform Transformers and LSTMs across a wide range of tasks, with dynamic-halting UTs achieving a new state of the art on LAMBADA.The reported improvements span algorithmic tasks, bAbI language understanding, and machine translation.

2 MODEL DESCRIPTION

The Universal Transformer repeatedly updates all sequence positions in parallel using self-attention and a shared recurrent transition. Its encoder-decoder architecture supports autoregressive generation and optional position-specific dynamic halting.

  • 2.1 THE UNIVERSAL TRANSFORMER: The UT applies recurrent neural networks across consecutive representation revisions rather than across sequence positions.Both encoder and decoder recurrently process representations for their respective positions.
  • 2.1 THE UNIVERSAL TRANSFORMER: At each step, self-attention exchanges information across all positions before a shared transition function independently updates each position.The transition can be applied any number of times.
  • 2.1 THE UNIVERSAL TRANSFORMER: The encoder initializes H0 with symbol embeddings and produces HT after T parallel updates, yielding d-dimensional representations for m input symbols.The final encoder output is HT ∈ R^{m×d}.
  • 2.1 THE UNIVERSAL TRANSFORMER: The recurrent transition is either a separable convolution or a position-wise fully connected network.The fully connected option uses one rectified-linear activation between two affine transformations.
  • 2.1 THE UNIVERSAL TRANSFORMER: The decoder adds attention to the final encoder representation and generates outputs autoregressively while masking future positions.During generation, it consumes previously produced symbols and emits one next-position distribution per iteration.
  • 2.2 DYNAMIC HALTING: Adaptive Computation Time dynamically modulates the number of recurrent steps allocated to each input symbol.A scalar halting probability determines each symbol’s ponder time, with halted states copied forward until all blocks halt or the maximum step count is reached.

3 EXPERIMENTS AND ANALYSIS

Across algorithmic, language-understanding, and translation evaluations, Universal Transformers achieve strong results, often surpassing standard Transformers and other recurrent baselines. Dynamic halting adapts computation to task demands and remains beneficial beyond simply increasing fixed recurrent steps.

  • bAbI question-answering: Universal Transformers achieve state-of-the-art results across all bAbI tasks in both 10K and 1K training regimes.Both standard UT and UT with dynamic halting lead in average error and number of failed tasks.
  • bAbI question-answering: 3.8±2.2 average ponder time for three supporting facts exceeds 3.1±1.1 for two and 2.3±0.8 for one.The model therefore varies recurrent processing with the number of supporting facts required.
  • Subject-verb agreement: UTs significantly outperform standard Transformers on subject-verb agreement and approach the current state of the art at 99.2% average accuracy.Their advantage increases as the number of agreement attractors grows.
  • LAMBADA language modeling: Universal Transformer achieves state-of-the-art results in both LAMBADA language modeling and reading comprehension, outperforming LSTMs and vanilla Transformers.The control set additionally tests standard language-modeling ability without LAMBADA filtering.
  • LAMBADA language modeling: Dynamic halting outperforms fixed UTs using 8 or 9 steps despite averaging 8.2±2.1 steps, suggesting computation is allocated selectively across symbols.The best fixed UT used 6 steps; fixed 8- and 9-step models improved over it but did not surpass dynamic halting.
  • Algorithmic tasks: On sequences trained at length 40 and evaluated at length 400, UTs outperform LSTMs and vanilla Transformers by a wide margin on Copy, Reverse, and Addition.The Neural GPU achieved perfect results with a special curriculum unavailable to the other models.
  • Machine translation: On WMT14 English-German translation, a UT improves by 0.9 BLEU over a Transformer and 0.5 BLEU over a Weighted Transformer.The comparisons use approximately the same number of parameters.

4 DISCUSSION

The Universal Transformer extends Transformer computation with recurrent, shared transformations and whole-sequence memory access, giving it greater computational power while retaining attention-based processing.

  • Discussion: With fixed recurrent steps, the Universal Transformer is a Transformer whose parameters are tied across layers.Both self-attention and transition weights are shared across layers.
  • Discussion: Unlike time-recurrent RNNs, UTs can attend to the entire previous layer during recurrent computation.This gives each recurrent step access to sequence-wide memory rather than only a fixed-size recurrent state.
  • Discussion: Given sufficient memory, the Universal Transformer can simulate any Turing machine, unlike the standard Transformer.The paper characterizes this as computational universality and links the added expressivity to improved sequence-modeling accuracy.
  • Discussion: A Universal Transformer can be reduced to a Neural GPU by using convolutional transitions and setting recurrent steps equal to input length.The input-length-dependent recurrent computation is the step that distinguishes this construction from the vanilla Transformer.
  • Discussion: Unlike end-to-end Memory Networks, UTs store states aligned with individual input or output positions and use an encoder-decoder architecture.The paper also reports competitive performance on large-scale sequence-to-sequence tasks.

5 CONCLUSION

The conclusion presents the Universal Transformer as a Transformer generalization combining shared weights and conditional computation, while identifying further improvements as an open direction.

  • Conclusion: The Universal Transformer combines Transformer generalization with state-of-the-art results across language-understanding and algorithmic sequence-modeling tasks.The stated goal is to address a key shortcoming of the standard Transformer.
  • Conclusion: Weight sharing provides an inductive bias intended to balance model expressivity, and the paper evaluates it in small- and large-scale experiments.The design is motivated by weight sharing in CNNs and RNNs.
  • Conclusion: Conditional computation lets each Universal Transformer position halt or continue, and dynamic halting performs better than fixed-depth computation.The mechanism is presented as part of the model’s route toward computational universality.
  • Conclusion: The authors identify further improvements as necessary for learning algorithms that are more powerful, data efficient, and able to generalize beyond the current state of the art.This is framed as a future direction beyond the basic Universal Transformer.
  • Conclusion: Training and evaluation code is available in the Tensor2Tensor repository.The cited repository is given as https://github.com/tensorflow/tensor2tensor.

APPENDIX A DETAILED SCHEMA OF THE UNIVERSAL TRANSFORMER

The appendix characterizes the Universal Transformer’s computational advantage through recurrent depth: unlike fixed-depth Transformers, its sequential computation can depend on input length.

  • Detailed schema: Figure 4 depicts a Universal Transformer with position and step embeddings, dropout, and layer normalization.These components are identified in the figure description.
  • Detailed schema: A standard Transformer’s sequential operation count is fixed by its number of layers, independent of input size.The appendix discusses this property under finite-precision assumptions.
  • Detailed schema: For functions requiring sequential processing of every input element, an input longer than a fixed depth can defeat a standard Transformer.The appendix contrasts this with input-length-dependent recurrent steps in UTs, RNNs, and Neural GPUs.

APPENDIX C UT WITH DYNAMIC HALTING

Dynamic halting updates each position’s state iteratively and stops computation when halting probabilities reach a threshold or the maximum step count is reached.

  • UT with dynamic halting: At each step, the implementation computes halting probabilities from the current state and tracks running positions, remainders, and update counts.These quantities are initialized and updated throughout the dynamic-halting loop.
  • UT with dynamic halting: The loop continues while some positions remain below the halting threshold and have not reached the maximum number of updates.The continuation predicate combines the threshold and maximum-step conditions.
  • UT with dynamic halting: A position’s halting probability accumulates its current probability, while newly halted positions receive a remainder.The remainder is computed from the unfinished halting probability and then added to the halted position’s total.
  • UT with dynamic halting: The update count increases for positions still running or newly halted at the current step.This count supports the maximum-step stopping condition.
  • UT with dynamic halting: Each step transforms the state through self-attention followed by the transition function, then interpolates transformed and previous states using update weights.The update weights assign zero to halted inputs, p to continuing inputs, and the remainder to inputs halting at that step.
  • UT with dynamic halting: The new state is the weighted combination of transformed and previous states, and the loop returns the updated halting variables and state.The listed computation returns transformed_state, step, halting_probability, remainders, n_updates, and new_state.

APPENDIX D DESCRIPTION OF SOME OF THE TASKS/DATASETS

Appendix D provides additional details on the bAbI, subject-verb agreement, LAMBADA, and learning-to-execute tasks.

  • The appendix covers four evaluation areas: bAbI, subject-verb agreement, LAMBADA language modeling, and learning to execute.

D.1 BABI QUESTION-ANSWERING

The bAbI dataset evaluates language understanding and reasoning across 20 synthetic tasks, with experiments considering data efficiency and joint training. Subject-verb agreement tests whether models capture hierarchical syntactic structure, including under intervening attractors.

  • bAbI question-answering: bAbI contains 20 synthetic tasks covering reasoning abilities such as counting, negation, coreference, time, position, path-finding, and motivations.
  • bAbI question-answering: The dataset includes 1k- and 10k-example training regimes, while joint training evaluates one model across all tasks without per-task tuning.
  • Subject-verb agreement: Subject-verb agreement predicts whether an English verb agrees in number with its subject and serves as a proxy for hierarchical syntactic structure.
  • Subject-verb agreement: The paper trains subject-verb agreement models with next-word prediction and evaluates ranking accuracy between the correct and incorrect verb forms.
  • Subject-verb agreement: Agreement attractors are intervening nouns with the opposite number that require identifying the subject head while ignoring distractors.
  • Subject-verb agreement: Examples range from no attractor to three attractors, increasing the number of intervening nouns that can confuse the model.

D.3 LAMBADA LANGUAGE MODELING

LAMBADA tests prediction of a final target word from broad narrative context, contrasting standard language modeling with reading-comprehension evaluation. The supplied task descriptions also characterize LTE programs by length and nesting, which control numerical size and parse-tree depth.

  • LAMBADA language modeling: LAMBADA asks models to predict the last word of the last sentence from the entire narrative passage.
  • LAMBADA language modeling: The passages are selected so humans can infer the final word from long context but not from the preceding target sentence alone.
  • LAMBADA language modeling: A randomly sampled control set with matching shape and size checks standard language modeling before evaluation on LAMBADA.
  • LAMBADA language modeling: LAMBADA is evaluated both as standard next-word language modeling and as reading comprehension, and the paper reports Universal Transformer results in both settings.
  • Learning to execute: LTE program-evaluation difficulty is controlled by integer length and operation nesting, with greater nesting producing deeper parse trees.

APPENDIX E BABI DETAILED RESULTS

The detailed bAbI results are organized around best-seed runs and distinguish single-task from joint-task training. The supplied table fragments include failed-task counts but do not identify their column mapping.

  • Results are reported as the best seed run for each task out of 10 runs.
  • The table distinguishes train single and train joint settings.
  • One supplied result row lists failed-task counts of 0, 0, 3, and 5 across the displayed conditions.

APPENDIX F BABI ATTENTION VISUALIZATION

The appendix visualizes attention distributions on bAbI examples across time steps and four attention heads. The examples cover questions requiring one, two, or three supportive facts.

  • Attention visualization: Attention visualizations show distributions over story facts and questions across different time steps and four heads.Different color bars indicate attention weights for the four heads.
  • One supportive fact: The appendix includes a one-fact example asking where Mary is after she travels to the bathroom.The example story states that Mary journeyed to the bathroom.
  • Two supportive facts: A two-fact example asks where the apple is after Mary takes and drops it in the bathroom.The corresponding attention visualization is provided for the question “Where is the apple?”.
  • Two supportive facts: Another two-fact example asks where the milk is after John moves to the bathroom, picks up the milk, and drops it.The associated figure visualizes attention for the question “Where is the milk?”.
  • Three supportive facts: A three-fact example asks where the apple was before the bathroom using a longer story containing multiple movements and object interactions.The appendix labels this example as requiring three supportive facts and provides its corresponding attention visualization.
Loading 1807.03819v3…