Newton solver to find the mode of the Gaussian (for INLA)

[theorashid]

Part of INLA roadmap #340.

We want to do the Laplace approximation around $\mu$. We need to find it first.

Notes

• jax implementation
• GPy implementation (maybe somewhere in here)

[ColtAllen]

Specific optimizers used in the referenced implementations:

jax implementation

scipy.optimize.newton

GPy implementation (maybe somewhere in here)

scipy.optimize.fixed_point

[ColtAllen]

Implementing scipy.optimize.newton or scipy.optimize.fixed_point will require creating Ops in pytensor. I'm curious if we can get by with the Hessian-approximation algorithms in the recently-merged Ops for scipy.optimize.minimize instead.

[jessegrabowski]

You shouldn't be calling newton directly, use the root_scalar(method='newton') api. All fixed point iteration problems can be re-cast as root finding problems, I show this in a new pytensor example notebook here. So if you have one variable, you need root_scalar, if you need many variables, you need root. I still haven't actually seen the math of what you're trying to do, so it's hard for me to give more concrete advice.

Also you don't need to approximate hessians in pytensor, you should just be able to compute it, unless the thing is really massive. I've done models with ~100k variables and used symbolic inverse Jacobians successfully.

[Michal-Novomestsky]

Thanks for the advice! I'm essentially trying to find the mode of a function (that is, maximise log likelihood). The INLA notes all point to using a fixed-point iterator with Newton's method to find the root of the derivative (approximated up to second order), but given that we need to calculate jac and hess for this at each step anyway, I don't see why we shouldn't simply call minimize on the NLL directly? If anything, this gives us the flexibility to choose the optimisation scheme we prefer (I believe find_MAP does exactly this, and indeed it uses BFGS as Colt mentions, rather than Newton).

[Michal-Novomestsky]

At the moment I mainly need to figure out how to pass in the NLL inputs and specify the latent field in some neat way. I've been a bit busy these past two days so I haven't had a chance to play around with INLA/PyMC in jupyter notebooks to get a feel for it yet. I'm part way through the optimize.minimize notebook Colt suggested and that's been very instructive so far. I'm aiming to get this done by the end of this week and I'll be sure to reach out if I get stuck!

[jessegrabowski]

As you point out, any minimization problem can be transformed into a root-finding problem by using the system of equations defined by the gradients $\nabla_x \mathcal L(x, \theta) = 0$, since the gradients must be zero at the minimum (assuming the function is convex and unconstrained blah blah).

If you have access to a contraction operator $T(x_i, \theta)$ such that repeated application $x_{i+1} = T(x_i, \theta)$ converges to $x^\star$, you can also transform that into a root-finding problem by directly solving for the roots of $T(x_0, \theta) - x_0$.

And if you have a system of equations $f(x, \theta) =0$, you can always define $\epsilon = f(x_0, \theta)$ as a vector of errors given solution guess $x_0$ and minimize sum of squares $<\epsilon, \epsilon>$ (or whatever you favorite convex twice-differentiable objective function is)

Which method to use is entirely up to you.