shrinkGPR: Gaussian Process Regression with Shrinkage and Normalizing...

In shrinkGPR: Scalable Gaussian Process Regression with Hierarchical Shrinkage Priors

Gaussian Process Regression with Shrinkage and Normalizing Flows

Description

shrinkGPR implements Gaussian process regression (GPR) with a hierarchical shrinkage prior for hyperparameter estimation, incorporating normalizing flows to approximate the posterior distribution. The function facilitates model specification, optimization, and training, including support for early stopping, user-defined kernels, and flow-based transformations.

Usage

shrinkGPR(
  formula,
  data,
  a = 0.5,
  c = 0.5,
  formula_mean,
  a_mean = 0.5,
  c_mean = 0.5,
  sigma2_rate = 10,
  kernel_func = kernel_se,
  n_layers = 10,
  n_latent = 10,
  flow_func = sylvester,
  flow_args,
  n_epochs = 1000,
  auto_stop = TRUE,
  cont_model,
  device,
  display_progress = TRUE,
  optim_control
)

Arguments

formula	object of class "formula": a symbolic representation of the model for the covariance equation, as in lm. The response variable and covariates are specified here.
data	optional data frame containing the response variable and the covariates. If not found in data, the variables are taken from environment(formula). No NAs are allowed in the response variable or covariates.
a	positive real number controlling the shrinkage prior for the covariance structure. The default is 0.5.
c	positive real number controlling the tail behavior of the shrinkage prior for the covariance structure. The default is 0.5.
formula_mean	optional formula for the mean equation. If provided, the response variable and covariates for the mean structure are specified separately from the covariance structure.
a_mean	positive real number controlling the shrinkage prior for the mean structure. The default is 0.5.
c_mean	positive real number controlling the tail behavior of the shrinkage prior for the mean structure. The default is 0.5.
sigma2_rate	positive real number controlling the prior rate parameter for the residual variance. The default is 10.
kernel_func	function specifying the covariance kernel. The default is kernel_se, a squared exponential kernel. For guidance on how to provide a custom kernel function, see Details.
n_layers	positive integer specifying the number of flow layers in the normalizing flow. The default is 10.
n_latent	positive integer specifying the dimensionality of the latent space for the normalizing flow. The default is 10.
flow_func	function specifying the normalizing flow transformation. The default is sylvester. For guidance on how to provide a custom flow function, see Details.
flow_args	optional named list containing arguments for the flow function. If not provided, default arguments are used. For guidance on how to provide a custom flow function, see Details.
n_epochs	positive integer specifying the number of training epochs. The default is 1000.
auto_stop	logical value indicating whether to enable early stopping based on convergence. The default is TRUE.
cont_model	optional object returned from a previous shrinkGPR call, enabling continuation of training from the saved state.
device	optional device to run the model on, e.g., torch_device("cuda") for GPU or torch_device("cpu") for CPU. Defaults to GPU if available; otherwise, CPU.
display_progress	logical value indicating whether to display progress bars and messages during training. The default is TRUE.
optim_control	optional named list containing optimizer parameters. If not provided, default settings are used.

Details

This implementation provides a computationally efficient framework for GPR, enabling flexible modeling of mean and covariance structures. Users can specify custom kernel functions, flow transformations, and hyperparameter configurations to adapt the model to their data.

The shrinkGPR function combines Gaussian process regression with shrinkage priors and normalizing flows for efficient and flexible hyperparameter estimation. It supports custom kernels, hierarchical shrinkage priors for mean and covariance structures, and flow-based posterior approximations. The auto_stop option allows early stopping based on lack of improvement in ELBO.

Custom Kernel Functions

Users can define custom kernel functions for the covariance structure of the Gaussian process by passing them to the kernel_func argument. A valid kernel function must follow the same structure as the provided kernel_se (squared exponential kernel). The function should:

1. Accept the following arguments:

   • thetas: A torch_tensor of dimensions n_latent x d, representing latent length-scale parameters.

   • tau: A torch_tensor of length n_latent, representing latent scaling factors.

   • x: A torch_tensor of dimensions N x d, containing the input data points.

   • x_star: Either NULL or a torch_tensor of dimensions N_new x d. If NULL, the kernel is computed for x against itself. Otherwise, it computes the kernel between x and x_star.

2. Return:

   • If x_star = NULL, the function must return a torch_tensor of dimensions n_latent x N x N, representing pairwise covariances between all points in x.

   • If x_star is provided, the function must return a torch_tensor of dimensions n_latent x N_new x N, representing pairwise covariances between x_star and x.

3. Requirements:

   • The kernel must compute a valid positive semi-definite covariance matrix.

   • It should use efficient tensor operations from the Torch library (e.g., torch_bmm, torch_sum) to ensure compatibility with GPUs or CPUs.

Testing a Custom Kernel Function

To test a custom kernel function:

1. Verify Dimensions:

   • When x_star = NULL, ensure the output is n_latent x N x N.

   • When x_star is provided, ensure the output is n_latent x N_new x N.

2. Check Positive Semi-Definiteness: Validate that the kernel produces a positive semi-definite covariance matrix for valid inputs.

3. Integrate: Use the custom kernel with shrinkGPR to confirm its compatibility.

Examples of kernel functions can be found in the kernel_funcs.R file in the package source code, which are documented in the kernel_functions help file.

Custom Flow Functions

Users can define custom flow functions for use in Gaussian process regression models by following the structure and conventions of the provided sylvester function. A valid flow function should be implemented as a nn_module in torch and must meet the following requirements:

Structure of a Custom Flow Function

1. Initialization (initialize):

   • Include all required parameters as nn_parameter or nn_buffer, and initialize them appropriately.

   • Parameters may include matrices for transformations (e.g., triangular matrices), biases, or other learnable components.

2. Forward Pass (forward):

   • The forward method should accept an input tensor z of dimensions n_latent x D.

   • The method must:

     • Compute the transformed tensor z.

     • Compute the log determinant of the Jacobian (log|det J|).

   • The method should return a list containing:

     • zk: The transformed samples after applying the flow (n_latent x D).

     • log_diag_j: A tensor of size n_latent containing the log determinant of the Jacobian for each sample.

3. Output Dimensions:

   • Input tensor z: n_latent x D.

   • Outputs:

     • zk: n_latent x D.

     • log_diag_j: n_latent.

An example of a flow function can be found in the sylvester.R file in the package source code, which is documented in the sylvester help file.

Value

A list object of class shrinkGPR, containing:

model	The best-performing trained model.
loss	The best loss value (ELBO) achieved during training.
loss_stor	A numeric vector storing the ELBO values at each training iteration.
last_model	The model state at the final iteration.
optimizer	The optimizer object used during training.
model_internals	Internal objects required for predictions and further training, such as model matrices and formulas.

Author(s)

Peter Knaus [e-mail]

Examples

if (torch::torch_is_installed()) {
  # Simulate data
  set.seed(123)
  torch::torch_manual_seed(123)
  n <- 100
  x <- matrix(runif(n * 2), n, 2)
  y <- sin(2 * pi * x[, 1]) + rnorm(n, sd = 0.1)
  data <- data.frame(y = y, x1 = x[, 1], x2 = x[, 2])

  # Fit GPR model
  res <- shrinkGPR(y ~ x1 + x2, data = data)

  # Check convergence
  plot(res$loss_stor, type = "l", main = "Loss Over Iterations")

  # Check posterior
  samps <- gen_posterior_samples(res, nsamp = 1000)
  boxplot(samps$thetas) # Second theta is pulled towards zero

  # Predict
  x1_new <- seq(from = 0, to = 1, length.out = 100)
  x2_new <- runif(100)
  y_new <- predict(res, newdata = data.frame(x1 = x1_new, x2 = x2_new), nsamp = 2000)

  # Plot
  quants <- apply(y_new, 2, quantile, c(0.025, 0.5, 0.975))
  plot(x1_new, quants[2, ], type = "l", ylim = c(-1.5, 1.5),
        xlab = "x1", ylab = "y", lwd = 2)
  polygon(c(x1_new, rev(x1_new)), c(quants[1, ], rev(quants[3, ])),
        col = adjustcolor("skyblue", alpha.f = 0.5), border = NA)
  points(x[,1], y)
  curve(sin(2 * pi * x), add = TRUE, col = "forestgreen", lwd = 2, lty = 2)
  }

shrinkGPR documentation built on April 4, 2025, 3:07 a.m.