Dynamic Automatic Differentiation of GPU Broadcast Kernels

Abstract

We show how forward-mode automatic differentiation (AD) can be employed within larger reverse-mode computations to dynamically differentiate broadcast operations in a GPU-friendly manner. Our technique fully exploits the broadcast Jacobian's inherent sparsity structure, and unlike a pure reverse-mode approach, this "mixed-mode" approach does not require a backwards pass over the broadcasted operation's subgraph, obviating the need for several reverse-mode-specific programmability restrictions on user-authored broadcast operations. Most notably, this approach allows broadcast fusion in primal code despite the presence of data-dependent control flow. We discuss an experiment in which a Julia implementation of our technique outperformed pure reverse-mode TensorFlow and Julia implementations for differentiating through broadcast operations within an HM-LSTM cell update calculation.

References in corpus (8)

• TensorFlow: A system for large-scale machine learning
• Gradient-based Hyperparameter Optimization through Reversible Learning
• Hierarchical Multiscale Recurrent Neural Networks
• Don't Unroll Adjoint: Differentiating SSA-Form Programs
• Confusion of Tagged Perturbations in Forward Automatic Differentiation of Higher-Order Functions
• Tangent: Automatic Differentiation Using Source Code Transformation in Python
• Demystifying Differentiable Programming: Shift/Reset the Penultimate Backpropagator
• Automatic Full Compilation of Julia Programs and ML Models to Cloud TPUs