Source-linked AI summary
Lingvo: a Modular and Scalable Framework for Sequence-to-Sequence Modeling
Jonathan Shen, Patrick Nguyen, Yonghui Wu, Zhifeng Chen, Mia X. Chen, Ye Jia, Anjuli Kannan, Tara Sainath, Yuan Cao, Chung-Cheng Chiu, Yanzhang He, Jan Chorowski, Smit Hinsu, Stella Laurenzo, James Qin, Orhan Firat, Wolfgang Macherey, Suyog Gupta, Ankur Bapna, Shuyuan Zhang, Ruoming Pang, Ron J. Weiss, Rohit Prabhavalkar, Qiao Liang, Benoit Jacob, Bowen Liang, HyoukJoong Lee, Ciprian Chelba, Sébastien Jean, Bo Li, Melvin Johnson, Rohan Anil, Rajat Tibrewal, Xiaobing Liu, Akiko Eriguchi, Navdeep Jaitly, Naveen Ari, Colin Cherry, Parisa Haghani, Otavio Good, Youlong Cheng, Raziel Alvarez, Isaac Caswell, Wei-Ning Hsu, Zongheng Yang, Kuan-Chieh Wang, Ekaterina Gonina, Katrin Tomanek, Ben Vanik, Zelin Wu, Llion Jones, Mike Schuster, Yanping Huang, Dehao Chen, Kazuki Irie, George Foster, John Richardson, Klaus Macherey, Antoine Bruguier, Heiga Zen, Colin Raffel, Shankar Kumar, Kanishka Rao, David Rybach, Matthew Murray, Vijayaditya Peddinti, Maxim Krikun, Michiel A. U. Bacchiani, Thomas B. Jablin, Rob Suderman, Ian Williams, Benjamin Lee, Deepti Bhatia, Justin Carlson, Semih Yavuz, Yu Zhang, Ian McGraw, Max Galkin, Qi Ge, Golan Pundak, Chad Whipkey, Todd Wang, Uri Alon, Dmitry Lepikhin, Ye Tian, Sara Sabour, William Chan, Shubham Toshniwal, Baohua Liao, Michael Nirschl, Pat Rondon
TL;DR
Collaborative deep-learning research needs rapid experimentation alongside reusable, reproducible, and scalable shared code. Lingvo addresses this with modular sequence-modeling components, centralized configurations, and integrated distributed-training and quantized-inference support. The framework has been used by dozens of researchers and has produced state-of-the-art results across multiple speech and language tasks.
Problem
Collaborative research needs code that can be rapidly prototyped, reused, documented, and reproduced across researchers and experiments.
Method
Lingvo provides modular shared components, centralized experiment configurations, and integrated support for distributed training and quantized inference.
Results
Lingvo has produced state-of-the-art results across machine translation, speech recognition, speech synthesis, and speech translation, and is used by dozens of researchers.
Takeaways & Limitations
Shared layers and configurations make experiments comparable and reproducible while allowing research code and ideas to be reused across tasks and deployment settings.
Takeaways & Limitations
These collaboration and scalability benefits require more discipline and boilerplate, and distributed job runners assume a shared checkpoint directory.
Abstract
from arXiv · showhide
Lingvo is a Tensorflow framework offering a complete solution for collaborative deep learning research, with a particular focus towards sequence-to-sequence models. Lingvo models are composed of modular building blocks that are flexible and easily extensible, and experiment configurations are centralized and highly customizable. Distributed training and quantized inference are supported directly within the framework, and it contains existing implementations of a large number of utilities, helper functions, and the newest research ideas. Lingvo has been used in collaboration by dozens of researchers in more than 20 papers over the last two years. This document outlines the underlying design of Lingvo and serves as an introduction to the various pieces of the framework, while also offering examples of advanced features that showcase the capabilities of the framework.
1 Introduction
Lingvo is an open-source Google framework for deep-neural-network sequence modeling, used across several speech and language tasks. The paper introduces its design, core components, training flow, and advanced distributed-training and inference capabilities.
- Lingvo is an open-source framework developed by Google for sequence modeling with deep neural networks.
- The framework has produced state-of-the-art results in machine translation, speech recognition, speech synthesis, and speech translation.
- The paper motivates Lingvo’s design, explains its core components and APIs, and walks through training-run logic from model construction to parameter updates.
- Advanced usage includes synchronous and asynchronous distributed training, multi-task models, and inference.
2 Design
Lingvo is designed for collaborative research through modular, reusable components, centralized experiment configuration, scalable performance, and shared research-to-production code. These benefits support reproducibility and deployment while requiring additional discipline and boilerplate.
- Motivation: Lingvo evolved to support many applied researchers working on speech and natural-language problems in one shared codebase.
- Guiding principles: Its guiding principles are modular extensibility, reproducible experiments, production-scale performance, and code reuse between research and production.
- Modular building blocks: Common interfaces and reusable building blocks let algorithmic improvements transfer across tasks and make existing models easier to adapt to new datasets.
- Experiment design: Dedicated, version-controlled hyperparameter configurations and shared layers make experiments easier to document, reproduce, compare, and understand.
- Performance and deployment: Lingvo supports production-scale datasets, synchronous and asynchronous distributed training, shared training and inference code, device-specific overrides, and built-in quantization.
- Trade-off: These scalability and collaboration benefits come at the cost of more discipline and boilerplate than fast prototyping typically requires.
- Core components: The framework organizes functionality into components including experiment registries, job runners, NestedMap data structures, custom operations, Models, Tasks, Layers, input generators, Params, and experiment configurations.
3 Implementation
Lingvo’s implementation centers on explicitly configured Layers and hierarchical Params objects. These APIs support construction, composition, extension, and experiment-specific customization of models.
- Params: The Params class defines explicit configuration keys, while experiment classes override defaults with experiment-specific values.
- Layer construction: Constructing a Layer requires Params containing its class, name, and variable-initialization specification.
- Layer construction: Because the class is stored in Params, a Layer can be constructed directly from its class or through the Params object.
- Layer composition: Layers execute forward computation through FProp() and can contain named child Layers created from child parameters.
3.3 Variable Management
Lingvo gives each Layer responsibility for creating and managing its variables, supporting transformations, distributed execution, placement policies, and research features such as weight noise.
- Each Layer creates and manages its own variables.
- Variables are created during layer initialization, registered in self.vars, and made available through self.theta, potentially after transformations such as variational noise.
- During FProp(), variables should be accessed through the theta argument because computation may execute on different devices in distributed training.
- Variable placement defaults to the least-allocated parameter server, while explicit policies can support model parallelism.
- Explicit variable management supports research ideas such as weight noise and simplifies weight sharing for synchronous replica training.
3.4 Input Processing
Lingvo supports text and TFRecord inputs, configurable tokenization and length-based bucketing, and extensible input-processing pipelines implemented in C++ or Python.
- Lingvo accepts inputs in either plain text or TFRecord format.
- Sequence inputs can be bucketed by length using bucket_upper_bound and bucket_batch_limit parameters.
- Text inputs can use VocabFileTokenizer, BpeTokenizer, or WpmTokenizer.These provide file-based lookup, byte pair encoding, and word-piece tokenization, respectively.
- Input processors read files through file_pattern and _DataSourceFromFilePattern(), then populate input batches and optionally preprocess them.The data source may use a custom C++ RecordProcessor op.
- Input processing can alternatively be defined directly in Python through the generic_input op.
3.5 Model Registration
Lingvo registers single-task model configurations through annotated parameter classes that define task and dataset settings, with the framework wrapping the task into a model.
- Configuration classes are annotated with @model_registry and registered under keys such as image.mnist.LeNet5.RegisterSingleTaskModel handles the typical single-task case.
- A single-task configuration subclasses SingleTaskModelParams and implements Task(), which returns parameters configuring a Task.Registration automatically wraps the Task into a SingleTaskModel.
- Train(), Test(), and optionally Dev() define input-generator parameters for different datasets.
- Figure 1 illustrates the registration process for a single-task model.
3.6 Overriding Params from the Command Line
Lingvo allows specific runs to override hyperparameter values from the command line or an override file, supporting similar-job hyperparameter tuning.
- Run-specific hyperparameters can be overridden with --model_params_override or --model_params_file_override.This simplifies starting similar jobs for hyperparameter tuning.
3.7 Assertions
Lingvo provides runtime assertions for values and shapes and numerical checks for NaNs, with command-line flags to disable either mechanism.
- py_utils.py provides runtime assertions for values and shapes and CheckNumerics() for detecting NaNs.
- Assertions can be disabled with --enable_asserts=false, while numerical checks can be disabled with --enable_check_numerics=false.
3.8 Code Layout
Lingvo’s code layout separates model registration, task-specific parameters, reusable layers and utilities, input processing, summaries, and custom operations.
- The global registry imports and registers model parameters, while task folders contain domain-specific projects and model parameter definitions.
- cluster.py specifies the policy for placing operations across devices.
- Reusable neural-network components include attention, layers, recurrent cells, recurrent layers, and a functional RNN implementation.
- Lingvo centralizes general-purpose utilities in py_utils.py and summary-handling utilities in summary_utils.py.
- The input pipeline is implemented through record_*.* files, while py_x_ops.py provides Python bindings for custom C++ operations defined in x_ops.cc.
4 Life of a Training Run
A Lingvo training run resolves a registered model configuration, independently builds job-specific graphs, initializes or restores variables, and iterates through checkpointed evaluation.
- Training starts by resolving a model name through the registry, obtaining its Params, and creating the job runners.
- Each runner independently instantiates the model and builds graphs for its job: training uses forward and backward propagation, whereas evaluation and decoding use evaluation metrics.Multiple Evalers and Decoders may serve different evaluation datasets.
- The model constructor recursively creates child-layer parameters and variables through CreateChild and CreateVariable.Child parameters may be exposed for configuration or constructed inside the model constructor.
- After graph construction, the Trainer waits for initialization or checkpoint restoration, while evaluation runners wait for a new checkpoint.
- Training then runs train_op repeatedly while the Controller writes summaries and checkpoints that Evalers and Decoders detect and evaluate until termination or max_steps.
5 Advanced Usage
Lingvo supports distributed training, multi-task models, and deployment-oriented inference features, including quantized layers and configurable execution strategies.
- 5.1 Distributed Training: Both synchronous and asynchronous distributed training are supported, with independent worker loops in asynchronous mode and a trainer-client-driven loop in synchronous mode.
- 5.1 Distributed Training: Asynchronous training uses Controller, Trainer, Parameter Server, and optional Data Processor jobs with distinct checkpointing, updating, storage, and preprocessing roles.
- 5.1 Distributed Training: Synchronous training uses workers and a trainer client that aggregates results and updates variables without parameter servers.
- 5.2 Multi-task Models: Multi-task models combine tasks that share variables, with sharing configurable from a shared encoder to fine-grained regular-expression-based selection.
- 5.2 Multi-task Models: Multi-task configurations define task parameters and relative task weights, and support knowledge distillation through a checkpoint-loaded teacher model.
- 5.3 Inference and Quantization: Quantized layers wrap training and inference computations to control dynamic range for deployment on servers or embedded devices using fixed-point arithmetic.
- 5.3 Inference and Quantization: Inference execution can use smaller batches, beam search, or timestep-by-timestep computation depending on operations, latency, parallelizability, memory, and power.