Recorded-Mesh Neural ODE Backpropagation
Implementation & benchmark of arXiv:2609.02876 — GRADSOLVE: fast exact gradients for ODE ensembles on GPUs
Source paper: GRADSOLVE: fast exact gradients for ODE ensembles on GPUs arXiv:2609.02876 ⓘ · analyzed Sep 3, 2026
AI-generated research hypothesis, automatically tested. Not peer-reviewed.
Idea description
Run an adaptive neural ODE solver once to determine accepted step sizes, then train using a regular fixed-length replay of those steps rather than differentiating through adaptive accept/reject logic. The replay can be fused across a batch of trajectories and differentiated with an ordinary reverse sweep, giving the exact discrete gradient of the replayed solver and predictable GPU work.
Formulas
Mathematical statement
The adaptive solver uses an embedded Runge–Kutta pair. For state y at time t, step size h, parameters theta, and s stages, it computes k_j = f(t+c_j h, y+h sum_{l<j} a_{jl} k_l, theta), with fixed coefficients a_{jl}, c_j, followed by y_new = y+h sum_{j=1}^s b_j k_j. The embedded estimate is \hat{y}_{new}=y+h\sum_{j=1}^s \hat{b}_j k_j and the local error estimate is e=y_new-\hat{y}_{new}; the adaptive solver accepts or rejects h using a normalized RMS error based on e/(tau_abs+tau_rel max(|y|,|y_new|)). After accepted steps are recorded, trajectory q is represented as a fixed chain y_i=\Phi_i(y_{i-1};h_i,theta), i=1,...,S, where \Phi_i is the chosen one-step map and h_i is the recorded accepted step. For a scalar loss L(y_S), the exact discrete adjoint is \lambda_S=
abla_{y_S}L and \lambda_{i-1}=(\partial_{y_{i-1}}\Phi_i)^T\lambda_i; parameter gradients accumulate as
abla_ heta L=\sum_{i=1}^S(\partial_ heta\Phi_i)^T\lambda_i. The fixed chain has no adaptive branches, so reverse-mode autodiff computes the exact derivative of the numerical map being replayed.
Implementation notes
1. Integration point: replace the differentiable solver inside a neural ODE or probability-flow diffusion training loop. The vector field is f_theta(t,y,c), where c is conditioning and y contains a batch or ensemble of initial states. In a mesh-recording pass, run an embedded adaptive RK solver with tolerances tau_abs and tau_rel, store each trajectory's accepted step sizes h_i, and stop gradients through this pass. Pad trajectories to S_max and attach an active mask; padded steps use h=0 or an identity map. 2. Replay pseudocode: record = stop_gradient(adaptive_solve(f_theta,y0,t0,t1)); h,mask = pad(record.steps); y=y0; t=t0; for i in range(S_max): y_old=y; y=rk_step(f_theta,y,t,h[:,i],theta); y=where(mask[:,i],y,y_old); t=t+h[:,i]; loss=objective(y,target); backward(loss). Implement rk_step with the displayed stage equations and fixed coefficients, using JAX lax.scan, XLA fusion, or a custom CUDA kernel. 3. The mathematics determines the accepted-step record, fixed RK map, and discrete-adjoint semantics. Empirically monitor replay-versus-adaptive state error, gradient cosine similarity, and trajectories exceeding S_max. Refresh the mesh every K updates or when a validation adaptive solve reports RMS error above twice the target tolerance. 4. First experiment: train a 2-layer MLP neural ODE on a 2D spiral or downsampled MNIST, comparing Diffrax adaptive backpropagation, fixed-step RK4, and recorded-mesh replay at matched forward error and batch size. Measure wall-clock time per update, peak memory, loss at equal GPU time, terminal-state error, and gradient cosine similarity. Success means at least 2x lower gradient-step time or lower memory at equal final accuracy without unstable gradients.
Verification
This idea has not been verified yet.
Verification happens in two stages: Stage 1 — a mechanism check on a toy system confirms the claimed mathematical phenomenon reproduces; Stage 2 — a benchmark implements the idea on a real (small) neural network task and compares it against a tuned baseline over 8 paired seeds with a permutation test.
Artifacts
Artifacts unavailable.