Tolerance schemes when using iterative solvers in least squares

[jpbrodrick89]

I wonder if anyone had a good, simple, ready-made recipe for dynamically tightening tolerance when using an iterative inner solver such as LSMR in Levenberg-Marquadt (e.g. schemes such as Wright and Holt's from 1985 or the more recent Eisenstat-Walker) before I try building one from scratch. My two ideas right now are either to write a new AbstractDampedNewton solver or just break up the solve and update the linear solver attribute of the nonlinear solver with eqx.tree_at.

@GJBoth I think you mentioned on lineax you had been using iterative solvers in LM?

[GJBoth]

Hi @jpbrodrick89

Yes I'm running LSMR currently as the linear solver inside an optimistix Levenberg Marquardt. It seems to work pretty well (as in I'm getting good solutions) but I haven't fully stress-tested it, so can't really help on the dynamic tolerance here. I'm using 1e-6 overall, although dynamic would probably be useful for me. Are you using LSMR as the main solver or just as a linear solver for some other solver (i.e. like me)?

On the more general side, I'm wondering if you can do this without recompiling the while loop? Did you test that?

[jpbrodrick89]

Thanks for coming back to me! Yes, just doing LSMR as the linear solver for LM least_squares, I'm trying to do PINN training with Levenberg-Marquadt and variable projection. When I tried SVD, I got massive out of memory errors because even though the Jacobian was small, the compiled jax function for the Jacobian was huge. Tried LSMR today with constant tolerance and worked a charm for low frequency problems, (1GB memory use instead 36GB, and converged in a few hundred iterations as opposed tens of thousands with Adam+L-BFGS), but Adam in initial phase still seems to be necessary for high frequency problems where loss landscape is noiser/sharper so the adaptive tolerance will likely be less useful to me now.

Good question about recompiling, I think if you just used eqx.filter_jit you probably would have to recompile as the tolerances which LSMR use are currently marked as primitive float's which are static by default. However, it doesn't complain or type-check if you pass in scalar arrays which would be considered static (e.g. lx.LSMR(atol=jnp.array(1e-6), rtol=jnp.array(1e-6))), so I THINK that will avoid recompilation. Let me know if you give this a go!

[GJBoth]

I'm trying to do PINN training with Levenberg-Marquadt and variable projection

Working with the Deepmind PINN paper? Sounds fun!

but Adam in initial phase still seems to be necessary for high frequency problems where loss landscape is noiser/sharper so the adaptive tolerance will likely be less useful to me now.

That makes sense; I wouldn't expect convex optimisers to work well in these non-convex landscapes. Although in my experience (thats a few years ago though) LBFGS always did remarkably well.

A few more comments:

• I think the current Lineax implementation of LSMR does not implement damping. That means the LM system gets solved as coupled system, but I thinking adding in the damping could help even more with memory? Although my knowledge of inner LSMR is a little fuzzy so could be wrong there.
• Optimistix LM does not scale the parameters (i.e. the damping term is lambda I). You can solve the problem in a reparametrised system to get the same result, and if it's very big I have a hutchinson-trace estimator to do this matrix-free too.

[patrick-kidger]

I think if you just used eqx.filter_jit you probably would have to recompile as the tolerances which LSMR use are currently marked as primitive float's which are static by default. However, it doesn't complain or type-check if you pass in scalar arrays which would be considered static (e.g. lx.LSMR(atol=jnp.array(1e-6), rtol=jnp.array(1e-6))), so I THINK that will avoid recompilation. Let me know if you give this a go!

Taking a look at our implementation I think this should be a constraint we can completely relax; these could be Real[ArrayLike, ""] instead. Happy to take a PR on that!

Optimistix LM does not scale the parameters (i.e. the damping term is lambda I)

Has been on my to-do list for literally years 😅 Another good candidate for being happy to take a PR on it!

[johannahaffner]

x-ref: Scaled LM in the works here: #189

[michael-0brien]

@GJBoth would be happy to have your feedback on #189 to make sure the proposed changes are sufficient for your usage!