Federated Generalised Variational Inference: A Robust Probabilistic Federated Learning Framework

We thank the reviewer for their appreciation of the clarity of our paper and the theoretical results, as well as their important suggestion for an ablation study that we have now conducted and which will significantly improve the empirical study of this work. We will include these in the revised version of the main paper.

Hyperparameters of FedGVI

In particular, we have carried out an ablation study on the selection of the hyperparameters of FedGVI in the training of a Bayesian Neural Network on 10% contaminated MNIST data, split across 5 clients, extending the results discussed in Section 5.5. We vary the $\delta$ parameter of the Generalised Cross Entropy, and the $\alpha$ parameter of the Alpha-Rényi divergence in the client optimisation step (Eq. 4), and record the predictive accuracies on the uncontaminated test data for the FedGVI posteriors found with the respective hyperparameters:

where we note that $\delta=0$ implies the negative log likelihood since the loss converges to it as $\delta$ tends down to zero (Zhang & Sabuncu, 2018). For the Alpha-Rényi divergence, $\alpha=1$ implies the Kullback-Leibler divergence, and $\alpha=0$ implies the reverse Kullback-Leibler divergence, i.e. $D_{AR}^{(0)}(q:q^{\backslash m})=D_{RKL}(q:q^{\backslash m}):=D_{KL}(q^{\backslash m}:q)$ (Amari, 2016). This means that for $\alpha=1.0$ and $\delta=0.0$, we recover PVI.

We propose to add these results on hyperparameter selection in FedGVI as an annotated heat map to Section 5.5 in the revised version of the paper; the figure can be viewed here: https://anonymous.4open.science/r/Resources-3CEF/ablation_study.png.

In the figure, we present the maximum result achieved across all server iterations and plot the percentage errors on uncontaminated test data for the different hyperparameter combinations. As evident from the table most of the combinations of $\alpha$ and $\delta$ explored show stability of the posterior regarding classification accuracy with different hyperparameters, however some care should be placed in selecting these since the wrong choice could significantly degrade model performance, e.g. $\alpha=5.0$ and $\delta=0.2$, where we do not sufficiently filter out outliers but place increased weight on the cavity distribution through the high alpha. Nevertheless, the majority of settings outperform PVI, especially for $\delta=0.6$.

Learning Rate Selection

Furthermore, we also present additional results on learning rate selection for the ADAM optimiser (Kingma & Ba, 2015) in the 3 Client FedGVI and PVI regimes presented in Tab. 1 of the paper. So far, we have fixed the learning rate to be 5e-4 in the BNN experiments of the paper, which we now vary while keeping the divergence and loss parameters fixed for the respective method. We highlight the best result for each learning rate in bold.

Here, FedGVI with a robust loss outperforms PVI in every scenario.

We also want to point out that by not carefully selecting the hyperparameters of FedGVI, as well as the learning rate, and keeping these constant across the BNN experiments, we have shown that you don't require extensive knowledge to adapt existing PVI approaches to FedGVI and outperform. For instance, FedGVI performs even better for the robust losses with a higher learning rate, but we have shown in Tab. 1 in the paper that it outperforms even when not carefully selecting a learning rate. Furthermore, choosing $\delta=0.6$ and $\alpha=2.5$ would have performed better when varying only the robustness parameters of FedGVI.

References not in paper

Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In 3rd International Conference on Learning Representations, 2015.

Thank you for appreciating our theoretical results and especially for taking the time to review our proofs.

Claims & Evidence

the examples provided are [...] specifically chosen to highlight FEDGVI's strengths over PVI

FedGVI with the negative log likelihood and KL divergence is equivalent to PVI, i.e. FedGVI(NLL,KLD)=PVI. In the idealised setting where the DGP $p_0=p_\theta$ is known, PVI with a correctly specified likelihood is preferable to FedGVI(RobustLoss, KLD), i.e. using a robust loss will degrade performance; see Knoblauch et al. (2022), Jewson et al. (2018), Bissiri et al. (2016). Arguably, we rarely have accurate likelihood functions, since these are either too complex to work with or practitioners opt for standard likelihoods without integrating domain expertise.

We will amend Sec. 3.2 and add:

It is important to highlight that GVI and FedGVI may underperform when using robust losses in the case of correct likelihood specification.

Relation to Literature

When viewing GVI/GBI as an optimisation problem on the space of probability distributions $\mathcal{P}(\Theta)$, Bayesian inference, VI, hierarchical Bayes/VI, all target a single element of this space. These methods either target the standard Bayesian posterior explicitly, or the posterior within some variational family with closest Kullback-Leibler distance to the Bayesian one. Through GBI and GVI we are able to target different elements of a subspace of $\mathcal{P}(\Theta)$, then simply a single point; in that regard, these approaches do 'generalise' Bayes. In the paper, 'generalised' is inherited from GVI and GBI. We should note that in the FedGVI setting, GBI and GVI allow us to generalise PVI or FedAvg to a broader subspace of possible posteriors. We will clarify this naming distinction in the main paper.

We have made a figure to highlight this, see https://anonymous.4open.science/r/Resources-3CEF/FedGVI.png, which we will include in the paper.

Question 2, Missed References, Hierarchical Bayes

Following your suggestion we will add and discuss the proposed references in relation to FedGVI and GVI, giving an additional discussion on personalised and hierarchical Bayesian FL in Sec 1.1. Furthermore, we propose to add the following to Sec. 6:

An interesting future direction is to extend FedGVI within personalised FL settings (Kotelevskii et al., 2020) and hierarchical Bayesian FL through latent variables (Kim & Hospedales, 2023), as well as the use of a structured posterior approximation (Hassan et al., 2023; 2024), in order to incorporate client level variations. Incorporating the hierarchical model structures and additional inductive biases from such settings, while maintaining conjugacy and favourable computational complexity, remain as open challenges.

Weakness 1, Comments, Experimental Design

We have now carried out further empirical and ablation studies, https://anonymous.4open.science/r/Resources-3CEF, see also our responses to reviewers 4k22 and gKSi, which we will include in the revised paper. Nevertheless, as we have shown even the 'simplistic' MNIST setting already provides challenges to traditional Bayesian FL making it worth studying.

We also agree with the reviewer that extending FedGVI to more complex data sets and scenarios, such as investigating different aspects of FL for instance privacy constraints, the cross-device setting, and massively distributed scenarios, present intriguing future directions for FedGVI. We thank the reviewer for pointing this out and will be included in Sec. 6.

For 'computational constraints' see our next answer.

Weakness 2, Question 1

unique methodological contributions that the use of [GBI] could specifically enable in [FL]

To the best of our knowledge, we have proposed the first theoretically grounded probabilistic FL framework that deals with model misspecification. Our framework brings the following methodological advances:

FedGVI can be more computationally efficient than PVI through the use of GBI, as we have shown in Prop. 4.9. The conjugacy with exponential family likelihoods through the robust score matching loss enables faster computation than PVI since it does not require sampling at the local clients. See also our response on computational complexity to Reviewer gKSi.

Through the guaranteed posterior robustness by Theorem 4.11, our prediction accuracy is less affected by outliers which would be detrimental for instance in medical settings; see e.g. Jonker et al. (2024) for FL in medicine.

The server divergence and the potential for robust aggregation at the servers (Sec. 6) without fundamentally changing our approach allows for unique optimisation on a global level not provided by Bayesian FL approaches.

We thank the reviewer for encouraging us to flesh these points out and we will include these in the revised paper.

We thank the reviewer for their in-depth, constructive, and positive review.

Question 1, Weakness 1, and Experiments

...method’s sensitivity to the choice of robust hyperparameters [...] how stable...

How sensitive is FEDGVI to [...] hyperparameters [...]? Would small changes degrade performance significantly?

We presented results on FashionMNIST while varying $\delta$ with different contaminations in Fig. 2. We now extend these to further contamination levels (0, 0.1, 0.2, 0.4) and $\delta$ (0, 0.4, 0.5, 0.8, 1), as well as FedAvg, FedPA, and $\beta$-PredBayes, to demonstrate hyperparameter sensitivity:

https://anonymous.4open.science/r/Resources-3CEF/FashionMNIST.png

We also now offer an ablation study on MNIST for hyperparameter selection: https://anonymous.4open.science/r/Resources-3CEF/ablation_study.png (see also response to reviewer 4k22).

To show the stability of FedGVI under small perturbations of $\delta$, we fix $\alpha=2.5$, vary $\delta$ around $0.8$, and report accuracies on uncontaminated test data after training on 10% contaminated MNIST data split across 5 clients:

https://anonymous.4open.science/r/Resources-3CEF/Perturbations.png

Question 2, Weakness 2 Computational Complexity/Overhead

W2

Communication cost or runtime overhead is not directly benchmarked...

Q2

do [...] loss functions [...] incur any significant computational overhead [...] updates?

As the number of clients increases, the comp. time at each client is reduced due to less data per client. By Proposition 4.9 we can choose a robust loss at each client that enables conjugate client updates whereas the negative log likelihood might not have closed form. These include intractable likelihood models (Matsubara et al., 2022) where the normalising constant is intractable but we can tractably solve.

The comp. complexity of the divergence is solely dependent on the choice of variational family, prior, and divergence without dependence on number of clients or data per client. The divergence may be more comp. expensive in the case where no closed form solutions exist, e.g. for the Total Variation Distance. However, for exponential family distributions, many robust divergences with closed form solutions exist (Pardo Llorente, 2006; Knoblauch et al., 2022).

The KL and A-R divergences between Gaussians (Sec. 5.5) are available in closed form, scaling only in the number of parameters. The A-R can be computed in $\mathcal{O}(1)$ times the comp. complexity of the KL, driven by the determinant of the covariances; assuming $\Theta\subset\mathbb{R}^d$, and using a Cholesky decomposition this naively takes $\mathcal{O}(d^3)$ for both; assume these are (block)-diagonal, makes this cheaper. The discrete losses in the classification setting are non-differentiable. We cannot apply the conjugate score matching loss, and require sampling from the approximation. We can transform the NLL in $\mathcal{O}(1)$ to the GCE, thereby, as the likelihood is used for the GCE and it's log for the NLL, these are of the same order of magnitude.

We will add in Section 5.5:

FedGVI incurs no additional computational complexity when compared to PVI. This is due to the KL and Alpha-Rényi divergences having closed form solutions between Multivariate Gaussians with complexity of $\mathcal{O}(1)$ in each other, and because we require $\mathcal{O}(1)$ additional, constant operations to get the GCE from the NLL.

To empirically validate this, we now report average wall-clock times at each client for FedGVI and PVI, see https://anonymous.4open.science/r/Resources-3CEF/runtime.png.

Comments

Heuristics for selecting robustness parameters

These include Knoblauch et al. (2018) and Yonekura & Sugasawa (2023) for the Density Power loss, with the latter a theoretically principled Sequential Monte Carlo sampler for selecting $\beta$, or Altamirano et al. (2024) for weighted score matching, through cross validation.

We propose to add this in Sections 3.2 and 3.3:

3.2: We can use a Sequential Monte Carlo sampler to estimate the $\beta$ or $\gamma$ hyperparameters in $\mathcal{L}_B$ and $\mathcal{L}_G$ (Yonekura & Sugasawa, 2023) or use cross validation to select optimal parameters (Altamirano et al., 2024).

3.3: Similarly to the losses, we can perform cross validation to select the $\alpha$ parameter, however as demonstrated in the ablation study (Figure TBD) FedGVI performs favourably under a range of $\alpha$ values.

Robust Aggregator strategies at the server

This is an area we are actively working on. We are exploring ways in which the summation in Eq. 6 can be replaced by a robust aggregator, such as Nearest Neighbour Mixing (Allouah et al., 2024) or indeed coordinate-wise median and trimmed mean, to achieve Byzantine robustness. We will discuss this in Sec. 6.

References

Yonekura, S., Sugasawa, S. Adaptation of the tuning parameter in general Bayesian inference with robust divergence. Stat Comput 33, 39, 2023.

Thank you for your constructive and positive review that highlights the clarity of our writing, the theoretical results on FedGVI, and its provable robustness. By addressing below the weaknesses you mentioned, we notably improve the clarity of the paper and its ease of understanding for the reader.

Weakness 1 $\beta$ parameter in Eq. 2 undefined:

Thank you for spotting this typo, it has now been corrected, and we will add the following clarification:

Here, $\beta\in\mathbb{R}_{>0}$ is a learning rate parameter that determines how much weight we place on the observed data, similar to power posteriors in VI (Grünwald, 2012; Kallioinen et al., 2024).

The $\beta$ parameter comes from the power/cold/tempered posteriors of e.g. Grünwald (2012), where the likelihood in Bayesian posteriors is raised to some power of $\beta>0$. This was originally done to add some robustness to the posterior, down-weighting observations if $\beta< 1$ and up weighting these for $\beta > 1$. Through a known result (Knoblauch et al., 2022), which we highlight in Lemma B.1 in the Appendix, this is equivalent to having a weighted Kullback-Leibler divergence, $\frac1\beta\mathrm{KL}$. This also allows us to define if we want to trust the prior more $\beta<1$ or less $\beta>1$, since up weighting the data means down weighting the prior and vice versa.

We will add this discussion to Appendix B.1 to give further intuition behind GBI.

Weakness 2 Intuition behind Definition 4.10 and Conditions 2 and 3:

The three conditions combined allow us to say whether the client posterior (or simply the posterior in a global, 1 Client, GBI setting) derived from such a robust loss is provably robust to Huber contamination.

From Condition 1 we are able to bound an infinitesimal change in the loss with the contaminating data point $z$ by some auxiliary function $\gamma$, possibly infinite for some values of $\theta$.

Condition 2 states that the product function, $\gamma(\theta)\pi(\theta)$ has finite uniform norm. This ensures that this product under the worst case contamination and the worst parameter $\theta$, is finite and hence it cannot be made arbitrarily bad, which does not hold for the negative log likelihood in general. Alternatively, the prior decays to zero faster than the auxiliary function can diverge to infinity in $\theta$.

Condition 3 further says that $\gamma(\theta)\pi(\theta)$ is finitely integrable, i.e. that this is in $L^1(\Theta)$. This, in effect, bounds the normalising constant of the contaminated posterior and will ensure that this is finite.

Taking all these conditions together tells us that the product function $\pi(\theta)\gamma(\theta)$ is in $L^1(\Theta)$ and that it is finite everywhere; two conditions that are mutually independent, i.e. one doesn't follow form the other.

These conditions characterise the notion of robustness we use for Theorem 4.11, a derivation of this notion is shown in Appendix B.7, by considering the worst choice for the contamination $z$ and the parameter $\theta$ with respect to small perturbations of the resulting posterior through $\epsilon$. The influence of the contamination $z$ and parameter $\theta$ on the posterior is defined as $\frac d{d\epsilon}q_m^{(t)}(\theta;\mathbb{P}_{n_m,\epsilon, z})|_0$ (evaluated at $\epsilon=0$), which is bounded through the conditions. This then implies that the local posterior is globally bias robust, i.e. robust to Huber contamination of all $z$ for all parameters $\theta\in\Theta$.

Thank you for pointing this out, we will provide the following clarification between Definition 4.10 and Theorem 4.11.

These conditions ensure that the influence of arbitrary contamination on the local posterior is not arbitrarily bad. In particular the auxiliary function $\gamma_m^{(t)}$ ensures that the influence of an adversarial data point $z$ on the posterior over infinitesimal contaminations, $\frac d{d\epsilon} q_m^{(t)}(\theta;\mathbb{P}_{n_m,\epsilon, z}) {\Big|}_0$ evaluated at $\epsilon=0$, are finite over all $\theta$ and $z$. Condition 2 ensures the loss increases slowly enough for the local posterior to concentrate around the data, and condition 3 ensures the resulting posterior will be normalisable.

References not in Paper

Kallioinen, N., Paananen, T., Bürkner, P.-C., and Vehtari, A. Detecting and diagnosing prior and likelihood sensitivity with power-scaling. Statistics and Computing, 34(1):57, 2024.