Fast Matrix Square Roots with Applications to Gaussian Processes and Bayesian Optimization

Abstract

Matrix square roots and their inverses arise frequently in machine learning, e.g., when sampling from high-dimensional Gaussians $\mathcal{N}(0,\mathbf{K})$ or whitening a vector $\mathbf{b}$ against covariance matrix $\mathbf{K}$. While existing methods typically require $O(N^3)$ computation, we introduce a highly-efficient quadratic-time algorithm for computing ${\mathbf{K}}^{1/2}\mathbf{b}$, ${\mathbf{K}}^{-1/2}\mathbf{b}$, and their derivatives through matrix-vector multiplication (MVMs). Our method combines Krylov subspace methods with a rational approximation and typically achieves $4$ decimal places of accuracy with fewer than $100$ MVMs. Moreover, the backward pass requires little additional computation. We demonstrate our method's applicability on matrices as large as $50,000\times 50,000$ - well beyond traditional methods - with little approximation error. Applying this increased scalability to variational Gaussian processes, Bayesian optimization, and Gibbs sampling results in more powerful models with higher accuracy.

arXiv commentNeurIPS 2020

References in corpus (19)

• Practical Bayesian Optimization of Machine Learning Algorithms
• Gaussian Processes for Big Data
• Predictive Entropy Search for Efficient Global Optimization of Black-box Functions
• Scalable Variational Gaussian Process Classification
• Gaussian Process Kernels for Pattern Discovery and Extrapolation
• Kernel Interpolation for Scalable Structured Gaussian Processes (KISS-GP)
• Max-value Entropy Search for Efficient Bayesian Optimization
• Scalable Global Optimization via Local Bayesian Optimization
• Krylov space solvers for shifted linear systems
• Variational Fourier features for Gaussian processes
• Kernel Operations on the GPU, with Autodiff, without Memory Overflows
• Fast Variational Inference in the Conjugate Exponential Family
• A Framework for Individualizing Predictions of Disease Trajectories by Exploiting Multi-Resolution Structure
• Learning to Detect Sepsis with a Multitask Gaussian Process RNN Classifier
• Finding structure with randomness: Probabilistic algorithms for constructing approximate matrix decompositions
• Product Kernel Interpolation for Scalable Gaussian Processes
• Numerical Methods for the QCD Overlap Operator IV: Hybrid Monte Carlo
• Preconditioning Kernel Matrices
• Scalable iterative methods for sampling from massive Gaussian random vectors

Cited by in corpus (5)

• Parallel Bayesian Optimization of Multiple Noisy Objectives with Expected Hypervolume Improvement
• Overcoming Failures of Imagination in AI Infused System Development and Deployment
• Multi-Objective Bayesian Optimization over High-Dimensional Search Spaces
• Error bounds for Lanczos-based matrix function approximation
• Hierarchical Inducing Point Gaussian Process for Inter-domain Observations