JuliaGaussianProcesses / KernelFunctions.jl

Julia package for kernel functions for machine learning
https://juliagaussianprocesses.github.io/KernelFunctions.jl/stable/
MIT License
267 stars 32 forks source link

Fix Matern Grads #478

Closed willtebbutt closed 1 year ago

willtebbutt commented 1 year ago

Summary

We've been seeing test failures for AD with the MaternKernel. This fixes these.

Proposed changes

Modifies the implementation of the kappa function for the MaternKernel to utilise if ... else ... end instead of ifelse. The reason that this fixes the problem is that, if iszero(d), the branch which calls the bessel function etc (rather that the one that just returns 1) has some NaNs and some Infs floating around in it. The ifelse branch evaluates this branch-with-NaNs, and when Zygote is performing the reverse pass, this NaN somehow gets propagated backwards. Conversely, if ... else ... end entirely avoids computing any NaNs, thereby avoiding the problem.

This has the potential downside of yielding worse performance under Zygote, the MaternKernel performance was already quite bad, so I don't view this as a problem (it might not even make it much worse tbh).

What alternatives have you considered?

The solution that's defnitely going to work + be performant here would be a custom rule, but I don't have the time, and I don't know that this can be improved without sorting out the besselk rrule anyway.

Breaking changes

Not breaking.

codecov[bot] commented 1 year ago

Codecov Report

Base: 68.87% // Head: 93.72% // Increases project coverage by +24.85% :tada:

Coverage data is based on head (43748e5) compared to base (ba7b37c). Patch coverage: 100.00% of modified lines in pull request are covered.

Additional details and impacted files ```diff @@ Coverage Diff @@ ## master #478 +/- ## =========================================== + Coverage 68.87% 93.72% +24.85% =========================================== Files 52 52 Lines 1346 1354 +8 =========================================== + Hits 927 1269 +342 + Misses 419 85 -334 ``` | [Impacted Files](https://codecov.io/gh/JuliaGaussianProcesses/KernelFunctions.jl/pull/478?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=JuliaGaussianProcesses) | Coverage Δ | | |---|---|---| | [src/basekernels/matern.jl](https://codecov.io/gh/JuliaGaussianProcesses/KernelFunctions.jl/pull/478/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=JuliaGaussianProcesses#diff-c3JjL2Jhc2VrZXJuZWxzL21hdGVybi5qbA==) | `100.00% <100.00%> (+75.00%)` | :arrow_up: | | [src/utils.jl](https://codecov.io/gh/JuliaGaussianProcesses/KernelFunctions.jl/pull/478/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=JuliaGaussianProcesses#diff-c3JjL3V0aWxzLmps) | `91.46% <0.00%> (+3.65%)` | :arrow_up: | | [src/kernels/transformedkernel.jl](https://codecov.io/gh/JuliaGaussianProcesses/KernelFunctions.jl/pull/478/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=JuliaGaussianProcesses#diff-c3JjL2tlcm5lbHMvdHJhbnNmb3JtZWRrZXJuZWwuamw=) | `72.72% <0.00%> (+6.06%)` | :arrow_up: | | [src/chainrules.jl](https://codecov.io/gh/JuliaGaussianProcesses/KernelFunctions.jl/pull/478/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=JuliaGaussianProcesses#diff-c3JjL2NoYWlucnVsZXMuamw=) | `87.65% <0.00%> (+18.51%)` | :arrow_up: | | [src/basekernels/periodic.jl](https://codecov.io/gh/JuliaGaussianProcesses/KernelFunctions.jl/pull/478/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=JuliaGaussianProcesses#diff-c3JjL2Jhc2VrZXJuZWxzL3BlcmlvZGljLmps) | `100.00% <0.00%> (+44.44%)` | :arrow_up: | | [src/basekernels/exponential.jl](https://codecov.io/gh/JuliaGaussianProcesses/KernelFunctions.jl/pull/478/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=JuliaGaussianProcesses#diff-c3JjL2Jhc2VrZXJuZWxzL2V4cG9uZW50aWFsLmps) | `100.00% <0.00%> (+50.00%)` | :arrow_up: | | [src/basekernels/polynomial.jl](https://codecov.io/gh/JuliaGaussianProcesses/KernelFunctions.jl/pull/478/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=JuliaGaussianProcesses#diff-c3JjL2Jhc2VrZXJuZWxzL3BvbHlub21pYWwuamw=) | `100.00% <0.00%> (+57.89%)` | :arrow_up: | | [src/basekernels/fbm.jl](https://codecov.io/gh/JuliaGaussianProcesses/KernelFunctions.jl/pull/478/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=JuliaGaussianProcesses#diff-c3JjL2Jhc2VrZXJuZWxzL2ZibS5qbA==) | `100.00% <0.00%> (+73.33%)` | :arrow_up: | | [src/basekernels/exponentiated.jl](https://codecov.io/gh/JuliaGaussianProcesses/KernelFunctions.jl/pull/478/diff?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=JuliaGaussianProcesses#diff-c3JjL2Jhc2VrZXJuZWxzL2V4cG9uZW50aWF0ZWQuamw=) | `80.00% <0.00%> (+80.00%)` | :arrow_up: | | ... and [8 more](https://codecov.io/gh/JuliaGaussianProcesses/KernelFunctions.jl/pull/478/diff?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=JuliaGaussianProcesses) | | Help us with your feedback. Take ten seconds to tell us [how you rate us](https://about.codecov.io/nps?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=JuliaGaussianProcesses). Have a feature suggestion? [Share it here.](https://app.codecov.io/gh/feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=JuliaGaussianProcesses)

:umbrella: View full report at Codecov.
:loudspeaker: Do you have feedback about the report comment? Let us know in this issue.

willtebbutt commented 1 year ago

Codecov is taking a little while to update. This is ready for review.

willtebbutt commented 1 year ago

Looks good to me, I just think it would be good to add a test for it that fails on the master branch.

Well I implemented this change to fix a test on master, so is it not sufficient that CI now passes?