Flexible convergence criteria: an AbstractConvergence interface

[michael-0brien]

Welp, brace yourself: here's an initial attempt for addressing #182. I propose creating an interface for handling terminations:

class AbstractTermination(eqx.Module):

    @abc.abstractmethod
    def __call__(self, y, y_diff, f, f_diff):
        ...

class CauchyTermination(AbstractTermination[Y]):

    rtol: float
    atol: float
    norm: Callable[[PyTree], Scalar] = max_norm

    @override
    def __call__(self, y, y_diff, f, f_diff) -> Bool[Array, ""]:
        ...

I've put together a very rough draft with this proposal, where tests are running for minimizers and least squares solvers. Halfway into working on this I realized that rtol and atol and pretty baked into the code and it's a pretty large refactor. However, I think this would be a significant enhancement to optimistix as I don't think that changing termination criteria are edge cases (for example, scipy.optimize.least_squares has different tolerances and norms for function vs. parameter convergence). This is particularly needed for messy problems on noisy scientific data. There are a few current issues implementing this in practice:

• Currently I've clobbered the root finders and the ability to convert problems between one another (e.g. least squares -> minimize or root find -> least squares). This is because I needed to remove the atol, rtol, and norm from the AbstractIterativeSolver interface. This could potentially be replaced by the AbstractTermination interface, but I didn't want to do this without discussion (also, it is a little out of my depth).
• I've written this as a non-API breaking change, where advanced users of optimistix will be able to take advantage of the new interface. It is particularly easy for users to do this for cases where implementing a new solver requires only writing an __init__ (i.e. the mixing and matching approach), but it is not-so-elegant to take advantage of the new changes for cases where the AbstractMinimiser, etc are subclassed directly to a concrete class (e.g. in the case of the OptaxMinimiser). In these cases, the user can resort to a tree_at call. This isn’t so bad but would be good to talk about what the right approach is.

[patrick-kidger]

Thanks for this! So:

1. does this just end up being a way to factor out the terminate method? If so, then this sounds like a fairly straightforward addition, and it sounds good to me!
2. however 'Halfway into working on this I realized that rtol and atol and pretty baked into the code and it's a pretty large refactor.' does give me pause :D Where are you finding that this is baked in? (I can completely believe that it is.) If this is going to be something more complicated then I think this might be out-of-scope.

(Don't forget that you can always already override the terminate method by subclassing/composing.)

[michael-0brien]

Thanks for this! So:

1. does this just end up being a way to factor out the terminate method? If so, then this sounds like a fairly straightforward addition, and it sounds good to me!

The idea is to replace lines like:

terminate = cauchy_termination(
    self.rtol, self.atol, self.norm, state.y_eval, y_diff, f_eval, f_diff
)

with:

terminate = self.termination(state.y_eval, y_diff, f_eval, f_diff
)

I keep API the same, so for example the BFGS.__init__ would replace:

self.atol = atol
self.rtol = rtol
self.norm = norm

with:

self.termination = CauchyTermination(rtol, atol, norm)

The essence of this is to indeed factor out the terminate method, but in practice it is slightly challenging to do this. Consider the AbstractGaussNewton.terminate. Here the termination was evaluated during the step, and the terminate method forwards along the value of the termination stored in the solver state. There are a few examples like this.

2. however 'Halfway into working on this I realized that rtol and atol and pretty baked into the code and it's a pretty large refactor.' does give me pause :D Where are you finding that this is baked in? (I can completely believe that it is.) If this is going to be something more complicated then I think this might be out-of-scope.

The issue is the AbstractIterativeSolver interface declares rtol, atol, and norm as AbstractVars. I think it should be relatively straightforward to remove this, but there are still challenges (for example, I still need to fix the code where root finders are converted to other problems). I wouldn’t be surprised if it were necessary to replace it with another AbstractVar or two somewhere, but I haven’t quite sorted this out yet.

Note that for this change, minimisers and least squares solvers would get the new interface, but I haven’t messed with things like the FixedPointIteration class, which still has the rtol, atol, and norm directly as fields. There are a few other classes that would maintain these as fields, but it would no longer be abstract API. Philosophically, this makes sense to me; I think it is too hard to unite all problems with this API.

There are a few other things to potentially discuss as well—for example, I could see there being a better naming pattern for the new interface.

[jpbrodrick89]

Just wanted to weigh in that this does sound really useful in principle as could help address #165 and allow for super cheap termination routines that just always return False so I can deterministically run for a certain number of steps.

[jpbrodrick89]

Just took an every so slightly closer look at this, I really like the idea, I would lean towards trying to support this with root finders too, particularly NewtonChord, and offer the HairerWanner termination as its own termination class. Offering one more standard termination class where it only checks for f convergence and not y might also be desirable in some cases where performance matters more than accuracy/robustness (of course just leaving it to user's to implement themselves to prevent offering a footgun seems sensible). If you need a hand with Newton Chord I'd be happy to help as I've just been getting very deep with that this week.

[michael-0brien]

@jpbrodrick89 this sounds reasonable to me! If I had to guess I think this PR is probably most likely to be pushed through if a comprehensive alternative is offered compared to the current model for terminations, so extending this to root finders sounds great. I could use the help as I don’t work on root finders. If you also want to take a whack at fixing where root finds are converted to minimizations and least squares solves please go for it as well!

[michael-0brien]

Just wanted to weigh in that this does sound really useful in principle as could help address #165 and allow for super cheap termination routines that just always return False so I can deterministically run for a certain number of steps.

Keep in mind with this it can be tricky what ends up being “cheap” termination at the end of the day. This PR would only change termination criteria, not the JAX control flow in a given iterative solver (which is expensive on GPU). Unless the compiler figures out it should remove that control flow as well, I suspect there won’t necessarily be performance increases.

[michael-0brien]

@patrick-kidger @johannahaffner when you have a moment could you help me come to a solution on this? I can work on this but want to make sure it’s the right direction, and if not what would be. There are two approaches to resolution in my mind:

1. Something like the current approach (add an API).
2. Some solution where flexible termination can be solved by overwriting terminate methods (keep API the same but allow for hacking). While this already works for some solvers, it does not for others like AbstractQuasiNewton classes

I would lean towards 1. because at the least optimistix should support the ability to have different norms and tolerances for y and f as this is quickly broken by many problems, but curious if there is a feasible path for 2. as well.

[michael-0brien]

To present an alternative to the current version of things, it would be possible to preserve atol and rtol as a part of the AbstractIterativeSolver API. This could look like the following:

class AbstractTerminationCriteria(eqx.Module):

    atol: eqx.AbstractVar[float]  # these are forwarded to the `AbstractIterativeSolver` API, when appropriate
    rtol: eqx.AbstractVar[float]
    norm: eqx.AbstractVar[Callable]

    @abc.abstractmethod
    def __call__(...):
        ...

class CustomTerminationCriteria1(AbstractTerminationCriteria):

    atol: float  # atol and rtol are used for function tolerance
    rtol: float
    yatol: PyTree[float]  # pytree-valued atol and rtol for parameter tolerance 
    yrtol: PyTree[float]
    norm: Callable

    def __call__(...):
        ...

class CustomTerminationCriteria2(AbstractTerminationCriteria):

    atol: float
    rtol: float
    norm: Callable

    def __call__(...):
        #  atol and rtol only used for function tolerance; no parameter tolerance considered
        ...
 

Under this idea I am still not sure how to handle the issue of problem conversion (i.e. using a minimiser to solve a least squares problem). In either case this is the only thing that breaks the proposed API. I personally think this is a bit messier, but it is less API breaking in case users are accessing solver.atol/rtol/norm.

[jpbrodrick89]

Apologies for the lack of update on my side, the reason I've not attempted this for root finders yet is we're still trying to determine the correct non-Cauchy termination conditions for Newton.

[michael-0brien]

@jpbrodrick89 no problem! There's no rush, I mostly want to discuss direction so I can start using this PR in my own fork. I also quickly want to correct myself:

Under this idea I am still not sure how to handle the issue of problem conversion (i.e. using a minimiser to solve a least squares problem). In either case this is the only thing that breaks the proposed API. I personally think this is a bit messier, but it is less API breaking in case users are accessing solver.atol/rtol/norm.

Looking through the code a little bit more closely, actually root finder problem conversion (i.e. using least squares solver or a minimiser to solve a root find) is the only case where this PR breaks. This is good news actually! I bet we can find a simple way to address this. Further, the current method of handling terminations here looks slightly hacky so I bet there's an elegant solution.

[johannahaffner]

Sorry for the slow replies, @michael-0brien! I was mulling this over a few days ago, and indeed we should also think about @jpbrodrick89's change to the termination criteria for the root finders in this context.

I'll try to sit down with pen and paper in the coming days and map this out.

I prefer an API break with a clean solution over some hybrid, and since concrete solvers already have an __init__, forwarding to wherever we think is appropriate is easy. The fraction of power users that define solver attributes for custom solvers is, I think, competent enough to adjust their code and upgrade to declaring new attributes.

[michael-0brien]

Thanks so much @johannahaffner! Let me know how I can help—maybe there is a way to adapt the current PR to encompass changes to the root finder behavior. This would make things much more complete, the current draft really only focuses minimizers and lstsq solvers!

One argument for something like the current proposal, which I haven’t had a chance to write a draft of, is the AbstractTermination API may fit nicely into the advanced API docs that currently has searches and descents. Thematically it is very similar—there is default behavior for every optimizer in optimistix, but it is extremely straightforward to create a concrete class with custom searches, descents, and now possibly termination criteria as well.

[michael-0brien]

Hi, ping! Wanted to check in. Wanted to share how I happen to be using this interface downstream:

class CustomTermination(optx.AbstractTermination[Y]):
    ftol: float | tuple[float, float]
    ytol: float
    y_scale: float | PyTree[float]
    norm: Callable = optx.max_norm

    @override
    def __call__(self, y: Y, y_diff: Y, f: F, f_diff: F) -> Bool[Array, ""]:
        eps = jnp.finfo(jax.dtypes.canonicalize_dtype(float)).eps
        if isinstance(self.ftol, Sequence):
            assert len(self.ftol) == 2
            frtol, fatol = self.ftol
        else:
            frtol, fatol = self.ftol, eps
        y_scale, f_scale = self.y_scale, ω(f).call(jnp.abs).ω
        if isinstance(y_scale, float):
            y_scale = _tree_full_like(y_diff, y_scale)
        y_converged = (
            self.norm((ω(y_diff).call(jnp.abs) / (eps + self.ytol * y_scale**ω)).ω) < 1
        )
        f_converged = (
            self.norm((ω(f_diff).call(jnp.abs) / (fatol + frtol * f_scale**ω)).ω) < 1
        )
        return y_converged & f_converged

This a) separates ytol and ftol criteria, b) allows per-parameter convergence criteria via a PyTree-valued y_scale, and c) overrides optimistix's default behavior of y_scale = (ω(y).call(jnp.abs)).ω, which for some optimization problems, like what I am doing, would be a bug / is an incorrect thing to do.

Let me know if anyone has an updated thoughts on this PR (@patrick-kidger & @johannahaffner?). Thanks so much in advance and no rush, I understand how busy the maintainers must be.

[patrick-kidger]

'Busy' is the right word these days... @johannahaffner and I recently made the questionably-wise decision of founding a startup :D This does mean that we're picking-and-choosing our battles a little more carefully on which PRs to work through and accept.

For that reason then right now I think – realistically – this PR is probably one of the lower-priority ones. Termination can anyway be overriden through subclasing/composing, so IIUC this is really a bit of polish on what is already a fairly advanced feature? Please let me know if you think I've misunderstood something there.

[michael-0brien]

'Busy' is the right word these days... @johannahaffner and I recently made the questionably-wise decision of founding a startup :D This does mean that we're picking-and-choosing our battles a little more carefully on which PRs to work through and accept.

Makes sense!

For that reason then right now I think – realistically – this PR is probably one of the lower-priority ones. Termination can anyway be overriden through subclasing/composing, so IIUC this is really a bit of polish on what is already a fairly advanced feature? Please let me know if you think I've misunderstood something there.

This is true in some cases, but not all. E.g. in AbstractBFGS termination cannot be overwritten without overwriting the hessian update. I am forgetting right now if this is an exception or if there are other cases as well, but I recall there are other cases as well where the termination method ends up simply returning a boolean passed from the state pytree.

I would also argue that there is a reason frameworks like scipy.optimize have more built in control of termination (i.e. scipy.optimize.minimize). There it is top-level API, not an advanced feature; my guess being that it is an important part of making the algorithms widely applicable. In making optimistix the scipy.optimize of JAX, I think this issue while small is pretty significant. Others I have spoken to also are in the world of “inference of physical parameters on experimental data in the presence of noise” seem to have similar experience.

Here’s an example I have not yet already mentioned in this thread: consider optimizing the translation of an object; i.e. there is noisy data with an object in frame, and you want to find how to shift a model to the reference. Let’s say the domain is [-100, 100], for the sake of demonstration. The object can be anywhere in frame, therefore whether it is located at x = 90.0004 or x = 0.0004, say, my localization precision does not change in terms of absolute x. Let’s say my noise threshold is such that I can only get three decimals of precision. In either case my x_scale should be ~ 10^-3, where my optimizer loses sensitivity. What is currently rtol when applied to x does not have meaning.

In this toy example, one could mess around with rtol and try to only converge on atol. But introduce two parameters, each with their own properties and sensitivities, and optimistix can’t accommodate. In any case, this is just one aspect added onto previous discussion in this thread; perhaps the biggest limitation in practice is still having the same rtol and atol for function and parameter convergence.

I understand things are quite busy and thank you for making the time to respond! I do feel passionately that this issue is an important one, and I suppose the only reason not to include it is that it is slightly non-trivial to engineer into the current code. Another route would be to consider making it truly possible to factor out each terminate method in the library, but I think I may have found that more complicated when I wrote this PR.

[michael-0brien]

Hi all, I've resolved outstanding issues and gotten this PR to the point of review. There has been a lot of discussion in this thread (mostly from me, haha), but to TL;DR:

1. To factor out termination criteria, we cannot easily overwrite terminate. As one example, AbstractQuasiNewton stores the convergence check in the state and passes it through terminate.
2. Therefore, I have proposed that we add an interface for the convergence check called AbstractConvergence with a single method called check.
3. To achieve this, it is necessary to remove atol, rtol, and norm on the AbstractIterativeSolver interface. Aside from adding this feature, I think this was an over-constrained design for the library and is generally a healthy thing to do.

Backwards compatibility: AbstractRootFinder and AbstractFixedPointSolver are fully backwards compatible as I have added atol, rtol, and norm to their interfaces. If there are changes to this in the future, I think this is a separate PR. AbstractMinimiser and AbstractLeastSquaresSolver lose their atol, rtol, and norm pytree leaves and gain a convergence abstract field. These could be declared as properties for more backwards compatibility.

I would have liked the convergence abstract field to be declared at the same level of abstraction as searches and descents (e.g. in AbstractQuasiNewton, rather than AbstractMinimiser). It was necessary to have the new interface on AbstractMinimiser to preserve the problem-conversion behavior in _root_find.py.

@patrick-kidger @johannahaffner @jpbrodrick89 no rush but when you get a moment, please do take a good look and let me know if you have any questions!