Maximum Likelihood Demonstration

[charlesknipp]

Overview

I added a quick little MLE demonstration, which should work for almost every AD backend (more on this later). This particular example was taken from Kalman.jl using Optimisers.jl to create a custom Newton's method.

A Note on Automatic Differentiation

All the AD here is done via DifferentiationInterface.jl to support something more universal. This allows us to quickly test various backends, switching to whichever suits the users' needs. Luckily, in this instance, every relevant backend can at least evaluate gradients. Unfortunately, since this is a Newton's method, the Hessian fails when using Enzyme. This is likely due to a bug in DifferentiationInterface, so it's not our problem.

Speaking of Enzyme, it requires type stability to efficiently compute gradients on the lowered code, which allowed me to catch an instability in the Kalman filter. See my post here where one of the Enzyme devs caught the issue.

Requests

• I don't have the environment properly set up since GeneralisedFilters is not on the official registry, and I can't figure out how to add it to Project.toml
• The SSMProblems dependency is not yet up to date, and I needed to set pkg> dev SSMProblems in order for it to execute
• If you have any suggestions, feel free to tear it to shreds

[charlesknipp]

I have been testing various backends and there seem to be some problems. I'm not entirely sure it's our fault, but it seems like Zygote, Enzyme, and Mooncake all yield some sort of error.

Zygote fails rather spectacularly because the observation coefficient matrix is non-square. I'm not sure how this even factors into the computation of the gradient; and, if I'm being honest, if Zygote is the one we can't get working then I'm okay with that.

Both Enzyme and Mooncake error specifically when computing the Hessian. For Mooncake, this is due to some recursive tangent type error; I'm not entirely sure why this wasn't an issue with the gradient, but I think there may be some problems. Enzyme just has issues with the additional context (the Constant seen in the objective), which is an underlying bug in the way it interacts with DifferentiationInterface. Lastly, the gradient computation for both is incredibly slow, even for shorter series. Which leads me to believe there might be something else under the hood.

Takeaways

Of all the working interfaces, I'm glad it's ForwardDiff since it's a stable and widely adopted framework. Although, the future of AD is clearly in the direction of both Mooncake and Enzyme. In the future, I'd really like to get those two backends to work consistently.

[THargreaves]

I don't have the environment properly set up since GeneralisedFilters is not on the official registry, and I can't figure out how to add it to Project.toml

I've just been deving it using the relative path. But yes, should get it registered soon. Was just waiting to tidy things up a touch more but think it should be there by the end of the week.

[THargreaves]

All sounds very strange, especially with it not working for non-square H. Glad the ForwardDiff version is working though.

I'll see if Will (Mooncake.jl author) can have a quick look and see if anything stands out to him.