Preconditioner and initial guess for backward pass in Krylov solvers

[gaspardbb]

Thanks a lot for this super nice library.

I'm using Lineax for the GMRES iterative solver, which I use with a preconditioner and a good initial guess.

    solver_out = lx.linear_solve(
        operator,
        rhs,
        lx.GRMES(...),
        options={'y0': x0, 'preconditioner': precond},
    )

I've been experiencing issues with convergence. Forward pass works fine, so I suspect issue comes from gradient computation.

My questions:

1. I'm not 100% sure issue comes from the VJP computation since I don't know how to capture the variable of the backward pass. If you have pointers on that I'm happy to look into.
2. Lineax is not using a preconditioner for the backward pass. Writing this issue, I realized this has been fixed in Fixes for preconditioned krylov solvers #162. Though, wouldn't x0.T be a good initial guess for the backward pass too?

[johannahaffner]

Forward pass works fine, so I suspect issue comes from gradient computation.

1. I'm not 100% sure issue comes from the VJP computation since I don't know how to capture the variable of the backward pass. If you have pointers on that I'm happy to look into.

We're having a similar discussion over in #176 (comment), where I think that issue is coming from worse conditioning on the backward pass as well.

What exactly do you mean by issues with convergence? If it comes from gradient computation, does that mean that the iterative linear solve is part of some other overall program that you want a gradient for?

3. Lineax is not using a preconditioner for the backward pass. Writing this issue, I realized this has been fixed in Fixes for preconditioned krylov solvers #162. Though, wouldn't x0.T be a good initial guess for the backward pass too?

Did you install from Github (current main) or from PyPi? We're preparing a release and the PR you mention will be in it.

[gaspardbb]

What exactly do you mean by issues with convergence? If it comes from gradient computation, does that mean that the iterative linear solve is part of some other overall program that you want a gradient for?

Exactly. Basically, it looks like:

def cost_fn(params):
  solver_out = lx.linear_solve(linear_operator(params))
  return ((solver_out.value.dot(op) - y) ** 2).sum()

i.e. we're minimizing |x(params) ⋅ op - y|² where A(params) x = b. This outer optimization loop (on params) is done with gradient descent.

I think that issue is coming from worse conditioning on the backward pass as well.

Can you elaborate on that?
The backward pass solves something like M.T A.T x = M.T b, with M the preconditioner. This has the same condition number as M A. No?

I don't know how to capture the variable of the backward pass.

Would you have pointers on that? I'd love to be able to catch the values which cause the solve to fail on the backward pass. I've tried os.environ['EQX_ON_ERROR'] = 'breakpoint' but the whole thing runs indefinitely.

Did you install from Github (current main) or from PyPi? We're preparing a release and the PR you mention will be in it.

Indeed, now I install from main to benefit from #162

[johannahaffner]

Alright, got it! Thanks for the extra context.

I think that issue is coming from worse conditioning on the backward pass as well.

Can you elaborate on that? The backward pass solves something like M.T A.T x = M.T b, with M the preconditioner. This has the same condition number as M A. No?

You're right on that, I did not recall that GMRES requires nonsingular matrices. If the solution of the linear solve is a pseudoinverse then the backward pass is a bit more complicated, but that is not the case here.

Are you sure that your operator is nonsingular / remains nonsingular throughout your outer optimisation? You could try to add a small multiple of the identity (Tikhonov regularisation) if not.

I don't know how to capture the variable of the backward pass.

Would you have pointers on that? I'd love to be able to catch the values which cause the solve to fail on the backward pass. I've tried os.environ['EQX_ON_ERROR'] = 'breakpoint' but the whole thing runs indefinitely.

The linear solve on the backward pass gets computed here:

https://github.com/patrick-kidger/lineax/blob/main/lineax/_solve.py#L243

I'd install it in editable mode and put some jax.debug.print statements in there. I can't help much with the environment variables, I don't really use this way of debugging things.

[gaspardbb]

Thanks.

I'll investigate this and will come back here if I found the solution. I'll start by calling jax.custom_linear_solve to experiment a bit with the transpose solve, see if it's really the culprit.

[johannahaffner]

I'll investigate this and will come back here if I found the solution. I'll start by calling jax.custom_linear_solve to experiment a bit with the transpose solve, see if it's really the culprit.

That would be great, thanks! I also appreciate that my answer isn't super satisfactory. If you have an MWE that reproduces the issue, then I'd be happy to give it a spin myself - to check if there is something about our error handling that could make potential issues more transparent here. I probably won't have time until mid December, though.