Kernel Operations on the GPU, with Autodiff, without Memory Overflows

Abstract

The KeOps library provides a fast and memory-efficient GPU support for tensors whose entries are given by a mathematical formula, such as kernel and distance matrices. KeOps alleviates the major bottleneck of tensor-centric libraries for kernel and geometric applications: memory consumption. It also supports automatic differentiation and outperforms standard GPU baselines, including PyTorch CUDA tensors or the Halide and TVM libraries. KeOps combines optimized C++/CUDA schemes with binders for high-level languages: Python (Numpy and PyTorch), Matlab and GNU R. As a result, high-level "quadratic" codes can now scale up to large data sets with millions of samples processed in seconds. KeOps brings graphics-like performances for kernel methods and is freely available on standard repositories (PyPi, CRAN). To showcase its versatility, we provide tutorials in a wide range of settings online at \url{www.kernel-operations.io}.

6 pages

References in corpus (3)

• GPyTorch: Blackbox Matrix-Matrix Gaussian Process Inference with GPU Acceleration
• TVM: An Automated End-to-End Optimizing Compiler for Deep Learning
• The Fast and Free Memory Method for the efficient computation of convolution kernels

Cited by in corpus (17)

• Differentiable Graph Module (DGM) for Graph Convolutional Networks
• Sinkhorn Divergences for Unbalanced Optimal Transport
• Elastic shape analysis of surfaces with second-order Sobolev metrics: a comprehensive numerical framework
• Pure Tensor Program Rewriting via Access Patterns (Representation Pearl)
• Sliced Gromov-Wasserstein
• A Generative Model for Texture Synthesis based on Optimal Transport between Feature Distributions
• Low-rank statistical finite elements for scalable model-data synthesis
• On UMAP's true loss function
• Towards a mathematical theory of trajectory inference
• Fast Matrix Square Roots with Applications to Gaussian Processes and Bayesian Optimization
• Targeted Likelihood-Free Inference of Dark Matter Substructure in Strongly-Lensed Galaxies
• Bulk topological states in a new collective dynamics model
• SKIing on Simplices: Kernel Interpolation on the Permutohedral Lattice for Scalable Gaussian Processes
• Bayesian Optimization of Sampling Densities in MRI
• Tighter Bounds on the Log Marginal Likelihood of Gaussian Process Regression Using Conjugate Gradients
• Learning landmark geodesics using Kalman ensembles
• ChebLieNet: Invariant Spectral Graph NNs Turned Equivariant by Riemannian Geometry on Lie Groups