Suggestions for implementing a composition-based optimization (i.e. fractional portion of ingredients)

[sgbaird]

For starters, my experience with Ax is running the Loop tutorial once and reading through some of the documentation such as the parameter types (i.e. fairly new). Also, I have some familiarity with Bayesian optimization.

The actual use-case is slightly different and more complicated, but I think the following is a suitable toy example. I go over the problem statement, some setup code, and possible solutions. Would love to hear some feedback.

Problem Statement

Take a composite material with the following class: ingredient combinations:

• Filler: Colloidal Silica (filler_A)
• Filler: Milled Glass Fiber (filler_B)
• Resin: Polyurethane (resin_A)
• Resin: Silicone (resin_B)
• Resin: Epoxy (resin_C)

Take some toy data of components and their fractional prevalences (various combinations of fillers and resins, and various numbers of components) along with their objective (training data), and some model which takes arbitrary input parameters and predicts the objective (strength) which we wish to maximize.

For constraints, I'm thinking:

• limit the total number of components in any given "formula" (e.g. max of 3 components)
• naturally, that the compositions sum to 1 (or that abs(1-sum(composition)) <= tol)
• there has to be at least one filler and at least one resin (if feasible)

Setup Code

To make it more concrete, it might look like the following:

choices = ["filler_A", "filler_B", "resin_A", "resin_B", "resin_C", "dummy"]

data = [
        [["filler_A", "filler_B", "resin_C"], [0.4, 0.4, 0.2]],
        [["filler_A", "resin_A", "resin_B"], [0.6, 0.2, 0.2]],
        [["filler_A", "filler_B", "resin_B"], [0.5, 0.3, 0.2]],
        [["filler_A", "resin_B", "dummy"], [0.5, 0.5, 0.0]],
        [["filler_B", "resin_C", "dummy"], [0.6, 0.4, 0.0]],
        [["filler_A", "filler_B", "resin_A"], [0.2, 0.2, 0.6]],
        [["filler_B", "resin_A", "resin_B"], [0.6, 0.2, 0.2]],
        ] # made-up data

def predict(objects, composition):
    ...
    return obj

Possible Solutions

One-hot-like prevalence encoding and components/composition

One-hot-like prevalence encoding

I've thought about trying to do a sort of "one-hot encoding" (assuming I'm using this term correctly), such that each component gets its own composition as a variable:

filler_A	filler_B	resin_A	resin_B	resin_C
0.4	0.4	--	--	0.2
0.6	0.0	0.2	0.2	--
0.5	0.3	--	0.2	--
0.5	--	--	0.5	--
--	0.6	--	--	0.4
0.2	0.2	0.6	--	--
--	0.6	0.2	0.2	--

which I think would look like the following:

best_parameters, values, experiment, model = optimize(
    parameters=[
        {
            "name": "filler_A",
            "type": "range",
            "bounds": [0.0, 1.0],
        },
        {
            "name": "filler_B",
            "type": "range",
            "bounds": [0.0, 1.0],
        },
        {
            "name": "resin_A",
            "type": "range",
            "bounds": [0.0, 1.0],
        },
        {
            "name": "resin_B",
            "type": "range",
            "bounds": [0.0, 1.0],
        },
        {
            "name": "resin_C",
            "type": "range",
            "bounds": [0.0, 1.0],
        },
    ],
    experiment_name="composition_test",
    objective_name="strength",
    evaluation_function=predict,
    parameter_constraints=["abs(1 - (filler_A + filler_B + resin_A + resin_B + resin_C)) <= 1e-6", "filler_A + filler_B > 0", "resin_A + resin_B + resin_C > 0"], # not sure if I can use `abs` here
    total_trials=30,
)

However, this could easily lead to compositions where all of the components have a finite prevalence and can be problematic from an experimental perspective.

components/composition

As I mentioned in the constraints, I've also thought about setting an upper limit to the number of components in a formula, which I think might look something like the following:

best_parameters, values, experiment, model = optimize(
    parameters=[
        {
            "name": "object1",
            "type": "choice",
            "bounds": choices,
        },
        {
            "name": "object2",
            "type": "choice",
            "bounds": choices,
        },
        {
            "name": "object3",
            "type": "choice",
            "bounds": choices,
        },
        {
            "name": "composition1",
            "type": "range",
            "bounds": [0.0, 1.0],
        },
        {
            "name": "composition2",
            "type": "range",
            "bounds": [0.0, 1.0],
        },
        {
            "name": "composition3",
            "type": "range",
            "bounds": [0.0, 1.0],
        },
    ],
    experiment_name="composition_test",
    objective_name="strength",
    evaluation_function=predict,
    parameter_constraints=["abs(1 - (composition1 + composition2 + composition3)) <= 1e-6"],
    total_trials=30,
)

How would you suggest implementing this use-case in Ax? If it would help, I'd be happy to flesh this out into a full MWE or try out any suggestions. The real use-case involves ~100 different components across 4 different classes, and the idea is to (eventually) use this in an experimental adaptive design scheme.

(tag @ramz-i who is the individual in charge of this project in our research group, post here if you have anything to add)

#706

[lena-kashtelyan]

Wow @sgbaird, this is one exemplary Github issue! : )

I think your idea about representing parameters as [0.0, 1.0] ranges is exactly right for this case. There are two challenges though:

1. Constraints are complicated to represent here.

I think the types of constraints you are looking for are not currently supported by Ax, at least not out-of-the box.

Just to make sure, my understanding is that the three constraints you'd like to have are:

1. Limit the total number of components in any given "formula" (e.g. max of 3 components),
2. The prevalences in compositions sum to 1,
3. There has to be at least one filler and at least one resin (if feasible).

Constraint 3) here is easy: it will in fact be two constraints: "filler_A + filler_B + ... >= 1e-6" and "resin_A + resin_B + ... >= 1e-6".

Constraints 1) and 2) are challenging though. For constraint 2), which is exact equality constraint, we have this wishlist issue open currently: #510.

2. The real use-case involving around 100 components and 4 classes of them.

This is too many parameters and constraints for standard Ax models, so this will also not work out-of-the-box.

Let me discuss this with @dme65, @Balandat, and @eytan, and get back to you with whether we can find a way to represent your problem in Ax.

Before we delve into this further, high-level question: how expensive will each predict function call be in your case? How many total trials could you run and how important running as few as possible vs. finding optimal outcome? I think depending on that, we might want to break up the search space into sub-search spaces with fewer components (as total of 100 is a lot).

[sgbaird]

Hi @lena-kashtelyan, I'm glad you like the issue 😄 I really appreciate your thorough and timely response. This is great!

I have some more thoughts on the constraints, I'm curious if the "components/composition" version above seems problematic (esp. with respect to symmetry concerns), and last, I tried to answer the higher-level questions you asked.

Constraints

Constraint 3 doesn't seem bad at all. Thank you! I think I have a workaround for 2, and I have a few more thoughts for 1.

Prevalence sums to 1

Perhaps a workaround for constraint 2 would be the following two constraints:

"filler_A + filler_B + resin_A + resin_B + resin_C <= 1.000001"
"filler_A + filler_B + resin_A + resin_B + resin_C >= 0.999999"

as a substitute for "abs(1-sum(...)) <= 1e-6"
Does this seem like it would work OK? I'm worried it might be a bit hacky and have some unintended consequences.

EDIT:
After reading #510 (comment), I see that the tighter the tolerance in the "hack" above, the tougher it is on the rejection sampling and would either result in an error or (likely) bloat the runtime.

raise SearchSpaceExhausted(
ax.exceptions.core.SearchSpaceExhausted: Rejection sampling error (specified maximum draws (10000) exhausted, without finding sufficiently many (1) candidates). This likely means that there are no new points left in the search space.

Limit the total number of components in any given formula

For constraint 1, maybe the L2 norm constraint could be used as a proxy for "complexity" of the suggested formula. For example:
norm([0.2 0.2 0.2 0.2 0.2]) --> 0.45 (higher complexity, lower score)
norm([0.4 0.4 0.2 0.0 0.0]) --> 0.60 (lower complexity, higher score)
If I'm thinking about this correctly, incorporating this into the objective function (predict) could bias the model towards less complex systems, but this probably pushes the problem into a multi-objective optimization (MOO) scheme and doesn't actually enforce an upper bound on the number of components. Rather than fiddle with the L2-norm constraint, at that point, it might make more sense to just use MOO and have the second objective be the number of components that exhibit a prevalence above some tolerance, e.g. count_components([0.4 0.4 0.2 1e-8 1e-7]) --> 3

Components/Composition version

Do you see any potential issues with using the "components/composition" version (i.e. sort of a hard-coded upper limit of 3 components)? Symmetry might be an issue (A B C --> 0.5 0.3 0.2 is degenerate relative to A C B --> 0.5 0.2 0.3), so I wondered if using an ordering constraint such as composition1 > composition2 > composition3 would remove most of the symmetry, but then there's still the issue of (A B C --> 0.41 0.39 0.2) being distant in parameter space to (B A C --> 0.41 0.39 0.2) despite actually being quite close. This seems very relevant to #710

e.g. when parameters (1,2) actually means the same as (2,1), we want to disallow one of them.

Is the main issue with this kind of symmetry just that it takes longer to search the space? Or is the worry that the search space might become unsmooth?

Additional Context

@ramseyissa may wish to comment more about the priorities of the project, but here are some thoughts to the best of my understanding:

• A single predict function call should be very inexpensive: i.e. millions of calls should be trivial
• It shouldn't be a problem to select a subset of the higher priority components (e.g. 10-50 components) rather than the full ~100 (curse of dimensionality - makes sense), I think that's what he's doing right now just to simplify the problem
• My guess is that finding the optimal outcome is more important than running as few as possible since it's a few months into a 4-year project. IIRC, the experimental throughput is ~10/week and the current dataset is on the order of a few hundred. Right now @ramseyissa has been using a custom script that suggests 10 new formulas to try each week by looking at 1D slices (using a subset of components) in composition space where the data is sparse. I think he also incorporated some chemical intuition into the search, but he's looking for something more sophisticated that considers the full parameter space and the target property.

[ramseyissa]

Thanks @sgbaird for raising the issue and thank you @lena-kashtelyan for the quick response and feedback!

Hopefully I can clear up some of the project goals / issues that I am facing.
To begin, I am fairly new to Ax and its complete capabilities.

Project Overview

• The dataset is close to 80 dimensions of different types of items that can be combined to make a single formulation. This high dimensional dataset is divided into 4 sections, SectionA/SectionB/SectionC/SectionD. We can think of these sections as divided up evenly for simplicity, so 20 items per section. From each section, 5 items are selected in various percentages summing to 1. To reduce the dimensionality in the dataset, I am holding SectionB/C/D constant and only looking to optimize Section A. So ideally what I would like to do, is of the 20 items in SectionA, pick 5 items at a time at various percentages and predict for this combination (the 5 items always summing to 1). Each prediction can be any combination of the 20 in SectionA (as long it chooses 5).

Useful Information

• As @sgbaird stated, the predict function call is not of concern as it is inexpensive to run.
• I am mainly concerned with finding the optimal prediction.
• One dimension search through the dataset in regions that are sparse has already produced a higher performing sample. So I believe a higher level search through the dataset using BO would be ideal and can definitely help optimize this problem.

Thank you both again and I am open to any suggestions or references.

[Balandat]

So if the predict function is cheap to evaluate and you can potentially do millions of evaluations I think that Ax might just not be the right tool for this problem, as it is primarily designed for very costly settings where you may be able to get maybe a few hundred evaluations.

There are also quite a few challenges regarding the setup (in terms of the dimensionality of the inputs, limiting the number of components, etc.) that would not be straightforward to resolve. But even if we could resolve these, you might actually end up getting better results by using an optimization algorithm that can just evaluate tons of different configurations. Maybe some kind of evolutionary algorithm could work well here?

[bernardbeckerman]

Brief drive-by comment regarding the constraints in this problem: @sgbaird correct me if i'm wrong, but you really only have 4 free parameters in this case, since the requirement that they sum to 1 removes a degree of freedom. For this reason, it might be useful for you to set a constraint:

"filler_A + filler_B + resin_A + resin_B <= 1"

and then for each parameterization resin_C = 1 - (filler_A + filler_B + resin_A + resin_B).

Of course this may end up being moot per @Balandat's comment above, but if you were to end up using Ax this might be a better way to compose your search space. Does that make sense?

[sgbaird]

@Balandat that's a good point. A brute force method or an EA might be more suitable. Do you have any suggestions on implementing the acquisition function (i.e. expected improvement) with a custom model (i.e. predict)?

@bernardbeckerman that's exactly right. I hadn't thought of that. If we try something out in Ax, I think we'll use that parameterization.

Thanks for the suggestions!

[Balandat]

A brute force method or an EA might be more suitable. Do you have any suggestions on implementing the acquisition function (i.e. expected improvement) with a custom model (i.e. predict)?

I am not sure what exactly you're getting at here. If you use a brute force or EA, why would you need to implement an acquisition function?

[sgbaird]

In the adaptive design scheme (even if this involves a brute force or EA approach), we are still looking to find a happy medium between exploitation and exploration. With brute force or EA, it would be straightforward to try to maximize the objective function, strength for example, but it's less clear to me how we would incorporate the "exploration" component which is why I mentioned an acquisition function. Perhaps a better phrasing of the question is: with brute force or EA, how would you suggest incorporating a trade-off between exploration and exploitation similar to Bayesian optimization? My questions here could very well stem from a misunderstanding on my part about the theory or not explaining the problem clearly 😅.

It might also be worth mentioning that while we can run the predict function millions of times (e.g. random forest or neural network regression), we are probably limited to a few hundred adaptive design iterations per year (make a new sample, characterize it, retrain the predict model, determine next experiment, repeat).

[lena-kashtelyan]

cc @Balandat and @dme65, who I know was also thinking through this

[Balandat]

It might also be worth mentioning that while we can run the predict function millions of times (e.g. random forest or neural network regression), we are probably limited to a few hundred adaptive design iterations per year (make a new sample, characterize it, retrain the predict model, determine next experiment, repeat).

I see, maybe we misunderstood the problem to some extent. It seems to me that the predict function really is your surrogate model, and the thing you're actually trying to optimize is the "outer loop" of optimizing the design where an evaluation is the fabrication and characterization of a sample, is that correct?

The way Ax is typically used is by building its own (typically Gaussian Process) surrogate model internally. It's generaly possible to swap out that model (though that will be some work). My main question is whether your predict function have any uncertainty quantification (i.e. how confident it is about the prediction and what the prediction variance is)?