Second derivative of Matern in zero is wrong

[FelixBenning]

julia> using KernelFunctions: MaternKernel

julia> k = MaternKernel(ν=5)
Matern Kernel (ν = 5, metric = Distances.Euclidean(0.0))

julia> import ForwardDiff as FD

julia> kx(x,y) = FD.derivative(t -> k(x+t, y), 0)
kx (generic function with 1 method)

julia> dk(x,y) = FD.derivative(t -> kx(x, y+t), 0)
dk (generic function with 1 method)

julia> dk(0,0)
0.0

This is wrong, because for a centered GP $Z$ with covariance function $k$

$$dk(x,y) = \partial_x \partial_y k(x,y) = \partial_x \partial_y \mathbb{E}[Z(x),Z(y)] = \mathbb{E}[\partial_x Z(x) \partial_y Z(y)] = \text{Cov}(Z'(x), Z'(y))$$

And $\text{Cov}(Z'(0), Z'(0)) >0$.

$\nu=5$ should be plenty of space for numerical errors since this implies the GP is 5 times differentiable.

[FelixBenning]

This is most likely due to this

function _matern(ν::Real, d::Real)
    if iszero(d)
        return one(d)
    else
        y = sqrt(2ν) * d
        b = log(besselk(ν, y))
        return exp((one(d) - ν) * oftype(y, logtwo) - loggamma(ν) + ν * log(y) + b)
    end
end

https://github.com/JuliaGaussianProcesses/KernelFunctions.jl/blob/master/src/basekernels/matern.jl#L43-L45

The if block should not be constant but rather a taylor polynomial so that autodiff in this branch works. I am looking for a reference...

[devmotion]

Might be due to the hardcoded (constant) value for x = y. Hence possibly the problem doesn't exist with ForwardDiff 0.11.

[devmotion]

The if block should not be constant but rather a taylor polynomial so that autodiff in this branch works.

Shouldn't be necessary in ForwardDiff 0.11: It skips measure zero branches (which fixes some problems but broke existing code).

[FelixBenning]

https://github.com/cgeoga/BesselK.jl

https://arxiv.org/pdf/2201.00090.pdf

would be a possibility

[FelixBenning]

Shouldn't be necessary in ForwardDiff 0.11: It skips measure zero branches (which fixes some problems but broke existing code).

wow neat - when is that going to be released?

[devmotion]

The package and paper is mainly concerned with derivatives wrt the order, which does not seem to be the issue in the OP.

when is that going to be released?

I don't know. Initially the change was released in a non-breaking 0.10.X version but it broke a lot of downstream packages that rely on the current behaviour. So it was reverted and re-applied to the master branch but releases are only made in a separate 0.10.X branch without this change recently. I don't think anyone plans to release a 0.11 any time soon because the same thing will happen and nobody wants to invest the time to fix all the broken downstream code.

[FelixBenning]

@devmotion you are probably right about the references - damn I thought I have seen that somehwere. Can't find it at the moment. Given that the variance of the derivatives is really important for everything, this is not really a corner case and will break working with derivatives of GPs...

[FelixBenning]

Are you interested in a taylor expansion for the if x==y branch? Given the AD promise

Automatic Differentiation compatibility: all kernel functions which ought to be differentiable using AD packages like ForwardDiff.jl or Zygote.jl should be.

Same issue with the rationalQuadratic btw (although that does not appear to branch in the same way)

julia> using KernelFunctions: RationalQuadraticKernel

julia> k = RationalQuadraticKernel()
Rational Quadratic Kernel (α = 2.0, metric = Distances.Euclidean(0.0))

julia> kx(x,y) = FD.derivative(t->k(x+t, y), 0)
kx (generic function with 1 method)

julia> dk(x,y) = FD.derivative(t->k(x,y+t), 0)
dx (generic function with 1 method)

julia> dk(0,0)
0.0

EDIT: Wikipedia https://en.wikipedia.org/wiki/Mat%C3%A9rn_covariance_function#Taylor_series_at_zero_and_spectral_moments

[Crown421]

I think the Matern kernel is in fact correct:
Made a mistake here, Matern is incorrect,but RationalQuadratic is fine

using KernelFunctions, Enzyme, Plots
k = Matern52Kernel()
mk(d) = k(1.0, 1.0 + d)

mk(0.1)

dr = range(0.0, 0.1; length=200)

p = plot(layout=(1, 2))
plot!(p[1], dr, mk.(dr), label="matern", legend=:topright)

dmkd(x, y) = only(autodiff(
    Forward,
    yt -> only(autodiff_deferred(Forward, k, DuplicatedNoNeed, Duplicated(x, 1.0), yt)),
    DuplicatedNoNeed,
    Duplicated(y, 1.0)))
dmkd(1.0, 1.0 + 0.1)

plot!(p[2], dr, dmkd.(1.0, 1.0 .+ dr), label="matern", legend=:topright, title="d^2/(dx1 dx2) ker")

but for RationalQuadratic

using KernelFunctions, Enzyme, Plots
k = RationalQuadraticKernel()
mk(d) = k(1.0, 1.0 + d)

mk(0.1)

dr = range(0.0, 0.1; length=200)

p = plot(layout=(1, 2))
plot!(p[1], dr, mk.(dr), label="RQ", legend=:topright)

dmkd(x, y) = only(autodiff(
    Forward,
    yt -> only(autodiff_deferred(Forward, k, DuplicatedNoNeed, Duplicated(x, 1.0), yt)),
    DuplicatedNoNeed,
    Duplicated(y, 1.0)))
dmkd(1.0, 1.0 + 0.1)

plot!(p[2], dr, dmkd.(1.0, 1.0 .+ dr), label="RQ", legend=:right, title="d^2/(dx1 dx2) ker")

[FelixBenning]

I think the Matern kernel is in fact correct: Made a mistake here, Matern is incorrect,but RationalQuadratic is fine

neat so switching to enzyme would at least fix the rational quadratic.

[Crown421]

I have been working on this for a bit, and it seems the issue is not the matern kernel, or that the Taylor Expansion is needed, but instead it seems to related to the Euclidean distance. In the following code I did it all by hand, and get

using KernelFunctions
using Enzyme
using Plots

k = Matern52Kernel()
mk(d) = k(1.0, 1.0 + d)

kappa(d) = (1 + sqrt(5) * d + 5 * d^2 / 3) * exp(-sqrt(5) * d)

r = range(0, 0.2, length=30)

begin
    dist1(x, y) = sqrt((x - y)^2)
    dist2(x, y) = abs(x - y)
    ck1(x, y) = kappa(dist1(x, y))
    ck2(x, y) = kappa(dist2(x, y))

    p = plot(layout=(2, 1))
    plot!(p[1], r, ck1.(1.0, 1.0 .+ r), label="dist1")
    plot!(p[1], r, ck2.(1.0, 1.0 .+ r), label="dist2")
    plot!(p[1], r, mk.(r), label="ref")

    dmkd1(x, y) = only(autodiff(
        Forward,
        yt -> only(autodiff_deferred(Forward, ck1, DuplicatedNoNeed, Duplicated(x, 1.0), yt)),
        DuplicatedNoNeed,
        Duplicated(y, 1.0)))
    dmkd2(x, y) = only(autodiff(
        Forward,
        yt -> only(autodiff_deferred(Forward, ck2, DuplicatedNoNeed, Duplicated(x, 1.0), yt)),
        DuplicatedNoNeed,
        Duplicated(y, 1.0)))

    plot!(p[2], r, dmkd1.(1.0, 1.0 .+ r), label="dist1")
    plot!(p[2], r, dmkd2.(1.0, 1.0 .+ r), label="dist2")
end

Repeating the same with the distance works correctly, so there seems to be some weird issue with the chain rule.

[FelixBenning]

If I understand your code correctly, you reimplemented kappa (in the first plot you make sure that it results in the same function by comparing it to the reference implementation) and then you take the derivative of this kappa.

But your kappa does not have an if iszero(d) like the general matern implementation does cf.

https://github.com/JuliaGaussianProcesses/KernelFunctions.jl/blob/master/src/basekernels/matern.jl#L43-L45

As autodiff simply takes the derivative of the branch it finds itself in, it will take the derivative of this if case (if d is zero). And the derivative of a constant is zero.