[Paper Review] TensorFlow: A system for large-scale machine learning

Martı́n Abadi, Paul Barham+20 more|arXiv (Cornell University)|May 27, 2016

Parallel Computing and Optimization Techniques56 references8,791 citations

TL;DR

TensorFlow presents a dataflow-graph based system for large-scale machine learning with mutable state, distributed execution, and extensibility for research and production use.

ABSTRACT

TensorFlow is a machine learning system that operates at large scale and in heterogeneous environments. TensorFlow uses dataflow graphs to represent computation, shared state, and the operations that mutate that state. It maps the nodes of a dataflow graph across many machines in a cluster, and within a machine across multiple computational devices, including multicore CPUs, general-purpose GPUs, and custom designed ASICs known as Tensor Processing Units (TPUs). This architecture gives flexibility to the application developer: whereas in previous "parameter server" designs the management of shared state is built into the system, TensorFlow enables developers to experiment with novel optimizations and training algorithms. TensorFlow supports a variety of applications, with particularly strong support for training and inference on deep neural networks. Several Google services use TensorFlow in production, we have released it as an open-source project, and it has become widely used for machine learning research. In this paper, we describe the TensorFlow dataflow model in contrast to existing systems, and demonstrate the compelling performance that TensorFlow achieves for several real-world applications.

Motivation & Objective

• Motivate the need for a scalable ML system capable of training on large datasets and models.
• Introduce a unified dataflow graph model that captures computation and mutable state.
• Demonstrate distributed execution and device placement across CPUs, GPUs, and TPUs.
• Showcase extensibility for research through built-in differentiation, large-model support, and fault tolerance.

Proposed method

• Define a single dataflow graph to represent computation and mutable state.
• Use variables and queues to enable in-graph state and coordination across concurrent executions.
• Implement distributed execution with per-device subgraphs and Send/Recv communication across devices.
• Provide dynamic control flow via Switch and Merge to support non-strict evaluation.
• Allow user-level extensibility for differentiation, optimization algorithms, and large sparse embeddings.
• Incorporate fault-tolerance through in-graph checkpointing and flexible synchronization schemes.

Experimental results

Research questions

• RQ1 How can a single dataflow graph model support both training and inference at large scale?
• RQ2 What mechanisms enable efficient distributed training and device placement across heterogeneous architectures?
• RQ3 How can mutable state and coordination be exposed to researchers to enable new optimization and model architectures?
• RQ4 What strategies (e.g., synchronous vs asynchronous updates, backup workers) improve training throughput and robustness?

Key findings

• TensorFlow enables data-parallel training across large clusters with step times suitable for large-scale learning (as shown in Subsection 6.3).
• The system supports both asynchronous and synchronous replica coordination, with backup workers improving throughput by up to 15%.
• Dynamic control flow and in-graph state allow flexible experimentation with new models and optimization algorithms without modifying the runtime.
• Sparse embedding and large-model handling are facilitated by colocating computation with parameters and by shard-based strategies.
• Checkpointing and fault-tolerance are implemented at the graph level, enabling flexible policy choices and transfer learning workflows.
• The dataflow model provides portability across devices, including CPUs, GPUs, and TPUs, and across production/mobile inference.