Connecting Federated ADMM to Bayes

We thank the reviewer for their comments and review. We are glad that the reviewer appreciates the connection between ADMM and Bayes, and appreciates the numerical and methodological parts of the new methods. We respond to the reviewer’s questions individually below. We have updated our paper on openreview, with changes highlighted in blue.

Q1: “FedLap-Func only provides marginal improvements over FedLap-Cov and sometimes even makes it worse.” We expect FedLap-Func to improve upon FedLap, not also improve upon FedLap-Cov. This is because FedLap-Func is an orthogonal way to improve FedLap compared to FedLap-Cov: FedLap-Func adds function-space information only, while FedLap-Cov adds covariance information. It is possible to combine the two (use both covariance information and function-space information in a single algorithm), but we do not test this in the experiments. We did not explore problems outside the benchmarks listed, because already FedLap-Func is outperforming FedLap and other existing baselines.

Q2: “Include a paper outline at the end of Section 1.” Thanks, we have done so in the new version of the paper.

Q3: “Could you bring the definition of $\hat{t}^\*_k$ from (5) to (4)”. Note that $\hat{t}^*_k$ used in Eq 4 are different than $\hat{t}_k$ used in Eq 5, therefore it is not possible to move them or use them interchangeably. Eq 5 gives a way to estimate the optimal sites in Eq 4, that is, as the PVI algorithm progresses, $\hat{t}_k$ converges to $\hat{t}^*_k$.

Q4: “FedLap seemingly indicates the Laplace variant of the PVI. But it is still Gaussian used as the variational distribution. Why is it called "Lap"?” Good question. FedLap is indeed a Laplace variant obtained through the variational PVI algorithm. Essentially, by taking a delta approximation in the PVI, we recover the Laplace variant. We also show how FedLap’s fixed points coincide with the Laplace solution of the global problem (old version lines 201-204, 274-275 and 341-342; new version lines 210-212, 291-292, 355-356).

We thank the reviewer for their comments and review. We are glad that the reviewer appreciates the connection between ADMM and Bayes, and finds the empirical results overall to be convincing. We respond to the reviewer’s questions individually below. We have updated our paper on openreview, with changes highlighted in blue.

Q1: “What is the complexity of the extended VB approaches against the baselines, in particular when full covariance estimates are considered.” We have added analysis as requested in the new submitted version, please see the response to all reviewers for the lines changed. To summarise, FedDyn and FedLap have the same computation and communication cost, while Fedlap-Cov and FedLap-Func increase the cost (but only slightly) due to the use of additional variables.

Q2: “FedLap-Func is used in the setting where some inputs are shared globally. Can you discuss the practical implication of that?” Sharing inputs between local clients and the central server is common in Federated Distillation (FD) work (see eg Li and Wang (2019)). The focus of our work is to show theoretical links between the derived methods and FD algorithms generally. FD works normally randomly split an existing benchmark dataset (eg CIFAR-10) into a publicly-available part (that all clients and the server has access to) and a private part; in contrast, we assume that all parties have access to a random 1 input per class per client (which is significantly fewer datapoints than FD settings). FedLap-Func still performs well in our low-shared-data setting because, unlike FD algorithms, it also sends/uses weights. We discuss the limitations of this in lines 422-424 (new version lines 437-439): it introduces privacy issues, but as we use very few datapoints, it might be reasonable if we obtain prior permission from the owner(s) of the data. We hope that future work can build upon these insights further and show that FedLap-Func is competitive in full FD settings.

Q3: “the empirical comparison needs more statistical rigorousness.” Currently, we simply bold the top two (mean) values, and do not bold if multiple (mean) values are the same (we have made this clearer in the new paper version’s Table 1 caption). We can alternatively bold methods that are within one standard deviation of the top-performing method. We note that previous papers usually do not run multiple random seeds at all (eg FedAvg, FedProx, FedDyn papers).

We thank the reviewer for their comments and review. We respond to the reviewer’s questions individually below. We have updated our paper on openreview, with changes highlighted in blue.

Q1: “Whether the federated locals can be used to recover the global solution is not justified both theoretically and numerically.” We do have a theoretical justification in lines 201-204, 274-275 and 341-342 (new version lines 210-212, 291-292, 355-356) where we show that the stationary points of FedLap (and its variants) are the same as the global solution. In addition, in Appendix D2 we show empirically how FedLap correctly targets its global loss. We also show this numerically, for example, our experiments section (Section 4) shows how FedLap performs at least as well as FedDyn, and FedLap-Cov and FedLap-Func are even quicker to converge.

Q2: “Why do you want to connect ADMM to Bayesian? Bayesian learning is known to be the go-to method for conducting uncertainty quantification and non-convex optimization. The algorithm is only evaluated in the optimization perspective.” … “the connections to Bayes are not fully evaluated empirically. I recommend the authors to see how the uncertainty experiments are evaluated in [1,2,3].” We disagree with the reviewer. Our focus on the paper is to draw methodological links between ADMM and Bayes (two historically-different fields), and show how this can be used to improve the optimisation itself. This is an important contribution. Also note that the extension of ADMM to preconditioned updates is not straightforward and often requires versions that use Bregman divergence. Our connection here shows that it can also be achieved by using Bayes where using covariances automatically leads to preconditioning by uncertainty estimates. This is a new result and has implications for learning with heterogenous data. For example, it suggests that the learning rates of clients should be modified according to their uncertainty.

Q3: “Missing important literatures in Bayes federated learning”. We thank the reviewer for these references, and have added them to our paper (new version lines 113-115). All three papers [1,2,3] use MCMC algorithms in a federated learning setting, and so are orthogonal to the Variational Bayes approach we use in this paper.

Q7: “I didn't understand very well what is indicated in L309.” We add and subtract two terms that are equal at the point $w_k$: the true loss over inputs in $M_k$, and the Gaussian contribution over those points (as derived in Khan et al. (2019)). The detailed derivation is in Appendix C1: we show in Eq 23 how terms cancel so that we do not need access to true labels, but only to soft labels. In the new version of the paper, we removed some of the detail in Sec 3.4 to avoid confusion (the details explaining the derivation require more explanation, as is in App C1). Please let us know if this did not make sense!

Q8: “L342: M_k = null space, Isn’t this obvious?” Yes, when $M_k$ is the empty set, then FedLap-Func recovers FedLap. Existing Federated Distillation works do not have this property: they are stand-alone algorithms. Please let us know if you meant something different here.

Q9: “For Algorithm in Eq. 17. How many training steps per round and updates?” We have now added how many training steps per round (and learning rate) for the global server optimisation in App E, thank you for pointing this out. In summary, we use 5000 updates at the server (500 for CIFAR), but did not hyperparameter sweep over this value: we stress that this will usually not be a problem because in most federated learning settings, the global server has significantly more compute resources than the local clients; also, the number of input points at the server much smaller than the data at local clients (as discussed in the computation complexity section, new version lines 364-372).

We thank the reviewer for their comments and review. We are glad that the reviewer found the paper well-written, and found the link between ADMM and PVI to be elegant and novel. We respond to the reviewer’s questions individually below. We have updated our paper on openreview, with changes highlighted in blue.

Q1: “Computational cost … I am just missing the analysis and the clarity.” We have added analysis as requested in the new submitted version, please see the response to all reviewers for the lines changed. To summarise, FedDyn and FedLap have the same computation and communication cost, while Fedlap-Cov and FedLap-Func increase the cost (but only slightly) due to the use of additional variables.

Q2: “I do not really know if sharing some soft labels or inputs between local clients and the central server is worth it.. It kind of modifies a bit the spirit of federated learning and introduces extra issues like privacy, etc, that are ignored imho.” Sharing inputs between local clients and the central server is common in Federated Distillation (FD) work (see eg Li and Wang (2019)). The focus of our work is to show theoretical links between the derived methods and FD algorithms generally. FD works normally randomly split an existing benchmark dataset (eg CIFAR-10) into a publicly-available part (that all clients and the server has access to) and a private part; in contrast, we assume that all parties have access to a random 1 input per class per client (which is significantly fewer datapoints than FD settings). FedLap-Func still performs well in our low-shared-data setting because, unlike FD algorithms, it also sends/uses weights. We discuss the limitations of this in lines 422-424 (new version lines 437-439): as you say, it introduces privacy issues, but as we use very few datapoints, it might be reasonable if we obtain prior permission from the owner(s) of the data. We hope that future work can build upon these insights further and show that FedLap-Func is competitive in more regular FD settings.

Q3: “scalability … I feel that all of this between Eq. 10 and Eq. 14 is ignored, or at least not being highlighted for the reader (the dimensionality problem, or an indication of what are we considering here).” … “What is the dimension of the second auxiliary/dual variable? I'm always afraid of the use of preconditioning.. could the authors add a bit more details in that regard”. In general, the second dual variable $V_k$ has dimension PxP, where P is the number of parameters of the model. However, as indicated in line 224 (new version line 238), we only consider diagonal approximations in our experiments, which we already show performs very well in FedLap-Cov (and as is common for both Laplace and variational neural networks), although other approximations are possible (eg block-diagonal K-FAC). In the diagonal case, $V_k$ has dimension P (same as the dimension of $w$ and of the first dual variable $v_k$). As requested, we have added/highlighted this point in lines 273-274 of the new paper version.

Q4: “Experiments are fine, but Table 1 is somehow limited. I basically missed some larger datasets, where there is not that much similarity or structure in the data as in MNIST or FMNIST … and augmenting a bit the number of clients.” We understand the point. We use similar experiments to other works in this area (especially cross-silo works as opposed to cross-device). Other reviewers (8nbj, E4hs, kAKo) also found the experiments to be sufficient, but we agree that larger experiments may support the paper's argument even more.

Q5: “I'm curious to know where the term ‘sites’ in L112 is coming from” ‘Sites’ are used in Bayesian message passing works to describe ‘locations’ where there is relevant information, and messages are passed to/from these sites. For example, this is seen in Belief Propagation and Expectation Propagation (Minka, 2001), where the $t_k$ are referred to as the ‘sites’.

Q6: “How does this delta approximation work in L185? How does this approximation affect the general problem?” As described in line 193 (new version line 201), this approximates an expectation of a function by the value at the mean: $\mathbb{E}_q[g(w)] \approx g(m)$. In our setting, this approximation allows us to recover (a local) Laplace approximation as a special case of (the more global) Variational-Bayes; see Opper and Archambeau (2009) for the local-global connection.

We thank the reviewer for their comments and review. We are glad that the reviewer found the paper well-written and intuitive to understand, and that there are extensive benchmarks in the paper. We respond to the reviewer’s questions individually below. We have updated our paper on openreview, with changes highlighted in blue.

Q1: “limitations do not seem to be clearly discussed”. A limitation is that the computation cost increases slightly when using more expressive posterior approximations, however this is accompanied by improvements, therefore may not be considered a direct limitation rather a feature of the Bayesian framework.

Q2: “computational complexity”. We have added analysis as requested in the new submitted version, please see the response to all reviewers for the lines changed. To summarise, FedDyn and FedLap have the same computation and communication cost, while Fedlap-Cov and FedLap-Func increase the cost (but only slightly) due to the use of additional variables.

Q3: “final model is a point estimate” … “the weights used in to make predictions were just the MAP estimate?” Yes, we use a point estimate as the final model in all experiments. We already see the benefit of adding a Bayesian viewpoint in the improved convergence using FedLap-Cov and FedLap-Func.

Q4: “Equation 13 … the first $m_k$ may be a typo” … “Equation 13 uses gradients from a point in space that is not the centre of the Taylor approximation, is this principled or just a necessary hotfix?” Yes, the first $m_k$ is a typo, it should be $m$. Thank you for pointing this out (we have fixed it in the new version). The derivation in Eq 13 avoids higher-order derivatives during optimisation. For instance, another way to do this would be to use a Taylor series approximation around $q(w) = N(w; m, S^{-1})$ as usual, but assume that $H_k(m)$ does not depend on $m$. We have updated the text around Eq 13 to make this clear in the new version.

Q5: “Equation 16: can the authors provide some context as to why the new weight space term is subtracted?” The intuition is that the contribution of the memory points is already added in the function-space term (first term in Eq 16), and so to avoid double-counting the contribution from the memory points, we need to subtract their contribution, which we do in weight-space (second-term in Eq 16). We have added a sentence after Eq 16 in the new version to add this intuition.

Typos. We thank the reviewer for pointing these out! We have (i) fixed the missing log in the equation after Eq 9, (ii) changed Eq 12’s parentheses to be square brackets, and (iii) changed $m_k$ to $m$ in Eq 13.

Thank you for the response and thank you for adding clear computational cost descriptions.

In the revision, on line 295, it is stated that the cost of FedLapCov with covariance would be $2P + P^2/2$, I wander if it should it be $P + P^2/ 2$? Just a mean vector and half a cov matrix?

After a read of the other reviews, I will keep my score as is.