Experiment Design

In Part One, we took a broad overview of some potentially useful parts of an AI scientist. We introduced simple conceptual objects to discuss them, and implemented them for a toy example.

We found that we could define model weights, and update our Beliefs over the models, and their internal parameters. We did not discuss in huge detail how to select a model as a result (or to prioritise them while doing hypothesis refinement and generation), an in the first part of this post we'll talk about one possible measure.

We will then go on to discuss experiment choices. Notably, our experimental Designs were chosen by us. A true AI scientist applies the scientific method for themselves: they have access to a system not just to observe, but also to probe. This would be particularly useful when trying to uncover a causal mechanism, for instance, or when given access to a fast emulator, or even access to a lab. In fact, nearly everything a human scientist does ultimately amounts to this kind of activity.

As mentioned, this is heavily inspired by a mix of recent works including Model Discovery Agent, where three key distinctions are made:

• Model discrimination: perhaps the most obvious follow on from Part 1 -- a good experiment is the one that helps us become more certain about the structure of the model that describes the world. In Part 1, this would mean an experiment that allowed us to identify hypothesis $h$. There is of course more depth to how one can identify a model in a pool, and the fact that pool size itself matters. More below.
• Joint learning: a good experiment aims to reduce uncertainty about both structure and parameters. This is an expected benefit, not a guarantee that every realised outcome makes us more certain or correct.
• Task aware design: Ultimately, the resulting model should be better at downstream forecasting on a new design.

For now, we will continue with our ridiculously simple toy example, and come back to how to deal with everything else later.

This post has two stages: compare explanations using existing observations, then choose another observation to obtain. The first connects to program-synthesis for SBI, which refines candidate programs against given observations. Its simulated training data are not extra measurements from the unknown world.

Both papers represent model and parameter uncertainty. Here we are just going to keep doing exact inference for now, while we learn about experiment design. The question is not just where our models disagree, but what we want to learn from the next observation. We'll start with one uncertain line: predict possible outcomes, imagine updating on each, then average how much we would learn. We can then bring the other hypotheses back in. The worked loop uses expected information gain about both structure and parameters; cheaper shortcuts and task-aware choices are in an expandable section below.

1. Model discrimination with no further information gain#

Before diving in to experiments, we can look a little more carefully at how we choose among models when beliefs have already been updated on the basis of past observations. Of course, this does not explicitly address the 'information gain' from doing a new experiment, but if we don't understand now to distinguish models anyway, that seems less useful.

For now, we stay with our exact inference case but note that general formulae still hold.

Start with a fixed observed dataset $\mathcal D$ and a candidate pool of hypotheses $\mathcal H$. as before, $Z_h=p(\mathcal D\mid h)$, is the model evidence after integrating over its coefficient prior, and $\pi_h=p(h)$ is its separate prior probability over models. This results in the posterior over models, or alterantively the weight of each hypothesis:

The Bayes factor compares two candidates' evidences:

Note this is in general intractable. However, the fact that the evidence integrates over the parameter prior for a given hypothesis reduces the weight of models with a larger number of parameters, since the prior contribution to the evidence becomes everywhere lower, even if the likelihood does not.

This is not a more principled alternative to posterior probabilities. It is part of the same calculation, the weight-ratio is the prior-adjusted Bayes factor:

In other words, if model $i$ has significantly larger evidence than $j$, it may still be granted lower weight if the prior probability is low. With a uniform prior over hypotheses, $i$ would be most weighted.

For the toy polynomial/Gaussian example (unnecessary to actually implement):

Here $\mathbf y$ contains the $n$ measured outputs, $X_h$ is that candidate's polynomial feature (or design) matrix, $\tau_h$ is its coefficient prior's standard deviation, and $\sigma$ is the s.d. of noise in the measurement. We compute log evidences and subtract them; we do not explicitly invert matrices, even though in this case we could (it was included for completeness).

Neural Ratio Estimation (NRE) for model comparison#

In program-synthesis for SBI it is of course not assumed you can do exact inference. That particular paper trains a neural network on examples drawn from simulation of the two models of interest $h_{i}$ and $h_{j}$, where parameters are drawn from the respective parameter priors, $\pi_{i}$ and $\pi_{k}$. The network consumes simulation examples and outputs a scalar, the optimised loss results in the output of the network itself being the log Bayes Factor, $\approx \log BF_{ji}$. This is following Evidence networks: simple losses for fast, amortized, neural Bayesian model comparison approach.

A whole pile of pairwise estimators can be trained, $\hat{s}_{ij}\approx BF_{ij}$, and ultimately used to get all the log evidences in such a way the models can be compared (I am omitting details here for brevity, and I'm a little sleepy).

SIGMA = 0.2
INITIAL_INPUTS = (-2.0, -1.0, 0.0, 1.0)  # Current Part 1: four observations.
ALLOWED_INPUTS = (-3.0, -2.0, -1.0, 0.0, 1.0, 2.0, 3.0)
world = LineWorld(a=2.0, b=1.0, sigma=SIGMA, grid=ALLOWED_INPUTS)
noise = GaussianNoise(SIGMA)
preparation_rng = np.random.default_rng(7)
initial_history = tuple(
    world.run(Design.from_dict({"x": x}, name=f"x={x:g}"), preparation_rng)
    for x in INITIAL_INPUTS
)
models = [PolynomialHypothesis(d, prior_scale=3.0, name=name)
          for d, name in enumerate(("constant", "linear", "quadratic"))]
updater = ExactGaussianUpdater(SIGMA)
beliefs = updater.update(initial_history, models)
print("Observed (x, y):", [(o.design["x"], float(o.y[0])) for o in initial_history])
for h in models:
    print(f"{h.name:10} log evidence={beliefs.log_evidences[h.name]:10.5f} "
          f"model probability={beliefs.model_probabilities[h.name]:.6g}")
log_bf = log_bayes_factor(beliefs.log_evidences, "linear", "quadratic")
odds = beliefs.model_probabilities["linear"] / beliefs.model_probabilities["quadratic"]
print("Line/quadratic log Bayes factor:", log_bf)
print("Line/quadratic Bayes factor:", np.exp(log_bf))
print("Posterior odds with equal model priors:", odds)
np.testing.assert_allclose(odds, np.exp(log_bf))

Observed (x, y): [(-2.0, -2.9997539693285034), (-1.0, -0.9402508924983061), (0.0, 0.9451724289275565), (1.0, 2.8218816322485454)]
constant   log evidence=-234.78341 model probability=9.68874e-101
linear     log evidence=  -4.52973 model probability=0.964209
quadratic  log evidence=  -7.82334 model probability=0.035791
Line/quadratic log Bayes factor: 3.293611966782885
Line/quadratic Bayes factor: 26.939994498870536
Posterior odds with equal model priors: 26.939994498870536

The line receives about 96.4% of the pool probability and the quadratic 3.6%. The Bayes factor is about 27 in favour of the line.

We know exact log evidences here; Part 3 will address replacing their computation by one of a number of inference methods.

Other checks using only existing observations#

Model entropy $H(w)=-\sum_h w_h\log w_h$ summarizes uncertainty in nats. Weight effective sample size, $1/\sum_h w_h^2$, ranges from 1 for a concentrated pool to the pool size for equal weights. They are useful but they don't capture whether the pool or the "best" model is adequate.

Parameter covariance allows us to access uncertainty conditional on a model. Training residuals diagnose fitted error, not independent test performance. Their scaling by measurement SD below does not include parameter uncertainty.

print("Model entropy:", model_entropy(beliefs.model_probabilities), "nats")
print("Weight effective sample size:", effective_model_count(beliefs.model_probabilities))

constant_only = updater.update(initial_history, [models[0]])
print("Now we reduce available model hypotheses to only the constant one.")
print("\nConstant-only entropy:", model_entropy(constant_only.model_probabilities))
print("Constant-only weight effective model count:", effective_model_count(constant_only.model_probabilities))
print("Constant residuals of noise standard deviations, at designs x:",
      residuals_in_noise_sds(constant_only.parameters["constant"], initial_history, SIGMA))
assert model_entropy(constant_only.model_probabilities) == 0
different_prior = {"constant": 0.05, "linear": 0.05, "quadratic": 0.90}
print("\nSame evidence, different model prior:",
      model_probabilities(beliefs.log_evidences, different_prior))
assert len(initial_history) == 4

Model entropy: 0.15432881561196404 nats
Weight effective sample size: 1.0741369146335678
Now we reduce available model hypotheses to only the constant one.

Constant-only entropy: -0.0
Constant-only weight effective model count: 1.0
Constant residuals of noise standard deviations, at designs x: [-14.78282  -4.48531   4.94181  14.32536]

Same evidence, different model prior: {'constant': 6.023662530331525e-101, 'linear': 0.5994659055765488, 'quadratic': 0.4005340944234515}

2. Choosing experiments#

Above, no data were added. We could revise the candidate pool using these observations, but now we will hold the coefficient priors and exact updater fixed. Only the next experimental input changes.

First, what are we actually trying to learn? Distinguishing a line from a quadratic is one question. Finding the slope of a line is another. An experiment can be useful even if we have only one hypothesis, because we can still be uncertain about its parameters.

Start with just the line#

For a moment, assume the structure is $h=\text{linear}$, with

Here $a$ is the experimental setting, which in our toy world just means choosing $x$. The intercept $b$ and slope $m$ are unknown. $\mathcal D$ is still the four observations from Part 1, not a simulated training set. We condition on the line, rather than pretending that its model probability has proved it correct.

For now, our objective is to learn both coefficients, measured by the expected reduction in their joint posterior entropy. Entropy measures the spread of this distribution; we will compare its value before and after a measurement. This is not yet the objective of making the best forecast at a particular input.

Here we are going to assume a fixed or bounded set of inputs; you can't pick arbitrary experimental settings. You can imagine this would often be the case, and trying to suggest extreme values would be a way to potentially maximise objectives in a non-useful way (this is worth some more thought). Repeated measurements are allowed, each giving a separate noisy record. We'll also assume measurements cost the same, and choose one step ahead.

Predict possible outcomes before running an experiment#

To assess the quality of an experiment, we can imagine running it if our hypothesis were true. But we need the distribution of possible outcomes, not just one best prediction:

The first density inside the integral is the observation likelihood under fixed parameters; the second is our current parameter posterior. Integrating gives the posterior predictive density, which includes uncertainty about the coefficients and measurement noise.

For this Gaussian example it is $\mathcal N(y;\mu_h(a),s_h^2(a)+\sigma^2)$. Here $\mu_h(a)$ is the mean noise-free prediction and $s_h^2(a)$ is its variance due to uncertain coefficients. $\sigma^2$ is the separate measurement-noise variance.

What would we believe after each possible outcome?#

Pick a possible output $y$. We can calculate

This is the update we would make if we observed $y$. It is not a reason to append invented measurements to our actual history. Each imagined outcome starts from the same $\mathcal D$, not from the previous imagined outcome.

Below, we compare $x=0$ and $x=3$. For each setting, take a low, middle and high possible output from its own predictive distribution. What would each do to our beliefs about $b$ and $m$?

line_hypothesis = models[1]
single_beliefs = updater.update(initial_history, [line_hypothesis])
line_before = single_beliefs.parameters[line_hypothesis.name]
candidate_settings = [Design.from_dict({"x": x}) for x in (0.0, 3.0)]
quantiles = np.array([0.1, 0.5, 0.9])
imagined = {}
history_before_lookahead = tuple(
    (obs.design, obs.y.copy()) for obs in initial_history
)

for candidate in candidate_settings:
    prediction_mean, parameter_variance = line_before.predict(np.array([candidate["x"]]))
    outcome_sd = np.sqrt(parameter_variance[0] + SIGMA**2)
    possible_outputs = norm.ppf(quantiles, loc=prediction_mean[0], scale=outcome_sd)
    alternatives = []
    for possible_y in possible_outputs:
        hypothetical_history = initial_history + (
            Observation(candidate, np.array([possible_y])),
        )
        # Same prior + full hypothetical history: the real data are used once.
        hypothetical = updater.update(hypothetical_history, [line_hypothesis])
        alternatives.append(hypothetical.parameters[line_hypothesis.name])
    imagined[candidate["x"]] = (possible_outputs, alternatives)
    print(f"x={candidate['x']:g}: possible y at 10%, 50%, 90% quantiles = {possible_outputs}")

assert len(initial_history) == 4
for obs, (design_before, y_before) in zip(initial_history, history_before_lookahead):
    assert obs.design == design_before
    np.testing.assert_array_equal(obs.y, y_before)

x=0: possible y at 10%, 50%, 90% quantiles = [0.63    0.92219 1.21438]
x=3: possible y at 10%, 50%, 90% quantiles = [6.22813 6.72091 7.21368]

Average the benefit, not just the possible outputs#

The coloured curves are alternative futures. We do not get all three measurements, and they are not three equally weighted cases covering the whole distribution. Their purpose is to show what updating would do.

At $x=0$, the new observation measures the intercept directly; at $x=3$, it measures the combination $b+3m$. The slope can also change after measuring at zero because our existing observations have made the coefficients correlated. The plots show each coefficient separately, but our objective concerns their joint uncertainty, including that correlation.

Let $\mathsf H[p]=-\int p(\theta)\log p(\theta)\,d\theta$ be the entropy of a parameter density. For candidate $a$ and possible outcome $y$, the reduction would be

We don't know which $y$ will happen, so average using the predictive density:

This is expected information gain: current uncertainty minus expected uncertainty after observing. We use natural logarithms, so the result is in nats. These are continuous entropies, not model entropies; their absolute values can be negative. It is their difference that matters here.

There is a convenient simplification in our ridiculously simple example. For a fixed design, the Gaussian posterior's covariance does not depend on the measured value. Different outcomes move its mean, but give the same entropy reduction. We can therefore do the average analytically:

This is a result of the Gaussian setup, not permission to ignore possible outcomes in general. The code below checks the formula against the hypothetical updates above. The details of the simplification can stay in a dropdown.

single_eig_values = []
print(f"{'x':>4} {'parameter prediction variance':>30} {'expected gain (nats)':>23}")
for candidate in candidate_settings:
    _, variance = line_before.predict(np.array([candidate["x"]]))
    expected_gain = 0.5*np.log1p(variance[0]/SIGMA**2)
    gains_from_updates = [
        0.5*(np.linalg.slogdet(line_before.covariance)[1]
             - np.linalg.slogdet(fit.covariance)[1])
        for fit in imagined[candidate["x"]][1]
    ]
    np.testing.assert_allclose(gains_from_updates, expected_gain, atol=1e-12)
    np.testing.assert_allclose(
        joint_eig(candidate, single_beliefs, noise=noise).value,
        expected_gain, atol=1e-10,
    )
    assert model_disagreement(candidate, single_beliefs, noise=noise).value == 0
    single_eig_values.append(expected_gain)
    print(f"{candidate['x']:4.0f} {variance[0]:30.6f} {expected_gain:23.6f}")
assert single_eig_values[1] > single_eig_values[0] > 0

   x  parameter prediction variance    expected gain (nats)
   0                       0.011982                0.131011
   3                       0.107849                0.653656

Both settings can teach us something, despite there being no competing model structures. With these observations and this objective, $x=3$ gives about 0.65 nats of expected information, compared with 0.13 at $x=0$. This does not mean that measuring at an extreme input is always best, or that a single measurement identifies both coefficients.

A useful check: if our target were only the intercept, would you necessarily choose the same experiment? We come back to this in the task-aware dropdown.

Now bring back the other hypotheses#

So far we assumed the line. Returning to the constant, linear and quadratic pool, an observation can now change both the model weights and the parameter posteriors. For prediction we average the candidate-specific predictive densities using their current weights:

The next figure shows the line and quadratic separately, to see how their possible outcomes compare. Overlap matters when trying to distinguish those structures, but even identical predictive distributions across models do not rule out learning parameters within each one.

Scoring experimental designs: what do we expect to learn?#

So, now we know how to look ahead at the possible output of our hypotheses, we need a way to give values to different experiments, also known as Designs. This is our design_scorer. In the code it returns a DesignScore (which is just the evaluation of whatever scoring function we use at that value of model inputs).

We'll keep the same idea: value the expected reduction in uncertainty, now about both the structure and its parameters. Let $H$ be the unknown model label, $\Theta$ its parameter vector, and $Y_a$ the as-yet-unknown output. The information chain rule separates their joint expected information gain:

The first term is the expected reduction in model entropy. The second is parameter information under each assumed model, averaged using the current model weights. For our Gaussian candidates, each parameter term is exactly the single-line formula above.

This is the objective we call joint_eig. With one hypothesis, the model term is zero but the parameter term remains. With several hypotheses, both can matter. It is still an expectation under our current pool and priors, not knowledge of what the world will actually do.

Appendix A.9 of Murphy's paper distinguishes model, joint and task-aware objectives. For this small example we can afford to compute joint EIG directly: the model term uses numerical integration over possible outputs and the parameter term is analytic. This is our teaching choice, not a claim that the paper's cheaper disagreement score computes the same quantity.

The more general value-of-information question is how much better a decision could be after seeing the outcome. Information gain chooses a particular objective for that question; improving a forecast or making a costly downstream decision need not favour the same experiment. The alternatives, including disagreement as a shortcut, are below.

designer = EnumeratingDesigner(
    scorer=partial(joint_eig, noise=noise),
    name="joint_eig",
)
decision = designer.choose(menu, beliefs, np.random.default_rng(100))
print(f"{'x':>6} {'model EIG':>14} {'parameter EIG':>16} {'joint EIG':>14}  (nats)")
for candidate, score in decision.scores:
    model_gain = score.components["model_eig"]
    parameter_gain = score.components["conditional_parameter_eig"]
    np.testing.assert_allclose(score.value, model_gain + parameter_gain)
    print(f"{candidate['x']:6.1f} {model_gain:14.6f} {parameter_gain:16.6f} {score.value:14.6f}")
print("Selected:", decision.design.as_dict())
assert len(initial_history) == 4

     x      model EIG    parameter EIG      joint EIG  (nats)
  -3.0       0.022414         0.480317       0.502730
  -2.0       0.000720         0.267640       0.268360
  -1.0       0.000992         0.134228       0.135220
   0.0       0.000993         0.134145       0.135139
   1.0       0.000718         0.267439       0.268157
   2.0       0.022404         0.480089       0.502494
   3.0       0.052836         0.693285       0.746120
Selected: {'x': 3.0}

As we can see, the selection score (here joint expected information gain) is computed for all possible designs. The highest score is selected. This is a choice about what we expect to learn before measuring; it is not the model evidence or a score of an observation we have already obtained.

The Designer is the object calling the scorer, and it records every score, not just the winning design, in its output Decision.

While scoring, $\mathcal D$ is fixed, $a$ is a candidate action and $Y_a$ is an unknown outcome. Only after observing an actual value $y_{\mathrm{new}}$ do we extend the dataset.

Run the selected experiment once#

1. Score the allowed experiments using current beliefs.
2. Select one and record its predictive mean and observed-output variance.
3. Execute it once and append the actual observation.
4. Update from priors on the complete updated history.

The Designer receives no world or held-out answers. There is as yet no touching of the hypothesis set, because we're getting to that... I know it's taking a while, but it is worth it :)

print(inspect.getsource(run_experiments))
updated_history, trace = run_experiments(
    world, initial_history, models, updater, designer, budget=1,
    measurement_rng=np.random.default_rng(1007),
    selection_rng=np.random.default_rng(2007),
)
step = trace[0]
print("Selected action:", step.decision.design.as_dict())
print("Forecast recorded BEFORE observation:", step.forecast_mean)
print("Forecast SD including noise:", np.sqrt(step.forecast_variance))
print("Actual observation:", step.observation.y[0])
print("Probabilities before:", step.before.model_probabilities)
print("Probabilities after: ", step.after.model_probabilities)
print("Entropy before/after:",
      model_entropy(step.before.model_probabilities), model_entropy(step.after.model_probabilities))
print("Records before/after:", len(initial_history), len(updated_history))
assert step.decision.design == decision.design
assert len(initial_history) == 4 and len(updated_history) == 5

def run_experiments(world: World, initial_data: Dataset, hypotheses, updater: Updater,
                    designer: Designer, *, budget: int, measurement_rng: RNG,
                    selection_rng: RNG) -> tuple[tuple[Observation, ...], list[ExperimentalRound]]:
    """Run a fixed-pool loop with exactly one world call per selected experiment.

    Repeated settings are allowed; each measurement is a separate record. The
    menu comes from world.design_space(), but only designs and beliefs reach the
    designer. This is ordinary information separation, not a Python sandbox.
    Full-history inference begins from priors each time. No model proposal or
    evaluation occurs here. The initial dataset is copied, never extended in place.

    This controller requires scalar polynomial posterior moments and known
    Gaussian measurement noise to record forecasts. Text-only agents need a forecast
    capability later; they must not silently receive fabricated numerical beliefs.
    """
    if isinstance(budget, bool) or not isinstance(budget, int) or budget < 0:
        raise ValueError("Experiment budget must be a non-negative integer.")
    data = list(deepcopy(initial_data))
    pool = tuple(hypotheses)
    beliefs = updater.update(tuple(data), pool)
    trace = []
    for number in range(1, budget + 1):
        menu = tuple(world.design_space())
        decision = designer.choose(menu, deepcopy(beliefs), selection_rng)
        if decision.design not in menu:
            raise ValueError("Designer selected an experiment outside the allowed menu.")
        parts = predictive_components(decision.design, beliefs, world.noise)
        weights, means = parts["weights"], parts["means"]
        mean = float(weights @ means)
        variance = float(weights @ (parts["observation_variances"] + (means - mean)**2))
        before = deepcopy(beliefs)
        observation = world.run(decision.design, measurement_rng)
        if observation.design != decision.design:
            raise ValueError("World returned an observation for a different design.")
        data.append(observation)
        beliefs = updater.update(tuple(data), pool)
        trace.append(ExperimentalRound(number, len(data)-1, decision, before, mean,
                                        variance, deepcopy(observation), deepcopy(beliefs)))
    return tuple(data), trace

Selected action: {'x': 3.0}
Forecast recorded BEFORE observation: 6.702533487149083
Forecast SD including noise: 0.4473272561144739
Actual observation: 6.785032673972183
Probabilities before: {'constant': 9.688740668402366e-101, 'linear': 0.9642090122802123, 'quadratic': 0.03579098771978728}
Probabilities after:  {'constant': 5.0786115157859765e-303, 'linear': 0.9891219522501264, 'quadratic': 0.010878047749873788}
Entropy before/after: 0.15432881561196404 0.05999841229168654
Records before/after: 4 5

An informative experiment is useful in expectation. Technically, we can make a measurement that upon updating the beliefs, make neglected alternatives more plausible and increase entropy, so we can't expect every observation as guaranteed to make us more certain or correct.

Running the loop: compare policies over several runs#

We can do a tiny experiment: keeping our starting data distribution, hypothesis pool and priors, and set of allowed actions all the same, what happens if we run a six-measurement budget where in one case we select using joint expected information gain, and in the other case, our designs are chosen uniformly at random. We have a RandomDesigner implemented, which will be helpful as a later baseline.

To test performance, we have held-out experiments at x=-4 and x=5, rerunning the toy problem from scratch. We can report model entropy and mean square error on our test points as successive experiments are performed.

Model entropy and prediction error need not move together, they represent different things; one measures uncertainty over structures and the other evaluates the model-averaged prediction. Neither is the full joint-information objective we used for design. A useful parameter-learning experiment need not reduce model entropy much. The initial data already constrain this simple line, so a large advantage for design is not guaranteed. These are paired runs on one toy world, not a benchmark claim about any of the papers I am discussing.

3. Other experiment-design objectives and shortcuts#

What next?#

Experiment design can improve which observations become available, and once we have updated our beliefs over models in some way, we need to select a model. Here, we have looked at those two points in more detail. The design calculation is: decide what we want to learn, predict possible outcomes, imagine the update, then average the benefit. It does not require several competing hypotheses, though we can learn about those too.

The question is, how can we have a method of updating those beliefs when our exact inference has become intractable? Of course, there are the usual suspects so the next post walks through possible approaches.