User-friendly syntax for block-sparsity masks, triangular matrices, etc.

[jeanfeydy]

Hi all,

With the presentation of our geometric demos at NeurIPS in a few days, December 2020 is probably a good time to plan ahead for the next stages of the library.

As you know, some work remains to be done on the low-level "KeOps++" engine: on top of the transition to CUDA 11 and the acceleration of compilation times, we should add support for Tensor cores (#100), a better interface for multi-GPU computing (#99) and some other features such as integer formulas. But overall, the foundations of the library are pretty solid now, which should allow us to start working on higher-level features - closer to data scientists and researchers, not just library developers and experienced programmers.

As a target/roadmap for 2021, I have created several projects through the GitHub interface, with detailed to-do lists ("cards") to help new contributors to find their feet:

• a scikit-learn interface,
• better integration with geometric deep learning libraries,
• support for approximate KNN queries and sum reductions.

All these new features will be implemented in Python (probably in the pykeops.common folder, just like the LazyTensor API) and rely extensively on block-sparse reductions. To ease their development (and help advanced users in general!), I believe that we should start thinking about a better interface for "masks" and "partial evaluation" of symbolic tensors in general. "ranges" will always be there as the low-level encoding for block-sparsity masks, but a user-friendly syntax would certainly be appreciated too (#80, #120).

As a starting point for a discussion, here are some ideas for a new API:

1. PyTorch_Geometric-like batch vectors. As is now standard for geometric deep learning, we should allow users to specify block-diagonal ranges using integer arrays of the form batch = [0, ..., 0, 1, ..., 1, ..., B-1, ..., B-1], where B is the batch size and batch is an integer Tensor of shape (N,), where N is a number of points.
   For instance, a batched K-NN search with non-uniform batches would read:

# x is (N, D), y is (M, D),
# batch_x is (N,) with non-decreasing integer values that range between 0 and B-1
# batch_y is (M,) with non-decreasing integer values that range between 0 and B-1

x_i = LazyTensor(x.view(N, 1, D), batch = batch_x)
y_j = LazyTensor(y.view(1, M, D), batch = batch_y)

# At this point, x_i.batch = (batch_x, None), y_i.batch = (None, batch_y).
# Operations automatically propagate the batch vectors
# (with appropriate compatibility checks), just like with
# array dimensions (N, M, etc.).
D_ij = ((x_i - y_j)**2).sum(-1)  # (N, M, 1) squared distances

# At this point, D_ij.batch = (batch_x, batch_y).
# At reduction time, we turn the "batches" into a
# ".ranges" that is used if D_ij.ranges is None.
ind_i = D_ij.argKmin(K=5, dim=1)  # (N, 5)

2. In the same vein, but at a finer level: cluster vectors that allow users to specify a "coarse" hierarchical representation of their data. This will be handy for everything that is related to kernel truncation or fast NN queries. The cluster tensors should have the same shape as the data arrays, minus the final vector dimension. Supporting batch processing shouldn't be too difficult, even though full broadcasting could be a little bit tricky. If "batch" vector and "cluster" vectors are both provided, clusters from different "sub-batches" should not interact with each other - i.e. we generate a "mask" which is thinner than block-diagonal.

# x is (N, D), y is (M, D),
# cluster_x is (N,) with integer values that range between 0 and Ci-1
# cluster_y is (M,) with integer values that range between 0 and Cj-1

# By default, we sort the points at reduction time to ensure that
# clusters are all contiguous in memory:
x_i = LazyTensor(x.view(N, 1, D), cluster = cluster_x)

# Advanced users may also skip this pre-processing if they re-use the same operation 
# several times and provide pre-sorted arrays.
# If "sorted" = True, we assume that the vector of cluster labels is non-decreasing.
y_j = LazyTensor(y.view(1, M, D), cluster = cluster_y, sorted = True)

# At this point, x_i.cluster = (cluster_x, None), y_i.cluster = (None, cluster_y).
# Operations automatically propagate the cluster vectors
# (with appropriate compatibility checks), just like with
# array dimensions (N, M, etc.).
D_ij = ((x_i - y_j)**2).sum(-1)  # (N, M, 1) squared distances
K_ij = (-D_ij).exp()             # (N, M, 1) kernel matrix

# At this point, K_ij.batch = (cluster_x, cluster_y).
# The user can now specify a "block-sparsity matrix":
K_ij.cluster_mask = mask  # (A, B) boolean matrix

# Typically, A = Ci and B = Cj, but we don't really have to check.
# We only keep cluster-cluster interactions such that mask[i,j] == True.

# At reduction time, we turn the "cluster" vectors and
# "cluster_mask" matrix into a ".ranges" 6-uple that 
# is used if D_ij.ranges is None.
a_i = K_ij.sum(dim=1)  # (N, 1)

3. Add support for triangular masks. Under the hood, these should be handled by a combination of block-sparsity masks (to specify a rough outline with blocks of 64-192 lines/columns at a time) and integer comparisons on "i" and "j" (only for the blocks around the diagonal). Prior to a reduction, users could specify something like:

K_ij.triangle = "upper"  # [None | "upper" | "lower"]

4. Add support for band-diagonal masks. Under the hood, these could be handled using either the same strategy as above (masks + integer checks on "i" and "j") or a specific "offset" in the KeOps++ reduction, that would be more efficient. Prior to a reduction, users could specify something like:

K_ij.bands = [-3, -2, 0, 1, 2, 5]  # offsets of the "useful" diagonals

5. Add support for the specific types of sparsity that come from working with times series, 2D images and 3D volumes. This could be useful for e.g. bi-lateral filtering, superpixels and other types of local operations on feature vectors that are defined on a grid domain.

# x is (N, D), y is (M, D),
x_i = LazyTensor(x.view(N, 1, D), grid = [Nx, Ny, Nz])  # N = Nx*Ny*Nz
y_j = LazyTensor(y.view(1, M, D), grid = [Mx, My, Mz])  # M = Mx*My*Mz

# At this point, x_i.grid = ([Nx, Ny, Nz], None), y_i.grid = (None, [Mx, My, Mz]).
# Operations automatically propagate the grid dimensions
# (with appropriate compatibility checks), just like with
# array dimensions (N, M, etc.).
D_ij = ((x_i - y_j)**2).sum(-1)  # (N, M, 1) squared distances
K_ij = (-D_ij).exp()             # (N, M, 1) kernel matrix

# At this point, K_ij.grid = ([Nx, Ny, Nz], [Mx, My, Mz]).
# The user can now specify a maximum interaction radius:
K_ij.grid_radius = 5  # pixels
# This is fine if [Nx, Ny, Nz] = [Mx, My, Mz],
# but the user should also provide some stride and offset otherwise...

# At reduction time, we turn the "grids" and radius into a
# ".ranges" that is used if D_ij.ranges is None.
a_i = K_ij.sum(dim=1)  # (N, 1)

Of course, as discussed in #120, we should also support "broadcasting over batch dimensions" for sparsity masks whenever possible.
What do you think?

See you soon,
Jean