def laplace_sample_U(mll: ExactMarginalLogLikelihood,
nsamp: int) -> Tuple[Tensor, Tensor, Tensor]:
"""Draw posterior samples of kernel hyperparameters using Laplace
approximation.
Only the Mahalanobis distance matrix is sampled.
The diagonal of the Hessian is estimated using finite differences of the
autograd gradients. The Laplace approximation is then N(p_map, inv(-H)).
We construct a set of nsamp kernel hyperparameters by drawing nsamp-1
values from this distribution, and prepending as the first sample the MAP
parameters.
Args:
mll: MLL object of MAP ALEBO GP.
nsamp: Number of samples to return.
Returns: Batch tensors of the kernel hyperparameters Uvec, mean constant,
and output scale.
"""
# Estimate diagonal of the Hessian
mll.train()
x0, property_dict, bounds = module_to_array(module=mll)
x0 = x0.astype(np.float64) # This is the MAP parameters
H = np.zeros((len(x0), len(x0)))
epsilon = 1e-4 + 1e-3 * np.abs(x0)
for i, _ in enumerate(x0):
# Compute gradient of df/dx_i wrt x_i
def f(x):
x_all = x0.copy()
x_all[i] = x[0]
return -_scipy_objective_and_grad(x_all, mll, property_dict)[1][i]
H[i, i] = approx_fprime(np.array([x0[i]]), f,
epsilon=epsilon[i]) # pyre-ignore
# Sample only Uvec; leave mean and output scale fixed.
assert list(property_dict.keys()) == [
"model.mean_module.constant",
"model.covar_module.raw_outputscale",
"model.covar_module.base_kernel.Uvec",
]
H = H[2:, 2:]
H += np.diag(-1e-3 *
np.ones(H.shape[0])) # Add a nugget for inverse stability
Sigma = np.linalg.inv(-H)
samples = np.random.multivariate_normal(mean=x0[2:],
cov=Sigma,
size=(nsamp - 1))
# Include the MAP estimate
samples = np.vstack((x0[2:], samples))
# Reshape
attrs = property_dict["model.covar_module.base_kernel.Uvec"]
Uvec_batch = torch.tensor(samples, dtype=attrs.dtype,
device=attrs.device).reshape(
nsamp, *attrs.shape)
# Get the other properties into batch mode
mean_constant_batch = mll.model.mean_module.constant.repeat(nsamp, 1)
output_scale_batch = mll.model.covar_module.raw_outputscale.repeat(nsamp)
return Uvec_batch, mean_constant_batch, output_scale_batch
def get_map_model(
B: Tensor,
train_X: Tensor,
train_Y: Tensor,
train_Yvar: Tensor,
restarts: int,
init_state_dict: Optional[Dict[str, Tensor]],
) -> ExactMarginalLogLikelihood:
"""Do random-restart optimization for MAP fitting of an ALEBO GP model.
Args:
B: Projection matrix.
train_X: X training data.
train_Y: Y training data.
train_Yvar: Noise variances of each training Y.
restarts: Number of restarts for MAP estimation.
init_state_dict: Optionally begin MAP estimation with this state dict.
Returns: non-batch ALEBO GP with MAP kernel hyperparameters.
"""
f_best = 1e8
sd_best = {}
# Fit with random restarts
for _ in range(restarts):
m = ALEBOGP(B=B,
train_X=train_X,
train_Y=train_Y,
train_Yvar=train_Yvar)
if init_state_dict is not None:
# pyre-fixme[6]: Expected `OrderedDict[typing.Any, typing.Any]` for 1st
# param but got `Dict[str, Tensor]`.
m.load_state_dict(init_state_dict)
mll = ExactMarginalLogLikelihood(m.likelihood, m)
mll.train()
mll, info_dict = fit_gpytorch_scipy(mll,
track_iterations=False,
method="tnc")
logger.debug(info_dict)
# pyre-fixme[58]: `<` is not supported for operand types
# `Union[List[botorch.optim.fit.OptimizationIteration], float]` and `float`.
if info_dict["fopt"] < f_best:
f_best = float(info_dict["fopt"]) # pyre-ignore
sd_best = m.state_dict()
# Set the final value
m = ALEBOGP(B=B, train_X=train_X, train_Y=train_Y, train_Yvar=train_Yvar)
# pyre-fixme[6]: Expected `OrderedDict[str, Tensor]` for 1st param but got
# `Dict[typing.Any, typing.Any]`.
m.load_state_dict(sd_best)
mll = ExactMarginalLogLikelihood(m.likelihood, m)
return mll
def get_map_model(
B: Tensor,
train_X: Tensor,
train_Y: Tensor,
train_Yvar: Tensor,
restarts: int,
init_state_dict: Optional[Dict[str, Tensor]],
) -> ExactMarginalLogLikelihood:
"""Do random-restart optimization for MAP fitting of an ALEBO GP model.
Args:
B: Projection matrix.
train_X: X training data.
train_Y: Y training data.
train_Yvar: Noise variances of each training Y.
restarts: Number of restarts for MAP estimation.
init_state_dict: Optionally begin MAP estimation with this state dict.
Returns: non-batch ALEBO GP with MAP kernel hyperparameters.
"""
f_best = 1e8
sd_best = {}
# Fit with random restarts
for _ in range(restarts):
m = ALEBOGP(B=B,
train_X=train_X,
train_Y=train_Y,
train_Yvar=train_Yvar)
if init_state_dict is not None:
m.load_state_dict(init_state_dict)
mll = ExactMarginalLogLikelihood(m.likelihood, m)
mll.train()
mll, info_dict = fit_gpytorch_scipy(mll,
track_iterations=False,
method="tnc")
logger.debug(info_dict)
if info_dict["fopt"] < f_best:
f_best = float(info_dict["fopt"]) # pyre-ignore
sd_best = m.state_dict()
# Set the final value
m = ALEBOGP(B=B, train_X=train_X, train_Y=train_Y, train_Yvar=train_Yvar)
m.load_state_dict(sd_best)
mll = ExactMarginalLogLikelihood(m.likelihood, m)
return mll