Is it possible to specify an outcome_constraint on ax_client.get_next_trial() with an exactly known (non-objective) outcome?

[sgbaird]

#273 seems very relevant to this. For my use-case #727, I've thought about having:

import pandas as pd
outcome_constraints = ["n_components <= 8.0"]

to constrain a compositional formula to have no more than 8 components above some low threshold (e.g. 1e-3). In my case, n_components is calculated very simply (and exactly) using only the parameters. For example:

def count_nonzero_components(parameters, tol=1e-3):
    df = pd.DataFrame(parameters)
    df[df < tol] = 0.0
    n_components = np.count_nonzero(df, axis=1)
    return n_components

However, the first outcome (the objective) is not known a-priori and can only really be sampled via wet-lab synthesis. In #273, it seems like the outcome constraint metric is being estimated from the GP model (I could be wrong on this). Is there a way to "tell" ax_client.get_next_trial() that n_components = count_nonzero_components(parameters) and that n_components <= 8.0?

I've been digging through custom generation strategies and custom acquisition functions #278, but without much luck so far. Similar to #278, I'm also using the Service API.

[sgbaird]

I'm open to switching to the Developer API as well, just been getting a bit lost/overwhelmed 😅. Let me know if I should post a MWE using the Developer API. Thanks again for all the help so far!

[lena-kashtelyan]

Hi @sgbaird, you should definitely be able to use Service API for this. Will get back to you with some help soon! These questions require some thought, but we're not ignoring them.

[sgbaird]

@lena-kashtelyan, thank you! That is good to know. I will hold off on the conversion to the Developer API. Thanks for the patience as I've been bombarding Ax Issues over the last couple weeks!

[lena-kashtelyan]

I will hold off on the conversion to the Developer API

I'd recommend that, yes!

[lena-kashtelyan]

Is there a way to "tell" ax_client.get_next_trial() that n_components = count_nonzero_components(parameters) and that n_components <= 8.0?

I think the issue is that what you are looking to do is really a parameter constraint, hence "n_components is calculated very simply (and exactly) using only the parameters". However, it's not a linear constraint, so it cannot be easily expressed in Ax right now (non-linear constraints are on our wish list: #566, but with uncertain status). I think something like this was also brought up on the original issue you posted: #727, and @dme65 and myself brainstormed that without an easy solution, unfortunately.

What you are doing right now (by expressing n_components as an outcome constraint) is working around the requirement of linearity that we have for parameter constraints; expressing the constraint as an outcome constraint indeed results in it being modeled using a GP during candidate generation, as you point out: "it seems like the outcome constraint metric is being estimated from the GP model." This kind of constraint will be tricky for the model to learn, so most likely the generated trials will not be observing the constraint well, and it might be throwing off the model.

We are still brainstorming how exactly you could express the constraints you need in Ax (any API); in the meantime, if you'd like to post an up-to-date version of your problem (is it still the one in #727 or have things changed?), that would help us. If this issue and the one in #727 are one and the same, feel free to close whichever one is less up-to-date and just keep one open, so we have the discussion centralized, @sgbaird!

[sgbaird]

@lena-kashtelyan that helps out and I think resolves this issue. It sounds like it may not be feasible in Ax's current implementation to swap out the GP model with a custom model in order to estimate the outcome constraint during trial generation, and that there may be some fundamental issues with this approach (i.e. trying to use a non-linear parameter constraint). EDIT unless @Balandat has an idea of how to use a custom generation strategy with a custom model that returns two outcomes, i.e.:

GenerationStrategy(
    steps=[
        GenerationStep(
            model=<my_custom_model_that_returns_two_outcomes>,
            ...
        )
    ]
)

in which case it seems like I could evaluate the second outcome exactly.

If I were to try to force a workaround using outcome_constraint without replacing the GP model, I imagine I could add tons of synthetic trials (i.e. random parameterizations) that return (NaN, n_components > 8). This would help the GP model to estimate n_components reasonably well; however, it could also throw off the model as @lena-kashtelyan mentioned. I don't know enough about the following two points to say how or to what extent:

1. how the GP model treats NaN values, e.g.
   1. ignoring NaN values entirely
   2. "blacklisting" nearby areas
   3. creating "breaks" in the model space (i.e. close points no longer affect each other as much because there is a NaN value in-between)
2. whether the GP model prediction for one outcome affects the prediction for another (e.g. strength predictions affecting n_components predictions and vice-versa)

I think I'm OK with "pruning" this approach for now in favor of something else. #727 (comment)

[lena-kashtelyan]

@sgbaird, this is definitely a tricky issue, and I think you are right that this approach would likely have issues:

If I were to try to force a workaround using outcome_constraint without replacing the GP model, I imagine I could add tons of synthetic trials (i.e. random parameterizations) that return (NaN, n_components > 8). This would help the GP model to estimate n_components reasonably well; however, it could also throw off the model as @lena-kashtelyan mentioned

cc @Balandat to check me on that one and also to overall check the thinking here

[Balandat]

Late to the party, but Lena is spot on with saying that

I think the issue is that what you are looking to do is really a parameter constraint, hence "n_components is calculated very simply (and exactly) using only the parameters". However, it's not a linear constraint, so it cannot be easily expressed in Ax right now

We could try to model this constraint as a black box function, but that will cause all kinds of issues and will also be quite inefficient (why do we learn the constraint if we already know it in closed form?).

I think that, as with your other issue, the correct approach here would be to use a custom generation strategy that is able to take in nonlinear constraints. If you look at your constraint \sum_i 1{x_i > 1e-3} < 9, you'll see that it's not just nonlinear but also non-convex (which I guess isn't really too big of a problem since the acquisition function is usually non-convex also).

So the proper solution here would be to change the API and allow Ax to take in some callable that evaluates the constraint and that we can pass to the optimizer. This is not too hard in principle, but b/c of the various transformations and normalizations that we apply in the modelbridge layer to both parameters and data, this can cause a bunch of headaches in practice. Essentially, allowing this would provide an excellent way for people to shoot themselves in the foot. That said, there clearly seems to be a need for this functionality, so maybe the right thing to do would be to throw together a proof of concept and just put slap a big red warning sign on it?