Sobol sampling isn't independent across multiple posteriors

[pvasu]

First off, thanks to the Botorch contributors for developing, maintaining, and supporting this incredible resource!!

Summary

I recently updated from Botorch 0.6.6 + GPyTorch 1.8.1 (I know, very out of date) to Botorch 0.8.5 + GPyTorch 1.10 and noticed some odd behavior: when evaluating a UCB acquisition function on a Botorch Objective that includes outcome constraints, I saw big differences across repeated runs.

After digging into it, I found that the issue lies in sampling: in Botorch 0.8.5, when I sample from multiple posteriors (wrapped in a PosteriorList), I don't get independent samples for each posterior!

Digging further, the issue seems to be in the "base samples": when constructing a separate SobolQMCNormalSampler for each posterior (as Botorch 0.8.5 does in get_sampler), the base samples are not independent for some reason.

(NOTE: I verified that the same problematic behavior also exists in Botorch 0.12.0 + GPyTorch 1.13 + torch 2.0.1. But this description will continue to compare Botorch 0.6.6 vs 0.8.5.)

Full details below. Please let me know if I'm missing something vs. whether you think this is a genuine bug. Thank you in advance!

Minimal example

Below is a minimal example to illustrate the problem. Call this file botorch_sampler.py:

"""
Illustrate a problem where in more recent versions of Botorch, `get_sampler` does not
independently draw samples from multiple posteriors in a `PosteriorList`.

Whereas in Botorch 0.6.6, constructing a `SobolQMCNormalSampler` for a `PosteriorList`
doesn't have this problem.
"""

import argparse
import os

import botorch
import matplotlib.pyplot as plt
import pandas as pd
import seaborn as sns
import torch
from botorch.posteriors import GPyTorchPosterior, PosteriorList
from gpytorch.distributions.multivariate_normal import MultivariateNormal


def main(manual: bool):
    # Define two normal distribution posteriors with mean=0 std=1.
    mvns = [
        MultivariateNormal(
            torch.tensor([0.0]),
            torch.tensor([[1.0]]),
        )
        for _ in range(2)
    ]
    posterior = PosteriorList(*[GPyTorchPosterior(mvn) for mvn in mvns])

    # Sample from the two distributions.
    if botorch.__version__.startswith("0.6"):
        from botorch.sampling import SobolQMCNormalSampler

        if manual:
            samplers = [
                SobolQMCNormalSampler(num_samples=512) for _ in posterior.posteriors
            ]
            samples = torch.stack(
                [
                    sampler(posterior).squeeze()
                    for posterior, sampler in zip(posterior.posteriors, samplers)
                ],
                dim=-1,
            )
        else:
            sampler = SobolQMCNormalSampler(num_samples=512)
            samples = sampler(posterior).squeeze()

    else:
        if manual:
            raise ValueError("--manual is meant to be used on 0.6")

        from botorch.sampling import get_sampler

        sampler = get_sampler(posterior, torch.Size([512]))
        samples = sampler(posterior).squeeze()

    assert samples.shape == torch.Size([512, 2])

    # Plot 2D distribution of samples from the two posteriors.
    df = pd.DataFrame(samples.detach().numpy(), columns=["posterior1", "posterior2"])
    sns.scatterplot(data=df, x="posterior1", y="posterior2")

    # Save figure to file.
    title = f"Samples. botorch@{botorch.__version__}, manual={manual}"
    plt.title(title)
    os.makedirs("tmp/", exist_ok=True)
    plt.savefig(f"tmp/{title}.png")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "--manual",
        action="store_true",
        required=False,
        help="Meant for Botorch 0.6: construct a separate SobolQMCNormalSampler for each posterior.",
    )

    args = parser.parse_args()

    main(args.manual)

Running the minimal example

✅ On Botorch 0.6.6, running this script produces a 2D distribution where each posterior is sampled independently, as expected:

⚠️ However, on Botorch 0.8.5, the two posteriors are not sampled independently:

⚠️ The problem also occurs on Botorch 0.6.6 when manually constructing a SobolQMCNormalSampler for each posterior (run the script with the --manual CLI argument), which is similar to what Botorch 0.8.5's get_sampler() does on the PosteriorList:

[saitcakmak]

Interesting finding! This pretty much boils down to two Sobol sequences not being independent of each other. We can produce similar plots using

from matplotlib import pyplot as plt
from botorch.utils.sampling import draw_sobol_normal_samples
seq_1 = draw_sobol_normal_samples(d=1, n=1024).view(-1)
seq_2 = draw_sobol_normal_samples(d=1, n=1024).view(-1)
plt.scatter(seq_1, seq_2)

If I were to be lazy and try to produce two separate sequences by splitting one in half, it's even worse!

from matplotlib import pyplot as plt
from botorch.utils.sampling import draw_sobol_normal_samples
seq = draw_sobol_normal_samples(d=1, n=2048).view(-1)
plt.scatter(seq[:1024], seq[1024:])

The right way to do it is to use a full dimensional Sobol sequence, which is what was being done prior to ListSampler.

from matplotlib import pyplot as plt
from botorch.utils.sampling import draw_sobol_normal_samples
seq = draw_sobol_normal_samples(d=2, n=1024).view(-1, 2)
plt.scatter(seq[..., 0], seq[..., 1])

I haven't looked into the code to see if there is an easy fix for this, but here are a couple thoughts that may help in the short term:

• Do you have a PosteriorList of GPyTorchPosterior because you're using ModelList rather than a ModelListGP? If so, ModelListGP would be more efficient and would produce a single GPyTorchPosterior. This will automatically dispatch to a single sampler.
• You can pass in sampler=SobolQMCNormalSampler(num_samples=512) (SobolQMCNormalSampler(sample_shape=torch.Size([512]) in latest versions) to the acquisition function constructor to reproduce the old behavior with the new code.

cc @Balandat -- maybe we should use IIDNormalSampler with ListSampler rather than SobolQMCNormalSampler?

[pvasu]

Hey folks, bumping this thread: do you think it's possible to restore proper Sobol sampling for a PosteriorList, aka without having to fallback to random sampling? I'm happy to try and work on this fix, but could use some guidance on the right path to take.

[pvasu]

Should I start by converting this into an Issue?

[Balandat]

@pvasu Yes please go ahead and create an issue for this as well since this is undesired behavior. I think the most reasonable approach here would be as Sait suggested to update the dispatcher to apply my suggestion for the special case of ListPosterior[GPyTorchPosterior, ...]. Does that make sense to you?

[pvasu]

Cool, submitted an issue 👍

I think the most reasonable approach here would be as Sait suggested to update the dispatcher to apply my suggestion for the special case of ListPosterior[GPyTorchPosterior, ...]. Does that make sense to you?

Sure that sounds like an improvement 👌 But what about the case of a posterior list containing a mix of gaussian posteriors and deterministic posteriors?

[Balandat]

Sure that sounds like an improvement 👌 But what about the case of a posterior list containing a mix of gaussian posteriors and deterministic posteriors?

For deterministic posteriors I guess there isn't anything to sample. If there are k Gaussian posteriors and n - k deterministic ones, the natural thing to do would be to have one Sobol sampler for the k Gaussians and then distribute the samples accordingly. But that will induce quite a bit of complexity in the code so not sure if that's a great idea.