Source-linked AI summary
Tensor2Tensor for Neural Machine Translation
Ashish Vaswani, Samy Bengio, Eugene Brevdo, Francois Chollet, Aidan N. Gomez, Stephan Gouws, Llion Jones, Łukasz Kaiser, Nal Kalchbrenner, Niki Parmar, Ryan Sepassi, Noam Shazeer, Jakob Uszkoreit
TL;DR
Neural machine translation needed alternatives to recurrent fixed-size bottlenecks and sequential computation. The paper presents Tensor2Tensor around Transformer-style self-attention, reporting state-of-the-art WMT translation results while supporting standardized, scalable experimentation.
Problem
Recurrent translation models compress entire input sentences into fixed-size vectors, degrading on longer sentences, while recurrent computation requires sequential state propagation.
Method
Tensor2Tensor provides TensorFlow-based models, datasets, and abstractions, including Transformer self-attention that directly accesses sequence history and supports training across CPU, GPU, and TPU.
Results
The Transformer achieved more than 2.0 BLEU over previous reported WMT 2014 English-to-German models and 41.8 BLEU on English-to-French, at lower training cost.
Takeaways & Limitations
Standardized models, datasets, and procedures make it easier to compare models across problems and iterate on experiments at scale.
Takeaways & Limitations
Fully mitigating training randomness is effectively impossible in multithreaded, distributed, floating-point systems.
Abstract
from arXiv · showhide
Tensor2Tensor is a library for deep learning models that is well-suited for neural machine translation and includes the reference implementation of the state-of-the-art Transformer model.
1 Neural Machine Translation Background
Neural machine translation progressed from recurrent sequence-to-sequence models, but fixed-size sentence representations degrade on longer inputs. Convolutional alternatives improved efficiency, while recurrent bottlenecks remained a concern.
- RNN encoder-decoder models compress the source sentence into a fixed-sized state vector before generating the target sentence.
- Longer sentences expose quality degradation when the entire input must be encoded into one fixed-size vector.Attention partially addressed this limitation by providing a neural mechanism for focusing on input information.
- Early convolutional translation models placed a standard RNN above the convolution, creating a bottleneck that hurt performance.
- Fully convolutional models removed the recurrent bottleneck, using recurrent gated convolutions or nonrecursive left-padded decoder convolutions.The latter approach was introduced in WaveNet and significantly improved model efficiency.
2 Self-Attention
The Transformer replaces recurrent sequence processing with stacked self-attention and feed-forward layers, giving each timestep direct access to sequence history. This improves long-range modeling and training speed, with memory cost scaling quadratically in sequence length, while achieving strong translation results.
- Self-Attention: Self-attention gives every timestep direct access to sequence history, enabling distant temporal relationships without hidden-state propagation across time.This also speeds training, but attention memory scales with t^2 for sequence length t.
- Architecture: The Transformer uses stacked self-attention and point-wise fully connected layers in both its encoder and decoder.The encoder and decoder are shown in the left and right halves of the architecture, respectively.
- Architecture: Each encoder layer combines multi-head self-attention with a positionwise feed-forward network.
- Architecture: The decoder adds multi-head attention over the encoder stack to the two sub-layers used by the encoder.
- Computational Performance: Self-attention uses a constant number of sequential operations across positions, whereas recurrent layers require O(n) sequential operations.It is computationally faster when sequence length n is smaller than representation dimensionality d, as commonly occurs in sentence representations.
- Machine Translation: On WMT 2014 English-to-German, Transformer (big) exceeded previous reported models, including ensembles, by more than 2.0 BLEU and reached 28.4 BLEU.Training took 3.5 days on 8 P100 GPUs, and the base model also surpassed published models at a fraction of their training cost.
- Machine Translation: On WMT 2014 English-to-French, the big model achieved 41.8 BLEU and outperformed previously published single models at less than 1/4 the prior state-of-the-art training cost.
3 Tensor2Tensor
Tensor2Tensor is a TensorFlow-based library intended to make deep learning research faster and more accessible. It standardizes models, datasets, and workflows across tasks and devices, supporting rapid experimentation with minimal device-specific configuration.
- Tensor2Tensor provides deep learning models and datasets through TensorFlow, supporting CPU, GPU, and TPU training locally or in the cloud.The library emphasizes both performance and usability while requiring minimal device-specific code or configuration.
- The library began with neural machine translation and includes successful NMT models and standard datasets before expanding across text, images, video, and audio.The number of models and datasets has grown significantly.
- Standardized usage across models and problems makes it easy to try one model on multiple problems or multiple models on one problem.Commands and procedures are unified across datasets, models, training, evaluation, and decoding.
- Open development on GitHub includes contributors from inside and outside Google.
4 System Overview
Tensor2Tensor organizes training runs around standardized dataset, device, hyperparameter, model, and runtime components. These abstractions let users change a problem or model while leaving other components available for reuse.
- System components: A training run is specified by five components: datasets, device configuration, hyperparameters, model, and Estimator/Experiment.Together, these components define data handling, hardware, model instantiation, optimization, evaluation, and runtime support services.
- System components: Problems encapsulate dataset generation, vocabulary construction, encoded samples, input pipelines, and feature-specific information.They support both training and evaluation and can download and prepare data from public sources.
- System components: TensorFlow and Tensor2Tensor support CPU, GPU, and TPU execution in single- and multi-device configurations, including synchronous and asynchronous data-parallel training.Device configuration specifies the type, number, and location of devices.
- Runtime: Estimator and Experiment manage runtime instantiation, training loops, checkpointing, logging, and alternation between training and evaluation.Appendix A outlines the code, while Appendix B provides example usage.
- Usability: Standardized abstractions let users define a new problem or modify a model without changing the other training components.This reduces mental load and supports faster iteration at scale.
5 Library of research components
Tensor2Tensor serves as a vehicle for rapidly trying, sharing, and reusing research components. Its reusable building blocks support extensions across models and media, while experiments in the library have produced components adopted elsewhere.
- Research components: Tensor2Tensor enables research ideas to be quickly tried out and shared, with useful components potentially committed to widely used libraries such as TensorFlow.The library contains standard layers, optimizers, and other higher-level components through this development pathway.
- Research components: Users can reuse specific Tensor2Tensor components in their own models or systems through library and script usage.Attention building blocks support ongoing extensions and variations of the Transformer model.
- Examples: The Image Transformer extends the Transformer model to images and relies heavily on Tensor2Tensor attention building blocks.It also adds many of its own components.
- Examples: The reversible-layer block and sequence-length bucketing functionality were first implemented and exercised in Tensor2Tensor.Bucketing enables efficient processing of sequence inputs on GPUs in the tf.data.Dataset input pipeline API.
- Examples: Adafactor was developed within Tensor2Tensor and tried on various models and problems, reducing memory requirements for second-moment estimates.The passage describes Adafactor as pending publication.
6 Reproducibility and Continuing Development
Tensor2Tensor addresses the difficulty of maintaining model quality during ongoing development while preserving historical reproducibility. It combines versioned configurations, regression tests, seeded randomness, and recoverable code, though full determinism remains impossible in distributed settings.
- Motivation: Ongoing development threatens historical reproducibility because model training is expensive and random, while freezing a codebase imposes substantial usability and development costs.The section frames reproducibility as a maintenance problem rather than a reason to stop development.
- Mechanisms: Tensor2Tensor mitigates ongoing-development effects through named versioned hyperparameters, recurring end-to-end regression tests, and random seeds at multiple levels.The tests verify that important model-problem pairs achieve specified quality metrics.
- Mechanisms: Random seeds for Python, numpy, and TensorFlow mitigate randomness but cannot fully control a multithreaded, distributed, floating-point system.The paper explicitly qualifies full reproducibility as effectively impossible in that setting.
- Recovery: Version-controlled GitHub code allows recovery of the exact code that produced particular experiment results.This provides a fallback when reproducing historical experiments requires the original implementation.
A Tensor2Tensor Code Outline
The Tensor2Tensor code outline separates input preparation, model execution, optimization, evaluation, and runtime support. It also configures parallel execution and standardizes transformations between feature representations and model targets.
- Configuration: RunConfig specifies devices, while Parallelism enables data-parallel duplication of the model across multiple devices.The passage gives multi-GPU synchronous training as an example.
- Runtime: Experiment provides training and evaluation hooks for support services such as logging and checkpointing, while Estimator encapsulates the model function.These classes organize runtime execution around the model function.
- Model pipeline: model.bottom transforms input features into representations consumable by the model body, such as dense embeddings of integer token ids.The transformation uses feature type information from the Problem.
- Model pipeline: model.body contains the core model computation, and model.top transforms its output into the target space using Problem information.Together they connect the model’s internal computation with the task’s output representation.
- Execution: During training, Tensor2Tensor runs the training operation for the specified number of steps; during evaluation, it accumulates metrics across the specified evaluation steps.The code outline distinguishes optimization from metric computation by execution mode.
- Input pipeline: Problem.input_fn uses TensorFlow’s tf.data.Dataset API to produce mode-specific input pipelines, stream examples, and pad and batch them for efficient processing.The pipeline prepares examples in a form ready for model execution.
B Example Usage
Tensor2Tensor standardizes the workflow for training and evaluating models, making experiments reproducible through problem, model, and hyperparameter choices. The example applies this workflow to Transformer translation on WMT English–German data.
- Example Usage: Tensor2Tensor experiments can typically be reproduced using a problem, model, and hyperparameter-set triple.The example identifies these as translate_ende_wmt32k, transformer, and transformer_base.
- Example Usage: The workflow installs Tensor2Tensor, generates data, trains and evaluates a model, and produces decodes from the trained model.These stages are represented by the installation, data-generation, training, evaluation, and decoding commands.
- Example Usage: The example trains the attention-based Transformer model on WMT data translating from English to German.