Potential precision loss in AutoDiff through the loss function

[IvanBioli]

Description
I have been testing the numerical accuracy of automatic differentiation in NeuralPDEs.jl and observed significant numerical errors when differentiating through the loss function defined by NeuralPDEs. Specifically, the forward-mode Jacobian-vector products (JVPs) with the residual vector defining the loss function exhibit errors in the range of 1e-8, while other related computations (such as direct evaluations of the model) remain at the expected machine precision (~1e-16).

This suggests that some internal operations in NeuralPDEs' loss formulation might be inadvertently using lower-precision arithmetic (Float32 instead of Float64), or otherwise introducing unexpected numerical instability.

Additional context and expected behavior
To investigate this issue, I implemented functions that explicitly define the loss as $L(\theta) = \sum_{i=1}^n (r_i(\theta))^2$, where $\theta$ are the neural network parameters, and the residuals are defined as:

• For internal points: $r_i(\theta) = \mathcal{D}u(x_i) - f(x_i)$ where $\mathcal{D}$ is the differential operator.
• For boundary points: $r_i(\theta) = u(x_i) - g(x_i)$ (with optional weighting constants).

I have attached a minimal working example below that demonstrates the issue. The key observations are:

1. As a sanity check of my implementations, the relative error in evaluating loss(θ) versus loss_neuralpdes(θ) is on the order of 1e-16
2. The error when computing JVPs via AutoForwardDiff for direct neural network evaluations is also on the order of 1e-16.
3. However, the error in JVPs computed for the residual function (which contributes to the loss) is 1e-8, indicating a significant loss of numerical precision.

Minimal Reproducible Example 👇

using NeuralPDE, Lux, LuxCUDA, Random, ComponentArrays
using Optimization
using OptimizationOptimisers
import ModelingToolkit: Interval
using Plots
using Printf

################################# AUXILIARY FUNCTIONS ####################################
using Optimization: OptimizationProblem
using NeuralPDE: NeuralPDE, PINNRepresentation, recursive_eltype, EltypeAdaptor, safe_get_device, GridTraining
using Statistics: Statistics, mean

# only for PhysicsInformedNN
function merge_strategy_with_residual_vector(pinnrep::PINNRepresentation,
        strategy::GridTraining, datafree_pde_loss_function, datafree_bc_loss_function)
    (; domains, eqs, bcs, dict_indvars, dict_depvars) = pinnrep
    eltypeθ = recursive_eltype(pinnrep.flat_init_params)
    adaptor = EltypeAdaptor{eltypeθ}()

    train_sets = generate_training_sets(domains, strategy.dx, eqs, bcs, eltypeθ,
        dict_indvars, dict_depvars)

    # the points in the domain and on the boundary
    pde_train_sets, bcs_train_sets = train_sets |> adaptor
    pde_loss_functions = [get_residual_vector(pinnrep, _loss, _set, eltypeθ, strategy)
                        for (_loss, _set) in zip(
        datafree_pde_loss_function, pde_train_sets)]

    bc_loss_functions = [get_residual_vector(pinnrep, _loss, _set, eltypeθ, strategy)
                        for (_loss, _set) in zip(datafree_bc_loss_function, bcs_train_sets)]

    return pde_loss_functions, bc_loss_functions
end

function get_residual_vector(
        init_params, loss_function, train_set, eltype0, ::GridTraining; τ = nothing)
    init_params = init_params isa PINNRepresentation ? init_params.init_params : init_params
    train_set = train_set |> safe_get_device(init_params) |> EltypeAdaptor{eltype0}()
    return θ -> loss_function(train_set, θ)
end


function get_full_residual(prob::OptimizationProblem, symprob::PINNRepresentation)
    # Get PDE and BC residuals
    pde_residuals, bc_residuals = merge_strategy_with_residual_vector(symprob,
        symprob.strategy, symprob.loss_functions.datafree_pde_loss_functions, symprob.loss_functions.datafree_bc_loss_functions)

    # Setup weights for PDE and BCs
    flat_init_params = prob.u0
    adaloss = discretization.adaptive_loss
    @assert isnothing(adaloss) # FIXME: Assuming no adaloss
    num_additional_loss = 0

    adaloss === nothing && (adaloss = NonAdaptiveLoss{eltype(flat_init_params)}())
    
    # setup for all adaptive losses
    num_pde_losses = length(pde_residuals)
    num_bc_losses = length(bc_residuals)
    adaloss_T = eltype(adaloss.pde_loss_weights)

    # this will error if the user has provided a number of initial weights that is more than 1 and doesn't match the number of loss functions
    adaloss.pde_loss_weights = ones(adaloss_T, num_pde_losses) .* adaloss.pde_loss_weights
    adaloss.bc_loss_weights = ones(adaloss_T, num_bc_losses) .* adaloss.bc_loss_weights
    adaloss.additional_loss_weights = ones(adaloss_T, num_additional_loss) .*
                                    adaloss.additional_loss_weights

    function full_residual(θ)
        pde_losses = [pde_residual(θ) for pde_residual in pde_residuals]
        bc_losses = [bc_residual(θ) for bc_residual in bc_residuals]

        weighted_pde_losses = sqrt.(adaloss.pde_loss_weights) .* pde_losses ./ sqrt.(length.(pde_losses))
        weighted_bc_losses = sqrt.(adaloss.bc_loss_weights) .* bc_losses ./ sqrt.(length.(bc_losses))

        # full_res = hcat(Iterators.flatten((weighted_pde_losses, weighted_bc_losses))...)
        full_res = hcat(hcat(weighted_pde_losses...), hcat(weighted_bc_losses...))
        return full_res
    end

    return full_residual
end

function get_quadpoints(symprob::PINNRepresentation, strategy::GridTraining)
    (; domains, eqs, dict_indvars, dict_depvars) = symprob
    eltypeθ = recursive_eltype(symprob.flat_init_params)

    train_sets = hcat(generate_training_sets(domains, strategy.dx, eqs, [], eltypeθ,
        dict_indvars, dict_depvars)[1]...)
    return train_sets
end

################################# NEURALPDES TUTORIAL ####################################
@parameters t x y
@variables u(..)
Dxx = Differential(x)^2
Dyy = Differential(y)^2
Dt = Differential(t)
t_min = 0.0
t_max = 2.0
x_min = 0.0
x_max = 2.0
y_min = 0.0
y_max = 2.0

# 2D PDE
eq = Dt(u(t, x, y)) ~ Dxx(u(t, x, y)) + Dyy(u(t, x, y))

analytic_sol_func(t, x, y) = exp(x + y) * cos(x + y + 4t)
# Initial and boundary conditions
bcs = [u(t_min, x, y) ~ analytic_sol_func(t_min, x, y),
    u(t, x_min, y) ~ analytic_sol_func(t, x_min, y),
    u(t, x_max, y) ~ analytic_sol_func(t, x_max, y),
    u(t, x, y_min) ~ analytic_sol_func(t, x, y_min),
    u(t, x, y_max) ~ analytic_sol_func(t, x, y_max)]

# Space and time domains
domains = [t ∈ Interval(t_min, t_max),
    x ∈ Interval(x_min, x_max),
    y ∈ Interval(y_min, y_max)]

# Neural network
inner = 25
chain = Chain(Dense(3, inner, σ), Dense(inner, 1))

strategy = GridTraining(0.1)
ps, st = Lux.setup(Random.default_rng(), chain)
ps = ps |> ComponentArray .|> Float64
discretization = PhysicsInformedNN(chain, strategy; init_params = ps)

@named pde_system = PDESystem(eq, bcs, domains, [t, x, y], [u(t, x, y)])
prob = discretize(pde_system, discretization)
symprob = symbolic_discretize(pde_system, discretization)

callback = function (p, l)
    println("Current loss is: $l")
    return false
end

# Definition of the residual vector
residual = get_full_residual(prob, symprob)
loss = θ -> sum(abs2, residual(θ))
loss_neuralpdes = θ -> prob.f(θ, prob.p)

################################# TESTS ON THE ACCURACY ####################################
using ForwardDiff, Zygote, DifferentiationInterface, LinearAlgebra

# Sanity check
θ = prob.u0
rel_err = (loss_neuralpdes(θ) - loss(θ)) / loss_neuralpdes(θ)
println("Error on the loss: \t\t\t $rel_err") # In the order of 1e-16

# Test of AutoForwardDiff for differentiation of model evaluations
x = get_quadpoints(symprob, strategy)
fun = ps -> chain(x, ps, st)[1]
v = randn(length(θ))
J_fwd = ForwardDiff.jacobian(fun, θ)
jvp_explicit = J_fwd * v
jvp_pushforward = DifferentiationInterface.pushforward(
    fun,
    AutoForwardDiff(),
    θ,
    (v,),
)[1]
println("AutoForwardDiff error on model jvp:\t $(norm(jvp_explicit - jvp_pushforward[:]) / norm(jvp_explicit))") # In the order of 1e-16

# Check with JVPs
v = randn(length(θ))
J_fwd = ForwardDiff.jacobian(residual, θ)
jvp_explicit = J_fwd * v
jvp_pushforward = DifferentiationInterface.pushforward(
    residual,
    AutoForwardDiff(),
    θ,
    (v,),
)[1]
println("AutoForwardDiff error on residual jvp:\t $(norm(jvp_explicit - jvp_pushforward[:]) / norm(jvp_explicit))") # In the order of 1e-8!!!

[lkwargs]

TL;DR

This is not a bug, and there is no incorrect use of FP32/FP64 arithmetic in NeuralPDE.
When calculating the JVP of a function:

• DifferentiationInterface.pushforward offers higher speed and lower memory usage, but suffers from lower precision due to directional accumulation.

• If you need higher JVP accuracy, you should call ForwardDiff.jacobian directly, as done in ode_solve.jl.

• Alternatively, you can implement your own derivative function (e.g., Chebyshev differentiation) and pass it to NeuralPDE:

  discretization = PhysicsInformedNN(chain, strategy; init_params = ps, derivative = cheb_numeric_derivative)

However, this is outside the main design goal of NeuralPDE.

---

What causes the precision loss?

The precision loss only happens in residuals.
Let's check where the loss comes from:

function get_single_residual(prob::OptimizationProblem, symprob::PINNRepresentation)
    # Get PDE and BC residuals
    pde_residuals, bc_residuals = merge_strategy_with_residual_vector(symprob,
        symprob.strategy, symprob.loss_functions.datafree_pde_loss_functions,
        symprob.loss_functions.datafree_bc_loss_functions)

    # Setup weights for PDE and BCs
    flat_init_params = prob.u0
    adaloss = discretization.adaptive_loss
    @assert isnothing(adaloss)  # FIXME: Assuming no adaloss
    num_additional_loss = 0

    adaloss === nothing && (adaloss = NonAdaptiveLoss{eltype(flat_init_params)}())

    num_pde_losses = length(pde_residuals)
    num_bc_losses = length(bc_residuals)
    adaloss_T = eltype(adaloss.pde_loss_weights)

    adaloss.pde_loss_weights = ones(adaloss_T, num_pde_losses) .* adaloss.pde_loss_weights
    adaloss.bc_loss_weights = ones(adaloss_T, num_bc_losses) .* adaloss.bc_loss_weights
    adaloss.additional_loss_weights = ones(adaloss_T, num_additional_loss) .* adaloss.additional_loss_weights

    # Residual selector: returns an individual unweighted residual function
    function select_residual(θ, type::Symbol, idx::Int)
        if type == :pde
            @assert 1 ≤ idx ≤ length(pde_residuals)
            pde_residual = pde_residuals[idx](θ)
            return sqrt(adaloss.pde_loss_weights[idx]) * pde_residual / sqrt(length(pde_residual))
        elseif type == :bc
            @assert 1 ≤ idx ≤ length(bc_residuals)
            bc_residual = bc_residuals[idx](θ)
            return sqrt(adaloss.bc_loss_weights[idx]) * bc_residual / sqrt(length(bc_residual))
        else
            error("Invalid residual type: use :pde or :bc")
        end
    end

    return select_residual
end

function test_jvp_accuracy(f, θ, rname; v=nothing, differentiation_backend=AutoForwardDiff())
    if isnothing(v)
        v = randn(length(θ))
    end
    
    J_fwd = ForwardDiff.jacobian(f, θ)
    jvp_explicit = J_fwd * v
    jvp_pushforward = DifferentiationInterface.pushforward(
        f,
        differentiation_backend,
        θ,
        (v,),
    )[1]
    rel_error = norm(jvp_explicit - jvp_pushforward[:]) / norm(jvp_explicit)
    println("AutoForwardDiff JVP relative error of $rname:\t $rel_error")
    return (rel_error=rel_error, 
            jvp_explicit=jvp_explicit, 
            jvp_pushforward=jvp_pushforward)
end
# Check with JVPs
test_jvp_accuracy(residual, θ, "full residual")

# Check single residual
residual_single = get_single_residual(prob, symprob)
pde_residual_1 = θ -> residual_single(θ, :pde, 1)
test_jvp_accuracy(pde_residual_1, θ, "pde residual 1")

for i in 1:5
    bc_residual_i = θ -> residual_single(θ, :bc, i)
    test_jvp_accuracy(bc_residual_i, θ, "bc residual $i")
end

Result:

AutoForwardDiff JVP relative error of full residual:     2.8512364355099622e-8
AutoForwardDiff JVP relative error of pde residual 1:    5.397813367809174e-7
AutoForwardDiff JVP relative error of bc residual 1:     6.956875052383842e-16
AutoForwardDiff JVP relative error of bc residual 2:     1.578640092555188e-15
AutoForwardDiff JVP relative error of bc residual 3:     4.1790727967444575e-16
AutoForwardDiff JVP relative error of bc residual 4:     2.8546667789409537e-16
AutoForwardDiff JVP relative error of bc residual 5:     3.563848180799332e-16

✅ Boundary conditions (BC) are fine.
⚠️ PDE residuals cause the major precision loss.

---

Why PDE residuals?

Inspecting the expressions:

=== BUILDING PDE LOSS FUNCTIONS ===
=== EXPRESSION DUMP ===
(cord, var"##θ#230", phi, derivative, integral, u, p)->begin
        begin
            let (t, x, y) = (cord[[1], :], cord[[2], :], cord[[3], :])
                begin
                    cord1 = vcat(t, x, y)
                end
                derivative(phi, u, cord1, [[6.055454452393343e-6, 0.0, 0.0]], 1, var"##θ#230") .- (+).(derivative(phi, u, cord1, [[0.0, 0.0, 0.0001220703125], [0.0, 0.0, 0.0001220703125]], 2, var"##θ#230"), derivative(phi, u, cord1, [[0.0, 0.0001220703125, 0.0], [0.0, 0.0001220703125, 0.0]], 2, var"##θ#230"))
            end
        end
    end
=== BUILDING BC LOSS FUNCTIONS ===
=== EXPRESSION DUMP ===
(cord, var"##θ#230", phi, derivative, integral, u, p)->begin
        begin
            let (t, x, y) = (cord[[1], :], cord[[2], :], cord[[3], :])
                begin
                    cord1 = vcat(t, x, y)
                end
                u(cord1, var"##θ#230", phi) .- (*).((exp).((+).(x, y)), (cos).((+).(x, y)))
            end
        end
    end
... ...

• PDE residuals involve numeric derivatives with small step sizes (e.g., 6e-6, 1e-4).
• BC residuals involve basic algebraic operations (no small epsilon scaling).

In pinn_types.jl, when no derivative is specified, it defaults to a numeric finite difference method with tiny step sizes (numeric_derivative).

Since pushforward only tracks a directional derivative, small step sizes amplify numerical errors compared to building the full Jacobian.

---

Verifying the hypothesis

You can recompile the PDE loss function with different step sizes:

# Play with step sizes
function get_numeric_derivative_func(phi, step_size, order)
    get_u() = (cord, θ, phi) -> phi(cord, θ)
    u = get_u()
    _loss_function = new_numeric_derivative
    return θ -> begin
        cord = pde_train_sets[1]
        let (t, x, y) = (cord[[1], :], cord[[2], :], cord[[3], :])
            cord1 = vcat(t, x, y)
            _loss_function(phi, u, cord1, step_size, order, θ)
        end
    end
end

pde_derivative_o1_0 = get_numeric_derivative_func(symprob.phi, [[6.055454452393343e-6, 0.0, 0.0]], 1)
pde_derivative_o2_0 = get_numeric_derivative_func(symprob.phi, [[0.0, 0.0, 2^(-13)], [0.0, 0.0, 2^(-13)]], 2)
pde_derivative_o2_1 = get_numeric_derivative_func(symprob.phi, [[0.0, 2^(-13), 0.0], [0.0, 2^(-13), 0.0]], 2)
combination_pde_derivative = θ -> pde_derivative_o1_0(θ) .- (pde_derivative_o2_0(θ) .+ pde_derivative_o2_1(θ))

test_jvp_accuracy(pde_derivative_o1_0, θ, "pde derivative o1_0")
test_jvp_accuracy(pde_derivative_o2_0, θ, "pde derivative o2_0")
test_jvp_accuracy(pde_derivative_o2_1, θ, "pde derivative o2_1")
test_jvp_accuracy(combination_pde_derivative, θ, "combination pde derivative")

📈 Observation:
Increasing step size (e.g., from 2e-13 to 2e0) reduces the relative error from 1e-7 to 1e-16, so small epsilons are the root cause when using pushforward.

---

How to fix it properly?

You can implement a better numerical derivative to avoid this:

Example: Chebyshev differentiation

function cheb_numeric_derivative(phi, u, x, εs, order, θ)
    adjusted_εs = deepcopy(εs)

    # Replace all non-zero values with 2e-13
    for i in eachindex(adjusted_εs)
        for j in eachindex(adjusted_εs[i])
            if adjusted_εs[i][j] != 0.0
                adjusted_εs[i][j] = 2^(-13)
            end
        end
    end
    ε = adjusted_εs[order]
    _epsilon = inv(first(ε[ε .!= zero(ε)]))
    ε = ε |> safe_get_device(x)
    
    # Define number of Chebyshev points to use
    n = 5
    
    if order == 1
        cheb_points = [cos(π * j / (n-1)) for j in 0:(n-1)] .* _epsilon
        
        f_vals = []
        for j in 1:n
            push!(f_vals, u(x .+ cheb_points[j] .* ε, θ, phi))
        end
        
        derivative = zeros(eltype(f_vals[1]), size(f_vals[1]))
        for j in 1:n
            for k in 1:n
                if j != k
                    weight = (j == 1 || j == n ? 2.0 : 1.0) / (j == 1 || j == n ? 2.0 : 1.0)
                    weight *= (-1)^(j+k) / (cheb_points[j] - cheb_points[k])
                    derivative .+= weight .* f_vals[k]
                end
            end
        end
        
        return derivative
        
    elseif order == 2
        cheb_points = [cos(π * j / (n-1)) for j in 0:(n-1)] .* _epsilon
        f_vals = []
        for j in 1:n
            push!(f_vals, u(x .+ cheb_points[j] .* ε, θ, phi))
        end
        
        derivative = zeros(eltype(f_vals[1]), size(f_vals[1]))
        
        weights = [1.0, -2.0, 1.0]  # Simple second-derivative weights
        for j in 2:(n-1)
            derivative .+= weights[1] .* f_vals[j-1] .+ weights[2] .* f_vals[j] .+ weights[3] .* f_vals[j+1]
        end
        
        return derivative .* (_epsilon^2)
    else
        error("Only first and second order derivatives are implemented for Chebyshev differentiation")
    end
end

... ...

discretization = PhysicsInformedNN(chain, strategy; init_params = ps, derivative = cheb_numeric_derivative)

... ...

# Check with JVPs
test_jvp_accuracy(residual, θ, "full residual")

# Check single residual
residual_single = get_single_residual(prob, symprob)
pde_residual_1 = θ -> residual_single(θ, :pde, 1)
test_jvp_accuracy(pde_residual_1, θ, "pde residual 1")

for i in 1:5
    bc_residual_i = θ -> residual_single(θ, :bc, i)
    test_jvp_accuracy(bc_residual_i, θ, "bc residual $i")
end

Then use it:

discretization = PhysicsInformedNN(chain, strategy; init_params = ps, derivative = cheb_numeric_derivative)

New JVP errors:

AutoForwardDiff JVP relative error of full residual:     2.833858748121392e-14
AutoForwardDiff JVP relative error of pde residual 1:    1.342686535899662e-14
AutoForwardDiff JVP relative error of bc residual 1:     2.560819481617652e-16
AutoForwardDiff JVP relative error of bc residual 2:     2.729062539276571e-16
AutoForwardDiff JVP relative error of bc residual 3:     1.2480183244468969e-15
AutoForwardDiff JVP relative error of bc residual 4:     6.396024128620241e-16
AutoForwardDiff JVP relative error of bc residual 5:     4.100363082697556e-16
...

Precision significantly improved.

---

Summary

• This is not a bug in NeuralPDE.
• pushforward is fast but less precise, especially with tiny perturbations.
• You can get better precision by:
  • Using ForwardDiff.jacobian if acceptable.
  • Implementing a better custom derivative (like Chebyshev differentiation).

[ChrisRackauckas]

DifferentiationInterface.pushforward offers higher speed and lower memory usage, but suffers from lower precision due to directional accumulation.

pushforward is fast but less precise, especially with tiny perturbations.
You can get better precision by:
Using ForwardDiff.jacobian if acceptable.

That's not true. The ForwardDiff.jacobian function is calculated using the same exact pushforward function. It just does it on the basis functions. All forward mode is just J*v operations, and directly calculating J*v instead of J = J*e_i and then J*v will be less numerically stable than doing it in a single operation, but only by a very minor amount. So it's actually the other way around, but it generally shouldn't be noticeable since it's effectively the same operation.

This suggests that some internal operations in NeuralPDEs' loss formulation might be inadvertently using lower-precision arithmetic (Float32 instead of Float64), or otherwise introducing unexpected numerical instability.

Yes, that is what I'd suspect here. There must be a Float32 somewhere, 1e-8 is not an accident.

But anyways, we're rewriting the parser soon so this should change quite a bit.

[IvanBioli]

• PDE residuals involve numeric derivatives with small step sizes (e.g., 6e-6, 1e-4).
• BC residuals involve basic algebraic operations (no small epsilon scaling).

In pinn_types.jl, when no derivative is specified, it defaults to a numeric finite difference method with tiny step sizes (numeric_derivative).

Since pushforward only tracks a directional derivative, small step sizes amplify numerical errors compared to building the full Jacobian.

....

Summary

• This is not a bug in NeuralPDE.

• pushforward is fast but less precise, especially with tiny perturbations.

• You can get better precision by:

  • Using ForwardDiff.jacobian if acceptable.
  • Implementing a better custom derivative (like Chebyshev differentiation).

If my understanding is correct, by using finite differences to compute the derivatives with respect to the input one loses precision in the ouput (the residual in this case) compared to autodiff or others methods like Chebichev, but this should not affect the precision in the JVPs and VJPs.

Yes, that is what I'd suspect here. There must be a Float32 somewhere, 1e-8 is not an accident.

But anyways, we're rewriting the parser soon so this should change quite a bit.

Any plan about when the parser should be ready?

[ChrisRackauckas]

End of summer?