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.

Second Order Differentiation

Mooncake is still working on their implementation, which is very exciting news; although that means it's functionality is limited to only basic gradient descent with appropriate hyperparametrs.

Enzyme on the other hand, fully supports the calculation of the Hamilton vector product. In our demonstration however, this computation will fail with both DI and base Enzyme.

Requests

• Make sure to dev both SSMProblems and GeneralisedFilters before resolving the toml
• 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.

[shravanngoswamii]

• 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 guess, this may help you:
https://pkgdocs.julialang.org/dev/toml-files/#The-[sources]-section

Monorepo example is what we need:

[sources]
Example = {url = "https://github.com/JuliaLang/Example.jl", rev = "custom_branch"}
WithinMonorepo = {url = "https://github.org/author/BigProject", subdir = "SubPackage"}
SomeDependency = {path = "deps/SomeDependency.jl"}

[gdalle]

Hi there! I'd be happy to help you figure out AD-related issues.

• Mooncake is not expected to handle second-order derivatives at the moment, we're working on it with @willtebbutt but it's a daunting task.
• Zygote's bug is rather disturbing, I can't really explain it.
• Enzyme even errors for the gradient when I try it on the latest version of everything. This makes it difficult for me to debug the Hessian. Since January, I have changed the DI.gradient bindings to forward as much as possible to Enzyme without meddling, so that might explain the difference wrt to @charlesknipp's first post. Could you maybe try it again to confirm that we're seeing the same things?

In general, when you encounter something you think is an issue with DI, I encourage you to open an issue on the DI repo with an MWE. It would help tremendously :)

[charlesknipp]

Enzyme even errors for the gradient when I try it on the latest version of everything. This makes it difficult for me to debug the Hessian. Since January, I have changed the DI.gradient bindings to forward as much as possible to Enzyme without meddling, so that might explain the difference wrt to @charlesknipp's first post. Could you maybe try it again to confirm that we're seeing the same things?

@gdalle I get an error for illegal type analysis with strict aliasing turned on. The computation succeeds with Enzyme.API.strictAliasing!(false), which is good but suboptimal. it works with the latest commit.

Just to clarify, I'm working with DI version 0.6.43 and Enzyme version 0.13.30 on Julia 1.10.4 (my work computer)

EDIT: I'm not sure why, but forcing symmetry on the innovation covariance caused Enzyme to report a type instability. I converted it to a symmetric type instead, and somehow that fixed the inconsistency.

[gdalle]

I'm afraid we'll have to concoct a pure-Enzyme MWE for the Hessian error, otherwise the Enzyme devs won't want to debug it

[yebai]

Second Order Differentiation

It is better to implement the basic ones first, i.e. first-order optimisation.

[charlesknipp]

@yebai good call. I'll focus on the non-linear methods from now on. The resampling step is likely to cause some trouble. In some variational algorithms here I was able to use Zygote.@ignore to stop the gradient calculation consistent with the literature.

@gdalle I added a minimum working example, isolated from the SSMProblems/GeneralisedFilters interface. The gradient works fine, but the Hessian needs some work. I created some sugar akin to their demo from here, which clearly doesn't work. If you have some ideas let me know. I'm not sure how DI interfaces with Enzyme for the Hessian.

[gdalle]

Your MWE looks simple enough to open an issue on Enzyme, so I'd say that's the right move