Kernel matrix reduction

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 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 2004.11127 §2. The reductions covered are sums, log-sum-exp, maxima and minima, and the indices of the K smallest entries per row 2004.11127 §2.

The paper's target is kernel and geometric problems with M and N from a hundred to a billion points 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 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.

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 2004.11127 §2. The engine then generates a GPU routine for that formula and reduction, compiles it once, and reuses the binary 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 2004.11127 §2. The paper claims differentiation to arbitrary order 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 2004.11127 §3. The Gaussian product is the paper's test case, with D = 3 and float32 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 2004.11127 §3, table. The text calls the speed-up "several orders of magnitude" 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

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 2004.11127 §2.

# 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,)

# 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 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: the KeOps library. It evaluates kernel and distance reductions on the GPU without storing the matrix, with autodiff.