SEBO (sparsity-aware BO): Tracking Issue

[lena-kashtelyan]

Motivation

Tracking issue for all SEBO-related discussion (this functionality has been in alpha-mode and not fully supported, but we have plans to add full support for it in the medium-term).

Issues we are tracking under this currently:

• https://github.com/facebook/Ax/issues/3794
• https://github.com/facebook/Ax/issues/3518

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

This issue will encapsulate SEBO-related requests.

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

• [x] I agree to follow Ax's Code of Conduct

[CompRhys]

With SEBO having been reaped from Ax (https://github.com/facebook/Ax/pull/4417) this should probably be closed.

can an example be provided for the intended replacement behavior with pruning? MOO to induce sparsity is important in many chemistry/materials problems and SEBO was useful for us. We can maintain independently of Ax so not a complaint about it's removal necessarily but more interest in what the alternative looks like. Pruning irrelevant parameters is qualitatively different than exploring the complexity performance trade-off for our purposes.

There is no longer tutorial example for how to do this and doc's search functionality comes up blank for "pruning"

@dme65 as the sebo author and reaper

[dme65]

Hi @CompRhys,

We ended up reaping SEBO since it had deviated from the base Acquisition class and was getting very hard to maintain. Attaching a short snippet below that shows how to use the recently added pruning functionality. @sdaulton and I are planning to put together a tutorial in the near future.

import numpy as np
from ax.api.client import Client
from ax.api.configs import RangeParameterConfig
from ax.core.arm import Arm
from ax.utils.measurement.synthetic_functions import Branin


DIM = 10
N_TRIALS = 50
branin = Branin()

# Initialize the Client.
client = Client()

# Configure where Ax will search (Branin embedded in DIM dimensions).
client.configure_experiment(
    name="embedded_branin_function",
    parameters=[
        RangeParameterConfig(
            name=f"x{i}",
            bounds=(-5.0, 10.0) if i == 0 else (0.0, 15.0),
            parameter_type="float",
        )
        for i in range(DIM)
    ],
)
client.configure_optimization(objective="-1 * branin")

# Turn on pruning
client.configure_generation_strategy(simplify_parameter_changes=True)

# Specify the pruning target parameterization.
optimization_config = client._maybe_experiment.optimization_config
optimization_config.pruning_target_parameterization = Arm(
    name="pruning", parameters={f"x{i}": 0.0 for i in range(DIM)}
)
client.set_optimization_config(optimization_config=optimization_config)

# Start the optimization.
for _ in range(N_TRIALS):
    for trial_index, parameters in client.get_next_trials(max_trials=1).items():
        client.complete_trial(
            trial_index=trial_index,
            raw_data={"branin": branin(np.array([parameters["x0"], parameters["x1"]]))},
        )

# Print the best point found.
client.get_best_parameterization(use_model_predictions=False)

[zhqrbitee]

Hi @dme65 , I come across this issue when searching SEBO (we have an optimization case shown below).

So how should I understand the pruning_target_parameterization? In your example, you set all the values to 0.0, is there any physical meaning for 0.0?

optimization_config.pruning_target_parameterization = Arm(
    name="pruning", parameters={f"x{i}": 0.0 for i in range(DIM)}
)