Add Local Regularization Support

[HengyeZhu]

Here are the latest changes

The modifications to SciMLBase are no longer needed now. I have implemented memory reuse to improve speed. Additionally, for the biased part, I used reservoir sampling to reduce memory consumption (this should be mathematically equivalent to the original implementation?)

This PR is to solve issue#817.

I have completed implementing the perform_step for an ODE solver — Tsit5, and I ran a 3000-step MNIST ODE experiment on the AMD Radeon RX 7900 XT following the parameter settings from the paper. The results are as follows:

Method	Train Accuracy (%)	Test Accuracy (%)	Training Time (hr)	Prediction Time (s / batch)	Testing NFE
none	99.4317	97.35	0.1655	0.0734	362.6736
unbiased	99.2800	97.78	0.1636	0.0630	286.8240
biased	99.3783	97.89	0.2114	0.0664	281.6016

I placed the implementation of perform_step in OrdinaryDiffEq.jl/lib/OrdinaryDiffEqTsit5/src/tsit_perform_step.jl. I am not sure if it would be appropriate to put it directly in DiffEqFlux.

@muladd function _tsit5_reg_step(k1, uprev, f, p, t, dt, abstol, reltol)
    T = constvalue(recursive_unitless_bottom_eltype(uprev))
    T2 = constvalue(typeof(one(t)))
    @OnDemandTableauExtract Tsit5ConstantCacheActual T T2
    a = dt * a21
    k2 = f(uprev + a * k1, p, t + c1 * dt)
    k3 = f(uprev + dt * (a31 * k1 + a32 * k2), p, t + c2 * dt)
    k4 = f(uprev + dt * (a41 * k1 + a42 * k2 + a43 * k3), p, t + c3 * dt)
    k5 = f(uprev + dt * (a51 * k1 + a52 * k2 + a53 * k3 + a54 * k4), p, t + c4 * dt)
    g6 = uprev + dt * (a61 * k1 + a62 * k2 + a63 * k3 + a64 * k4 + a65 * k5)
    k6 = f(g6, p, t + dt)
    unew = uprev + dt * (a71 * k1 + a72 * k2 + a73 * k3 + a74 * k4 + a75 * k5 + a76 * k6)
    k7 = f(unew, p, t + dt)
    utilde = dt *
        (
        btilde1 * k1 + btilde2 * k2 + btilde3 * k3 + btilde4 * k4 + btilde5 * k5 +
            btilde6 * k6 + btilde7 * k7
    )
    residuals = utilde ./ (abstol .+ max.(abs.(uprev), abs.(unew)) .* reltol)
    reg_val = sqrt(sum(abs2, residuals) / length(unew)) * dt
    return unew, reg_val
end

@muladd function perform_step(integrator, cache::Tsit5ConstantCache, p)
    (; t, dt, uprev, f) = integrator
    unew, reg_val = _tsit5_reg_step(
        integrator.fsalfirst,
        uprev,
        f,
        p,
        t,
        dt,
        integrator.opts.abstol,
        integrator.opts.reltol,
    )
    return unew, reg_val, 6, dt
end

function perform_step(integrator, p)
    return perform_step(integrator, integrator.cache, p)
end

Checklist

• Appropriate tests were added
• Any code changes were done in a way that does not break public API
• All documentation related to code changes were updated
• The new code follows the
  contributor guidelines, in particular the SciML Style Guide and
  COLPRAC.
• Any new documentation only uses public API

Additional context

Add any other context about the problem here.

[ChrisRackauckas]

Missing tests?

I placed the implementation of perform_step in OrdinaryDiffEq.jl/lib/OrdinaryDiffEqTsit5/src/tsit_perform_step.jl. I am not sure if it would be appropriate to put it directly in DiffEqFlux.

You can instead step! and use integrator.EEst

[HengyeZhu]

Sorry for the delayed response. I was going through final exams over the past three weeks.

I rebased and resubmitted the PR on top of the latest DiffEqFlux master.

I updated the implementation so that it no longer requires any changes to OrdinaryDiffEq.jl. All changes are now contained in this DiffEqFlux PR.

Following your suggestion, the forward local regularization path now calls step!(integrator) and uses integrator.EEst to obtain the local error estimate. This removes the previous need to add a custom perform_step implementation to OrdinaryDiffEqTsit5.

For the backward pass, I added a custom rrule so that the pullback does not differentiate through the mutable integrator state updated by step!(integrator). Instead, the pullback computes the parameter cotangent by differentiating a pure Tsit5 error-estimate helper that mirrors the local error estimate used in the forward pass.

I will add tests soon. For now, I ran a 3000-step MNIST ODE experiment on the AMD Radeon RX 7900 XT with Julia 1.12.5, following the parameter settings from the paper. The results are:

Method	Train Accuracy (%)	Test Accuracy (%)	Training Time (hr)	Prediction Time (s / batch)	Testing NFE
none	99.5883	97.43	0.1824	0.0872	362.9808
unbiased	99.1167	97.70	0.2075	0.0675	287.4384
biased	99.1600	97.64	0.2334	0.0724	281.4384