GRADSOLVE: GPU上でのODEアンサンブルのための高速な正確な勾配

GRADSOLVE: fast exact gradients for ODE ensembles on GPUs

https://export.arxiv.org/api/query·2026/9/2

AI要約

GPU上でのODEアンサンブルの正確な勾配計算を高速化する。現在のGPUソフトウェアでは、最速のソルバーは微分できず、微分用に構築されたソルバーは遅いというトレードオフがある。この手法は、両方の利点を提供する。

AI要点

GPU上で常微分方程式(ODE)アンサンブルの正確な勾配を高速に計算するJAXライブラリ「GRADSOLVE」を開発した。
従来のGPUソルバーは、速度と微分可能性のトレードオフがあったが、GRADSOLVEは適応的ソルバーのステップを記録し、固定ステップで微分することでこの問題を解決する。
Diffraxと比較して、フォワードパスは2.8倍速く、勾配計算は5.6〜14.1倍高速であることを複数のGPU世代で実証した。
明示的およびRosenbrock積分器をサポートし、既存のDiffraxライブラリと互換性がある。

なぜ重要か

GRADSOLVEは、科学技術計算におけるODEモデルのパラメータ感度解析をGPUで高速化する画期的なライブラリであり、計算コストを大幅に削減し、より複雑なモデルの最適化や感度分析を現実的な時間で可能にします。