On-the-fly Operation Batching in Dynamic Computation Graphs
Dynamic neural network toolkits such as PyTorch, DyNet, and Chainer offer more flexibility for implementing models that cope with data of varying dimensions and structure, relative to toolkits that operate on statically declared computations (e.g., TensorFlow, CNTK, and Theano). However, existing toolkits-both static and dynamic-require that the developer organize the computations into the batches necessary for exploiting high-performance algorithms and hardware. This batching task is generally difficult, but it becomes a major hurdle as architectures become complex. In this paper, we present an algorithm, and its implementation in the DyNet toolkit, for automatically batching operations. Developers simply write minibatch computations as aggregations of single instance computations, and the batching algorithm seamlessly executes them, on the fly, using computationally efficient batched operations. On a variety of tasks, we obtain throughput similar to that obtained with manual batches, as well as comparable speedups over singleinstance learning on architectures that are impractical to batch manually.
Summary:
The authors of this paper extend neural network toolkit DyNet with automatic operation batching. Batching enables efficient utilization of CPUs and GPUs by turning matrix-vector products into matrix-matrix products and reducing kernel launch overhead (for GPUs) but it is commonly done manually. Manual batching is manageable for simple feed-forward-networks but it becomes increasingly a headache as we explore more flexible models that take variable-length input, tree-structured input, or networks that perform dynamic control decisions.
Chainer, DyNet, and PyTorch are recently proposed neural network toolkits that allow user to dynamically define the computation graph using the syntax of the host language (if, while, etc in python). This is desirable as it avoids tookit specific constructions (e.g., cond in TensorFlow) and make the network definition intuitive but it tends to limit performance because the network construction and computation happens at the same time. Thus although it would be straightforward to program a control flow that supports variable length or tree-structured input in these toolkits, in practice it would be too inefficient to process single instance at a time.
The  key contribution of this paper is to delay and automatically batch the computation so that a user can still define an arbitrary complex control flow using the host language as in original DyNet without worrying about operation batching. The approach is similar to TensorFlow Fold (TFF) but differs in two places: first, the computation graph is defined in a dynamic manner using the control flow of the host language and the user does not need to learn any new language (note that TFF introduces many new operations that are not in TensorFlow); second, it employs an agenda-based approach for operation batching in contrast to the depth-based approach employed in TFF. The authors empirically show that agenda-based batching is slightly more efficient than depth-based batching.
Detailed comments:
It wasn't clear to me if there is a stronger argument for dynamic definition of the computation graph other than the simplicity of coding compared to static definition of the graph in TFF, for example.
Would the lazy evaluation still work when the control decision depends on the computed value (this is the case for the training with exploration scenario [2, 4, 9]. In this case, static graph definition may have an advantage.