Learnable prior mean

This example uses the empirical Bayes method to learn prior mean function.

import os
import torch
import matplotlib.pyplot as plt

torch.manual_seed(42)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

n_epochs = int(os.getenv("GPYTORCHQR_N_EPOCHS", 10000))

Input data

def mean(x):
    # x: (N, 1)
    return 2 * x.squeeze(-1)


def std(x):
    # x: (N, 1)
    return x.squeeze(-1) + 0.1


x_range = torch.linspace(0, 1, 100, device=device).reshape(-1, 1)
x = x_range.repeat(2, 1)
y = mean(x) + torch.randn(len(x), device=device).mul(std(x))
q = torch.tensor([0.1, 0.25, 0.5, 0.75, 0.9], device=device)
standard_normal = torch.distributions.Normal(
    torch.zeros((), device=device), torch.ones((), device=device)
)
icdf = standard_normal.icdf(q)
true_quantiles = mean(x_range).unsqueeze(1) + std(x_range).unsqueeze(1) * icdf

x_pred = torch.linspace(0, 1.5, 100, device=device).reshape(-1, 1)

plt.scatter(x.cpu(), y.cpu(), c="k", marker=".")
plt.plot(x_range.cpu(), true_quantiles.cpu(), "--", c="gray")
plt.show()

Prior mean

from gpytorch.means import Mean
import torch.nn as nn


class PriorMean(Mean):
    def __init__(self):
        super().__init__()
        self.linear = nn.Linear(in_features=1, out_features=1)
        with torch.no_grad():
            self.linear.weight.fill_(0.0)
            self.linear.bias.fill_(0.0)

    def forward(self, x):
        return self.linear(x)


prior_mean = PriorMean().to(device)

plt.scatter(x.cpu(), y.cpu(), c="k", marker=".")
plt.plot(x_pred.cpu(), prior_mean(x_pred).detach().cpu(), ls="--", c="k")
plt.show()

Model

from gpytorch.means import ConstantMean
from gpytorch.kernels import ScaleKernel, RBFKernel
from gpytorch.variational import (
    CholeskyVariationalDistribution,
    VariationalStrategy,
    IndependentMultitaskVariationalStrategy,
)
from gpytorch_qr.models import CenterGapQuantileGP
from gpytorch_qr.likelihoods import CenterGapQuantilesLikelihood


class QuantileGPModel(CenterGapQuantileGP):
    def __init__(
        self,
        inducing_points,
        num_quantiles,
        num_lower_quantiles,
    ):
        N, D = inducing_points.size()
        variational_distribution = CholeskyVariationalDistribution(
            N,
            batch_shape=torch.Size([num_quantiles]),
        )
        variational_strategy = IndependentMultitaskVariationalStrategy(
            VariationalStrategy(
                self,
                inducing_points,
                variational_distribution,
                learn_inducing_locations=True,
            ),
            num_tasks=num_quantiles,
        )

        mean = ConstantMean(batch_shape=torch.Size([num_quantiles]))
        covar = ScaleKernel(
            RBFKernel(ard_num_dims=D, batch_shape=torch.Size([num_quantiles])),
            batch_shape=torch.Size([num_quantiles]),
        )
        super().__init__(
            variational_strategy, mean, covar, [num_quantiles], [num_lower_quantiles]
        )


inducing_points = torch.linspace(0, 1, 10, device=device).reshape(-1, 1)
central_q_index = 2

likelihood = CenterGapQuantilesLikelihood(q, central_q_index).to(device)
model = QuantileGPModel(
    inducing_points,
    len(q),
    central_q_index,
).to(device)

from gpytorch.mlls import VariationalELBO

prior_mean.train()
model.train()
likelihood.train()

optimizer = torch.optim.Adam(
    list(prior_mean.parameters())
    + list(model.parameters())
    + list(likelihood.parameters()),
    lr=0.001,
)
mll = VariationalELBO(likelihood, model, num_data=y.numel())

for _ in range(n_epochs):
    train_mean = prior_mean(x)
    res = y - train_mean.squeeze(1)

    output = model(x)
    loss = -mll(output, res)
    loss.backward()
    optimizer.step()
    optimizer.zero_grad()

prior_mean.eval()
model.eval()
likelihood.eval()

with torch.no_grad():
    pred_mean = prior_mean(x_pred)
    mean_q = model.mean_quantiles_mc(x_pred)
    lower_q, upper_q = model.quantile_quantiles_mc(
        x_pred, torch.tensor([0.025, 0.975], device=device)
    )

pred_mean_q = (pred_mean + mean_q).detach().cpu()
pred_lower_q = (pred_mean + lower_q).detach().cpu()
pred_upper_q = (pred_mean + upper_q).detach().cpu()

colors = plt.cm.tab10.colors

plt.scatter(x.cpu(), y.cpu(), c="k", marker=".")

plt.plot(x_pred.cpu(), pred_mean.detach().cpu(), ls="--", c="k")
for i in range(len(q)):
    plt.plot(x_pred.cpu(), pred_mean_q[:, i], color=colors[i])
    plt.fill_between(
        x_pred.cpu().squeeze(),
        pred_lower_q[:, i],
        pred_upper_q[:, i],
        color=colors[i],
        alpha=0.3,
    )
plt.show()