一句话定位
DiffTaichi 是一门专为「构建高性能可微物理仿真器」定制的可微编程语言（嵌入 Python、基于 taichi 编译器），用一套两尺度自动微分系统（kernel 内做源码变换 SCT + kernel 间用轻量 tape 反向重放）让梯度端到端穿过接触、流体、弹性等物理过程；作者用它在 3 小时内各自实现了 10 个可微仿真器，可微弹性体仿真器比 TensorFlow 实现快 188×、代码短 4.2×（vs 手写 CUDA）——由此把「用暴力梯度下降优化神经网络控制器」变得比 model-free RL 快 1~4 个数量级（ICLR 2020）。
背景与定位
这是一篇仿真基础设施 / 工具论文，不是模型论文。它要解决的痛点是：可微物理仿真器已被证明是机器学习系统里高效的组件——de Avila Belbute-Peres 等 (2018) 和 ChainQueen (Hu et al. 2019b) 都显示，用可微仿真器做控制器优化，比 model-free 强化学习收敛快 1~4 个数量级；但用现有工具实现一个高性能可微仿真器很难。
问题的根源在于「深度学习工作负载」与「物理仿真工作负载」的错配：TensorFlow / PyTorch 这类函数式数组编程框架围绕大 data blob（如 AlexNet 第二卷积层 27×27×128×128）设计，提供的是粗粒度算子「原子」，用户只能把它们拼成复杂物理算子，导致大量临时 buffer 与过多 GPU kernel 启动、缺乏生产者-消费者局部性——论文引 ChainQueen 数据，纯 TensorFlow 实现比 CUDA 慢 132×。物理仿真真正需要的是三样现有可微框架普遍缺失的能力：
- Megakernel（大核）：把多阶段计算融进单个 kernel，提高算术强度（arithmetic intensity，每字节内存访问对应的 FLOP 数），中间结果缓存在寄存器而非主存。
- 命令式并行编程：传统物理仿真是 Fortran/C++ 命令式风格（并行 for 循环 + if 控制流处理碰撞/边界条件/迭代求解器），而现代 DL 是函数式数组风格；命令式让移植现有仿真代码更容易。
- 灵活索引：粒子-网格交互、数值 stencil（如
y[p[i]*2, j] = x[q[i+j]]）不是逐元素操作，在数组框架里只能用低效难维护的 scatter/gather 表达；DiffTaichi 允许直接按任意索引读写全局数组元素。
范式上它接续两条线：一是可微数组编程语言（Halide、Autograd、JAX、Enoki），DiffTaichi 用命令式 + 灵活索引替代不可变数组以方便移植仿真算法；二是可微物理仿真器（Degrave 2016、de Avila Belbute-Peres 2018、chainqueen、Liang 2019 的可微布料）——这些此前多靠手写 CUDA 梯度或用 Theano/PyTorch 拼，DiffTaichi 把「自动生成梯度仿真器」变成语言原语。它建立在 SIGGRAPH Asia 2019 的 taichi 语言之上（Taichi 的关键设计是计算与数据结构解耦）。
模型架构
本节即「系统 / 编译器架构」。DiffTaichi 前端嵌入 Python，用 Python AST transformer 把代码编译成 Taichi 中间表示（IR），再扩展 Taichi 编译器把 IR 自动微分成前向与反向可执行文件。核心是两尺度自动微分系统——SCT 与 tracing 两种经典 AD 方案各有短板（对上千时间步整体用 SCT 性能高但灵活性差、编译慢；纯 tracing 灵活但会破坏 megakernel 结构导致性能差），DiffTaichi 取两者之长：
- Local AD（kernel 内，源码变换 SCT）：对单个 Taichi kernel 做反向模式 AD。先经两步 IR 预处理简化 loop body——① Flatten Branching：把
if替换成三元select(cond, a, b)（梯度定义清晰，是程序向量化里的常见变换）；② Eliminate Mutable Local Variables：做 store-forwarding，把可变局部变量消成真正的单赋值（SSA）。之后对直线代码跑标准 reverse-mode SCT（make_adjointpass）：为每条 SSA 指令分配 adjoint 局部变量，逆序遍历累加梯度，梯度贡献通过 atomic add 累加到 adjoint 张量。并行 for 循环结构在 AD 中保留（每次迭代映射到 CPU/GPU 线程）；非并行 loop 在 AD 时逆转循环顺序。不支持携带可变局部变量的 loop（否则需要昂贵的运行时栈维护历史），要求用户改用满足全局数据访问规则的全局变量。 - Global AD（kernel 间，轻量 tape）：构造一条 tape 记录 kernel 执行，反向时逆序重放梯度 kernel。tape 极轻量——因为中间结果都存在全局张量里，前向时 tape 只记录 kernel 函数指针与（标量）输入参数，不像其他函数式数组系统要把所有中间 buffer 都记进 tape。求梯度时逆序遍历 tape、用记录的参数调用梯度 kernel。全局张量天然是梯度求值的 checkpoint。
全局数据访问规则（让命令式下的 AD 良定义的两条假设）：① 若某全局张量元素被写多于一次，从第二次写起必须是 atomic add（累加）形式；② 在累加完成前，不得对该张量元素发生读访问。为满足规则，前向仿真里要记录 x、v 的完整历史而非只留最新值——由此带来的内存开销用 checkpointing 缓解。
- Complex kernel（覆写默认 AD）：提供
ti.complex_kernel/ti.complex_kernel_grad两个装饰器，让用户手写梯度覆盖编译器默认 AD——典型用途是 3×3 矩阵 SVD 的迭代求解器（直接微分迭代过程数值不稳，改用手工推导的 SVD 导数子程序，如 Jiang 2015 §2.1.1.2），以及实现 checkpointing。 - Checkpointing（省内存）：① 时间步内重算——backward 时重做 p2g、grid_op 重算网格态，diffmpm 里只需分配一份网格拷贝而非 O(n) 份；② 分段重算——把仿真切成 S 步的段，前向只存每段第一个状态，反向段内重算，内存从 O(n) 降到 O(S + n/S)，取 S=O(√n) 则降到 O(√n)，时间复杂度仍 O(n)。
典型仿真规模：512~2048 时间步，每步最多约 1000 个并行操作。
数据
可微编程语言没有「训练数据」；这一维对应它验证所用的 10 个可微仿真器 / 工作负载，覆盖刚体、可变形体、流体三大类（论文 Appendix E）：
| 仿真器 | 物理模型 | 规模（分辨率 / 粒子 / 时间步） |
|---|---|---|
diffmpm 可微连续介质弹性体 | MLS-MPM（移动最小二乘物质点法，动量+质量守恒） | 2D 6.4K 粒子（benchmark）；2D 机器人 4 肌肉；3D 机器人 16 肌肉 30K 粒子 |
liquid 可微 3D 液体 | 弱可压缩流体（Tampubolon 2017），可与弹性体双向耦合 | 机器人 30K 粒子 + 液体 13K 粒子 |
smoke 可微不可压缩流体 | 半拉格朗日平流（Stam 1999）+ 隐式压力投影 | 110×110 网格，100 时间步，每步 6 次 Jacobi 压力投影（Table 2；Appendix G 讨论反传时另记为 10 次 Jacobi 迭代） |
wave 可微高度场浅水 | 波动方程 FDTD 离散（Wang 2018） | 128×128 网格，256 时间步，200 步梯度下降 |
mass_spring 可微弹簧质点 | Hooke 定律 + 半隐式 Euler + 地面碰撞 + TOI | 2048 时间步（README 载 682 步），3 个机器人，2 层 NN 控制器 |
rigid_body 可微刚体 | 冲量法（Catto 2009），支持碰撞/摩擦/弹簧/关节/驱动 | 2048 时间步 |
billiards 可微台球 | 前向 Euler + 动量与动能守恒解碰撞 | — |
water_renderer 可微（对抗）水面渲染 | 有限差分重建水面 + 光线折射 + 双线性插值 | 三级流水线（仿真→渲染→VGG16 识别）端到端可微 |
volume_renderer 可微体渲染 | ray marching 积分密度场 | 7 张目标视角图，重建密度场 |
electric 可微电场 | Coulomb 定律，8 电极 + 带电球 | 2 层 NN 控制器 |
数据来源说明：diffmpm 直接把 ChainQueen 的原始 CUDA 仿真器翻译成 DiffTaichi 语法；smoke 跟随 Autograd 库里的示例基线（用 numpy.roll 做周期边界以便 Autograd 能表达）；所有结果都能用仓库提供的单个 Python 脚本复现。
训练方法
这一维对应基于梯度的优化 / NN 控制器学习流程，而非模型预训练：
- 优化目标：为每个任务定义标量 loss（如 mass_spring 最大化前进距离、diffmpm 让软体机器人前进、wave 让高度场演化成目标 “Taichi” 图案的 L2 loss、water_renderer 用对抗目标骗 VGG16 把松鼠认成金鱼）。
- 优化循环：用
ti.Tape(loss)记录前向 kernel 启动 → 反向自动重放梯度 → 对参数（弹簧静止长度 / NN 控制器权重 / 初始速度场 / 初始高度场）做梯度下降。神经网络控制器通常在几十次迭代内优化完成（rigid_body 机器人 20 步梯度下降学会走路；spring rest-length 优化约 20 步收敛）。 - 关键工程 trick——Time of Impact (TOI) 修正梯度：这是全文最重要的方法学发现。刚体碰撞里，前向仿真即便用很小 ∆t 也能得到合理结果，但朴素时间积分器给出的梯度完全错误——弹性碰撞后小球该以 −1 的梯度响应初始高度变化，朴素积分器无论 ∆t 多小都给出 +1（因为碰撞检测只在 ∆t 整数倍发生，时间离散化本身没被编译器微分，loss 呈锯齿状）。解法是在前向程序里加连续碰撞检测（考虑精确碰撞时刻 TOI，把一个 ∆t 拆成 old_v 段 + new_v 段）：它几乎不改变前向结果，却能有效修正梯度，在控制器优化任务里显著提升梯度质量与优化效果（每个实验重复 5 次验证）。作者的 takeaway：即便仿真器前向做得对，微分它也不一定给出物理系统的有用梯度。
- 其他梯度陷阱（Appendix G）：初始化很关键（billiards 里大部分初始角度落在平坦区、梯度为零，还有局部极小）；理想物理模型自带不连续/奇异（刚体+摩擦时哪个角先着地会造成不连续；electric / mass_spring 的 1/r²、1/r 项在 r→0 时梯度数值不准，需给 r 加安全下界）。刚体碰撞在可数个不连续点不可微，但类比 ReLU 在 x=0 不可微仍广泛可用——除这些点外几乎处处可微，配合 TOI 梯度仍对优化有用。
- smoke 里穿过压力投影反传：跟随 Autograd 基线用 10 次 Jacobi 迭代做压力投影（不足以完全无散度但够用），成功地对展开的 10 次 Jacobi 迭代反传；更大规模下（如 MGPCG 5 层多重网格 + 50 次共轭梯度）AD 精度会不够，此时用 complex kernel 手工实现反向 Poisson solve（adjoint method）。
Infra（训练 / 推理工程）
- 硬件：diffmpm benchmark（Table 1）明示在单张 NVIDIA GTX 1080 Ti GPU 上；smoke（Table 2）仅写”在 GPU 上跑”、未指明型号。支持 CPU 与 GPU 双后端。
- 精度：主用 float32（Autograd 基线用 float64，运行时间约翻倍）。
- 编译 / JIT：smoke 例子整个程序在 GPU 上跑 10 秒，其中 2 秒花在 JIT；作为对比 JAX 的 JIT 编译要 2 分钟。
- 性能数字（diffmpm，2D 6.4K 粒子，Table 1）：
| 方案 | 前向 | 反向 | 总时间 | 有效代码行 |
|---|---|---|---|---|
| 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×、跑得几乎一样快（0.92× vs 1.00×）；比 TensorFlow 短 1.7×、快 188×。
- 性能数字（smoke，110×110 网格 100 时间步，Table 2）：
| 方案 | 前向 | 反向 | 总时间 | 有效代码行 |
|---|---|---|---|---|
| PyTorch (CPU, f32) | 405 ms | 328 ms | 733 ms（13.8×） | 74 |
| PyTorch (GPU, f32) | 254 ms | 457 ms | 711 ms（13.4×） | 74 |
| Autograd (CPU, f64) | 307 ms | 1197 ms | 1504 ms（28.4×） | 51 |
| JAX (GPU, f32) | 24 ms | 75 ms | 99 ms（1.9×） | 90 |
| DiffTaichi (CPU, f32) | 66 ms | 132 ms | 198 ms（3.7×） | 75 |
| DiffTaichi (GPU, f32) | 24 ms | 29 ms | 53 ms（1.0×） | 75 |
- 推理 / 控制频率、延迟、边端硬件：未披露（本文关注仿真-优化吞吐，非部署时的实时控制频率）。README 补充：多数示例不需要 GPU 即可运行，10 个仿真器多数可在 2~3 小时内实现。
评测 benchmark
论文没有 ML 模型意义上的 accuracy benchmark，评测即上节的性能与生产力对比（Table 1 diffmpm、Table 2 smoke），命名基线为 TensorFlow、手写 CUDA、PyTorch、Autograd、JAX。核心量化结论：
- diffmpm：vs TensorFlow 快 188×、短 1.7×；vs CUDA 短 4.2×、速度 0.92×~1.00× 持平。
- smoke：DiffTaichi (GPU) 比 PyTorch 快 13.4~13.8×、比 Autograd (f64) 快 28.4×、比 JAX (GPU) 快 1.9×；且 JIT 2 秒 vs JAX 2 分钟。
- 定性对比表（Table 3）：DiffTaichi 在「GPU Megakernels / 命令式 / 并行 / 灵活索引」四项全支持，PyTorch、TensorFlow、Enoki、JAX、Halide、Julia、Swift 各有缺失（如 PyTorch/TensorFlow 在 megakernel 与命令式上只是部分支持 ∆）。
- water_renderer：三级流水线端到端优化后，让 VGG16 把折射后的松鼠图以 99.91% 置信度认成金鱼，证明能与其他可微编程系统（VGG16）无缝拼接。
创新点与影响
- 核心贡献：把「自动生成高性能梯度仿真器」变成一门可微编程语言的原语。为此针对物理仿真的三大需求（megakernel、命令式并行、灵活索引）设计了定制的两尺度自动微分系统（kernel 内 SCT + kernel 间轻量 tape），既保性能又保灵活性；配套 complex kernel 覆写机制与 O(√n) checkpointing。
- 改变了什么：让研究者能用暴力梯度下降 + 可微仿真器在几十次迭代内优化 NN 控制器，替代样本效率低的 model-free RL（快 1
4 个数量级）；大幅降低可微物理仿真在 ML/机器人社区的实现门槛（10 个仿真器各约 23 小时、单脚本复现）。DiffTaichi 后来正式并入 taichi 主项目，成为可微仿真基础设施的代表工作，为后续 gradient-based 具身控制、可微软体机器人、sim2real 提供了工具底座。 - 重要方法学洞见（TOI）：揭示了「前向仿真正确 ≠ 微分它得到有用梯度」——时间离散化本身不被编译器微分会给出方向相反的错误梯度，须在前向程序里引入连续碰撞检测（精确碰撞时刻）来修正。这是可微仿真领域被反复引用的经验。
- 作者自陈的局限：① 刚体碰撞在可数个不连续点不可微、甚至不连续（比 ReLU 更复杂），梯度只是「几乎处处」有用；② 长迭代求解器（多重网格 + 大量共轭梯度）里 AD 数值精度不够，需手工写 adjoint；③ 记录完整状态历史带来 O(n) 内存压力，靠 checkpointing 缓解；④ 物理过程存在平坦区/局部极小/奇异（1/r 项），对初始化与数值安全下界敏感。
原始链接
- arXiv（含全部 Appendix，ICLR 2020 版 v3）：https://arxiv.org/abs/1910.00935
- PDF：https://arxiv.org/pdf/1910.00935
- GitHub（示例代码，DiffTaichi 本体已并入 Taichi）：https://github.com/taichi-dev/difftaichi
- 复现视频（含每个 demo 的复现说明）：https://www.youtube.com/watch?v=Z1xvAZve9aE
- 基础语言 Taichi：https://github.com/taichi-dev/taichi
一手源存档（sources/）
- difftaichi—github-readme — GitHub README 快照（
sources/embodied/2019/difftaichi--github-readme.md） - arXiv 原文 PDF（不入 git）：见上方 arXiv PDF 链接