[Paper Review] DiffTaichi: Differentiable Programming for Physical Simulation

Yuanming Hu, Luke Anderson+5 more|arXiv (Cornell University)|Oct 1, 2019

Model Reduction and Neural Networks35 references75 citations

TL;DR

DiffTaichi introduces a differentiable programming language tailored for high-performance physical simulation, enabling end-to-end gradient computation with a two-scale AD system and megakernel fusion.

ABSTRACT

We present DiffTaichi, a new differentiable programming language tailored for building high-performance differentiable physical simulators. Based on an imperative programming language, DiffTaichi generates gradients of simulation steps using source code transformations that preserve arithmetic intensity and parallelism. A light-weight tape is used to record the whole simulation program structure and replay the gradient kernels in a reversed order, for end-to-end backpropagation. We demonstrate the performance and productivity of our language in gradient-based learning and optimization tasks on 10 different physical simulators. For example, a differentiable elastic object simulator written in our language is 4.2x shorter than the hand-engineered CUDA version yet runs as fast, and is 188x faster than the TensorFlow implementation. Using our differentiable programs, neural network controllers are typically optimized within only tens of iterations.

Motivation & Objective

• Motivate the need for high-performance differentiable physical simulators in ML and robotics.
• Provide a language design that preserves arithmetic intensity and parallelism while enabling differentiability.
• Demonstrate automatic differentiation through a two-scale system and light-weight tape across 10 simulators.

Proposed method

• Extend Taichi with a Python frontend that compiles to a Taichi IR and differentiates forward kernels via source code transformations.
• Adopt a two-scale AD: local AD within kernels using SCT and global AD via a light-weight kernel tape.
• Impose global data access rules to ensure well-defined gradients in imperative, in-place computations.
• Use megakernel fusion to increase arithmetic intensity and map parallel loops efficiently to CPU/GPU.
• Provide tools to customize gradients for complex kernels via ti.complex_kernel decorators.
• Demonstrate differentiable simulators across continuum mechanics, fluids, and rigid bodies with performance benchmarks.

Experimental results

Research questions

• RQ1 How can differentiable physical simulators be built to preserve performance and parallelism on modern hardware?
• RQ2 Can a two-scale AD approach (within-kernel SCT plus end-to-end tape) provide both speed and flexibility for complex simulators?
• RQ3 What design patterns and rules (e.g., global data access rules) are needed to ensure correct gradient propagation in imperative, high-performance simulations?
• RQ4 How does DiffTaichi compare to existing differentiable programming tools in terms of code size, speed, and scalability across various simulators?

Key findings

• DiffTaichi-enabled differentiable simulators are substantially more productive, e.g., a differentiable elastic object simulator is 4.2× shorter than a hand-tuned CUDA version.
• The same simulator runs as fast as the hand-tuned CUDA version and 188× faster than TensorFlow for the elastic object example.
• Across 10 simulators, DiffTaichi achieves high performance, with gradients generated efficiently via the two-scale AD system.
• A lightweight tape records kernel launches and reverses them for end-to-end backpropagation, avoiding heavy intermediate buffering.
• TOI (time of impact) based continuous collision handling significantly improves gradient quality for controller optimization.
• Comparisons show improvements in code size and performance versus TensorFlow, Autograd, PyTorch, and JAX in representative tests.
• The approach enables learning-based control to converge within tens of iterations in typical scenarios.