[论文解读] torchode: A Parallel ODE Solver for PyTorch

Marten Lienen, Stephan Günnemann|arXiv (Cornell University)|Oct 22, 2022

Parallel Computing and Optimization Techniques被引用 6

一句话总结

torchode 是一个高性能、并行的 PyTorch ODE 求解器，可独立地对批量中的多个 ODE 进行积分，每个 ODE 拥有独立的求解器状态，从而实现每步最多 4.3 倍的加速，并有效缓解批量引起的步长膨胀问题。它支持 JIT 编译、可扩展的步长控制（包括 PID 控制器），并可收集详细的求解器统计信息，便于研究扩展与内部观察。

ABSTRACT

We introduce an ODE solver for the PyTorch ecosystem that can solve multiple ODEs in parallel independently from each other while achieving significant performance gains. Our implementation tracks each ODE's progress separately and is carefully optimized for GPUs and compatibility with PyTorch's JIT compiler. Its design lets researchers easily augment any aspect of the solver and collect and analyze internal solver statistics. In our experiments, our implementation is up to 4.3 times faster per step than other ODE solvers and it is robust against within-batch interactions that lead other solvers to take up to 4 times as many steps. Code available at https://github.com/martenlienen/torchode

研究动机与目标

• 解决 PyTorch ODE 求解器与 JAX 和 Julia 等其他框架相比存在的性能差距。
• 消除批量 ODE 之间因非预期交互而导致的步长膨胀问题，避免训练效率下降。
• 使研究人员能够轻松扩展求解器行为并收集内部统计信息，用于模型分析。
• 通过批量并行、状态隔离的积分方式优化 GPU 性能，支持高级步长控制器。
• 提供一个生产就绪、可扩展的 ODE 求解器，支持连续归一化流（CNF）及其他模型的独立与联合伴随反向传播。

提出的方法

• torchode 将批量中的每个 ODE 视为独立问题，拥有独立的初始条件、积分区间、步长和接受状态。
• 它为每个 ODE 实例维护独立的求解器状态，包括步长、误差估计以及接受/拒绝决策，避免跨批量干扰。
• 求解器采用基于控制理论的 PID 控制器进行自适应步长选择，提升对刚性变化的响应能力。
• 通过 PyTorch 的 JIT 编译器支持 JIT 编译，适用于性能关键的推理与训练场景。
• 实现支持组件插拔替换，如步长控制器和积分器，便于扩展自定义学习或分析逻辑。
• 为 CNF 模型提供独立的联合伴随反向传播，利用霍纳法则（Horner’s rule）和融合核函数，显著降低反向传播计算量。

实验结果

研究问题

• RQ1 在 PyTorch 中采用并行独立 ODE 求解是否能消除批量引起的步长膨胀，从而提升训练效率？
• RQ2 基于 PID 的步长控制器在求解刚性 ODE 时，与积分控制器相比是否能更有效地减少求解步数？
• RQ3 独立 ODE 求解在连续归一化流和时间序列建模中对模型性能的影响有多大？
• RQ4 原生 PyTorch ODE 求解器能否实现与 JAX 基础求解器（如 diffrax）相当的性能？
• RQ5 求解器可观察性与统计信息收集对模型调试与超参数调优有何影响？

主要发现

• torchode 相较于现有 PyTorch ODE 求解器，每步执行速度最高可提升 4.3 倍，尤其在刚性问题上表现突出。
• 该求解器避免了其他求解器因批量交互导致的 4 倍步长膨胀问题，在不同积分区间下保持一致的性能表现。
• 在高度刚性问题（如范德波尔振子中 μ=25）中，PID 控制使求解步数减少 3–5%；但在平滑问题上无明显优势。
• torchode 的联合伴随反向传播相比 torchdiffeq 和 TorchDyn 更快，得益于优化的核融合与霍纳法则的应用。
• 独立 ODE 求解不会降低模型性能——bits/dim 和 MAE 指标与联合求解器相当，表明对学习无负面影响。
• 详细求解器统计信息与可扩展性设计，使研究人员无需全局状态或代码注入即可监控与分析求解器行为。