shrinkMVGPR: Multivariate Gaussian Process Regression with Shrinkage and...

In shrinkGPR: Scalable Gaussian Process Regression with Hierarchical Shrinkage Priors

Multivariate Gaussian Process Regression with Shrinkage and Normalizing Flows

Description

Fits a multivariate Gaussian process regression (MVGPR) model to an N \times M response matrix Y. The joint distribution is matrix normal, Y \sim \mathcal{MN}(0,\, K + \sigma^2 I,\, \Omega), where K is the GP kernel matrix with triple-gamma shrinkage priors on the inverse length-scales, and \Omega is an M \times M output covariance matrix with an LKJ prior on its correlations and triple-gamma priors on its scale parameters. The joint posterior is approximated by normalizing flows trained to maximize the ELBO.

Usage

shrinkMVGPR(
  formula,
  data,
  a = 0.5,
  c = 0.5,
  eta = 4,
  a_Om = 0.5,
  c_Om = 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. Specifically, the response is created by binding the M response variables together with cbind() on the left-hand side of the formula, e.g., cbind(y1, y2) ~ x.
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 behavior at the origin of 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.
eta	positive real number controlling the concentration of the LKJ prior on the correlation matrix of the output covariance. Higher values push the prior towards the identity matrix. The default is 4.
a_Om	positive real number controlling the behavior at the origin of the shrinkage prior for the output covariance scale parameters. The default is 0.5.
c_Om	positive real number controlling the tail behavior of the shrinkage prior for the output covariance scale parameters. 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 shrinkMVGPR 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

Model Specification

Given N observations with d-dimensional covariates and M response variables, the response matrix Y \in \mathbb{R}^{N \times M} follows a matrix normal distribution:

Y \sim \mathcal{MN}_{N,M}(0,\; K(\theta, \tau) + \sigma^2 I_N,\; \Omega),

which is equivalent to

\mathrm{vec}(Y) \sim \mathcal{N}_{NM}(0,\; \Omega \otimes (K + \sigma^2 I_N)).

Here K_{ij} = k(x_i, x_j;\, \theta, \tau) is the kernel matrix and \Omega is the M \times M between-response covariance. The output covariance is parameterized as \Omega = S D S, where D is a correlation matrix and S = \mathrm{diag}(s_1, \ldots, s_M) contains the marginal standard deviations. The product of the diagonal elements of S is constrained to equal 1 to ensure identifiability. The default squared exponential kernel is

k(x, x';\, \theta, \tau) = \frac{1}{\tau} \exp\!\left(-\frac{1}{2} \sum_{j=1}^d \theta_j (x_j - x'_j)^2\right),

where \theta_j \ge 0 are inverse squared length-scales and \tau > 0 is the output scale. Users can specify custom kernels by following the guidelines below, or use one of the other provided kernel functions in kernel_functions.

Priors

\theta_j \mid \tau \sim \mathrm{TG}(a, c, \tau), \quad j = 1, \ldots, d,

\tau \sim F(2c, 2a),

\sigma^2 \sim \mathrm{Exp}(\sigma^2_\mathrm{rate}),

D \sim \mathrm{LKJ}(\eta),

s_m \mid \tau_\Omega \sim \mathrm{TG}(a_\Omega, c_\Omega, \tau_\Omega), \quad m = 1, \ldots, M,

\tau_\Omega \sim F(2c_\Omega, 2a_\Omega).

The LKJ(\eta) prior on the correlation matrix D is uniform over correlations when \eta = 1 and concentrates near the identity as \eta increases.

Inference

The posterior is approximated by a normalizing flow q_\phi trained to maximize the ELBO. auto_stop triggers early stopping when the ELBO shows no significant improvement over the last 100 iterations.

Custom Kernel Functions

Users can define custom kernel functions by passing them to the kernel_func argument. A valid kernel function must follow the same structure as kernel_se. The function must:

1. Accept arguments thetas (n_latent x d), tau (length n_latent), x (N x d), and optionally x_star (N_new x d).

2. Return a torch_tensor of dimensions n_latent x N x N (if x_star = NULL) or n_latent x N_new x N (if x_star is provided).

3. Produce a valid positive semi-definite covariance matrix using torch tensor operations.

See kernel_functions for documented examples.

Custom Flow Functions

Users can define custom flow functions by implementing an nn_module in torch. The module must have a forward method that accepts a tensor z of shape n_latent x D and returns a list with:

• zk: the transformed samples, shape n_latent x D.

• log_diag_j: log-absolute-determinant of the Jacobian, shape n_latent.

See sylvester for a documented example.

Value

A list object of class shrinkMVGPR, 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 multivariate data
  torch::torch_manual_seed(123)
  sim <- simMVGPR(N = 100, M = 2, d = 2)

  # Fit MVGPR model
  res <- shrinkMVGPR(cbind(y.1, y.2) ~ x.1 + x.2, data = sim$data)

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

  # Check posterior of length-scale parameters
  samps <- gen_posterior_samples(res, nsamp = 1000)
  boxplot(samps$thetas)

  # Predict at new covariate values
  newdata <- data.frame(x.1 = runif(10), x.2 = runif(10))
  y_new <- predict(res, newdata = newdata, nsamp = 500)
  # y_new is an array of shape nsamp x N_new x M
  }

shrinkGPR documentation built on March 30, 2026, 5:06 p.m.