Federated Wasserstein Distance

Thank you for your reply,

"Could FedWd yield any approximations for the gradient?"

Yes, I refer to the approximation of the gradient of the Wasserstein distance w.r.t. to the parameters. It is still unclear to me how the gradients of the coreset are computed. I assume that the measure $\varepsilon$ is found then the gradient is taken from $W(\varepsilon$, empirical measure over coreset). Is it correct?

We thank the reviewer for their positive evaluation of our work and for finding our contribution novel.

Could FedWd be extended to continuous cases?

• In theory, FedWad can be extended to the continuous distributions since the notion of geodesics is also well-defined in this case and if a transport map exists then the interpolating measure can be written as \mu_t = ( (1-t)Id + tT)_{\\#} \mu_0 Hence for a squared L2 cost and continuous distributions, situation for which an optimal map always exists, our approach can still hold. For instance, for general multivariate Gaussians with diagonal covariance matrices, one can compute the barycenter $\mu_t$ of 2 Gaussians in closed form. We provide in the Appendix D.2 an illustration on how the measure $\xi^{(k)}$ evolves in this situation and how the Federated WD converges towards the true WD.

Could FedWd yield any approximations for the gradient?

• We do not fully understand the question. The reviewer seems to ask if FedWad can be used to approximate the gradient of the Wasserstein distance w.r.t. to the parameters of $\mu$ or $\nu$, and if so, whether the approximation would be accurate enough. If so, then the answer is affirmative and the quality of the approximation depends on the error on the computed transport plan. This is the case for any numerical scheme that computes the Wasserstein distance since any scheme would approximate only the true optimal plan (we use the envelope theorem for computing the gradient). We note that in our experiments on coresets, the coresets learned by gradient descent on FedWad are visually similar to those obtained with exact WD (Figure 4), and they are useful for yielding a federated Knn that outperforms existing FL algorithms. While these empirical results are encouraging, a theoretical analysis might be needed to rigorously characterize the approximation error on the gradient; this is out of scope of our work.

For some ground costs that are not metric e.g., in domain adaptation application, the corresponding Wasserstein cost is not metric hence there is no triangle inequality. In this case, could we have any solutions?

• FedWad strongly relies on the fact that the true Wasserstein distance satisfies the triangle inequality. Hence, it is not straightforward to adapt our algorithm to divergences which do not verify the triangle inequality.

How does $t$ affect the performance?

• the choice of $t$ does not affect the performance of the algorithm neither in terms of running time nor in terms of error. At each iteration of the algorithm, we look for any element on the geodesic. Hence, any value of $t$ can be chosen. However a deterministic choice of $t$ might facilitate privacy leaks.

It is still unclear to me how the gradients of the coreset are computed. I assume that the measure $\varepsilon$ is found then the gradient is taken from $W(\varepsilon, empirical measure over coreset)$. Is it correct?

Let say, we want to minimize $W(f_\theta, \mu)$ with respect to $\theta$, $f_\theta$ being on server. In a FedWad notation, this mean that we want to solve $min_\theta~~W(f_\theta, \xi^K) + W(\mu, \xi^K)$ Once we have applied FedWad, we have $W(f_\theta, \xi^K) = \langle P^\star, M_{f,{\xi^K}} \rangle$ where $P^\star$ is the optimal plan and $M_{f,\xi^K}$ is the matrix of pairwise distances. Based on the envelope theorem the gradient $\nabla_\theta W$ is $\langle P^\star, \nabla_\theta M_{f,\xi^K} \rangle$ where the gradient is computed componentwise and this gradient is used for updating $f_\theta$ (in pytorch, the gradient is computed by autograding torch.sum(P.detach()*M))

We thank the reviewer for their positive rating and for finding the problem well motivated, our theoretical development sound and our experiments convincing.

Do we need the two distributions to be discrete and the samples of the two distributions to be the same? Theorem 1 seems to be less general than Theorem 2. If so, the motivation of Theorem 1 should be emphasized and highlighted.

• In theory, the two distributions do not need to be discrete, since the notion of geodesics is well-defined for continuous distributions. However, we need a closed-form expression of the transport map or a closed-form of the Wasserstein barycenter in order to be able to compute any element of a geodesic.

• Furthermore, the number of samples in each distribution can be different. There is no restriction on this and we have run experiments in which the datasets are imbalanced (one measure as $n$ samples and the other $3n$, see Figure 2)

• Theorem 1 is indeed less general than Theorem 2. Its goal is to formalize the soundness of the approximate interpolating measure given in Equation 10 compared to the exact formulation in Equation 5. We stated that under certain situations, Equation 10 provides the same result as Equation 5.

I think an experiment where the datasets are unbalanced would be super helpful; say what if a dataset is always a certain times of the other dataset.

We have run such an experiment on a toy problem and reported the results in Figure 2. We have shown there that the imbalance does not impact much the running time nor the approximation error of FedWad either using the exact interpolating measure or the approximated ones. In practice, using the approximated interpolating measure leads to better approximation and faster computation.

An experiment that compares the Gaussian & non-Gaussian assumptions to show the convergence speed could be better.

• We have included an example of convergence for Gaussians and non-Gaussian distributions. It is reported in the Appendix D.3 and Figure 11. We can note there that for non-Gaussians, the approximation error is reduced by an order of magnitude (vs $1.5$ order of magnitude for Gaussians) for every $2$ early iterations and the rate tends to decrease for the higher iterations.

We thank the reviewer for finding our work and theoretical developments valuable for the community.

The main criticism of the reviewer is related to privacy issue. While we acknolewdge that having privacy guarantees is an important desiderata in many applications, federated learning does not provide strong guarantee about this: we refer to our answer to all reviewers, as well as our answer to reviewer FgAh, for more explanations.

I think that the server can, after just two communication rounds, infer $\mu$ and $\nu$.

The attack proposed by the reviewer is based on finding the intersection between the lines joining $\xi^{(k)}$, $\xi_\mu^{(k)}$ and the one joining $\xi^{(k+1)}$, $\xi_\mu^{(k+1)}$.

• In the simplified setting described by the reviewer, we note that our approach needs only 1 communication round while the attack needs 2 rounds. Hence, arguably a simple defense for the client is to allow only 1 query of interpolating measure. In addition, we can also emphasize that the revealed information is an aggregate information on the distribution from which it is difficult to infer individual samples.

• In a more general setting (beyond Gaussians), the strategy proposed by the reviewer is supposed to reveal the location of the $\mu$ in the Wasserstein space. However, we believe that the key notion of lines $\xi^{(k)}$, $\xi_\mu^{(k)}$ needs a proper formal definition. Indeed, the definition of the geodesic between $\xi^{(k)}$ and $\xi_\mu^{(k)}$ is known but it is not clear for us how one should define the extension of this segment beyond $\xi_\mu^{(k)}$ towards $\mu$. Then, assuming that these lines are properly defined, and one can estimate their intersection (in the Wasserstein space) -- which we believe is an hard problem --, what we reveal is a global information (the distribution) but not individual samples.

• The approximated interpolating measure approach (given in Equation 10 and used in practice) not only facilitates favorable computational complexity and communication costs, but also encourages privacy. Indeed, as stated in Equation 10, the support of the interpolating measure depends on the data also through a matrix multiplication with the optimal plan. This plan is not revealed to the server, thus making the inference of $\mu$ not straighforward. In addition, when using the approximated interpolating measure, there exists situations in which several distributions $\mu$ lead to the same $\xi_\mu$, increasing the difficulty of inferring the true $\mu$.

In Algorithm 1, you use $d_{\mu,\xi^{(k)}}$ . While I can infer from the context what it probably means, this symbol has not been used before.

We thank the reviewer for spotting this: $d_{\mu,\xi^{(k)}}$ is the Wasserstein distance between $\mu$ and $\xi_{(k)}$, and we revised our manuscript to clarify this. We use the same notation as in the text $\mathcal{W}_p(\xi^{(k)},\nu)$.

In the experiments, for the Wasserstein Coreset it does not seem straight forward how your algorithm can be applied to this problem. Could you please explain it to me which agents hold which distributions and where your algorithm enters the computation?

We added the following discussion in our paper to clarify the setting of this experiment.

• Finding a coreset of $\mu$ aims at finding the support of an empirical distribution that minimizes the problem $\min_{\{x_i^\prime\}} W(\mu, \frac{1}{n}\sum_n \delta_{x_i^\prime})$ We propose to solve this problem in a federated learning manner, by computing the Wasserstein distance with FedWad. Hence, in practice, at each iteration $j$ of the coreset optimization algorithm, we update the coreset support based on the distance $d_{\xi^{(K)}, \frac{1}{n}\sum_n \delta_{x_i^j}}$, where $x_i^ĵ$ is the current estimation of the support.
• For the federated coreset classification model experiment, we have computed the coresets for each empirical distribution stored on clients. For instance, in Figure 5, we have 1000 clients with about 30 samples per client, and for each client, we learn a Wasserstein coreset based on 2 samples. This leads to 2000 learnt coresets stored on the server.

We thank the reviewer for acknowledging the importance of the problem and the novelty of our contribution. We are grateful for their very insightful questions.

Under an FL setting, is it in line to essentially reveal raw data to the server?

While we do not have formal guarantees on the robustness of FedWad to privacy leaks, FedWad is in line with the perspective on Federated Learning, as described in Kairouz et al., 2019: "federated learning (...) provides a level of privacy to participating users through data minimization: the raw user data never leaves the device, and only updates to models (e.g., gradient updates) are sent to the central server. (...) there is still no formal guarantee of privacy in this baseline federated learning model. For instance, it is possible to construct scenarios in which information about the raw data is leaked".

As we argue below, we believe it is not straightforward to infer the distribution $\mu$ from the sole information that navigates through the servers, especially when using approximated interpolating measures with random values of $t$.

The server sends a distribution $\xi$ to client with data $\mu$, then a $t$-barycenter $\xi_\mu$ is sent back to the server. By structure of WD geodesic, with knowledge of $\xi$ and $\xi_\mu$ it is already very immediate to reconstruct $\mu$.

The main question and weakness raised by the reviewer is based on the assumption that there exists a simple attack that reveals the raw data. However, we propose below a rebuttal to this attack and would be happy to further discuss this aspect with the reviewer.

• The attack proposed by the reviewer assumes the knowledge of some important information: $t$, the Wasserstein distances and the number of samples in $\mu$. As acknowledged by the reviewer themselves, none of these elements are supposed to be revealed to the server during the iterations.
• Another insight that we learn from the reviewer's question is that our approximated interpolating measure approach (given in Equation 10) not only helps improve computational complexity and communication costs, but also facilitate robustness to privacy leaks. Indeed, as stated in Equation 10, the support of the interpolating measure depends on the data through a matrix multiplication with the optimal plan. This plan is not revealed to the server, thus making the inference of $\mu$ not straighforward. In addition, when using the approximated interpolating measure, there exists situations in which several distributions $\mu$ lead to the same $\xi_\mu$, increasing the difficulty of inferring the true $\mu$.

Hence, while we do not have a formal result guaranteeing that FedWad does not reveal any information, we do not think that the situation described by the reviewer can lead to disclosure of the raw data. Furthermore, in practice, we use approximated interpolating measure which employs a important piece of information -- the transport plan -- that is not revealed by the client.

If a formal guarantee of privacy is needed, then we can resort to differential privacy and consider a differentially private Wasserstein distance instead.

[1] Kairouz, P., McMahan, H. B., Avent, B., Bellet, A., Bennis, M., Bhagoji, A. N., ... & Zhao, S. (2021). Advances and open problems in federated learning. Foundations and Trends® in Machine Learning, 14(1–2), 1-210.

We would like to thank all of the reviewers for their time and thoughtful efforts in providing detailed and constructive feedbacks. We have carefully considered their comments and have prepared a rebuttal addressing each one. We are open to further discussions and are more than willing to respond to any additional comments or queries they may have.

Thanks to reviewer xDsE for our discussions and for acknowledging the relevance of our work.

We would like to take this opportunity to provide another connection with some existing federated learning algorithms and our approaches in term of data leak. If we assume a FL that looks at solving a regression problem with quadratic loss in which clients share local gradients then, clients send to server

• $X^\top(y - Xw)$, where $X$ and $y$ are local data

In our case, clients send to server the support of the approximated interpolating measures which is

• $(1-t)X + t n PY$, where $X$ and $Y$ are the support of $\xi$ and $\mu$, and $P$ the transport plan (which is not revealed)

Hence, while we do not have a formal proof of data privacy, we can assume that our approach is at least as compliant with federated learning as some existing FL algorithms.