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.

[ColtAllen]

You shouldn't be calling newton directly, use the root_scalar(method='newton') api.

Awesome! Didn't realize the recent optimizer PR supports this out of the box. Good to know we don't have to write any more Ops for this.

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.

From what I've read, ~100k variables is pretty common for INLA models, but every working implementation I've seen also uses Newton methods for the Hessian rather than trying to approximate it. If the new root API supports it, I'm in favor of following their example and using it.

[Michal-Novomestsky]

I've opted to use minimize in the end, as I found that root needs the same amount of equations as variables, but since the equation we're solving is an inner product over vectors, there's technically only one equation, but multiple varibles given that we're optimising over a vector.

I don't believe there should be any difference between finding the root and just minimising the negative of the Taylor expansion of logp(x|y,params) explicitely (as it should be equivalent to maximising p(x|y,params) which is what we want).

To the best of my knowledge, the reason why that Newton algorithm is used is because it saves you one extra call to the model, as you use (almost) the same function to find the root as you do compute logp(x|y,params), but given that newton would happen internally inside root and we're compiling this as a symbolic function rather than doing it explicitely, I don't think we'd get this speedup anyway.

[jessegrabowski]

I've opted to use minimize in the end, as I found that root needs the same amount of equations as variables, but since the equation we're solving is an inner product over vectors, there's technically only one equation, but multiple varibles given that we're optimising over a vector.

If you want to use root, you will use the gradient vector of the loss function w.r.t the variables as the system of equations

Pytensor will always seek out and re-use any pieces of computation across your graph. If you compute the value logp(x | y, params) somewhere then use it in a downstream computation later, it will only compute it one. This is true even if you don't literally re-use the same python variable. If the underlying graphs are equivalent (e.g. you call a logp function on the same symbolic inputs twice, like from a pymc model), then this will only be computed once.

I can concoct some examples for you if it isn't clear how this works

[Michal-Novomestsky]

Oh yes you're absolutely right, sorry I forgot you mentioned that point about solving $\nabla_x L = 0$ in a previous post. Anyway, since that's just a proxy for doing the minimisation upfront, it's probably clearer and cleaner to just minimize explicitely.

Pytensor doing stuff like that automatically is really neat! I would appreciate some examples if you've got the time - doesn't that lead to some weirdness with PyTensor accessing variables which the dev doesn't intend it to if it can just take equivalent graphs? I guess it shouldn't matter since it's only using identical graphs so the output should be the same regardless?