Ax: Modeling a hierarchical search space using ParameterConstraint

[FelixNeutatz]

Hi everybody,

I saw that issue #140 already discusses the extension for arbitrary nesting. I was thinking about how to implement a hierarchical search space using the parameter constraints.
E.g.

A = ChoiceParameter('A', parameter_type=ParameterType.BOOL, values=[True, False], is_ordered=False)
if A == True:
    B = ChoiceParameter('B', parameter_type=ParameterType.BOOL, values=[True, False], is_ordered=False) 

That means that we do not allow the combination A = False and B = True. So, with constraints that would look like this:

SearchSpace(
    parameters=[
        ChoiceParameter('A', parameter_type=ParameterType.BOOL, values=[True, False], is_ordered=False),
        ChoiceParameter('B', parameter_type=ParameterType.BOOL, values=[True, False], is_ordered=False)
    ],
    parameter_constraints=[ParameterConstraint(constraint_dict={'A': -1.0, 'B': 1.0}, bound=0.0)]
)

However, it seems like that ChoiceParameters are not supported for ParameterConstraints:

Traceback (most recent call last):
  File "/home/neutatz/Software/DeclarativeAutoML/fastsklearnfeature/declarative_automl/optuna_package/myautoml/cbays/hierarchical_space.py", line 35, in <module>
    sobol = Models.SOBOL(exp.search_space)
  File "/root/anaconda3/envs/DeclarativeAutoML/lib/python3.8/site-packages/ax/modelbridge/registry.py", line 261, in __call__
    model_bridge = bridge_class(
  File "/root/anaconda3/envs/DeclarativeAutoML/lib/python3.8/site-packages/ax/modelbridge/base.py", line 151, in __init__
    obs_feats, obs_data, search_space = self._transform_data(
  File "/root/anaconda3/envs/DeclarativeAutoML/lib/python3.8/site-packages/ax/modelbridge/base.py", line 196, in _transform_data
    search_space = t_instance.transform_search_space(search_space)
  File "/root/anaconda3/envs/DeclarativeAutoML/lib/python3.8/site-packages/ax/modelbridge/transforms/one_hot.py", line 138, in transform_search_space
    return SearchSpace(
  File "/root/anaconda3/envs/DeclarativeAutoML/lib/python3.8/site-packages/ax/core/search_space.py", line 48, in __init__
    self.set_parameter_constraints(parameter_constraints or [])
  File "/root/anaconda3/envs/DeclarativeAutoML/lib/python3.8/site-packages/ax/core/search_space.py", line 77, in set_parameter_constraints
    self._validate_parameter_constraints(parameter_constraints)
  File "/root/anaconda3/envs/DeclarativeAutoML/lib/python3.8/site-packages/ax/core/search_space.py", line 309, in _validate_parameter_constraints
    raise ValueError(
ValueError: `A` does not exist in search space.

I can work around it like this:

def eval(parameterization):
    A = parameterization.get("A") > 0.5
    B = parameterization.get("B") > 0.5

SearchSpace(
    parameters=[
        RangeParameter(name="A", parameter_type=ParameterType.FLOAT, lower=0.0, upper=1.0),
        RangeParameter(name="B", parameter_type=ParameterType.FLOAT, lower=0.0, upper=1.0)
    ],
    parameter_constraints=[ParameterConstraint(constraint_dict={'A': -1.0, 'B': 1.0}, bound=0.0)]
)

However, it doesn't look so great. Does anybody have an idea how to make this nicer?

Best regards,
Felix

[Jakepodell]

Hey Felix-

Taking a look at this, will have a suggestion shortly!

[2timesjay]

@FelixNeutatz You're right that a continuous relaxation is the best option currently available.

https://ax.dev/api/core.html#module-ax.core.parameter_constraint OrderConstraint will represent your constraint more elegantly (though it is equivalent).

A = RangeParameter(name="A", parameter_type=ParameterType.FLOAT, lower=0.0, upper=1.0)
B = RangeParameter(name="B", parameter_type=ParameterType.FLOAT, lower=0.0, upper=1.0)
SearchSpace(
    parameters=[A, B],
    parameter_constraints=[OrderConstraint(lower_parameter=B, upper_parameter=A)]
)

similarly, a SumConstraint can enforce a continuous relaxation of XOR of two parameters.

Depending on your problem this may or may not model your problem accurately. Can you give us an idea of how many parameters and constraints you'll have?

Note: Our pseudorandom Sobol sampler enforces constraints by rejecting samples that violate them. If you have many constraints it may fail to find a valid sample before it hits a sampling limit. If this is your case let us know and we'll try to provide a solution.

For our BayesOpt models We have an improved solver that should do a decent job than before of sampling these linear-constraint relaxations, including nested constraints. @Balandat Any comments on what to expect performance-wise?

[Balandat]

The hit and run sampler should work decently for this - @sdaulton how far are we away from hooking this up?

Taking a step back, another option would be to just ignore the hierarchical structure on the modeling / optimization side and handle it in the evaluation code. Basically in your eval code you do B = B and A and on the Ax end you just don't impose a constraint. The model should eventually be able to figure out that the value of B is irrelevant when A=False. This is of course not ideal, but it might be a better option than having to deal with rather brittle workarounds on the constraints.

cc @dme65 who has some experience with this and can probably also speak to it.

[sdaulton]

The hit and run sampler should work decently for this - @sdaulton how far are we away from hooking this up?

The hit and run sampler is currently used by default in BO for generating initial conditions for acquisition optimization. Using the hit and run sampler as an Ax RandomModel does not really seem like a high priority at the moment, if the rejection sampling with Sobol is not prohibitively slow and we are not hitting the max sample limit.