shrinkGPR {shrinkGPR} | R Documentation |
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 |
data |
optional data frame containing the response variable and the covariates. If not found in |
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 |
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 |
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 |
cont_model |
optional object returned from a previous |
device |
optional device to run the model on, e.g., |
display_progress |
logical value indicating whether to display progress bars and messages during training. The default is |
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:
-
Accept the following arguments:
-
thetas
: Atorch_tensor
of dimensionsn_latent x d
, representing latent length-scale parameters. -
tau
: Atorch_tensor
of lengthn_latent
, representing latent scaling factors. -
x
: Atorch_tensor
of dimensionsN x d
, containing the input data points. -
x_star
: EitherNULL
or atorch_tensor
of dimensionsN_new x d
. IfNULL
, the kernel is computed forx
against itself. Otherwise, it computes the kernel betweenx
andx_star
.
-
-
Return:
If
x_star = NULL
, the function must return atorch_tensor
of dimensionsn_latent x N x N
, representing pairwise covariances between all points inx
.If
x_star
is provided, the function must return atorch_tensor
of dimensionsn_latent x N_new x N
, representing pairwise covariances betweenx_star
andx
.
-
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:
-
Verify Dimensions:
When
x_star = NULL
, ensure the output isn_latent x N x N
.When
x_star
is provided, ensure the output isn_latent x N_new x N
.
-
Check Positive Semi-Definiteness: Validate that the kernel produces a positive semi-definite covariance matrix for valid inputs.
-
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
-
Initialization (
initialize
):Include all required parameters as
nn_parameter
ornn_buffer
, and initialize them appropriately.Parameters may include matrices for transformations (e.g., triangular matrices), biases, or other learnable components.
-
Forward Pass (
forward
):The
forward
method should accept an input tensorz
of dimensionsn_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 sizen_latent
containing the log determinant of the Jacobian for each sample.
-
-
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 peter.knaus@wu.ac.at
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)
}