Differentiate the Solver, Not the Equation: Reverse-Sweep Adjoints for Block Implicit Simulation

TL;DR: When a physics simulator runs a fixed, small number of solver sweeps per timestep — as every real-time simulator does — the standard "implicit function theorem" gradient is the derivative of a program that never ran, and this paper measures the damage: 37% gradient error after one sweep. The fix is a shift in what you differentiate: not the converged equation, not an unrolled autodiff tape, but the executed solver itself. Because block-coordinate solvers like Vertex Block Descent are ordered compositions of tiny local updates, their exact adjoint is the same sweep run backwards, made entirely of local 3×3 solves — no global matrix ever exists. Same codebase, same GPU: 33× faster and 71× less memory than unrolled autodiff, exact to machine precision at every solver depth, and differentiable elastodynamics scaled to 8M vertices with a million contact-coupled bodies on one H100.

Two ways to differentiate an implicit step, and why both hurt

An implicit timestep in elastodynamics solves a stationary problem

where \bm{x} stacks the 3|V| vertex positions, M is the mass matrix, h the timestep, \bm{y} the inertial prediction, and E the elastic energy. To use simulation as a layer — system identification, control, training with a simulator in the loop — you need gradients through this solve.

Route one is unrolled automatic differentiation: record every operation the solver executed and backprop through the tape. It's exact by construction (the paper uses it as ground truth), but the tape grows with solver depth × rollout length × mesh size, and it's the first thing to run out of memory. Route two is equation-level implicit differentiation: throw the solver away, assert \nabla G(\bm{x}^\star)=\bm{0}, apply the implicit function theorem, and solve a global adjoint system against \nabla^2 G(\bm{x}^\star). Memory becomes O(1) in solver depth, but you've reintroduced exactly the global sparse linear algebra that modern GPU solvers were designed to avoid — and, more subtly, you're differentiating the converged solve.

That subtlety is where the paper plants its flag. Fast simulators don't run to convergence. Vertex Block Descent (VBD), the solver this paper builds on, sweeps over vertices one at a time, doing a single 3×3 Newton solve per vertex, with a graph coloring so non-adjacent vertices update in parallel. Each local update monotonically decreases G, so the method is unconditionally stable even when you cap the sweep count K at something tiny — and capping it tiny is the whole point, because it's fast. So the state after K=1 or K=5 sweeps satisfies \nabla G = \bm{0} not even approximately: the paper measures a residual of 1.6\times 10^{-1} after one sweep. The implicit-function-theorem gradient is the derivative of a fixed point the solver never reached. The forward pass computed a specific finite function of its inputs; the equation-level adjoint differentiates a different function that happens to agree in the K\to\infty limit.

The aha: the adjoint of a sweep is the sweep, backwards

Here's the structural observation that makes everything work. Write the executed solver as a composition of block updates, S = S_n \circ S_{n-1} \circ \cdots \circ S_1, where each S_i modifies only vertex i's three coordinates. By the chain rule, the Jacobian of the sweep is an ordered product, so its transpose — the thing a backward pass applies — is the product in reverse:

Two things fall out immediately. Reverse ordering isn't a design choice; it's what "transpose of a composition" means computationally. And because each \partial S_i touches only one block, so does each (\partial S_i)^\top. The exact adjoint of the executed sweep is therefore another sweep of the same shape: same color partition (traversed in reverse), same constant 3×3 block size, same memory layout, same parallelism. Nothing global is ever formed.

An analogy that actually fits: Gauss–Seidel-style sweeps are like an LU forward-substitution — an ordered cascade of small solves that together act as an approximate inverse. Transposing a triangular solve gives you back-substitution: the same cascade, in reverse. The backward pass here is exactly that, at the granularity of one vertex.

The mechanism, concretely

Each VBD update solves H_i \Delta\bm{x}_i = -\bm{g}_i where \bm{g}_i and H_i are the local gradient and 3×3 Hessian block of the incremental potential at vertex i, then applies \bm{x}_i \leftarrow \bm{x}_i + \Delta\bm{x}_i immediately. To differentiate one such update, the authors differentiate the local optimality condition F_i = H_i\Delta\bm{x}_i + \bm{g}_i = \bm{0}, not the internals of the linear solve. Given the incoming adjoint \bar{\bm{x}}_i^+ = \partial L / \partial \bm{x}_i^+ (accumulated by the reverse updates already executed), one local adjoint solve

produces a vector \bm{p}_i from which everything else follows by cheap contractions. For any upstream quantity z (neighboring vertex positions, Lamé parameters \mu, \lambda, external forces, rest shape):

The second term — the "Hessian tangent" — exists because H_i itself depends on state; it's evaluated as a directional derivative, never as an assembled third-order tensor, at the cost of a few 3×3 products. Crucially, state gradients and parameter gradients share the same \bm{p}_i, so adding material gradients to a rollout is nearly free.

The forward pass records just \Delta\bm{x}_i and the six independent entries of the symmetric H_i — 72 bytes per vertex per sweep. The backward walks sweeps, colors, and vertices in exact reverse order, undoing each recorded displacement to reconstruct the linearization state rather than storing it, and recomputing everything else. That's the entire "tape."

The paper is refreshingly honest that exactness at finite depth has traps. Three terms are easy to drop, and each produces gradients that are wrong at small K but self-correct at large K — invisible to the standard finite-difference-at-convergence validation:

1. Update saturation. Production solvers damp the raw step (\Delta\bm{x} \leftarrow \Delta\bm{x}\,c/(\|\Delta\bm{x}\|+c+\varepsilon)) to stabilize early sweeps. Its Jacobian is a rank-one correction; omitting it costs up to 26% error at K{=}1.
2. The Hessian tangent above; omitting it leaves error \propto \|\Delta\bm{x}\|, and it grows under mesh refinement.
3. Validity gates — disabling tangent terms on compressed elements (\det F > \tau) when the forward applies no such cutoff. One gated element out of 245 moved the whole gradient by 8.2\times10^{-3}, and the effect grows with refinement.

All three were bugs the authors actually shipped in earlier revisions of their own code and caught only against an autograd reference at K{=}1. That's a genuinely useful methodological warning for anyone hand-deriving adjoints.

The evidence

Exactness (Table 1, the central measurement). On a |V|{=}115 mesh, backpropagating through four rollout steps, the reverse-sweep gradient matches PyTorch autograd through the identical executed forward to \sim 3\times10^{-15} at every depth. The equation-level adjoint (verified in both assembled and matrix-free implementations, agreeing with each other to 2\times10^{-16}) is off by 37% at K{=}1, 2.3\times10^{-5} at K{=}5, and only converges by K{=}20:

This is the operational content of the title: the two routes are not two implementations of one gradient but exact gradients of two different functions, and only one of them is the function the forward pass evaluated.

Cost (Table 3). The comparison is unusually well controlled — one codebase, one GPU, one process, identical meshes, forward trajectories verified identical to 10^{-16} before any gradient is compared. At |V|{=}115, K{=}20, four rollout steps:

They deliberately include the matrix-free global adjoint (Hessian-vector products inside PCG, never forming the matrix) as the leanest possible opponent — it's exactly what you'd get by handing a custom vector–Jacobian product to an AD framework — and still win by 3× time and 20× memory. The regimes also separate cleanly with stiffness: at \mu{=}\lambda{=}5\times10^5 PCG needs 233 iterations and the equation-level backward climbs to 247–389 ms, while the reverse sweep doesn't move, because it never solves a system whose conditioning tracks the material.

Scaling. Across a K \times |V| grid up to 201k vertices, the backward stays within 0.96–1.14× the forward's runtime, and memory holds at ≈4 KB per vertex per sweep to ±13%. The headline demo optimizes per-body initial velocity and stiffness for 10^6 contact-coupled soft cubes (8M vertices, 6M tets) over 240 implicit steps with backpropagation through time, checkpointed every 5 steps, in ~74 GB on one H100. For reference, DiffIPC reports ~100 GB at 10k vertices; this paper measures 1,074 MB at 27k vertices — though the authors correctly flag this as a cross-paper comparison of published numbers, unlike Table 3. Inverse problems (material ID to <0.25%, cloth draping, contact-rich planning that drives an uncovered-area objective to 0.0%) confirm the gradients are useful, not merely correct.

The ablation errors grow with refinement — the opposite of what large-scale use needs — which is why the finite-depth exactness terms matter beyond pedantry:

What to be skeptical about

Does the "right" gradient help the optimizer? The paper's honesty cuts both ways here. Appendix G tests whether solver depth (and thus gradient fidelity of the equation-level view) changes what Adam achieves on a small task, and reports a null result: means are statistically indistinguishable across K \in \{5, 20, 100\}, though deeper solves halve the variance. So the 37% error is real and measured, but no experiment demonstrates the equation-level gradient actually fails an optimization task the executed-solver gradient passes. If your application genuinely wants the derivative of a converged equilibrium — some system-identification settings do — the paper itself notes you're better served equation-level.

Generality. The full evaluation covers one solver family. Theorem 1's hypotheses are mild (finite composition of C^1 block updates with invertible local systems), and the construction transfers to XPBD and projective dynamics at machine precision — but those verifications are "deliberately plain NumPy" on toy spring systems, not performance demonstrations. Whether the approach extends gracefully to solvers with global coupling per iteration (projective dynamics' global solve is handled by reusing a constant prefactorization; harder cases may not be so lucky) remains open. Coulomb friction with active-set changes is explicitly future work, and contact adjoints flow through the executed barrier evaluations rather than a converged contact configuration — a feature or a bug depending on what you're optimizing.

Small comparison meshes. The controlled four-way comparison lives at |V|{=}115; the large-scale numbers are for their method alone (unrolled AD literally doesn't fit at 201k vertices, which is part of the point, but it means the 33×/71× ratios are measured where all methods still run).

Why this matters

The framing — solver structure as the organizing principle for the backward pass — is the durable contribution, more than the VBD instantiation. The differentiable-simulation literature has treated the backward pass as a linear-algebra problem (assemble the adjoint system, pick a solver for it), inherited from the adjoint-PDE tradition where forward solves were themselves global. But GPU-era simulators won by abandoning global structure for ordered local work, and this paper shows the backward can inherit everything that made the forward fast: coloring, block size, memory layout, thread schedule. It closes what the authors call the gap that has held differentiable elastodynamics near ~30k DoF, turning any forward that scales to millions of vertices into a differentiable one for the price of one forward-equivalent reverse pass and 72 bytes per vertex per sweep.

The same lesson applies well beyond graphics: anywhere a fixed-budget iterative solver sits inside a learning loop — deep equilibrium models cut off early, ADMM layers, iterative decoders — "differentiate what ran, at the granularity it ran at" is a design principle with teeth, and this paper supplies the exactness conditions and a bug taxonomy for doing it right.

Where to spend your time: §3 (especially §3.5 on what finite-depth exactness requires — the most transferable content) and Tables 1–3 in §6. If you implement adjoints by hand for anything, read §3.5 twice.