Meaning of "Trace Observations" in qMultiFidelityMaxValueEntropy Based Acquisition Functions

[adrian-martens]

Hi there,

I am working with the qMultiFidelityLowerBoundMaxValueEntropy AF and was wondering what the meaning of "trace observations" is in this context.

They are mentioned when talking about the expand function one can define and I thought it means "expanding" the current point of a certain fidelity level to all the other possible fidelities in order to calculate the covariance between the fidelity levels.

Is it necessary to define it when using this acquisition function?

Here is an excerpt of the code:

class qMultiFidelityLowerBoundMaxValueEntropy(qMultiFidelityMaxValueEntropy):
    r"""Multi-fidelity acquisition function for General-purpose Information-Based
    Bayesian optimization (GIBBON).

    The acquisition function for multi-fidelity max-value entropy search
    with support for trace observations. See [Takeno2020mfmves]_
    for a detailed discussion of the basic ideas on multi-fidelity MES
    (note that this implementation is somewhat different). This acquisition function
    is similar to `qMultiFidelityMaxValueEntropy` but computes the information gain
    from the lower bound described in [Moss2021gibbon].

    The model must be single-outcome, unless using a PosteriorTransform.
    The batch case `q > 1` is supported through cyclic optimization and fantasies.

    Example:
        >>> model = SingleTaskGP(train_X, train_Y)
        >>> candidate_set = torch.rand(1000, bounds.size(1))
        >>> candidate_set = bounds[0] + (bounds[1] - bounds[0]) * candidate_set
        >>> MF_qGIBBON = qMultiFidelityLowerBoundMaxValueEntropy(model, candidate_set)
        >>> mf_gibbon = MF_qGIBBON(test_X)
    """

    def __init__(
        self,
        model: Model,
        candidate_set: Tensor,
        num_fantasies: int = 16,
        num_mv_samples: int = 10,
        num_y_samples: int = 128,
        posterior_transform: PosteriorTransform | None = None,
        use_gumbel: bool = True,
        maximize: bool = True,
        cost_aware_utility: CostAwareUtility | None = None,
        project: Callable[[Tensor], Tensor] = lambda X: X,
        expand: Callable[[Tensor], Tensor] = lambda X: X,
    ) -> None:
        r"""Single-outcome max-value entropy search acquisition function.

        Args:
            model: A fitted single-outcome model.
            candidate_set: A `n x d` Tensor including `n` candidate points to
                discretize the design space, which will be used to sample the
                max values from their posteriors.
            cost_aware_utility: A CostAwareUtility computing the cost-transformed
                utility from a candidate set and samples of increases in utility.
            num_fantasies: Number of fantasies to generate. The higher this
                number the more accurate the model (at the expense of model
                complexity and performance) and it's only used when `X_pending`
                is not `None`.
            num_mv_samples: Number of max value samples.
            num_y_samples: Number of posterior samples at specific design point `X`.
            posterior_transform: A PosteriorTransform. If using a multi-output model,
                a PosteriorTransform that transforms the multi-output posterior into a
                single-output posterior is required.
            use_gumbel: If True, use Gumbel approximation to sample the max values.
            maximize: If True, consider the problem a maximization problem.
            cost_aware_utility: A CostAwareUtility computing the cost-transformed
                utility from a candidate set and samples of increases in utility.
            project: A callable mapping a `batch_shape x q x d` tensor of design
                points to a tensor of the same shape projected to the desired
                target set (e.g. the target fidelities in case of multi-fidelity
                optimization).
            expand: A callable mapping a `batch_shape x q x d` input tensor to
                a `batch_shape x (q + q_e)' x d`-dim output tensor, where the
                `q_e` additional points in each q-batch correspond to
                additional ("trace") observations.
        """
        super().__init__(
            model=model,
            candidate_set=candidate_set,
            num_fantasies=num_fantasies,
            num_mv_samples=num_mv_samples,
            num_y_samples=num_y_samples,
            posterior_transform=posterior_transform,
            use_gumbel=use_gumbel,
            maximize=maximize,
            cost_aware_utility=cost_aware_utility,
            project=project,
            expand=expand,
        )

[Balandat]

@adrian-martens not sure what version of the code you're citing here, but he current main branch actually raises an error if you pass the expand argument for this acquisition function: https://github.com/pytorch/botorch/blob/4352bf3f2d94618cb2f4b80b8ba22c0753556440/botorch/acquisition/max_value_entropy_search.py#L815-L818 (this was added in #2769 since the acquisition function apparently did not handle this correctly). So it's definitely not necessary.

The trace observation setup is mostly used for the trace-aware Knowledge Gradient acquisition functions. You can find more on the meaning of trace observations in Section 2.2 of https://arxiv.org/pdf/1903.04703