tutorials/optimize_stochastic.ipynb

{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Optimize acquisition functions using torch.optim\n",
        "\n",
        "In this tutorial, we show how to use PyTorch's `optim` module for optimizing BoTorch MC acquisition functions. This is useful if the acquisition function is stochastic in nature (caused by re-sampling the base samples when using the reparameterization trick, or if the model posterior itself is stochastic).\n",
        "\n",
        "*Note:* A pre-packaged, more user-friendly version of the optimization loop we will develop below is contained in the `gen_candidates_torch` function in the `botorch.gen` module. This tutorial should be quite useful if you would like to implement custom optimizers beyond what is contained in `gen_candidates_torch`.\n",
        "\n",
        "As discussed in the [CMA-ES tutorial](./optimize_with_cmaes), for deterministic acquisition functions BoTorch uses quasi-second order methods (such as L-BFGS-B or SLSQP) by default, which provide superior convergence speed in this situation. "
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Set up a toy model\n",
        "\n",
        "We'll fit a `SingleTaskGP` model on noisy observations of the function $f(x) = 1 - \\|x\\|_2$ in `d=5` dimensions on the hypercube $[-1, 1]^d$."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 1,
      "metadata": {},
      "outputs": [],
      "source": [
        "import torch\n",
        "\n",
        "from botorch.fit import fit_gpytorch_model\n",
        "from botorch.models import SingleTaskGP\n",
        "from gpytorch.mlls import ExactMarginalLogLikelihood"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 2,
      "metadata": {},
      "outputs": [],
      "source": [
        "d = 5\n",
        "\n",
        "bounds = torch.stack([-torch.ones(d), torch.ones(d)])\n",
        "\n",
        "train_X = bounds[0] + (bounds[1] - bounds[0]) * torch.rand(50, d)\n",
        "train_Y = 1 - torch.norm(train_X, dim=-1, keepdim=True)\n",
        "\n",
        "model = SingleTaskGP(train_X, train_Y)\n",
        "mll = ExactMarginalLogLikelihood(model.likelihood, model)\n",
        "fit_gpytorch_model(mll);"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Define acquisition function\n",
        "\n",
        "We'll use `qExpectedImprovement` with a custom sampler that uses a small number of MC samples and re-samples upon each evaluation of the function. This results in a stochastic acquisition function that one should not attempt to optimize with the quasi-second order methods that are used by default in BoTorch's `optimize_acqf` function."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 3,
      "metadata": {
        "collapsed": true
      },
      "outputs": [],
      "source": [
        "from botorch.acquisition import qExpectedImprovement\n",
        "from botorch.sampling import IIDNormalSampler\n",
        "\n",
        "sampler = IIDNormalSampler(num_samples=100, resample=True)\n",
        "qEI = qExpectedImprovement(model, best_f=train_Y.max(), sampler=sampler)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Optimizing the acquisition function\n",
        "\n",
        "We will perform optimization over `N=5` random initial `q`-batches with `q=2` in parallel. We use `N` random restarts because the acquisition function is non-convex and as a result we may get stuck in local minima."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 4,
      "metadata": {
        "collapsed": true
      },
      "outputs": [],
      "source": [
        "N = 5\n",
        "q = 2"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Choosing initial conditions via a heuristic\n",
        "\n",
        "Using random initial conditions in conjunction with gradient-based optimizers can be problematic because qEI values and their corresponding gradients are often zero in large parts of the feature space. To mitigate this issue, BoTorch provides a heuristic for generating promising initial conditions (this dirty and not-so-little secret of Bayesian Optimization is actually very important for overall closed-loop performance).\n",
        "\n",
        "Given a set of `q`-batches $X'$ and associated acquisiton function values $Y'$, the `initialize_q_batch_nonneg` samples promising initial conditions $X$ (without replacement) from the multinomial distribution\n",
        "\n",
        "$$ \\mathbb{P}(X = X'_i) \\sim \\exp (\\eta \\tilde{Y}_i), \\qquad \\text{where} \\;\\; \\tilde{Y}_i = \\frac{Y'_i - \\mu(Y)}{\\sigma(Y)} \\;\\; \\text{if} \\;\\; Y'_i >0 $$\n",
        "\n",
        "and $\\mathbb{P}(X = X'_j) = 0$ for all $j$ such that $Y'_j = 0$. \n",
        "\n",
        "Fortunately, thanks to the high degree of parallelism in BoTorch, evaluating the acquisition function at a large number of randomly chosen points is quite cheap."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 5,
      "metadata": {
        "collapsed": true
      },
      "outputs": [],
      "source": [
        "from botorch.optim.initializers import initialize_q_batch_nonneg\n",
        "\n",
        "# generate a large number of random q-batches\n",
        "Xraw = bounds[0] + (bounds[1] - bounds[0]) * torch.rand(100 * N, q, d)\n",
        "Yraw = qEI(Xraw)  # evaluate the acquisition function on these q-batches\n",
        "\n",
        "# apply the heuristic for sampling promising initial conditions\n",
        "X = initialize_q_batch_nonneg(Xraw, Yraw, N)\n",
        "\n",
        "# we'll want gradients for the input\n",
        "X.requires_grad_(True);"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Optimizing the acquisition function\n",
        "\n",
        "If you have used PyTorch, the basic optimization loop should be quite familiar. However, it is important to note that there is a **key difference** here compared to training ML models: When training ML models, one typically computes the gradient of an empirical loss function w.r.t. the model's parameters, while here we take the gradient of the acquisition function w.r.t. to the candidate set.\n",
        "\n",
        "Thus, when setting the optimizer from `torch.optim`, we **do not** add the acquisition function's parameters as parameters to optimize (that would be quite bad!).\n",
        "\n",
        "In this example, we use a vanilla `Adam` optimizer with fixed learning rate for a fixed number of iterations in order to keep things simple. But you can get as fancy as you want with learning rate scheduling, early termination, etc.\n",
        "\n",
        "A couple of things to note:\n",
        "1. Evaluating the acquisition function on the `N x q x d`-dim inputs means evaluating `N` `q`-batches in `t`-batch mode. The result of this is an `N`-dim tensor of acquisition function values, evaluated independently. To compute the gradient of the full input `X` via back-propagation, we can for convenience just compute the gradient of the sum of the losses. \n",
        "2. `torch.optim` does not have good built in support for constraints (general constrained stochastic optimization is hard and still an open research area). Here we do something simple and project the value obtained after taking the gradient step to the feasible set - that is, we perform \"projected stochastic gradient descent\". Since the feasible set here is a hyperrectangle, this can be done by simple clamping. Another approach would be to transform the feasible interval for each dimension to the real line, e.g. by using a sigmoid function, and then optimizing in the unbounded transformed space. "
      ]
    },
    {
      "cell_type": "code",
      "execution_count": 6,
      "metadata": {},
      "outputs": [
        {
          "name": "stdout",
          "output_type": "stream",
          "text": [
            "Iteration  15/75 - Loss: -0.495\n",
            "Iteration  30/75 - Loss: -0.626\n",
            "Iteration  45/75 - Loss: -0.766\n",
            "Iteration  60/75 - Loss: -1.023\n",
            "Iteration  75/75 - Loss: -1.066\n"
          ]
        }
      ],
      "source": [
        "# set up the optimizer, make sure to only pass in the candidate set here\n",
        "optimizer = torch.optim.Adam([X], lr=0.01)\n",
        "X_traj = []  # we'll store the results\n",
        "\n",
        "# run a basic optimization loop\n",
        "for i in range(75):\n",
        "    optimizer.zero_grad()\n",
        "    # this performs batch evaluation, so this is an N-dim tensor\n",
        "    losses = - qEI(X)  # torch.optim minimizes\n",
        "    loss = losses.sum()\n",
        "    \n",
        "    loss.backward()  # perform backward pass\n",
        "    optimizer.step()  # take a step\n",
        "    \n",
        "    # clamp values to the feasible set\n",
        "    for j, (lb, ub) in enumerate(zip(*bounds)):\n",
        "        X.data[..., j].clamp_(lb, ub) # need to do this on the data not X itself\n",
        "    \n",
        "    # store the optimization trajecatory\n",
        "    X_traj.append(X.detach().clone())\n",
        "    \n",
        "    if (i + 1) % 15 == 0:\n",
        "        print(f\"Iteration {i+1:>3}/75 - Loss: {loss.item():>4.3f}\")\n",
        "    \n",
        "    # use your favorite convergence criterion here..."
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": true
      },
      "outputs": [],
      "source": []
    }
  ],
  "metadata": {
    "kernelspec": {
      "display_name": "python3",
      "language": "python",
      "name": "python3"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 2
}