Ax: How to allow multi-objective metrics to share results of expensive evaluation function to avoid recomputing it?

[du-210]

Hi,

I am using Ax for a multi-obj optimisation problem. Following the tutorial on Multi-Objective Optimization Ax API, I defined a MultiObjectiveOptimisationConfig containing two metrics that I want to optimise. These two metrics run an exactly the same computationally expensive objective function and extract a result value from a returned dictionary. The problem here is that every time that I fetch data for a new trial, I ended up running my objective function twice for these two metrics.

I also looked at Trial Evaluation tutorial where 3 different paradigms for evaluating trials are introduced. I think using Synchronous paradigm should solve my issue however, I have to use SimpleExperiment class for it that is not designed for multi-objective optimisation.

Is there any way around this problem within the Ax framework?

[lena-kashtelyan]

Hi, @du-210, this is an interesting problem! There are a few possible approaches:

1. Easiest: configure your metrics in a way that writes the data somewhere (e.g. to a local file) after computing the expensive function. This way, in your MyMetric.fetch_trial_data method (example in tutorial), you could first check that place and grab data from there if it already exists? You will need to use these metrics instead of NoisyEvaluationMetric-s used in the multi-objective tutorial.
2. Possible: use the attach_data flow and bypass fetching data through metrics. I elaborated on this below.
3. Probably the best, but not yet available path: use Service API. Stay tuned for multi-objective optimization support in the Service API!

---

Elaborating on option 2, if option 1 is not possible: to leverage attach_data flow instead of defining your metrics, you will need to configure your optimization config to use just base class Metric-s and to manually add data through Experiment.attach_data. Your main goal is to ensure that for any trial, data for which will be fetched (via experiment.fetch_data or trial.fetch_data), data is attached before fetching, and the steps below will hopefully clarify what that means.

Here is what you will need to change from the multi-objective optimization tutorial:

1. Set your metrics in MultiObjective to be just base Metric-s, not NoisyFunctionMetric-s or any other manually defined metric class. You can just define each metric like so: metric_a = Metric(name="my_metric_a", lower_is_better=True); you should not need to specify anything else or define any new metrics.
2. Create a get_data function that manually creates data for the metrics you need while only running the expensive evaluation function once. This function needs to output an Ax Data object. Example of such function is here in cell 9: https://ax.dev/tutorials/building_blocks.html#4.-Define-an-optimization-config-with-custom-metrics; here it is used to implement BoothMetric.fetch_data, but your function will be standalone instead. Something like this:

from ax.core.data import Data
from ax.core.trial import Trial

def get_data(trial: Trial) -> Data:
    expensive_function_output_dict = expensive_evaluation_function()
    df_dict = {
        "trial_index": trial.index,
        "metric_name": [metric_a.name, metric_b.name],   # Names of your metrics here
        "arm_name": trial.arm.name,
        # Below, pass metric mean values in same order as "metric_names":
        "mean": [expensive_function_output_dict["metric_a"], expensive_function_output_dict["metric_b"],
        "sem": [None, None],  # Can be actual SEM float values if known
    }
    return Data(df=pd.DataFrame.from_records(df_dict))

1. Set initialize_experiment to this:

def initialize_experiment(experiment):
    sobol = Models.SOBOL(search_space=experiment.search_space)

    for _ in range(N_INIT):
        trial = experiment.new_trial(sobol.gen(1)).run()
        experiment.attach_data(get_data(trial))

    # Because you are using `attach_data` and base `Metric`-s (which do not implement
    # fetching logic), this call will now just grab cached data you attached).
    return experiment.fetch_data()

1. Everywhere you do trial.run subsequently, also attach the data for that trial. Assuming you are just running the EHVI algorithm, that will just be in cell 14 in this section: https://ax.dev/tutorials/multiobjective_optimization.html#qEHVI. That cell will now look like this:

ehvi_data = initialize_experiment()
...
for i in range(N_BATCH):   
    ehvi_model = get_MOO_EHVI(
        experiment=ehvi_experiment, 
        data=ehvi_data,
    )
    generator_run = ehvi_model.gen(1)
    trial = ehvi_experiment.new_trial(generator_run=generator_run)
    trial.run()
    trial_data = get_data(trial)  # This line is new and uses `get_data` function you added! 
    experiment.attach_data(trial_data)  # This line is new!
    ehvi_data = Data.from_multiple_data([ehvi_data, trial_data])  # This line changed! 
    ...

Let me know if these steps do not make sense or if you run into any issues!

[Balandat]

Are you evaluating this function locally? If so another option would be to just cache the arguments like so:

from functools import lru_cache  # could use cache in py 3.9+

def expensive_function_wrapper(params):
    return expensive_function(**params)


@lru_cache(maxsize=maxsize)
def expensive_function(**params):
    # perform your computation here ....
    return {"metric_1": value_1, "metric_2": value_2}

That way calling expensive_function_wrapper with the same args twice will result in using the cached result rather than recomputing it. (the wrapping part is necessary b/c you can't hash the dict passed to expensive_function_wrapper directly). That way you won't have to make any changes to the Ax part of your setup. If you want to evaluate the same metric of different trials before going back to evaluate other metrics, you'll want maxsize of the chace to be larger than the total # of trials you run.

[du-210]

Hi @lena-kashtelyan and @Balandat,

Thanks for your quick answer! I have been using the proposed solution by Balandat but I was wondering if this can be sorted within Ax. I tried Lena solution and seems working with one small change. Trials should be marked completed otherwise fetch_data would not show anything.

Am I right to say that by using the proposed approach, we would not be able to model a noisy function anymore? If so, do you have any suggestions to address this issue?

[lena-kashtelyan]

There are a few things going on here, so I'll number my answers to make to easy to reply : )

I have been using the proposed solution by Balandat but I was wondering if this can be sorted within Ax.

1. I'm not sure I understand what "sorted within Ax" means here. Do you mean that Ax can expose some metric template for metrics that share a function that is expensive to compute?

---

I tried Lena solution and seems working with one small change. Trials should be marked completed otherwise fetch_data would not show anything.

1. You mean the solution #2, correct? That's right, sorry I forgot to mention that trials need to be marked as COMPLETED. That is something we should fix in the tutorial overall, we will make a note of it.

---

Am I right to say that by using the proposed approach, we would not be able to model a noisy function anymore? If so, do you have any suggestions to address this issue?

1. With approach #2 from my answer, as long as in your Data (returned from get_data), the SEM is returned as None and not as 0.0, Ax will be assuming that the function is noisy and using an optimization model that infers the noise levels. To specify that a function is noiseless, you would need to manually set the SEM to 0.0 –– we assume noise by default.

2. With @Balandat's approach (I'm assuming you apply it by setting the f in your metric to expensive_function_wrapper), the SEM given to Ax will depend on the noise_sd setting you pass to the metric (e.g. metric_a = MetricA("a", ["x1", "x2"], noise_sd=0.0, lower_is_better=False) means that SEM will be set to 0 and Ax will be modeling your metrics as noiseless. You will have to manually pass noise_sd=None (the default for that argument is 0) to indicate to Ax that your metrics are noisy.

[du-210]

Thanks @lena-kashtelyan for the detailed answers,

1. Yes, that's exactly what I meant. That would be really great if such template added to the Ax framework. I think it is quite standard to use a shared function among several metrics in multi-objective problems.
2. That's right. I tried the solution #2 as I didn't want to rely on an external cache.
   3, and 4. it is more clear now how to model noisy functions.

[lena-kashtelyan]

1. Got it! I think in general the Service API caters to this use case, and it will be extended with multi-objective optimization in the near future.
2. Awesome; let us know if you need more help! And feel free to reopen the issue if you have follow-ups.