Batching
BoTorch makes frequent use of "batching", both in the sense of batch acquisition functions for multiple candidates as well as in the sense of parallel or batch computation (neither of these should be confused with mini-batch training). Here we explain some of the common patterns you will see in BoTorch for exploiting parallelism, including common shapes and decorators for more conveniently handling these shapes.
Batch Acquisition Functions
BoTorch supports batch acquisition functions that assign a joint utility to a set of design points in the parameter space. These are, for obvious reasons, referred to as q-Acquisition Functions. For instance, BoTorch ships with support for q-EI, q-UCB, and a few others.
As discussed in the design philosophy, BoTorch has adopted the convention of referring to batches in the batch-acquisition sense as "q-batches", and to batches in the torch batch-evaluation sense as "t-batches".
Internally, q-batch acquisition functions operate on input tensors of shape , where is the number of t-batches, is the number of design points to be considered concurrently, and is the dimension of the parameter space. Their output is a one-dimensional tensor with elements, with the -th element corresponding to the -th t-batch. Always requiring explicit t-batch and q-batch dimensions makes it easier and less ambiguous to work with samples from the posterior in a consistent fashion.
Batch-Mode Decorator
In order to simplify the user-facing API for evaluating acquisition functions,
BoTorch implements the
@t_batch_mode_transform
decorator, which allows the use of non-batch mode inputs. If applied to an
instance method with a single Tensor
argument, an input tensor to that method
without a t-batch dimension (i.e. tensors of shape ) will automatically
be converted to a t-batch of size 1 (i.e. of batch_shape
torch.Size([1])
),
This is typically used on the forward
method of an AcquisitionFunction
.
The @t_batch_mode_transform
decorator takes an expected_q
argument that, if
specified, checks that the number of q-batches in the input is equal to the
one specified in the decorator. This is typically used for acquisition functions
that can only handle the case (e.g. any AnalyticAcqusitionFunction
).
Batching Sample Shapes
BoTorch evaluates Monte-Carlo acquisition functions using (quasi-) Monte-Carlo
sampling from the posterior at the input features . Hence, on top of the
existing q-batch and t-batch dimensions, we also end up with another batch
dimension corresponding to the MC samples we draw. We use the PyTorch notions of
sample_shape
and event_shape
.
event_shape
is the shape of a single sample drawn from the underlying
distribution:
- Evaluating a single-output model at a tensor,
representing data points in dimensions each, yields a posterior with
event_shape
being . Evaluating the same model at a tensor (representing a t-batch-shape of , with data points of -dimensions each in every batch) yields a posterior withevent_shape
being