Scale the JVP finite difference step by norm(x)/norm(v) - #230
Conversation
The JVP step multiplies `v`, so it must have units of `[x]/[v]`. The old rule `max(relstep*sqrt(|dot(x,v)|), absstep)` has units of `sqrt([x][v])`, which pins the step near the `absstep` floor for unit-norm directions and makes the result depend on the scaling of `v`. Use `max(relstep*norm(x), absstep)*dir/norm(v)` instead, so the perturbation `epsilon*v` is a `relstep` relative change of `x`. Co-Authored-By: Chris Rackauckas <accounts@chrisrackauckas.com>
|
CI status on the first push, with both red jobs traced to pre-existing breakage rather than to this PR:
I also ran that integration group locally against this branch ( |
|
Root cause of the red
Standalone reproducer with no FiniteDiff involved: using OrdinaryDiffEq, OrdinaryDiffEqRosenbrock, LinearAlgebra
const n = 10
f(du, u, p, t) = (du .= 0; du[2:(end-1)] .= p[1] .* (u[3:end] .- 2 .* u[2:(end-1)] .+ u[1:(end-2)]); nothing)
u0 = sin.(range(0, 1, length = n))
jp = Tridiagonal(similar(u0, n-1), similar(u0), similar(u0, n-1))
prob = ODEProblem(ODEFunction(f; jac_prototype = jp), u0, (0.0, 1.0), [0.42])
solve(prob, Rodas4P(), saveat = 0.1) # ArgumentError from fill! in build_J_WThree open draft PRs upstream already fix this |
2 participants
Please ignore until reviewed by @ChrisRackauckas.
What changed and why
finite_difference_jvp/finite_difference_jvp!computed their step asepsilon = max(relstep*sqrt(|dot(x,v)|), absstep)*dir. That expression is dimensionallywrong: the step multiplies
v, so it must carry units of[x]/[v], whilesqrt(x⋅v)carriessqrt([x][v]). This replaces it withepsilon = max(relstep*norm(x), absstep)*dir/norm(v), so theperturbation
epsilon*vis arelsteprelative change ofxfloored atabsstep— the same ruleSparseDiffTools.num_jacvec!used, and the same rule the colored-Jacobian path injacobians.jlalready uses for its own perturbation directions.
The dimensional argument
The JVP evaluates
f(x + h*v), sohis not a step inxand is not dimensionless:h*vmust bea perturbation of
x, henceh ~ [x]/[v].h = max(relstep*sqrt(|x⋅v|), absstep). For a unit-norm direction (the Krylov/Arnoldi case)and a random
v,|x⋅v| ~ norm(x)/sqrt(n), sosqrt(|x⋅v|)stays O(1) no matter how large thestate is, and
hsits at theabsstepfloor (~1.5e-8 forVal(:forward)with the defaultrelstep = absstep = sqrt(eps())) forever. It is also not invariant to rescalingv: theperturbation grows like
norm(v)^1.5instead of being independent ofnorm(v).h = max(relstep*norm(x), absstep)*dir/norm(v), i.e.norm(h*v) = max(relstep*norm(x), absstep).Scaling
vbysscaleshby1/sand leaves the computed JVP unchanged (up to the floor).norm(v) == 0falls back to the unscaled step (sov == 0returns exactly zero rather thanNaN),and a non-finite
norm(x)falls back to theabsstepfloor rather than making every outputNaN.Accuracy
2D Brusselator,
Val(:forward), 20 unit-norm random directions per size, exact reference from aForwardDiff dual;
relerr = norm(jvp - exact)/norm(exact). Old rule evaluated with the same codepath and the old
epsilonformula:The old rule's error grows linearly with
norm(x)because the step is stuck at the floor; the newrule's error is flat. The regime is roundoff-dominated (
relerr ∝ 1/h), so being 70× low inhcosts 70× in accuracy — h sweep at N=32:
Invariance to the scaling of
v(same problem, same direction,vscaled bys, result dividedby
s):Why it matters downstream
Matrix-free Krylov cannot see the inaccurate operator. GMRES (Krylov.jl) on the same Brusselator
with
A = J - I/gammaapplied through each JVP variant: identical iteration counts and identical"converged" status, but the residual measured against the exact (ForwardDiff) operator is 30–70×
worse with the old step.
In OrdinaryDiffEq's matrix-free Krylov path that inaccurate Newton direction makes Newton stagnate,
burn its iteration budget, record a convergence failure and collapse
dt.End-to-end solver impact
Brusselator N=32 (2048 unknowns),
KenCarp4(linsolve = KrylovJL_GMRES(), concrete_jac = false),abstol = reltol = 1e-8,tspan = (0.0, 11.5), one solve per process, wall clock including compile:AutoForwardDiff()AutoFiniteDiff()AutoFiniteDiff()With the corrected step,
AutoFiniteDiffreaches 1.03x the ForwardDiff wall clock with bit-identicalstep statistics. With the old step the same solve was still running when I killed it at a 30 minute
cap, i.e. >43x slower and not converging: Newton cannot make progress on an operator that is
100-5000x less accurate than the tolerance it is being asked to hit, so it exhausts its iteration
budget, records a convergence failure and
dtcollapses. (nfis higher forAutoFiniteDiffbecause DifferentiationInterface recomputes
f(x)per matvec — see the note at the bottom.)Tests: failing before, passing after
New file
test/jvp_accuracy_tests.jl:f_i(x) = x_i * x_{i+1}(cyclic) atnorm(x) ≈ 2164with adeterministic unit direction, checked against the analytically exact
J*v; plusv-scalinginvariance, degenerate
v == 0/x == 0/NaN-in-xbehaviour, and a complex-valued state.With
src/jvp.jlstashed (i.e. the old step rule, everything else identical):With the fix applied:
Existing suite
GROUP=Core julia --project -e 'using Pkg; Pkg.test()'on Julia 1.12.6, unchanged assertions andunchanged tolerances — nothing in the existing suite encoded the old step rule:
(The one
Brokenis pre-existing on master.)GROUP=Downstreamerrors, but identically on unmodified master — it is not caused by this PR:I ran
GROUP=Downstreamtwice, once on this branch and once on agit worktreeoforigin/master(9bae080) with no other differences: both fail with the same
fill!-on-Tridiagonalerror, fromsolve(prob, Rodas4P(), saveat=0.1)on line 20 — the defaultAutoForwardDiffpath, before anyAutoFiniteDiffsolve is reached. It is an upstreamOrdinaryDiffEqDifferentiationproblem and isbeing tracked separately.
Docs build:
julia --project=docs docs/make.jlexits 0 (only the pre-existing:missing_docswarning list, which
make.jlalready sets towarnonly).typosis clean over the diff. The repodeclares no formatter (
.JuliaFormatter.tomlabsent, no format CI), so nothing was reformatted;the new code follows the surrounding style.
Versioning
2.32.1 → 2.33.0. No exported name, signature or type changes, but this changes the numerical output
of
finite_difference_jvp/finite_difference_jvp!for every caller, so it is not a patch. I judgedit a minor bump rather than 3.0 because the change is an accuracy fix inside an approximation rather
than an API break, and a major bump would strand the ecosystem's
FiniteDiff = "2"bounds. Pushback if you would rather ship it as 3.0.
What I did not verify
norm(x)where the old usedsqrt(|dot(x,v)|); both mixabsstepwith a quantity carrying units ofx, so unitful statesare no better and no worse than before, but I did not test them.
InterfaceIIintegration group was run locally against this branch:AutoSparse Detection Testspassed 4/4, then
Enum Testserrored on an OrdinaryDiffEq master test-file bug(
UndefVarError: DiscreteProblem) which aborts the group.rather than letting it run to completion, so "how much slower" is unquantified beyond >43x.
Related, deliberately not in this PR
DifferentiationInterface'spushforward!forAutoFiniteDiffnever passesf_intofinite_difference_jvp!, sof(x)is recomputed on every matvec — a 2× cost in matrix-freeKrylov. Different repo, separate PR.
finite_difference_jvp!(jvp, f, x, v, fdtype, f_in)(the cache-less in-place method) builds itscache with
JVPCache(x, fx, fdtype), the non-allocating constructor, sox1 === xand theuser's
xis silently overwritten withx + epsilon*v. Reproduces on released 2.32.1(
x = [1.0,2.0,3.0]comes back as[1.0000000149011612, 2.0, 3.0]), so it is pre-existing andindependent of this change. Separate PR.