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

Introduction

Recent advances in machine learning have been driven by the diffusion of two pieces of software: automatic differentiation and GPU backends for tensor computations. Today, thanks to e.g. the TensorFlow or PyTorch libraries Abadi et al. (2015); Paszke et al. (2017), users routinely perform gradient descent on functions that involve millions of parameters.

These high-level Python frameworks unlock the use of massively parallel hardware for machine learning research. Under the hood, they rely on C++ routines that are often supported by hardware manufacturers: the cuBLAS and cuDNN libraries edited by Nvidia provide the binaries for linear algebra and convolutions that power a majority of deep learning models. In practice, the presence of a complete software stack (from low-level binaries to well-documented Python libraries) is a prerequisite for the widespread adoption of a research idea by the machine learning community. The KeOps library intends to provide such a solid numerical foundation for all methods that involve large distance or kernel matrices. A motivating example is the computation of pair-wise interactions of the form:

KeOps Purpose and Usage

A generic reduction framework. The workhorse of the KeOps library is a C++ engine for generic reductions on sampled data. Let us assume that we have at hand:

a reduction operation such as a sum, max, argmin, log-sum-exp, etc.

Then, a single call to the KeOps C++ engine allows users to evaluate the expression:

efficiently, with a linear memory footprint on GPUs and CPUs. As illustrated in our gallery of tutorials, this level of generality allows KeOps to handle off-grid convolutions, $k$-nearest neighbors classification, $k$-means clustering and many other tasks.

The LazyTensor abstraction. The “LazyTensor” wrapper for NumPy arrays and PyTorch tensors lets users specify computations along the lines of Eq. (2) with a tensor-like interface. For instance, we can specify the Gaussian matrix-vector product of Eq. (1) with:

Note that variable types ($i$-, $j$-variable or parameter) are inferred from the shapes of the input tensors at lines 2 and 3: in practice, symbolic tensors are as easy to use as sparse matrices. The LazyTensor wrapper turns a dense array into a symbolic matrix whose axes -3 and -2 are understood as “virtual” dimensions; a reduction on these axes is the signal that triggers a call to the KeOps C++ engine.

As showcased on our website and at the end of this paper, KeOps scripts for kernel and geometric applications generally outperform their Numpy and PyTorch counterparts by several orders of magnitude while keeping a linear memory footprint. LazyTensors support a wide range of mathematical operations that mimic the usual interface for NumPy arrays and PyTorch tensors. They fully support broadcasting and batch dimensions, as well as a decent collection of reduction operations: .sum(), .logsumexp(), .max() and .min() but also .argmin() or .argKmin(K=...) that return the indices of the smallest (or K-smallest) elements of the rows of a symbolic tensor.

Inner engine. Internally, KeOps creates an optimized C++ code for every new reduction and formula $F$ that it encounters. Binaries are then compiled and stored on the hard drive for later use: compilation relies on the standard CUDA stack (nvcc, gcc and/or clang compilers) and is only performed once per reduction.

Backpropagation. Crucially, KeOps supports automatic differentiation up to arbitrary orders of differentiation: a new binary is created automatically for every new partial derivative that is required by the user’s computations. This mechanism is fully integrated with the torch.autograd engine and lets users “backprop” through KeOps calls using the usual torch.autograd.grad() and .backward() methods.

Performance evaluation

As evidenced by this table, KeOps turns NumPy-like scripts into highly competitive binaries. Going further, it can be neatly interfaced with the iterative linear solvers of the Scipy Jones et al. (2001) or GPytorch Gardner et al. (2018) libraries and supports the specification of cluster-wise block-sparsity patterns: this allows users to solve large kernel linear systems efficiently, with applications to geology (Kriging), imaging (splines), statistics (Gaussian processes) and data sciences (kernel methods).

Intended Use, Limitations and Future Works

KeOps fills a specific but important niche in machine learning research. Unlike most other compilers for deep learning computations, such as Halide and TVM, it is meant to be used directly by theorists of the machine learning community. Our focus on the simple yet powerful concept of symbolic matrices allows us to keep a transparent interface, while being more efficient than the generalist PyTorch and XLA frameworks on a wide range of computations. In future works, we intend to add support for approximation schemes such as the Nyström and FFM methods Yang et al. (2012); Aussal and Bakry (2019), beyond the block-sparse truncation rule that is currently supported. These features will be valuable to many researchers in the field, while being out of scope for most deep learning libraries.

Acknowledgements

The three first authors are the project leaders: they contributed equally to the library and its documentation. The authors also thank Alain Trouvé, whose theoretical work in shape analysis was the first motivation for the development of the KeOps engine.

References