Factorial design with empirical Bayes and Thompson Sampling¶

This tutorial illustrates how to run a factorial experiment. In such an experiment, each parameter (factor) can be assigned one of multiple discrete values (levels). A full-factorial experiment design explores all possible combinations of factors and levels.

For instance, consider a banner with a title and an image. We are considering two different titles and three different images. A full-factorial experiment will compare all 2*3=6 possible combinations of title and image, to see which version of the banner performs the best.

In this example, we first run an exploratory batch to collect data on all possible combinations. Then we use empirical Bayes to model the data and shrink noisy estimates toward the mean. Next, we use Thompson Sampling to suggest a set of arms (combinations of factors and levels) on which to collect more data. We repeat the process until we have identified the best performing combination(s).

In [1]:

import numpy as np
import pandas as pd
import sklearn as skl
from typing import Dict, Optional, Tuple, Union
from ax import (
    Arm,
    ChoiceParameter,
    Models,
    ParameterType,
    SearchSpace,
    Experiment,
    OptimizationConfig,
    Objective,
)
from ax.plot.scatter import plot_fitted
from ax.utils.notebook.plotting import render, init_notebook_plotting
from ax.utils.stats.statstools import agresti_coull_sem

In [2]:

init_notebook_plotting()

[INFO 08-11 16:01:30] ax.utils.notebook.plotting: Injecting Plotly library into cell. Do not overwrite or delete cell.

1. Define the search space¶

First, we define our search space. A factorial search space contains a ChoiceParameter for each factor, where the values of the parameter are its levels.

In [3]:

search_space = SearchSpace(
    parameters=[
        ChoiceParameter(
            name="factor1",
            parameter_type=ParameterType.STRING,
            values=["level11", "level12", "level13"],
        ),
        ChoiceParameter(
            name="factor2",
            parameter_type=ParameterType.STRING,
            values=["level21", "level22"],
        ),
        ChoiceParameter(
            name="factor3",
            parameter_type=ParameterType.STRING,
            values=["level31", "level32", "level33", "level34"],
        ),
    ]
)

/Users/mpolson64/opt/anaconda3/envs/ax-dev/lib/python3.9/site-packages/ax/core/parameter.py:517: UserWarning:

`is_ordered` is not specified for `ChoiceParameter` "factor1". Defaulting to `False` for parameters of `ParameterType` STRING. To override this behavior (or avoid this warning), specify `is_ordered` during `ChoiceParameter` construction.

/Users/mpolson64/opt/anaconda3/envs/ax-dev/lib/python3.9/site-packages/ax/core/parameter.py:517: UserWarning:

`sort_values` is not specified for `ChoiceParameter` "factor1". Defaulting to `False` for parameters of `ParameterType` STRING. To override this behavior (or avoid this warning), specify `sort_values` during `ChoiceParameter` construction.

/Users/mpolson64/opt/anaconda3/envs/ax-dev/lib/python3.9/site-packages/ax/core/parameter.py:517: UserWarning:

`is_ordered` is not specified for `ChoiceParameter` "factor2". Defaulting to `False` for parameters of `ParameterType` STRING. To override this behavior (or avoid this warning), specify `is_ordered` during `ChoiceParameter` construction.

/Users/mpolson64/opt/anaconda3/envs/ax-dev/lib/python3.9/site-packages/ax/core/parameter.py:517: UserWarning:

`sort_values` is not specified for `ChoiceParameter` "factor2". Defaulting to `False` for parameters of `ParameterType` STRING. To override this behavior (or avoid this warning), specify `sort_values` during `ChoiceParameter` construction.

/Users/mpolson64/opt/anaconda3/envs/ax-dev/lib/python3.9/site-packages/ax/core/parameter.py:517: UserWarning:

`is_ordered` is not specified for `ChoiceParameter` "factor3". Defaulting to `False` for parameters of `ParameterType` STRING. To override this behavior (or avoid this warning), specify `is_ordered` during `ChoiceParameter` construction.

/Users/mpolson64/opt/anaconda3/envs/ax-dev/lib/python3.9/site-packages/ax/core/parameter.py:517: UserWarning:

`sort_values` is not specified for `ChoiceParameter` "factor3". Defaulting to `False` for parameters of `ParameterType` STRING. To override this behavior (or avoid this warning), specify `sort_values` during `ChoiceParameter` construction.

2. Define a custom metric¶

Second, we define a custom metric, which is responsible for computing the mean and standard error of a given arm.

In this example, each possible parameter value is given a coefficient. The higher the level, the higher the coefficient, and the higher the coefficients, the greater the mean.

The standard error of each arm is determined by the weight passed into the evaluation function, which represents the size of the population on which this arm was evaluated. The higher the weight, the greater the sample size, and thus the lower the standard error.

In [4]:

from ax import Data, Metric
from ax.utils.common.result import Ok
import pandas as pd
from random import random


one_hot_encoder = skl.preprocessing.OneHotEncoder(
    categories=[par.values for par in search_space.parameters.values()],
)


class FactorialMetric(Metric):
    def fetch_trial_data(self, trial):
        records = []
        for arm_name, arm in trial.arms_by_name.items():
            params = arm.parameters
            batch_size = 10000
            noise_level = 0.0
            weight = trial.normalized_arm_weights().get(arm, 1.0)
            coefficients = np.array([0.1, 0.2, 0.3, 0.1, 0.2, 0.1, 0.2, 0.3, 0.4])
            features = np.array(list(params.values())).reshape(1, -1)
            encoded_features = one_hot_encoder.fit_transform(features)
            z = (
                coefficients @ encoded_features.T
                + np.sqrt(noise_level) * np.random.randn()
            )
            p = np.exp(z) / (1 + np.exp(z))
            plays = np.random.binomial(batch_size, weight)
            successes = np.random.binomial(plays, p)
            records.append(
                {
                    "arm_name": arm_name,
                    "metric_name": self.name,
                    "trial_index": trial.index,
                    "mean": float(successes) / plays,
                    "sem": agresti_coull_sem(successes, plays),
                }
            )
        return Ok(value=Data(df=pd.DataFrame.from_records(records)))

3. Define the experiment¶

We now set up our experiment and define the status quo arm, in which each parameter is assigned to the lowest level.

In [5]:

from ax import Runner


class MyRunner(Runner):
    def run(self, trial):
        trial_metadata = {"name": str(trial.index)}
        return trial_metadata


exp = Experiment(
    name="my_factorial_closed_loop_experiment",
    search_space=search_space,
    optimization_config=OptimizationConfig(
        objective=Objective(metric=FactorialMetric(name="success_metric"))
    ),
    runner=MyRunner(),
)
exp.status_quo = Arm(
    parameters={"factor1": "level11", "factor2": "level21", "factor3": "level31"}
)

[WARNING 08-11 16:01:31] ax.core.experiment: Experiment's status_quo is updated. Generally the status_quo should not be changed after being set.

4. Run an exploratory batch¶

We then generate an a set of arms that covers the full space of the factorial design, including the status quo. There are three parameters, with two, three, and four values, respectively, so there are 24 possible arms.

In [6]:

factorial = Models.FACTORIAL(search_space=exp.search_space)
factorial_run = factorial.gen(
    n=-1
)  # Number of arms to generate is derived from the search space.
print(len(factorial_run.arms))

Now we create a trial including all of these arms, so that we can collect data and evaluate the performance of each.

In [7]:

trial = exp.new_batch_trial(optimize_for_power=True).add_generator_run(
    factorial_run, multiplier=1
)

By default, the weight of each arm in factorial_run will be 1. However, to optimize for power on the contrasts of k groups against the status quo, the status quo should be sqrt(k) larger than any of the treatment groups. Since we have 24 different arms in our search space, the status quo should be roughly five times larger. That larger weight is automatically set by Ax under the hood if optimize_for_power kwarg is set to True on new batched trial creation.

In [8]:

trial._status_quo_weight_override

Out[8]:

4.795831523312719

5. Iterate using Thompson Sampling¶

Next, we run multiple trials (iterations of the experiment) to hone in on the optimal arm(s).

In each iteration, we first collect data about all arms in that trial by calling trial.run() and trial.mark_complete(). Then we run Thompson Sampling, which assigns a weight to each arm that is proportional to the probability of that arm being the best. Arms whose weight exceed min_weight are added to the next trial, so that we can gather more data on their performance.

In [9]:

models = []
for i in range(4):
    print(f"Running trial {i+1}...")
    trial.run()
    trial.mark_completed()
    thompson = Models.THOMPSON(experiment=exp, data=trial.fetch_data(), min_weight=0.01)
    models.append(thompson)
    thompson_run = thompson.gen(n=-1)
    trial = exp.new_batch_trial(optimize_for_power=True).add_generator_run(thompson_run)

/var/folders/w_/g_40lr6j0s98g6921f3qxcmr0000gn/T/ipykernel_92917/2149564789.py:35: DeprecationWarning: