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