Paper: Federated Wasserstein Distance (ICLR-2024)

Reviewer comment (Theory):
Could FedWd be extended to continuous cases?

Author response:
* 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.

Reviewer comment (Theory):
Could FedWd yield any approximations for the gradient?

Author response:
* 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.

Reviewer comment (Experiments):
How does $t$ affect the performance?

Author response:
* 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.

Reviewer comment (Experiments):
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.

Author response:
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.

Reviewer comment (Evaluation):
Two main desiderata of federated learning are privacy and low communication cost. While the problem of communication cost is addressed with (10), the problem of privacy remains largely unanswered. If the authors convincingly address the problem of privacy leak, I am open to change my recommendation.

Author response:
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.

Reviewer comment (Theory):
If I correctly understand the reasoning at the bottom of page 4, the authors propose to randomly select $t$ such that the server cannot easily infer $d_{\mu,\xi^{(k)}}$. While I agree that this improves the privacy, I think that the server can, after just two communication rounds, infer $\mu$ and $\nu$. Consider the simplified setting of theorem 3. If the server knows $(m_{\xi^{(k)}}, m_{\xi^{(k)}\mu})$ and $(m_{\xi^{(k+1)}}, m_{\xi^{(k+1)}\mu})$, she knows that $\mu$ must lie on the line through $(m_{\xi^{(k)}}, m_{\xi^{(k)}\mu})$ and on the line through $(m_{\xi^{(k+1)}}, m_{\xi^{(k+1)}\mu})$. These lines must cross at $m_\mu$ by construction. (Here, $m_{\xi^{(k)}\mu}$ denotes the mean of $\xi^{(k)}_\mu$).

Author response:
> 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$.

Reviewer comment (Writing):
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. I recommend defining it somewhere.

Author response:
> 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)$.

Reviewer comment (Experiments):
In the experiments, for the Wasserstein Coreset it does not seem straightforward how your algorithm can be applied to this problem. Could you please explain to me which agents hold which distributions and where your algorithm enters the computation?

Author response:
> 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.