# Student-t Process Regression with Shrinkage and Normalizing Flows `shrinkTPR` implements Student-t process regression (TPR), as described in Shah et al. (2014), 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. ```r shrinkTPR( formula, data, a = 0.5, c = 0.5, formula_mean, a_mean = 0.5, c_mean = 0.5, sigma2_rate = 10, nu_alpha = 0.5, nu_beta = 2, 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 `NA`s 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. - `formula_mean`: **optional** formula for the linear mean equation. If provided, the covariates for the mean structure are specified separately from the covariance structure. A response variable is not required in this formula. - `a_mean`: positive real number controlling the behavior at the origin of the shrinkage 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. - `nu_alpha`: positive real number controlling the shape parameter of the shifted gamma prior for the degrees of freedom of the Student-t process. The default is 0.5. - `nu_beta`: positive real number controlling the scale parameter of the shifted gamma prior for the degrees of freedom of the Student-t process. The default is 2. - `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 `shrinkTPR` 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. ## Returns A list object of class `shrinkTPR`, 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. ## Details This implementation provides a computationally efficient framework for TPR, 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 `shrinkTPR` function combines Student-t 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 `shrinkTPR` 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. ## Examples ```r 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 TPR model res <- shrinkTPR(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) # Add mean equation res2 <- shrinkTPR(y ~ x1 + x2, formula_mean = ~ x1, data = data) } ``` ## References Shah, A., Wilson, A., & Ghahramani, Z. (2014, April). Student-t processes as alternatives to Gaussian processes. In Artificial intelligence and statistics (pp. 877-885). PMLR. ## Author(s) Peter Knaus peter.knaus@wu.ac.at

shrinkTPR function