[FEATURE REQUEST]: Query Max Expected Improvment Directrly from Ax Client

[avivgood]

Motivation

The convincing stopping rule for BO's original paper was that max EI < 1%. Therefore it would be great if you could query max EI directly and use it as a stopping rule.
Currently there are two ways to do it: Either calculate manually or get the next trial and eval EI there.
Calculating manually is a pain because it means re-implementing the max EI algorithm the same exact way it is implemented in the backend of Ax (and you also need to consider here complex multimodel generation strategy), and getting the next trial and evaling EI there is also not so good because are so Bothorch hacks needed for this and it again depends on the model, and it is quite verbose.

Describe the solution you'd like to see implemented in Ax.

It would be beneficial if you could just do client.max_expected_improvement (not exact API) and even better if you could do a botorch generation strategy that stop generating after a certain max EI threshold is reached.
It will also be a nice addition to compute_diagnostics()

Describe any alternatives you've considered to the above solution.

No response

Is this related to an existing issue in Ax or another repository? If so please include links to those Issues here.

No response

Code of Conduct

• I agree to follow Ax's Code of Conduct

[mgarrard]

Hi @avivgood thanks for reaching out -- cc @mpolson64 for a request to AxClient + analyses.

Curious if you've had a chance to check out our early stopping tutorial: https://ax.dev/docs/tutorials/early_stopping/ this may be useful w.r.t your idea about a GenerationStrategy that could stop at some threshold :)

[lena-kashtelyan]

Hi @avivgood ! We call this kind of "overall optimization stopping" once a certain threshold is reached, "global stopping" (as opposed to "early stopping", which refers to stopping individual trials that are currently running (typically to save capacity).

An EI threshold-based global stopping strategy relies on the fit of the BayesOpt model being good, which is not always guaranteed. If there is model mismatch, it's possible to be overconfident and stop too early since the computed expected improvement is lower than what improvement could be achieved on the actual problem.

At the end of the day, it's a tradeoff between performance and sample efficiency. We could implement an ExpectedImprovementGlobalStoppingStrategy (let us know if you are interested in taking a stab at implementing this!), but we'd be hesitant to enable it by default.

In the meantime, you can get the acquisition function values the Ax generators produced for each arm in a given trial, like so:

trial = client.experiment.trials[trial_index]
trial.generator_run.gen_metadata["expected_acquisition_value"]  # Should be a one-element list if using `Client` API

Let us know whether this helps!

[avivgood]

@lena-kashtelyan

1. Jones suggests multiple strategies when you want a EI stopping rule but model isn't fitted properly, such as getting EI < 1% twice, or reducing it to < 0.5%, I don't know what the modern EI based stopping rules are, I am sure they are more sophisticated. We can derive from this multiple generations strategies.
2. Of course not enabled by default but as a custom opt in stopping strategy.
3. As for the suggested way to query EI, I think it would be beneficial to add a computed getter directly under Ax client, as the current API is a bit complex.
4. What do you think about the other suggestions of adding it to diagnostics
5. I am ready to try to implement this feature myself. I am not sure how the process works for this repo tho, do I simply get open a PR and feedback is there? Or do you have a step where you discuss the exact API with a maintainer?