Source-linked AI summary
Kernel Operations on the GPU, with Autodiff, without Memory Overflows
Benjamin Charlier, Jean Feydy, Joan Alexis Glaunès, François-David Collin, Ghislain Durif
TL;DR
Large kernel and geometric computations are bottlenecked by the memory required to materialize pairwise matrices. KeOps represents these computations as symbolic formulas and executes reductions with linear memory, while supporting autodiff and high-level tensor interfaces. It generally outperforms NumPy and PyTorch counterparts by several orders of magnitude and targets large-scale kernel workflows.
Problem
Explicit M-by-N kernel matrices do not scale for large kernel and geometric computations because storing them requires substantial memory.
Method
KeOps uses symbolic formula-defined arrays, lazy reductions, optimized C++/CUDA execution, and high-level NumPy and PyTorch interfaces without materializing intermediate pairwise buffers.
Results
KeOps generally outperforms NumPy and PyTorch counterparts by several orders of magnitude while keeping a linear memory footprint.
Takeaways & Limitations
KeOps supports large kernel and geometric applications, including iterative kernel solvers and block-sparse patterns, through a transparent interface.
Takeaways & Limitations
KeOps is of limited use when a kernel matrix fits in memory and is reused many times or when evaluating F(p, x_i, y_j) takes significant time.
Abstract
from arXiv · showhide
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}.
1. Introduction
KeOps targets large kernel and geometric computations by avoiding explicit kernel-matrix storage while preserving high-level GPU and automatic-differentiation workflows. Its motivating operation is a Gaussian-kernel interaction over potentially billion-scale point sets.
- Context: Automatic differentiation and GPU tensor backends have made gradient-based computation with millions of parameters routine in modern machine learning.TensorFlow and PyTorch exemplify this software trend.
- Motivation: KeOps provides a numerical foundation for methods involving large distance or kernel matrices through high-level software supported by low-level C++ routines.The library is motivated by the need for a complete, usable software stack for massively parallel hardware.
- Target operation: The motivating computation evaluates a_i = Σ_j k(x_i, y_j)b_j for M and N ranging from hundreds to a billion.Here x_i and y_j lie in R^D, b_j are real signals, and k is a Gaussian kernel with deviation σ > 0.
- Target operation: The same operation can be viewed either as multiplication by a Gaussian kernel matrix or as convolution sampled on two point clouds.These equivalent interpretations connect kernel methods with geometric and convolutional computations.
- Memory bottleneck: Tensor-centric GPU workflows commonly materialize an explicit M-by-N kernel buffer, which fails to scale when M and N reach about 50k or more.Tiling reduces memory pressure but introduces GPU-memory transfer overhead, especially for D ⩽100.
2. KeOps Purpose and Usage
KeOps exposes formula-defined computations through symbolic arrays and executes reductions without materializing intermediate pairwise matrices. Its interfaces support common tensor workflows, automatic differentiation, and broad kernel or geometric applications with linear memory usage.
- Generic reduction framework: KeOps evaluates generic reductions of vector-valued symbolic formulas over parameter, i-variable, and j-variable inputs.The engine accepts reductions such as sum, max, argmin, and log-sum-exp.
- Generic reduction framework: The C++ engine computes formula-defined expressions with a linear memory footprint on GPUs and CPUs, supporting convolutions, nearest-neighbor classification, and k-means clustering.These applications are enabled by evaluating reductions directly rather than storing full pairwise arrays.
- LazyTensor abstraction: LazyTensor chains remain symbolic until a reduction triggers computation, so no intermediate N-by-M buffer is created in global device memory.The Gaussian distance and kernel formulas are assembled lazily before the matrix product performs the reduction.
- LazyTensor abstraction: LazyTensor provides a tensor-like interface in which array shapes infer variable types and virtual dimensions trigger calls to the KeOps C++ engine.This lets NumPy and PyTorch users specify symbolic computations with familiar operations.
- Performance and differentiation: KeOps scripts generally outperform NumPy and PyTorch counterparts by several orders of magnitude while retaining a linear memory footprint.LazyTensors also support broadcasting, batch dimensions, and reductions including sum, logsumexp, max, min, argmin, and argKmin.
- Performance and differentiation: KeOps supports automatic differentiation up to arbitrary orders through binaries integrated with torch.autograd.Users can differentiate KeOps calls with standard grad() and backward() methods.
3. Performance evaluation
KeOps is evaluated on Gaussian kernel matrix-vector products against PyTorch, TensorFlow, Halide, and TVM as the number of points increases. The results show that NumPy-like KeOps scripts produce highly competitive binaries and support large kernel-system applications.
- Benchmark setup: KeOps is compared with PyTorch, TensorFlow, Halide, and TVM on Gaussian kernel matrix-vector products with increasing point counts.The benchmark uses M = N points in dimension D = 3.
- Compared systems: The benchmark includes PyKeOps alongside PyTorch, PyTorch-TPU, TF-XLA, Halide, and TVM.
- Results: KeOps turns NumPy-like scripts into highly competitive binaries.
- Applications: KeOps interfaces with SciPy and GPytorch solvers and supports cluster-wise block-sparsity for large kernel linear systems.Applications include kriging, splines, Gaussian processes, and kernel methods.
4. Intended Use, Limitations and Future Works
KeOps targets machine-learning computations expressed through symbolic matrices, combining a transparent interface with efficient execution. Its main scope is large, moderate-dimensional problems, while approximation methods and deployment improvements remain future work.
- Intended use: KeOps targets a specific machine-learning niche through transparent symbolic matrices and efficiency on a wide range of computations.The paper contrasts this focus with generalist deep-learning compilers and frameworks.
- Future works: Future work includes Nyström and FFM approximation schemes beyond the currently supported block-sparse truncation rule.The paper also prioritizes easier deployment of pre-compiled binaries, shorter compilation times, and CUDA Tensor-core support.
- Intended use: KeOps is intended for computations involving 10k points or more in spaces of dimension D ⩽100.
- Limitations: The engine is limited when the kernel matrix fits in memory and is reused many times, or when evaluating F(p, x_i, y_j) is time-consuming.