---
kind: concept
slug: kernel-matrix-reduction
title: Kernel matrix reduction
summary: >-
  Reduce a pairwise kernel or distance formula over two point sets on the GPU without storing the full matrix, with autodiff through the formula. KeOps is the one paper so far: author-reported, a Gaussian product over 1M 3D points takes 1.39 s where PyTorch runs out of memory from 100k. Trade-off: a compile per new formula, and evidence from one GPU and one kernel.
tags: [gpu, kernel-methods, memory-efficiency, automatic-differentiation]
papers: ["2004.11127"]
related: [nearest-neighbor-search, gaussian-process, random-feature-approximation, fused-attention-kernel, arithmetic-intensity]
status: draft
updated_at: 2026-10-19
---

## Overview

Many kernel and geometric methods need, for each point `x_i`, a sum over all
points `y_j` of a kernel `k(x_i, y_j)` times a weight. Written densely, this
is an M-by-N matrix product with the kernel matrix. At 50k points or more the
matrix no longer fits in GPU memory [arxiv:2004.11127 §1]. The shared idea is
to evaluate the reduction `a_i = Reduction_j F(p, x_i, y_j)` as one fused GPU
computation, so the M-by-N table never exists in memory [arxiv:2004.11127 §2].
The reductions covered are sums, log-sum-exp, maxima and minima, and the
indices of the K smallest entries per row [arxiv:2004.11127 §2].

The paper's target is kernel and geometric problems with M and N from a
hundred to a billion points [arxiv:2004.11127 §1]. Splitting the product into
a loop over blocks avoids the memory limit but is slow for dimensions up to
100, because moving tiles between GPU memory levels is the bottleneck
[arxiv:2004.11127 §1]. Only one paper card covers this concept, and it
describes one library, so the page has one variant and no comparison between
papers yet. Card: [Kernel Operations on the GPU, with Autodiff, without Memory
Overflows](/papers/2004.11127).

## Variants

### Symbolic formula with fused reduction (KeOps)

The user writes the pairwise formula with ordinary tensor operations on
`LazyTensor` wrappers. Nothing is computed until a reduction is called. The
wrapper infers whether each input is a per-row variable or a per-column
variable from its shape, and the formula is stored as a string
[arxiv:2004.11127 §2]. The engine then generates a GPU routine for that formula
and reduction, compiles it once, and reuses the binary [arxiv:2004.11127 §2].

Gradients use the same machinery. Each partial derivative the computation
needs gets its own generated binary, and the result plugs into
`torch.autograd` [arxiv:2004.11127 §2]. The paper claims differentiation to
arbitrary order [arxiv:2004.11127 §2].

The paper reports two interfaces: a NumPy-like Python one (PyKeOps) and a C++
one (KeOps++). Both appear in the benchmark table [arxiv:2004.11127 §3].
The Gaussian product is the paper's test case, with D = 3 and float32
[arxiv:2004.11127 §3].

Reported timings, author-reported, on one Nvidia RTX 2080 Ti: at 10k points
PyKeOps takes 0.7 ms and PyTorch 9 ms. At 100k points PyTorch is out of
memory and PyKeOps takes 15.0 ms. At 1M points PyKeOps takes 1.39 s
[arxiv:2004.11127 §3, table]. The text calls the speed-up "several orders of
magnitude" [arxiv:2004.11127 §3]. The table is the only place the paper gives
numbers, and it does not show that size of gap at 10k points.

## When to use which

| Situation | Use | Why |
|---|---|---|
| The kernel matrix fits in memory and is reused many times | A stored dense matrix | The paper says the library is of limited use here [arxiv:2004.11127 §4] |
| Point sets of 10k or more, dimension up to 100, a GPU, and only reductions are needed | Symbolic formula with fused reduction | Best fit in the paper's regime [arxiv:2004.11127 §4]; memory stays linear [arxiv:2004.11127 §2] |
| The result needs the full matrix, not a reduction | Not covered by this concept | The reduction interface returns sums, minima, or indices, not the matrix |
| Dimension above 100 | Check the paper's regime first | The paper states that performance is good for dimension at most 100 [arxiv:2004.11127 §1] [arxiv:2004.11127 §4] |
| The formula is cheap to evaluate | Dense code, unless measured otherwise | The paper names cheap formulas as a case of limited use [arxiv:2004.11127 §4] |
| Formulas change often, or each call is short | Budget for compile time | Each new formula and reduction is compiled once [arxiv:2004.11127 §2]; the paper lists reducing this as future work [arxiv:2004.11127 §4] |
| Gradients through the kernel are needed | Symbolic formula with autodiff | The paper claims autodiff to any order [arxiv:2004.11127 §2]. Its timings are forward products only; see the card's paper-vs-code check |

## Implementation sketch

The reduction below is the shape of the Gaussian product in the paper's example.
The engine never builds the `M x N` table. The reference meaning of a generic
reduction is the double loop that follows it [arxiv:2004.11127 §2].

```pseudocode
# Gaussian kernel matrix-vector product  [arxiv:2004.11127 §1, eq. 1]
# Inputs:  x (M, D), y (N, D), b (N,), sigma (not fixed by the paper's runs)
# Output:  a (M,)

def gaussian_matvec(x, y, b, sigma):
    # Build the symbolic formula; no values are computed yet  [§2]
    sq_dist = squared_distance(x[:, None], y[None, :])   # (M, N) symbolic
    kernel = exp(-sq_dist / (2 * sigma ** 2))            # (M, N) symbolic
    # The reduction over j runs the fused GPU computation  [§2]
    return sum_over_j(kernel * b[None, :])               # (M,)
```

```pseudocode
# Reference meaning of a generic reduction  [arxiv:2004.11127 §2, eq. 2]
# Not how the engine runs it: the full table is never stored.

def generic_reduction(formula, params, x, y, reduction):
    # Row i reduces the formula over every column j
    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)
```

Choose `sigma` yourself. The paper's runs do not state a value
[arxiv:2004.11127 §3].

## Contradictions

None known. Only one paper is covered on this page. Within that paper, the
text's "several orders of magnitude" is not matched by the table at 10k points
(see Variants).

## Papers

- [Kernel Operations on the GPU, with Autodiff, without Memory Overflows](/papers/2004.11127): the KeOps library. It evaluates kernel and distance reductions on the GPU without storing the matrix, with autodiff.