Improving LM implementation

Hi @patrick-kidger, as promised I wanted to help improve some of the optimisation methods in optimistix. I'd like to start with the LM implementation.

Trust region acceptance

The implemented approach has two thresholds used for determining any improvement, and sufficient improvement to warrant taking larger steps. The damping parameter is then taken to be 1/step_size in the damped newton iteration.

There are several points here:

1. The point of the trust region approach is to determine when the function is locally acting quadratic, and to take more Newton-like steps in this case. That said, you cover the case where the actual reduction is less than predicted, but you miss the case where it's more than predicted. Intuitively you would think that if it improves more than predicted then you should take even bigger steps, but this is wrong. The function is not behaving quadratically.

Therefore, you should have a third cutoff sensing when actual_reduction/pred_reduction is sufficiently greater than one (1.1 is usually fine). In this case, accept but do not make the step more newton. Basically, only make the iterations more newton if the gain is within a region around 1.

2. The units of dampening parameter in damped Newton step are those of the hessian-like operator. For minimisation it's actual hessian with [f]/[x]^2. For normalised least-squares it's J^T.J it's [f]^2/[x]^2, which is consistent because we normalised the equations. Anyways, choosing lambda=1/step_size is not dimensionally correct. Much better is to let lambda = |grad(f)| / mu (or lambda = |J^T.F| / mu for LM). Note, the units are now correct when mu has units [x]. The intuition behind this is, in the asymptotic steepest descent case x -> x - mu * grad(f) / |grad(f)|, i.e. a step-size times the gradient unit-vector.

Therefore, you can improve the damping in two ways.

i. setting lambda = |grad(f)|/step_size for minimisation, and lambda = |J^T.F| / step_size for LM.

ii. Choosing the initial value of step_size can be done by line search for a value of mu that leads to a reduction in the objective. This only needs to be done once, and thereafter step_size is modified following the normal approach. A good approach is to start from mu = |grad(f)| and half until x - mu * grad(f) / |grad(f)| leads to an objective improvement. You don't need to satisfy any other particular conditions to accept the value of mu.

Reusing J/JVP

Multi-step "approximate" LM is easily implemented by first linearising the JVPop around the current parameter and then performing one exact LM step, followed by a number of approximate steps using same JVPop. In the dense J case this is really valuable as you only form the J matrix once per 1 + num_approx_steps. It's also still helpful in the sparse case, wherein using jax.linearize is helpful. It is shown in literature to significantly reduce the amount of computation and only require a few more iterations to converge. There are simple criteria to determine when J should be recomputed, however JAX precludes these dynamic decisions. Simplest is a fixed number of approximate steps per exact step.

I didn't have time to attach literature, but hopefully this gets the ball rolling. I also suggest that a suite of simple but difficult benchmarks be written first to assess an improvement to the algorithm.

Start with optimistix/_solver/trust_region.py, especially the linked trust-region acceptance logic and damped Newton iteration. Review the LM implementation and existing tests, then establish simple benchmarks before changing the acceptance thresholds, damping initialization, or reuse of J/JVP information. Done should include benchmark evidence and tests covering the revised LM behavior.