Question: Element-wise matrix addition

[kheyer]

Thanks for creating this library!

I'm trying to use keops to implement a formula that involves element-wise addition between two 2D matrices. This is the basic format:

# x = (N,F), y=(M,F)

x_i = LazyTensor(x[:, None, :])
y_j = LazyTensor(x[None, :, :])

C_ij = ((x_i - y_j)**2).sum(dim=2)**0.5 # (N,M)

# Z_ij = (N,M)
O_ij = C_ij + Z_ij # (N,M)

Z_ij is a pre-defined torch tensor. When I try to convert Z_ij to a LazyTensor (ie Z_ij = LazyTensor(Z_torch, 0)), the axis specification causes the addition to return a tensor of size (N,M,M) rather than a (N,M) tensor.

Is element-wise addition between two 2D matrices possible with LazyTensors?

[jeanfeydy]

Hi @kheyer ,

Thanks for you interest in the library :-)
Could you maybe give me some more context on the computation that you are trying to implement, e.g. where the matrix Z_ij comes from and what you intend to do with O_ij?

The main feature of KeOps LazyTensors is that they are lightweight: just like sparse matrices, they can be used to implement some operations very efficiently, without every allocating a large (N,M) buffer in GPU memory. As a consequence, "mixing" them with dense torch Tensors using pointwise operations to create dense (and heavy) (N,M) arrays is not really advised.

In this context, there can be two answers to your question:

• If you really need to define a dense (N,M) array O_ij in your program, as a "displayable output", then LazyTensors probably won't be of much use to you: there is no shortcut to take, and you'd better stick to regular torch Tensors to allocate and manage these large 2D arrays.

• If the array O_ij is "reduced" into a (N,) or (M,) vector or used to perform a matrix-vector or matrix-matrix product further down the line (which is the case in the vast majority of applications), then it is likely that we can refactor your code to leverage the speed of KeOps routines. For instance, write down the "expensive" computation O_ij.sum(dim=1) into a cheap reduction C_ij.sum(dim=1) plus an other term Z_ij.sum(dim=1).

With some more information on your problem, I should be able to help you to see whether or not KeOps is relevant to your work. What do you think?

Best regards,
Jean

[kheyer]

Thanks for the reply.

I'm implementing the algorithm from Geometric Dataset Distances via Optimal Transport which incorporates label-to-label distances in the OT calculation.

Here's the gist of it:

You have two feature vectors x = (N,F), y=(M,F) which have corresponding labels x_c = (N,), y_c = (N,). We first compute the OT costs from all classes in x to all classes in y. So for classes c1, c2, we have cost_c1_c2 = OT(x[x_c==c1], y[y_c==c2]).

We iterate over all classes to build up a matrix of class distances class_costs = (C1, C2).

From here we create the matrix Z_ij where each element in Z_ij indexes the class_costs matrix based on the classes of 'i' and 'j' respectively.

Then we augment our cost matrix C_ij to get the final transport matrix O_ij = (C_ij**2 + Z_ij**2)**0.5 and we solve the OT problem over O_ij instead of over C_ij.

You can see the implementation here. I started out building around the POT library, but I want to get a Pytorch implementation working. I've implemented the algorithm using the tensorized Sinkhorn cost from the Geomloss library here, and I started to work on implementing an online/keops version when I ran into the matrix addition issue.

[jeanfeydy]

Hi @kheyer ,

I see, thanks a lot for letting me know about these methods!

To summarize, I understand that you are trying to design a hierarchical cost matrix O_ij = |x_i - y_j|**2 + Z[class(x_i), class(y_j)], where Z[i,j] is a matrix of cluster-cluster distances; O_ij should then be "fed" to an optimal transport solver to compute a class-aware transport distance between two distributions.
To keep notations consistent with GeomLoss, I will assume that x_i is an (N,D) array, y_j is an (M,D) array and that Z is an (Nc,Mc) array - i.e. you are working with Nc << N and Mc << M classes for each distribution.

I see four main ways of doing this (ranked from the simplest to the hardest to implement):

1. Augment the points x_i and y_j with additional coordinates to encode the cluster-cluster distances. For instance, you could start with two point clouds xc of shape (Nc,D), yc of shape (Mc,D) initialized at the cluster barycenters, and optimize their positions by gradient descent to ensure that the matrix of squared distances between xc and yc is close enough to Z. If the fit is not good enough, you could increase the size of the embedding space until you reach a mean squared error of ~1%, for instance. Then, you could augment each point x_i with the coordinates of xc[class(i)], do the same for the y_j's, and solve a regular OT problem.
   This is a pragmatic approach, that will allow you to leverage fast and well-maintained OT solvers without having to spend time implementing them or keeping them up to date. It's a bit heavy though (moving around all these additional coordinates between CUDA registers takes time), and not super precise. I would implement this first as a good baseline - but you may have done so already :-)

2. Augment the points x_i and y_j with coordinates that directly come from the matrix Z[i,j]. More precisely, for a reduction over the j index, the lightest approach would be to add a line variable Z_i = LazyTensor(Z[label_i,None,:]) to every point x_i, an index L_j = LazyTensor(label_j[None,:,None]) encoded as a one-hot vector to every variable y_j and work with the symbolic LazyTensor O_ij = C_ij + ( Z_i | L_j ). This approach can be a bit heavy if you have many clusters though.

3. Do the same thing as above, but use dynamic indexing instead of dot products with one-hot vectors to select the appropriate coefficients of the lines Z_i and columns Z_j. E.g. work with a formula O_ij = C_ij + Z_i[label_j] when performing a reduction over j, and O_ij = C_ij + Z_j[label_i] for a reduction i - to minimize memory transfers, it's best to keep the "heavy" variables Z_i and Z_j fixed in the reduction loop. This would be both more efficient than 2. and more precise than 1.... But unfortunately, this is not possible with the current KeOps engine: we still have to add support for "integer" variables and "dynamic" indexing. This is definitely on our road-map for early 2021 though, so you can expect this to work around May-June :-)

4. Implement your own OT solver from scratch, with custom Sinkhorn updates that handle C_ij and Z_ij separately. This is very much do-able: I would refer to my PhD thesis (Chapter 3) and our NeurIPS 2020 paper (esp. the supplementary materials, Section E) for tips on efficient implementations... But honestly, this may be a little bit overkill. Relying on "standard" solvers is bound to be easier in the long run, especially with respect to maintenance issues.

What do you think?

Of course, don't forget to profile your routines to identify bottlenecks. If you have any question on this, feel free to ask! I now have several years of experience on fast OT solvers, and would be happy to help you to get the most out of these libraries: your project is exactly the type of work for which KeOps and GeomLoss have been designed.

We're still mid-development at the moment (KeOps is quite stable now but requires a lot of maintenance work to keep up with all the latest hardware improvements such as Tensor cores, and GeomLoss is still very much in its alpha-beta stage), and interacting with advanced users like you is really key to identify priority areas.

Best regards, and good luck,
Jean

[kheyer]

Thanks for the detailed reply!

I was able to implement both methods 1 and 2 in a way that worked with the GeomLoss routines.

One thing I'm trying to figure out though is how to get the transport plan as a LazyTensor. For the tensorized version, I've used the formula from here:

P_ij = ((F_i + G_j - C_ij) / eps).exp() * (a_i * b_j)

I've tried implementing this with LazyTensors and I keep getting a 3D tensor instead of a 2D output. Can this formula be computed with some general reduction, or do general reductions always need to return a vector?

[jeanfeydy]

Hi @kheyer ,

Perfect!
There should be no problem using this formula, if your arrays are reshaped correctly as in:

F_i = LazyTensor(F.view(N,1,1))
G_j = LazyTensor(G.view(1,M,1))
a_i = LazyTensor(a.view(N,1,1))
b_j = LazyTensor(b.view(1,M,1))

# ...
# C_ij is computed using e.g. some points x_i, y_j
# ...

# The line below defines a (N, M, 1) symbolic LazyTensor:
P_ij = ((F_i + G_j - C_ij) / eps).exp() * (a_i * b_j)

As detailed in the LazyTensor intro, LazyTensors never completely collapse their last "vector" dimension: your matrix P_ij will represent a symbolic (N,M) array, handled as a (N,M,1) symbolic tensor internally - but this really shouldn't make much difference. You can then use this transport plan P_ij any way you like, e.g. to transfer some labels or move points around: most of the useful quantities in OT theory can be computed as matrix-vector or matrix-matrix products P_ij @ b or P_ij.t() @ b, which are fully supported.

Does this answer your question? If you have any other question, feel free to ask!
I have been very busy these past few days, but am fairly available now :-)

Best regards,
Jean

[kheyer]

Thanks, that answers things. I was getting an extra axis in there somehow.

[jeanfeydy]

I'm closing the issue as the problem appears to be solved: feel free to re-open if needed!