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

In plain words

Many machine learning methods compare every point in one set with every point in another, for example to weight each point by how close it is to the others. The natural code builds the full table of pairwise values first. For large point sets that table no longer fits in GPU memory, so the program crashes.

KeOps is a library that lets you write the same formula as ordinary tensor code, but never builds the table. You describe each pairwise value as a formula, and the library computes the sums (or minimums, or nearest-neighbour lists) on the fly, one small block at a time, directly on the GPU. It also knows how to take gradients through these formulas, so the method can still be trained. The library writes and compiles a fast GPU routine for each formula the first time it is used.

The paper shows that this is both memory-light and faster than the usual PyTorch approach on the simple kernel benchmark it uses.

TL;DR

• KeOps evaluates reductions of formulas over two point sets without storing the kernel or distance matrix, so memory stays linear in the number of points §2.
• A LazyTensor wrapper for NumPy and PyTorch builds the formula lazily; the reduction call runs the computation §2.
• Gradients come from automatic differentiation over the formula, to any order §2.
• Author-reported: on a Gaussian matrix-vector product in 3D, KeOps is faster than PyTorch and the other compared libraries where those fit in memory, and PyTorch runs out of memory from 100k points §3, table.

Problem

A Gaussian matrix-vector product between point sets x (M points) and y (N points), with weights b, needs the M-by-N kernel matrix §1, eq. 1. Storing that matrix does not scale when M and N are around 50k or more. Splitting it into a loop of smaller products avoids the memory limit but stays slow in dimensions up to 100, because moving tiles between GPU memory layers is the bottleneck §1. The paper targets kernel and geometric methods where M and N range from a hundred to a billion §1. It notes that the approach is of limited use when the kernel matrix fits in memory and is reused many times, or when the formula is cheap to evaluate, and that it is most useful with 10k points or more in dimension at most 100 §4.

Method

Problem form. A KeOps reduction takes parameters p^1..p^P (vectors), i-variables x^1..x^X (matrices with M rows), j-variables y^1..y^Y (matrices with N rows), a vector-valued symbolic formula F, and a reduction such as sum, max, argmin or log-sum-exp §2, list of four items. It evaluates

a_i = Reduction over j of F(p, x_i, y_j), for each i   [eq. 2]

with a linear memory footprint on CPU and GPU §2, eq. 2.

Gaussian case. The paper's motivating example is a_i = sum_j exp(-||x_i - y_j||^2 / (2 sigma^2)) * b_j eq. 1, with k(x, y) a Gaussian kernel of deviation sigma > 0 §1.

Python interface. Inputs are wrapped as LazyTensors. A (M, D) array becomes an i-variable of shape (M, 1, D) and a (N, D) array becomes a j-variable of shape (1, N, D). Variable types are inferred from these shapes §2, listing; §2 text. Operations such as squared distance, exp and products only extend a symbolic formula string §2. Reductions over the virtual axes (-3 and -2) run the computation: .sum, .logsumexp, .max, .min, .argmin and .argKmin(K) return indices of the K smallest entries per row, and @ with a dense tensor is a matrix product §2.

Engine. The engine creates optimized C++ code for each new formula and reduction, compiles it with the CUDA toolchain (nvcc plus gcc or clang), and stores the binary for later use. Compilation is done once per reduction §2.

Autodiff. KeOps supports differentiation to arbitrary order. A new binary is generated for each partial derivative a computation needs. It plugs into torch.autograd, so .backward() and torch.autograd.grad() work §2.

Overview of the example call sequence §2, listing:

# Gaussian kernel matrix-vector product  [eq. 1]  [§2 listing]
# Inputs:  x (M, D), y (N, D), b (N,), sigma
# Output:  a (M,)

def gaussian_matvec(x, y, b, sigma):
    # Build the symbolic squared distance; nothing is computed yet  [§2]
    sq_dist = squared_distance(x[:, None], y[None, :])   # (M, N) symbolic
    # Symbolic Gaussian kernel  [eq. 1]
    kernel = exp(-sq_dist / (2 * sigma ** 2))             # (M, N) symbolic
    # Reduction over j triggers the fused GPU computation  [§2]
    return sum_over_j(kernel * b[None, :])                # (M,)

Reference meaning of the generic reduction (not how the engine runs it; the engine never builds the full M-by-N table) eq. 2:

# Reference semantics of a generic reduction  [eq. 2]
# Inputs:  formula F, params, x (M, d_x), y (N, d_y)
# Output:  a (M, d_out)

def generic_reduction(formula, params, x, y, reduction):
    # Row i of the output reduces the formula over every column j  [eq. 2]
    rows = []
    for i in range(M):
        values = [formula(params, x[i], y[j]) for j in range(N)]
        rows.append(reduction(values))           # e.g. sum, max, argmin
    return stack(rows)                            # (M, d_out)

Hyperparameters

KeOps has no training hyperparameters. The benchmark settings are:

Results

Author-reported. Setup: Gaussian kernel matrix-vector product, M = N, D = 3, float32, one RTX 2080 Ti (PyTorch-TPU: Google Colab) §3. The paper does not state the timing protocol. The table has no caption or number in the compiled paper, so it is cited as the unnumbered table in §3.

Source: unnumbered table in §3. The same numbers appear in the benchmark README at the pinned commit README.md.

Other claims in the paper, with no number attached: KeOps "outperforms standard GPU baselines" [abstract]; "several orders of magnitude" faster than NumPy and PyTorch counterparts while keeping linear memory §3, text; interfaces with Scipy and GPyTorch iterative solvers and supports block-sparse reductions [§3, text]; compared with JAX and FAISS in an online gallery, not in the paper §3.

Paper vs code

Checked against the official repository, getkeops/keops, at commit ef72bc907df78cd18bff866d3dc7cb6375f11a4d [code_commit]. This is a newer release than the arXiv v2 (2021-04-08) paper: the pinned tree reports version 2.4b keops_version and copyright 2023 licence.txt.

• Gaussian bandwidth. The paper's Eq. 1 uses exp(-||x - y||^2 / (2 sigma^2)) eq. 1. The KeOps benchmark computes Exp(-SqDist(X,Y)) * B with no sigma KeOps.py. The PyTorch benchmark computes exp(-D_xy) for the squared distance D_xy PyTorch_GPU.py. Both correspond to sigma^2 = 1/2; the paper does not state a value.
• Distance computation in the PyTorch baseline. The baseline uses the expansion ||x||^2 - 2 x.y + ||y||^2 PyTorch_GPU.py, not a direct difference of points. The paper does not say which form it timed.
• Timing protocol. The paper does not state the repetition count or the synchronisation. The KeOps script runs 10 timed calls and does not call torch.cuda.synchronize() before stopping the clock KeOps.py. The PyTorch script runs 100 timed calls and synchronises before stopping the clock PyTorch_GPU.py. The table compares times from these two protocols.
• Dataset and dimension. Both scripts draw x, y and b from standard normal noise with D = 3 and N in {10k, 100k, 1M} [code: benchmarks/KeOps.py#L11, L25-L27]. This matches the setup in §3.
• Autodiff. The compiled paper claims gradients to arbitrary order §2. The benchmark scripts only time forward products KeOps.py. No gradient benchmark was found in the repository's benchmarks/ folder.

Pitfalls

• Compile latency is a real cost. Each new formula and reduction is compiled once, with nvcc and a host compiler §2. The compiled paper gives no compile time, and §4 lists reducing it as a priority §4. Benchmarks that time the first call would include compilation; the scripts avoid this with a warm-up call.
• Speed-up wording vs the table. The text says "several orders of magnitude" §3, text. The only quantified comparison in the compiled paper is the table: at 10k points PyTorch takes 9 ms and PyKeOps 0.7 ms. The other rows have PyTorch out of memory. The text's wording is not backed by a number in the paper.
• Only one setup is measured. One GPU (RTX 2080 Ti), float32, D = 3, one kernel formula §3. §4 says the library is "of limited use" when the kernel matrix fits in memory and is reused, or when the formula is expensive to evaluate §4.
• Bandwidth is not in the paper's experiments. Implementers must choose sigma themselves. The benchmark's implicit value is sigma^2 = 1/2 KeOps.py.
• Sync in the KeOps timing loop. The KeOps script does not synchronise before reading the clock KeOps.py. The reported KeOps times may therefore be shorter than the true GPU time. This card did not re-run the benchmarks to check.
• "Lines of code" is a subjective measure. The row counts the benchmark code for each library, not the effort to build it §3, table.
• Source-only text is not in the paper. The LaTeX source (main.tex) has commented-out material: a benchmark figure (lines 364-374), an "Internal Engine" section with tiled GPU reduction details (lines 419-528), and older discussion of compile times. The compiled paper does not contain these, so do not cite them as the paper's claims.
• Unverified tool comparisons. The Halide and TVM columns come from separate scripts (gauss_conv_halide.cpp, TVM.py) that this card did not run README.md.

Lineage

• Builds on: GPU tiled map-reduce schemes from GPU Gems 3, chapter 31 [§1, cited as GpuGems3]. The paper treats them as the basis for the engine.
• Compared against: PyTorch, PyTorch-TPU, TF-XLA, Halide and TVM §3. Also named but not measured in the compiled paper: JAX and FAISS, in the online gallery §3.
• Future work named in the paper: Nyström and fast furious methods, both approximation schemes beyond the block-sparse truncation already supported §4. Superseded: not known from this wiki's content.