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

By Benjamin Charlier, Jean Feydy, Joan Alexis Glaunès, François-David Collin, G. Durif. 2020

222 citations recorded by the source corpus.

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{this http URL}.

Subfields: Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Optimization & Training

Citations in the corpus

FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness (2022, background): “The KeOps library [8, 26] is a successful example of how reducing memory reads/writes can speed up kernel operations.”
POT: Python Optimal Transport (2021, methodology): “It is based on the very elegant KeOps framework (Charlier et al., 2020), and is currently a better alternative to POT for using Sinkhorn losses in deep learning architectures.”
OpenFold: retraining AlphaFold2 yields new insights into its learning mechanisms and capacity for generalization (2023, unknown): “Predictions made by OpenFold and AlphaFold2 on the CAMEO validation set are assessed using the lDDT-C α 33 (Local Distance Difference Test with respect to the alpha carbon) metric (Fig.”
P HYSICS - INFORMED DEEP NEURAL NETWORK FOR RIGID - BODY PROTEIN DOCKING (2022, methodology): “In order to increase speed and reduce the memory footprint of our computations, we rely extensively on the KeOps library (Charlier et al., 2021; Feydy et al., 2020).”
Fast end-to-end learning on protein surfaces (2020, methodology): “We compute exact interactions between all points of a protein surface efficiently using the recent KeOps library [13, 19] for PyTorch [30] that optimizes a wide range of computations on generalized distance matrices.”
Diagonal State Spaces are as Effective as Structured State Spaces (2022, methodology): “To alleviate this issue, several Transformer variants have been proposed with reduced compute and memory requirements [QML ` 20, KKL20, RSVG21, BPC20, GB20, KVPF20, WLK ` 20, CLD ` 21, ZS21, GDG ` 21] (see [TDBM20] for a survey).”
Interpolating between Optimal Transport and MMD using Sinkhorn Divergences (2018, methodology): “To alleviate this problem, we leveraged the KeOps library (Charlier et al., 2018) that provides online map-reduce routines on the GPU with full PyTorch integration.”
Quasi Linear Kernel Attention with Infinite Capacity (2026, unknown): “• KEOPS : a GPU-friendly O ( N 2 ) reference (Charlier et al., 2021).”