[論文レビュー] DiffTaichi: Differentiable Programming for Physical Simulation

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

Model Reduction and Neural Networks参考文献 35被引用数 75

ひとこと要約

DiffTaichi は、高性能な物理シミュレーション用に特化した微分可能なプログラミング言語を導入し、二段階の AD システムと 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.

研究の動機と目的

• ML とロボティクスにおける高性能微分可能物理シミュレータの必要性を動機付ける。
• 算術強度と並列性を保持しつつ、微分可能性を可能にする言語設計を提供する。
• 10 のシミュレータにわたって、二段階システムと軽量テープを通じた自動微分をデモンストレーションする。

提案手法

• Taichi を Python フロントエンド付きに拡張し、Taichi IR にコンパイルし、ソースコード変換を通じてフォワード・カーネルを微分化する。
• 二段階 AD を採用する：カーネル内のローカル AD は SCT を用い、グローバルな AD を軽量カーネル・テープで実現する。
• 命令的でインプレース計算における勾配を適切に定義するため、グローバルなデータアクセス規則を課す。
• メガカーネル融合を用いて算術強度を高め、CPU/GPU への並列ループのマッピングを効率化する。
• ti.complex_kernel デコレータを通じて複雑なカーネルの勾配をカスタマイズするツールを提供する。
• 連続体力学、流体、剛体にわたる微分可能なシミュレータを、パフォーマンスベンチマークとともにデモンストレーションする。

実験結果

リサーチクエスチョン

• RQ1 現代のハードウェアで性能と並列性を維持しつつ、微分可能な物理シミュレータをどのように構築できるか？
• RQ2 二段階 AD のアプローチ（カーネル内 SCT とエンドツーエンドのテープ）は、複雑なシミュレータに対して速度と柔軟性の両方を提供できるか？
• RQ3 インプリティブで高性能なシミュレーションにおける勾配伝播を正しく保証するために、設計パターンと規則（例：グローバルデータアクセス規則）は何が必要か？
• RQ4 DiffTaichi はコードサイズ、速度、スケーラビリティの点で、既存の微分可能プログラミングツールと比較してどうか？

主な発見

Approach	Forward Time	Backward Time	Total Time	# Lines of Code
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 対応の微分可能シミュレータは大幅に生産性が高く、例として微分可能な弾性オブジェクトシミュレータは手動チューニングされた CUDA バージョンより 4.2× 短い。
• 同じシミュレータは手動チューニングされた CUDA バージョンとほぼ同等の速度で動作し、Elastic object の例では TensorFlow より 188× 速い。
• 10 のシミュレータ全体で、DiffTaichi は高性能を実現し、二段階 AD システムを介して効率的に勾配を生成する。
• 軽量テープがカーネル起動を記録し、エンドツーエンドのバックプロパゲーションのためにそれらを反転させることで、重い中間バッファリングを回避する。
• TOI（Time of Impact）ベースの連続衝突処理は、コントローラ最適化の勾配品質を大幅に改善する。
• 代表的なテストにおいて、TensorFlow、Autograd、PyTorch、JAX と比較してコードサイズと性能が改善されることを示す。
• このアプローチにより、典型的なシナリオで学習ベースの制御が数十回の反復で収束する。