How to tackle the challenge of calculating the similarity between two sets of embeddings, each held by a different party, while respecting privacy restrictions?

Paper: Federated Wasserstein Distance
Authors: Alain Rakotomamonjy (Criteo AI Lab), Kimia Nadjahi (MIT), Liva Ralaivola (Criteo AI Lab).
Category: Federated Wasserstein Distance
Revue: ICLR 2024

Why did we work on this topic, the problem we want to solve?

We delved into this subject to tackle the challenge of calculating the similarity between two sets of embeddings held by two distinct parties while ensuring that the embeddings are not shared between them. For example, under certain privacy restrictions, Criteo might be interested in comparing how the user embeddings of one party are similar to those of another party, and we would be capable of doing so by addressing this issue. Resolving this issue will also enable us to enhance federated learning algorithms by being able to pinpoint dataset similarities across different parties.

What did we find? What did we achieve?

We proposed an algorithm named FedWad, which is capable of calculating the Wasserstein distance between two sets of embeddings held by two distinct parties in a federated manner. We prove the convergence properties of FedWaD, and were able to exhibit the convergence rate of the algorithm for some specific cases (when the distributions are Gaussians).

We show the applicability of their approach and its usefulness by computing federated coresets and by proposing federated KNN that is able to achieve state of the art when the number of parties is very large. We also demonstrated that we can enhance the performance of well-known federated learning (FL) algorithms using the FedWaD algorithm. Indeed, by measuring the similarity between a party’s dataset in a federated manner, we suggested clustering the parties by dataset similarity and applying FL algorithms on clusters of parties.

How did we proceed?

We used the notion of shortest (based on the Wasserstein distance) path between distributions to be able to define a proxy of the true distributions and exchange those proxies instead. The algorithm iteratively approximates the Wasserstein distance by exchanging distributions (on the geodesic) while a central entity/server orchestrates the computations. We established the convergence properties of FedWaD and provided empirical results on federated coresets and optimal transport dataset distance. The algorithm has been shown to have convergence properties and has demonstrated improved performance in federated learning algorithms.

What is the originality here?

The originality of this research lies in the development of the FedWaD algorithm, which enables the computation of the Wasserstein distance between two distributions in a federated manner without directly accessing the samples. The approach based on the notion of shortest path is novel and original.

Check all our research cards 👇
https://medium.com/criteo-engineering/research-cards/home