[논문 리뷰] DiffTaichi: Differentiable Programming for Physical Simulation

Yuanming Hu, Luke Anderson외 5명|arXiv (Cornell University)|2019. 10. 01.

Model Reduction and Neural Networks참고 문헌 35인용 수 75

한 줄 요약

DiffTaichi는 고성능 물리 시뮬레이션을 위한 미분 가능 프로그래밍 언어를 도입하여 이원 스케일의 AD 시스템과 메가 커널 융합으로 엔드-투-엔드 그래디언트 계산을 가능하게 한다.

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.

연구 동기 및 목표

• ML과 로봇공학에서 고성능의 미분 가능 물리 시뮬레이터의 필요성을 동기부여한다.
• 미분 가능성을 가능하게 하면서 산술 강도와 병렬성을 보존하는 언어 설계를 제공한다.
• 10개의 시뮬레이터에 걸친 이중 스케일 시스템과 경량 테이프를 통해 자동 미분을 시연한다.

제안 방법

• Taichi를 Python 프런트엔드와 함께 확장하여 Taichi IR로 컴파일하고 소스 코드 변환을 통해 순방향 커널을 미분한다.
• SCT를 사용한 커널 내 로컬 AD와 경량 커널 테이프를 통한 전역 AD의 이중 스케일 AD를 채택한다.
• 명령형 대입 연산에서 잘 정의된 그래디언트를 보장하기 위해 글로벌 데이터 접근 규칙을 부과한다.
• 메가커널 융합을 사용하여 산술 강도를 높이고 병렬 루프를 CPU/GPU에 효율적으로 매핑한다.
• ti.complex_kernel 데코레이터를 통해 복잡한 커널에 대한 그래디언트를 커스터마이즈하는 도구를 제공한다.
• 연속체 역학, 유체 역학, 강체를 아우르는 미분 가능 시뮬레이터를 성능 벤치마크와 함께 시연한다.

실험 결과

연구 질문

• 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?

주요 결과

접근 방식	정방향 시간	역방향 시간	총 시간	코드 줄 수
TensorFlow	13.20 ms	35.70 ms	48.90 ms (188×)	190
CUDA	0.10 ms	0.14 ms	0.24 ms (0.92×)	460
DiffTaichi	0.11 ms	0.15 ms	0.26 ms (1.00×)	110

• 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.