Vectorized computation of squared distances

[ssydasheng]

Hi,

Thanks for your awesome work ! I am starting to use keops to make my program more memory-efficient, while I cannot get my code to work. So I wonder if anyone can help me with it.

Assume x1 ~ [n, d], x2 ~ [m, T, d] where n, m are big and T, d are small. I also have matrices V ~ [m, T] and Z ~ [T, T]. I want to compute the n dimensional vector: (SqDist(x1, x2) @ Z).mul(V).sum(1, 2). Though this formula seems easy to implement, I still cannot get it to work.

An example of my current failing attempt is as the follows,

import torch
from pykeops.torch import LazyTensor

N, M, T, d = 200, 100, 2, 3

x1 = torch.randn([N, d])
x2 = torch.randn([M, T, d])
v = torch.randn([M, T])
z = torch.randn([T, T])

lazy_x1 = LazyTensor(x1, axis=0)
lazy_x2 = LazyTensor(x2)
lazy_v = LazyTensor(v, axis=0)

dist = lazy_x1.sqdist(lazy_x2)   # [n, m, T]
dist_mat = dist.keops_tensordot(z, (T), (T, T))   # [n, m, T] from matrix multiplication
dist_mat_mul = dist_mat * lazy_v   # [n, m, T]
res = dist_mat_mul.sum(1).sum(1)   # [n]

However, even initializing x2 causes an error about the 3D tensor. Thanks you ahead of time for any advice.

[jeanfeydy]

Hi @ssydasheng ,

Thanks for your kind words :-) First of all: your question allowed me to notice that the full documentation of the LazyTensor API had disappeared from our website! This was due to a re-factoring of the code that we made for the v1.4.2, two weeks ago: we created a new abstract "GenericLazyTensor" class common to both NumPy and PyTorch LazyTensors, but forgot to add a link for it in our documentation template... All my apologies, this was really a stupid mistake: we always check the benchmarks when re-rendering the doc, but don't always pay attention to these detailed doc pages - which are actually much more important than a few more tables. I have fixed it and re-rendered the doc this afternoon, so this is fixed now. If you haven't seen it already, don't forget to check the API documentation for a full list and documentation of the supported operations.

Now coming back to your problem: as far as I can tell, you have run into one of the current limitations of our math engine. As of today, KeOps formulas can only be written using tensors of shape (B1, ..., BK, I, J, D) where B1, ...., BK are batch dimensions, I and J are "symbolic" or "virtual" axes and D is a single vector dimension.

Long term, we plan to add support for more general shapes and support more than two "symbolic" axes and a single "vector" dimension... But other features have higher priority, and we probably won't have time to work on this before the summer. Fortunately though, we have implemented a "minimal" workaround to support basic linear algebra with the `matvecmult' routine and other similar methods.

Using the fact that |x-y|^2 = |x|^2 - 2 * <x,y> + |y|^2, we can use these operations to implement your formula as follows:

import torch
from pykeops.torch import LazyTensor

N, M, T, D = 200, 100, 2, 3

# Data arrays:
x1 = torch.randn([N, D])
x2 = torch.randn([M, T, D])
v = torch.randn([M, T])
z = torch.randn([T, T])

# Pre-compute the squared norms of the points x1, x2:
n1 = (x1 ** 2).sum(-1)
n2 = (x2 ** 2).sum(-1)

# Encode as KeOps LazyTensors:
x1_i = LazyTensor(x1.view(N, 1, D))
n1_i = LazyTensor(n1.view(N, 1, 1))
x2_j = LazyTensor(x2.view(1, M, T * D))
n2_j = LazyTensor(n2.view(1, M, T))
v_j = LazyTensor(v.view(1, M, T))
z_ = LazyTensor(z.view(1, 1, T * T))  # You may want to transpose z here

# Actual computations:
D_ij = n1_i - 2 * x2_j.matvecmult(x1_i) + n2_j   # [N, M, T]
Dz_ij = z_.matvecmult(D_ij)  # [..., T, T] @ [..., T] -> [N, M, T] from matrix multiplication
Dzv_ij = Dz_ij * v_j   # [N, M, T]
res = Dzv_ij.sum(2).sum(1)   # [N, 1]

What do you think?

Best regards, Jean

[ssydasheng]

Hi @jeanfeydy ,

Thank you so much for the detailed reply. Your code works perfectly, which is exactly what I expected. Also, by trying to make sense of it, I think I understand how to use keops much better. It is REALLY helpful !

Thanks for putting the API back to the website. Having it is much more convenient for me, since I was wondering why there was no API introductions previously.

[jeanfeydy]

Hi @ssydasheng ,

Great to hear about it! If you end up using KeOps in your project, please feel free to let us know about it once you've made it publicly available: we're always happy to learn about interesting, unexpected use of these tools :-)

By the way, speaking about the documentation: I'm working on it right now, but we will probably have to wait until early January to have all of the new deep learning, clustering and non-Euclidean tutorials available online. Until then, you may be interested by the supplementary materials of our NeurIPS paper, if you haven't seem them already. We've put quite a few code snippets, that explain how to implement e.g. Multi-Layer Perceptrons as KeOps formulas.

Best regards, and good luck for your work, Jean

[ssydasheng]

@jeanfeydy

Thanks for the notice. I will let you know if there is anything interesting I make with KeOps ! The supplementary materials of your paper is a really good reference, thanks!