论文精选73°

GRADSOLVE:GPU上ODE集合的快速精确梯度计算

GRADSOLVE: fast exact gradients for ODE ensembles on GPUs

精选理由

GRADSOLVE解决了ODE集合在GPU上求解与反向微分速度不匹配的问题,大幅提升计算效率。

AI 摘要

GRADSOLVE是一个开源JAX库,专为在NVIDIA GPU上求解和反向微分低维ODE集合而设计。该库记录自适应求解器接受的步骤,并对固定步长重放进行微分,返回精确的离散伴随梯度。在三种GPU代际上,GRADSOLVE的梯度计算速度比Diffrax的检查点伴随方法快5.6-14.1倍,前向求解速度比DiffEqGPU.jl快2.8倍。

原文 · arXiv cs.LG

GRADSOLVE: fast exact gradients for ODE ensembles on GPUs

Ordinary differential equations (ODEs) underlie models in science and engineering, and many applications need derivatives of their solutions with respect to parameters. Ensembles of independent trajectories suit graphics processing units (GPUs), but current GPU software forces a trade-off: the fastest ensemble solvers cannot be differentiated in reverse mode at the speed they solve, and the solvers built for differentiation solve more slowly. No single tool has yet offered a reverse-mode gradient at the speed of a fused-kernel solve. We present GRADSOLVE, an open-source JAX library for solving and reverse-mode differentiating low-dimensional ODE ensembles on NVIDIA GPUs. It records the steps an adaptive solver accepts and differentiates a fixed-step replay of them; the returned gradient is the exact discrete adjoint of those steps, the same derivative Diffrax returns by default, obtained more cheaply from a fixed-length chain than from an adaptive loop. It targets ensembles differentiated many times against one recorded mesh, keeps Diffrax as a fallback, and supports explicit and Rosenbrock integrators. Used as a solver, GRADSOLVE's forward-only kernel ran 2.8x faster than DiffEqGPU.jl; used for gradients, once a record exists, it computed them 5.6-14.1x faster than Diffrax's checkpointed adjoint at matched forward-state accuracy across three GPU generations, the advantage narrowing on large ensembles and, on stiff systems, down to parity at tight accuracy. GRADSOLVE is released at https://github.com/ECLIPSE-AI4Science/gradsolve.