LSMR init overflow, pt 2.

[GJBoth]

If max_steps is not given to the LSMR solver (and I guess the other solvers too?), max steps gets initialised as:

if self.max_steps is None:
    max_steps = min_dim * 10  # for consistency with other iterative solvers

For large scale solving this results in an overflow - easily fixed by just setting a max_steps, but the error doesn't tell you where. I think adding an extra safety here so the int32 doesn't overflow would be helpful!

[johannahaffner]

Thank you for letting us know! I'll address this in #176.

Do you have an MWE that we could include as a test case?

Edit: I've implemented a fix.

[GJBoth]

Sorry, was held up!

And yes, I'm using Lineax for some pretty large scale optical simulations, say 30 x 2048 x 2048 -> 30 4 x 2048 x 2048. So that's why it overflowed :-).

[johannahaffner]

And yes, I'm using Lineax for some pretty large scale optical simulations, say 30 x 2048 x 2048 -> 30 4 x 2048 x 2048. So that's why it overflowed :-).

Super cool to know! I'm always interested in the use cases people have :)

[GJBoth]

Two questions, now that I have your attention :-)

1. I'm using LSMR as the linear solver inside optimistix LevenBergMarquardt. That means the operator has a jvp in it - and I think LSMR calls .conj() internally, which seems to use a vjp, right? So does that mean we're secretly doing second-order here? I think we should be able to write a custom Operator here which doesn't do that, but it's a bit more involved!
2. Because I'm so memory constrained, Lineax / Optimistix approach to keep around y_eval and y and f_eval and f really sucks up huge bits of memory. Is there an easy work around, or should I write a custom solver (I guess more of an optimistix issue, actually?).

[patrick-kidger]

1. I think .conj() is just for complex numbers. It would be .transpose() to transform jvp->vjp. (Not sure if this answers your question.)
2. Interesting point! Are you suggesting that we hold only the y-values in memory, and recompute the f-values as required?

[GJBoth]

1. Sorry, my bad, I meant .transpose indeed. But my questions stands: the operator going into LSMR has a jvp when using LevenbergMarquardt, and LSMR calls operator.transpose, meaning we're secretly doing a vjp over jvp, which is I don't think is necessary?
2. No, not entirely - the forward pass can be pretty expensive. Taking Levenberg-Marquardt as example, as I work with that one, the solver & state keeps
   1. 2 times the parameter to be optimised (once as the parameter, once as y_eval in the state)
3. the full blown output f in the state.

This approach is fine if you're not memory bound, but for me this not the case. From what I can see, both are mainly used for convergence criteria, which doesn't seem necessary: optimistix / lineax use the norm of the difference for convergence, rather than the difference of the norm (which I think is what scipy does?) - memory wise a huge difference!

I think this would require some pretty major changes to the structure so not suggesting you'd do that, but if you have some ideas how to most easily tackle this on my end that'd be much appreciated :-)

[johannahaffner]

From what I can see, both are mainly used for convergence criteria, which doesn't seem necessary: optimistix / lineax use the norm of the difference for convergence,
if you have some ideas how to most easily tackle this on my end that'd be much appreciated :-)

This is one reason for having both y and y_eval - the difference does need to be computed somewhere (and it would have the same dimension).

But first, let's unpack what this solver does, to see where we currently access y and y_eval:
ClassicalTrustRegion and DampedNewtonDescent are the search and descent paired up to make LevenbergMarquardt, on top of AbstractGaussNewton.

1. This is a trust region solver, which means that the search outputs the radius in which the Hessian approximation is deemed to be of good quality. It uses the difference of y and y_eval to do that, in addition to f and f_eval.
2. The step (direction) is computed using this radius, which makes the direction a function of the radius. Since the radius changes while we're searching for an acceptable step, the directions change as well. This is in contrast to a line search. In DampedNewtonDescent the direction is correspondingly computed in descent.query, and not in descent.step.

The search uses the Jacobian and vector of residuals at the already accepted y. The descent likewise uses just these to compute the next step. Once a trial step y_eval has been accepted, its Jacobian and vector of residuals comprise the new f_info. We're computing f_eval and linearise at the same time (here) because that is efficient in most cases (and great for compilation time), but you could potentially move the computation of residuals and the Jacobian into the accepted branch of a custom solver (tweaking this part of AbstractGaussNewton) if you're willing to bake in assumptions about the search and descent you're going to use. This would allow you to pass only the squared sum of the residuals to the search, which is all that is required.

This might also make your forward passes more efficient, depending on how your GPU is treating the cond / avoiding extraneous evaluations.

You could also mildly hack around with the descent API, to avoid writing f_info into the descent state - it technically exists in two places, the solver state and in the case of the DampedNewtonDescent it also gets written into the descent state. You could probably evict it from either of these - although this might be an optimisation that already happens under the hood, depending on how XLA parses your computation graph (there are some heuristics in there)).

I don't see how you can get rid of having both y and y_eval unless you're willing to make the search a LearningRate - you will need to compute new candidate y_evals until the search is satisfied with one, and the computed $\Delta y$ always gets added to y. If your search is a LearningRate then the DampedNewtonDescent isn't really needed anymore, unless you want to set the learning rate to something that you believe will yield a reasonable regulariser throughout the course of your optimisation. If you want to experiment with that, try GaussNewton - this solver takes a full Newton step at each iteration, but you could tweak the learning rate it uses here.

1. Sorry, my bad, I meant .transpose indeed. But my questions stands: the operator going into LSMR has a jvp when using LevenbergMarquardt, and LSMR calls operator.transpose, meaning we're secretly doing a vjp over jvp, which is I don't think is necessary?

The operator is computed with jax.linearize here. A linearised function can be transposed (with jax.linear_transpose) and that is indeed what we're doing with lx.FunctionLinearOperator, here.