Data collection and storage are becoming increasingly decentralized, both due to the proliferation of smart devices and privacy concerns stemming from transferring and storing data at a centralized location. Federated Learning (FL) is a framework designed to learn the parameters W ∈ R d of a model f (W , •) on decentralized data distributed across a network of clients under the supervision of a central server (Kairouz et al., 2021;Li et al., 2020a;Yang et al., 2019). Formally, we formulate FL as the following distributed optimization problem:
where
Here, M is the number of clients, D i is the i-th client's local dataset consisting of input-label pairs {(x ij , y ij )} n j=1 where x ∈ R p is the input and y ∈ R C is the label, and n = |D i | is the dataset size. For simplicity, we consider the case where clients have equal amounts of data; our algorithm and analysis can easily be extended to the case where client objectives are unequally weighted based on their local dataset sizes. The loss function ℓ(•, •) penalizes the difference between the prediction of the model f (W , x) and the true label y. We use D = {D i } M i=1 to denote the collection of data across all clients and N = M n to denote the total number of data samples across clients.
The fundamental challenge with learning in federated settings is data heterogeneity across clients (Karimireddy et al., 2020;Li et al., 2020a). In particular, while clients can independently train a model that fits their local data, it is unclear how the server can combine these local models to get a global model that works for all the clients. To tackle this problem, standard FL algorithms such as FedAvg (McMahan et al., 2017) and FedProx (Li et al., 2020b) resort to an iterative approach requiring multiple rounds of communication between clients and the server in order for the local models to reach a consensus. However, this multi-round approach has several drawbacks. First, clients need to frequently connect with the server to send and receive updates. Second, a client needs to invest computational resources repeatedly to update the global model every time it participates in training. Third, it increases susceptibility to attacks such as data and model poisoning since attackers can continuously modify and update their attacks based on the global model they receive from the server in each round.
To overcome these drawbacks, a recent line of work has focused on the paradigm of one-shot FL which aims to learn the parameters of the global model in a single round of communication between clients and the server. Existing work for one-shot FL can be broadly split into two categories: (i) knowledge distillation methods and (ii) neuron matching methods. Knowledge distillation methods treat the collection of client models as an ensemble and propose to distill the knowledge from this ensemble into a single global model. To perform the distillation step, some works assume that the server has access to an auxiliary public dataset (Gong et al., 2021;Li et al., 2021;Lin et al., 2020), which clearly simplifies the problem. Another set of works proposes training generative models such as GANs (Zhang et al., 2022;Zhu et al., 2021) or variational autoencoders (Heinbaugh et al., 2023) to artificially generate data using local models of the clients. This again raises privacy questions, since a sufficiently powerful generative model can simply reconstruct client data (Hitaj et al., 2017). Moreover, both the data generation and distillation steps impose a significant computation cost on the server and require careful hyperparameter tuning (Kurach et al., 2019), which is itself a challenge to implement in one-shot FL.
Neuron matching methods are based on the observation that neural networks (NNs) are permutation invariant, i.e., it is possible to create NNs that differ only in the ordering of weights while having the same output (Ainsworth et al., 2023;Akash et al., 2022;Entezari et al., 2022;Jordan et al., 2023;Liu et al., 2022;Singh and Jaggi, 2020). Based on this observation, these works propose to first align the weights of the client models according to a common ordering (called matching) and then average the aligned client models. Although this idea has been shown to work well when combining simple models like feedforward NNs, the performance drops considerably for more complex models such as CNNs (Wang et al., 2020a). Another line of work, termed as model fusion has looked at fusing the capabilities of multiple existing models into a single model (Choshen et al., 2022;Ilharco et al., 2023;Jin et al., 2023;Matena and Raffel, 2022;Yadav et al., 2023). While not explicitly designed for one-shot FL, some of these techniques, such as Jin et al. (2023); Matena and Raffel (2022); Yadav et al. (2023) can be considered as an improvement over vanilla averaging. However, we note that these methods focus mostly on fusing pre-trained models and do not consider the effect of data heterogeneity used in training the models. We defer a more extensive discussion of such methods and other related work to Appendix A. Finally, we note that none of these existing works provide any theoretical guarantees for their proposed methods. Thus, motivated by the limitations of multi-round FL and current approaches for one-shot FL, we ask the following questions: Can we devise a one-shot FL method that is simultaneously communication and computation efficient (at both the server and clients), privacypreserving, and has good practical performance? Furthermore, can we provide theoretical guarantees for such a method? Our Contributions. In this work, we take a step towards providing an affirmative answer to both of the questions formulated above. To do so, we use the idea that Equation (1) can alternatively be reformulated as a posterior inference problem, specifically finding the mode of a global posterior over the model parameters (Al-Shedivat et al., 2021;Guo et al., 2023). While Al-Shedivat et al. (2021) and Guo et al. (2023) use this idea to propose a multi-round algorithm, our contribution lies in showing how this reformulation can yield a novel one-shot FL algorithm, which we term as FedFisher. Some highlights of our contribution are as follows.
• We propose FedFisher and show how the problem of finding the mode of a global posterior can be solved in a one-shot manner, using the local models at clients and some approximation of the Fisher information matrices computed at these local models (Section 2).
• We theoretically analyze FedFisher for the case of over-parameterized two-layer neural networks.
In particular, we show that when we utilize the full Fisher information in FedFisher, the error of our one-shot global model becomes vanishingly small as the width of the neural networks and the amount of local training at the clients increase (Section 3).
• We propose practical variants of FedFisher using the diagonal and K-FAC approximation, which we term as FedFisher(Diag) and FedFisher(K-FAC), respectively, and highlight the communication and compute efficiency of these variants along with their compatibility with secure aggregation (Section 4).
• We evaluate FedFisher(Diag) and FedFisher(K-FAC) on a range of one-shot FL tasks using deep neural networks and show that they give a consistent 5 -10% improvement in global model accuracy compared to competing one-shot baselines involving knowledge distillation, neuron matching, or model fusion (Section 5).
To begin our discussion, we first state the following standard assumption on the loss function ℓ(•, •), which is true for most common loss functions such as the squared loss and the cross-entropy loss.
Assumption 1. Given z = f (W , x), we assume that ℓ(z, y) is proportional to the negative log likelihood of y under some exponential-family probabilistic model, i.e., ℓ(z, y) ∝ -log P(y|z) where P(y|z) = h(y) exp(z ⊤ T (y) -A(z)) and h(y), T (y), A(z) are some real-valued functions.
Let us define the likelihood for a data point (x, y) for a given W as P((x, y)|W ) = P(x)P(y|x, W ) ∝ q(x) exp(-ℓ(f (W , x ij ), y)), where q(•) is some prior on x, independent of W . Given this definition, we can adopt a Bayesian viewpoint and try to find the maximum a posteriori probability (MAP) estimate, i.e., find W where the posterior likelihood P(W |D) ∝ P(D|W )P(W ) is maximized with P(W ) being some prior belief over W . Our motivation to do so comes from the following proposition. Proposition 1 tells us that as long as clients compute and send their local posteriors P(W |D i ) to the server, no further server-client communication is needed to find the global MAP estimate or equivalently a minimizer to our FL objective, giving us a one-shot inference procedure. However, doing so is challenging since P(W |D i ) typically does not have an analytical expression. To get a tractable solution, we propose to use some approximate inference techniques, as discussed below.
Mode of Local Posterior. To apply the approximate inference techniques detailed below, clients first need to compute Wi ≈ arg max W P(W |D i ), an estimate of the mode of their local posterior under the flat prior. Note that this corresponds to a minimizer of L i (W ) under Assumption 1 and therefore Wi can be obtained using standard gradient-based optimizers.
Laplace Approximation for Local Posterior. Now using a second order Taylor expansion around Wi , we can get the following approximation for the log-posterior at the i-th client:
where we additionally use
Here
is the Hessian of the negative log-posterior at client i computed at Wi . Thus, we see that the local posterior of a client is now parameterized by Wi and H i .
Approximating Hessian with Fisher. The Fisher information matrix F i (for brevity, we refer to it as "the Fisher" in the rest of the paper) of the local model Wi at client i is defined as follows:
computed at W = Wi . Now, under Assumption 1 and assuming Wi fits the data at client i perfectly, i.e.,
it can be shown that the Hessian at Wi corresponds exactly to the Fisher, i.e., H i = F i (Martens, 2020;Singh and Alistarh, 2020). The latter condition usually holds true for modern overparameterized deep learning models when trained for sufficient epochs. Motivated by this observation, we can further approximate H i with F i . This approximation is useful since, unlike the Hessian, the Fisher is guaranteed to be positive semi-definite, a condition required for tractable inference at the global server. However, computing (and communicating) the full Fisher would require O(d 2 ) bits, which is infeasible when the models are neural networks with d in the order of millions. In practice, clients can replace the full Fisher F i with another computationally tractable approximation Fi such as the diagonal Fisher or K-FAC (Grosse and Martens, 2016), which preserves the positive semi-definiteness of F i (see discussion in Section 4). Thus we have
Computing Mode of Global Posterior. Now assuming clients compute and send back Wi and Fi to the server, we can use Proposition 1, Equation (2) and Equation (4) to approximate the logarithm of the global posterior as
Algorithm 1 FedFisher 1: Input: initial W 0 , no. of iterations K, T , client and server step sizes η and η S respectively 2: Global server does:
Communicate W 0 to all clients; 4: Clients i ∈ [M ] in parallel do:
For k = 0, . . . , K -1 iterations:
Compute Fi ; // Approximation to true Fisher 10:
Communicate Wi and Fi to the server; 11: Global server does:
For t = 0, . . . , T -1 iterations:
With this approximation, finding a mode of the global posterior can be written as the following optimization problem:
Since each Fi is positive semi-definite, a global minimizer of Equation ( 6) can be found by simply setting the derivative of the objective to zero. Doing so, we have the following proposition.
Proposition 2. Any W that satisfies
For over-parameterized models, i.e, d ≫ N , the rank of M i=1 Fi will be smaller than d and therefore the system of equations ( M i=1 Fi )W = M i=1 F i Wi will not have a unique solution. To resolve this, we propose to use the solution that minimizes the sum of distances from the local models Wi of each client, i.e., mini-
. Such a constraint ensures fairness of the FedFisher objective towards all clients and prevents the solution from drifting too far away from the local models of each client. We express this mathematically as follows.
FedFisher objective:
We refer to the minimizer W * of Equation ( 7) as the FedFisher global model in the rest of our discussion. We now show that W * can be easily computed using gradient descent (GD) with proper initialization and learning rate conditions via the following lemma.
Lemma 1. Let W (1) , W (2) , . . . be the iterates generated by running the following gradient descent (GD) procedure:
Proofs for all lemmas and theorems in this paper can be found in Appendix B. This concludes our discussion of the proposed algorithm FedFisher. An algorithm outline for FedFisher can be found in Algorithm 1, with practical implementation considerations discussed in Section 4. Next, we analyze the performance of FedFisher for over-parameterized two-layer networks.
Over-parameterized Neural Network
In this section, we characterize the performance of the FedFisher global model W * , on the true global objective L(W ) when f (W , •) is a two-layer overparameterized neural network.
FedFisher. We see that the mismatch between FedFisher objective L(W ) in Equation (7) and L(W ) can be characterized into the following sources of error: (i) the suboptimality of Wi , i.e, ∇L i ( Wi ) ̸ = 0, (ii) the Laplace approximation, i.e., neglecting higher order terms in the Taylor expansion in Equation (2), and (iii) the error introduced by approximating the true Fisher F i with Fi in Equation ( 4). The error introduced by (iii) is orthogonal to (i) and (ii) and depends on the specific approximation used. Moreover, commonly used approximations, such as the diagonal and K-FAC, are primarily empirically motivated and have limited theoretical guarantees. Therefore, to simplify our analysis, we assume that clients send their full Fisher and focus on bounding the error due to the first two sources.
Novelty of Analysis. While the error due to (i) can be thought of as a function of the number of local training steps K, it is less clear how the error due to (ii) can be bounded. In this analysis, we show that for over-parameterized two-layer ReLU NNs, the error due to (ii) can be controlled using the width of the network. Intuitively, the quality of the approximation in Equation (2) depends on the distance between the weights of the FedFisher global model W * and clients' local models Wi ; the larger the distance the greater will be the effect of the higher-order terms and hence worse the approximation. Our contribution lies in showing that for a sufficiently wide model, this distance decreases as O(1/m) where m is the width of the model. Doing so requires a careful analysis of the trajectory of iterates generated by the optimization procedure in Lemma 1 and novel bounds on the smallest and largest eigenvalues of the aggregated Fisher M i=1 F i /M (see proofs in Appendix B). We believe that our analysis can also be extended to deeper NNs as future work.
We begin by modeling a two-layer ReLU NN as follows:
Here, m is the number of hidden nodes in the first layer, {w r } M r=1 are the weights of the first layer, {a r } m r=1 are the weights of the second layer, and σ(x) = max{x, 0} is the ReLU function. We use W = vec(w 1 , w 2 , . . . , w m ) to parameterize the weights of the first layer. Similar to Du et al. (2019), we consider a r 's to be fixed beforehand (initialized to be +1 or -1 uniformly at random) and only consider the case where w r 's are trained. We also consider the squared loss function ℓ(z, y) = 1 2 (y -z) 2 . Recall that D is the collection of data across all the clients.
Define {(x k , y k )} N k=1 = {(x 11 , y 11 ), (x 12 , y 12 ), . . . , (x M n , y M n )} to be an ordering of the data in D. We state some definitions and assumptions that will be needed in our analysis.
Assumption 2. (Data normalization) For any (x, y) ∈ D, we have ∥x∥ 2 = 1 and |y| ≤ C, where C is some positive constant. Assumption 3. (Data uniqueness) For any (x, y), (x ′ , y ′ ) ∈ D, we have ∥x -x ′ ∥ 2 ≥ ϕ, where ϕ is some strictly positive constant.
Assumption 4. (Full Fisher) Clients compute and send their full Fisher, i.e, Fi = F i .
Note on Definition and Assumptions. The rate of convergence of GD for two-layer ReLU neural networks depends closely on H ∞ and λ 0 as shown in Du et al. (2019). Assumption 2 and Assumption 3 are standard in NN optimization literature (Allen-Zhu et al., 2019a,b;Du et al., 2019;Zou and Gu, 2019) and are needed to ensure bounded loss at initialization and λ 0 > 0, respectively. Assumption 4 is specific to our setting, as discussed at the start of this section.
Theorem 1. Under Assumptions 2, 3, 4, for m ≥ poly(N, λ -1 0 , δ -1 , κ -1 ), and i.i.d Gaussian initializa-tion weights of W 0 as w 0,r ∼ N (0, κ), and initializing the second layer weights a r = {-1, 1} with probability 1/2 for all r ∈ [m], for step sizes η = O(λ 0 /N 2 ), η S = O(λ 0 /N 2 ) and for a given failure probability δ ∈ (0, 1), the following is true with probability 1 -δ over the random initialization:
ηλ0 log(N/δϵ) , then for failure probability δ ∈ (0, 1) and target error ϵ ∈ (0, 1), the global loss L(W * ) ≤ ϵ with probability 1 -δ.
Takeaways. Theorem 1 shows that for sufficiently wide networks, L(W * ) can be decomposed into two distinct sources of error. The first term is the local optimization error, which measures how well Wi fits the data at the i-th client, and the second term corresponds to the error introduced by the Laplace approximation in Equation (2). While the overall error always decreases as the width m increases, there is a trade-off in the number of local optimization steps K. In particular, a larger K reduces the local optimization error but increases the Laplace error as each local model Wi drifts further away from W * . Corollary 1 shows that by setting m and K large enough, the total training error can be made vanishingly small in just one round.
Effect of Data Heterogeneity. Note that our result does not make any explicit assumptions on the heterogeneity of data across clients such as bounded dissimilarity (Karimireddy et al., 2020;Wang et al., 2020b). We find that the notion of heterogeneity that implicitly appears in our analysis is the distance between the weights of any two local models, i.e, ∥ wi,rwj,r ∥ 2 for any i ̸ = j. For two-layer neural networks this quantity can be bounded as follows: ∥ wi,rwj,r ∥ 2 ≤ ∥ wi,rw 0,r ∥ 2 + ∥ wj,r -w0,r ∥ 2 ≤ O(1/ √ m) (see first part of Lemma 6 in Appendix B), where w 0,r are the weights of the common initialization point across clients. This bound is agnostic to the correlation of data across clients and shows that ∥ wi,rwj,r ∥ 2 always decreases as the width of the model increases. Improving this bound by incorporating the effects of data correlation can be an interesting direction for future work.
Empirical Validation. We conduct experiments on a synthetic dataset with M = 2 clients, p = 2 dimensions, and n = 100 samples to verify our theoretical findings. This setting allows us to ensure that the conditions in Theorem 1, such as computing the full Fisher, are computationally tractable. The synthetic data is generated following a similar procedure as Li et al. (2020b). Further details of the experimental setup can be found in Appendix D. In Figure 1a we fix the number of local steps K = 2048 and vary the width of the model. As predicted by our theory, we see that the training error for FedFisher monotonically decreases as we increase the width of the model. In Figure 1b we fix the width of the model to m = 512 and vary the number of local steps. Here we see a trade-off in the effect of local steps -the error initially decreases as clients do more local steps and then rises again, as predicted by our theory. Figure 1: Empirical validation of Theorem 1 in a synthetic setting. For a fixed number of local steps (Figure 1a), the error for FedFisher decreases as the width of the model increases. For a fixed width (Figure 1b), the error first decreases and then increases as local steps increases.
Generalization Bounds. Motivated by Arora et al.
(2019), we can also provide generalization bounds on the performance of the FedFisher global model by bounding ∥W * -W 0 ∥ 2 , which is the total distance between the weights of W * and the initialization point W 0 . We provide these bounds and subsequent discussion in Appendix B.
Two of the most popular approximations for the Fisher are the diagonal Fisher and the Kronecker Factored Approximate Curvature (K-FAC). The diagonal Fisher simply approximates the full Fisher with its diagonal elements and has been used in previous work for applications such as network pruning (LeCun et al., 1989) and transfer learning (Kirkpatrick et al., 2017).
The K-FAC introduced by Grosse and Martens (2016); Martens and Grosse (2015) uses two approximations: (i) it assumes that the Fisher is block diagonal where each block corresponds to the Fisher of a particular layer, and (ii) it approximates the Fisher of each block as the Kronecker product of two smaller matrices, i.e, F ≈ A ⊗ B. Thus, storing (and communicating) the K-FAC only takes O (dim(A) + dim(B)) bits instead of O (dim(A) • dim(B)) bits. Using these approximations as a substitute for the true Fisher in Equation ( 4), we get two practical variants of FedFisher, which we term as FedFisher(Diag) and FedFisher(K-FAC), respectively. We highlight the computation and communication efficiency of these variations along with their compatibility with secure aggregation below.
Computation Efficiency. To compute their diagonal Fisher or K-FAC, clients need to perform an additional forward-backward pass over the data, i.e., an additional epoch of training, plus some small overhead cost. This is a reasonable cost for FL setups, since most of the computation cost goes into computing Wi , which requires multiple local epochs. This is demonstrated in Table 1, where we see that FedFisher(Diag) and FedFisher(K-FAC) add less than 14.5% of the total computational cost of FedAvg at clients for our FL setups described in Section 5. Communication Efficiency. The communication cost of FedFisher depends on the number of parameters in the approximate Fisher. For the diagonal Fisher this is exactly d, while for the K-FAC it is model specific, but usually less than 10d. We find that the total communication cost of both FedFisher(Diag) and FedFisher(K-FAC) (including local models) can be easily reduced to match that of FedAvg using a mix of standard compression techniques such as quantization and low-rank decomposition without significantly affecting the accuracy of the final one-shot model. We discuss more details about this compression in Appendix C. For all our experiments in Section 5, we ensure that the communication cost of FedFisher(Diag) and FedFisher(K-FAC) matches that of FedAvg.
Compatibility with Secure Aggregation. In theory, the server optimization in FedFisher just needs the aggregate quantities preventing the server from accessing individual Wi 's and Fi 's and improving privacy. For K-FAC, however, while the server does not need access to individual Wi 's, it does need access to individual Fi 's. This is because there does not exist a summation operation over K-FAC matrices that preserves the Kronecker factorization property. In other words, to compute matrix vector products of the form M i=1 Fi x, the server needs to store each Fi and individually compute and sum up Fi x. However, we argue that having access to Fi can be more private than having access to local models Wi , where the latter is needed for knowledge distillation and neuron matching baselines. We discuss this in more detail, including an experiment on measuring privacy using model inversion attacks (Zhang et al., 2020), in Appendix E.
We evaluate the performance of FedFisher in comparison to state-of-the-art (SOTA) neuron matching, knowledge distillation, and model fusion baselines across a range of image recognition tasks in a FL setting. The datasets that we use are (i) Fashion-MNIST (Xiao et al., 2017), (ii) SVHN (Netzer et al., 2011), (iii) CIFAR10 (Krizhevsky et al., 2009) and (iv) CINIC10 (Darlow et al., 2018). Our code is available at https://github.com/Divyansh03/FedFisher.
Baselines. We compare FedFisher with 6 other oneshot baselines. FedAvg is the de facto baseline in all our experiments. The other baselines that we compare with are (i) PFNM (Wang et al., 2020a;Yurochkin et al., 2019), (ii) OTFusion (Singh and Jaggi, 2020), (iii) DENSE (Zhang et al., 2022), (iv) FisherMerge (Matena and Raffel, 2022) and (v) RegMean (Jin et al., 2023). PFNM and OTFusion are popular neuron matching methods that first permute the weights of local models and then average them, rather than directly averaging the weights. DENSE is a SOTA method for one-shot FL based on data-free knowledge distillation that uses GANs to artificially generate data for distillation at the server. FisherMerge and RegMean are SOTA model fusion baselines. We note that, similar to FedFisher(Diag), FisherMerge also proposes using the diagonal Fisher when merging models. However, our application of the diagonal Fisher differs in how we compute the Fisher averaged model (Lemma 1) and additional regularization (Equation ( 7)). We defer a more detailed discussion on the difference between the algorithms to Appendix A. We avoid comparing with baselines that need auxiliary data or maintain an ensemble of local models (Diao et al., 2023;Guha et al., 2019) at the server to ensure fairness of comparison. We also avoid comparing with algorithms that are inherently multi-round in nature such as FedProx (Li et al., 2020b), SCAFFOLD (Karimireddy et al., 2020) and adapative variants of FedAvg such as FedAdam (Reddi et al., 2021) since their performance would be similar to FedAvg for one round.
Models and Other Details. For FashionMNIST we use the LeNet architecture (LeCun et al., 1998); for other datasets we use a CNN model proposed in Wang et al. (2020a). To simulate data heterogeneity among the client datasets, we split our original image dataset into M partitions using a Dirichlet sampling procedure with parameter α (Hsu et al., 2019;Reddi et al., 2021), with a lower value of α implying a more heterogeneous split. The local optimization procedure is the same across all algorithms. In particular, clients perform E = 30 epochs of local training using the SGD optimizer with local learning rate η = 0.01, batch size 64 and momentum factor 0.9. To compute the Fisher diagonal and Fisher K-FAC we use the nngeometry package (George, 2021). Further details, including how we tuned hyperparameters, can be found in Appendix D.
FedFisher outperforms baselines across varying heterogeneity. As a first step, we seek to understand the impact of data heterogeneity on the performance of FedFisher and other one-shot baselines. To do so, we fix the number of clients M = 5 and vary α in the range {0.05, 0.1, 0.2} which can be considered moderate to high data heterogeneity. Table 2 summarizes the results obtained by the algorithms across the various datasets. We see that FedFisher variants, especially FedFisher(K-FAC), consistently outperform baselines and give almost 10 -20% improvement over vanilla averaging in most cases. This highlights the effectiveness of FedFisher as a one-shot algorithm that can tackle data heterogeneity in FL settings while being computation and communication efficient.
FedFisher outperforms baselines across varying number of clients. In Appendix D, we provide experimental results in which we fix heterogeneity α = 0.3 and vary the number of clients as M = {10, 20, 30}.
Our results show that FedFisher(K-FAC) continues to outperform baselines across different M with up to 10% improvement in some cases like CIFAR10.
Extension to Few-Shot Settings. A natural question to consider is if we can gain additional utility from FedFisher by extending it to the few shot setting, i.e., using a few more rounds of local training and aggregation. Figure 2 presents the results of performing 5 rounds of aggregation with different algorithms with α = 0.1 and M = 5. We omit comparison with OTFusion here because we did not find the performance to be competitive with FedAvg after the second round.
We also omit comparison with PFNM as it modifies the global model architecture, leading to increased client computation and communication in every round. We see that while the use of multiple rounds improves the performance of all algorithms, FedFisher variants continue to show the greatest improvement, especially for CINIC10 which can be considered the hardest dataset in our experiments. This highlights the additional utility offered by FedFisher in few-shot settings.
Using Pre-trained Models. As motivated by recent literature in FL (Chen et al., 2022;Nguyen et al., 2022;Tan et al., 2022), in many cases the server might have access to a model pre-trained on a large public dataset. We consider a setting where the server has a ResNet-18 (He et al., 2016) pre-trained on Tiny-ImageNet (Le and Yang, 2015) and wants to fine-tune this model. The fine-tuning datasets that we consider are CIFAR10, CINIC10, GTSRB Stallkamp et al. (2011) and CIFAR100. In all cases we split the fine-tuning dataset across M = 5 clients with α = 0.1 heterogeneity. We focus on full fine-tuning, i.e., clients update all weights in the model, for E = 30 local epochs for CIFAR100 and E = 10 local epochs for the rest. We also use a smaller step size of η = 0.001 in this setting. We omit comparison with PFNM since it does not support ResNet-like architectures. Table 3 summarizes the results of one-shot aggregation in this setting. We see that algorithms that use Fisher information, including FisherMerge, improve on FedAvg and other baselines by almost 20%, with our methods achieving the highest accuracy. We attribute this large improvement to the reduced distance between weights of the local models when starting from a pre-trained model, which in turn reduces the approximation error in Fisher averaging as discussed in Section 3.
In 2023) discusses how we can improve the prediction of this ensemble when the label distribution across client data is highly skewed. Another line of work proposes that clients send some distilled form of their data to the server (Shin et al., 2020;Zhou et al., 2020). However, the privacy guarantees of such methods is unclear. The work of Garin et al. (2022) proposes techniques to optimize the weights of given to the local models when aggregating them at the server to improve one-shot performance. However, their analysis is limited to simple linear models and does not consider combing neural networks. We also note the existence of works that propose to perform clustering in a one-shot manner in FL setting (Armacki et al., 2022;Dennis et al., 2021); these approaches our orthogonal to our problem of finding a global minimizer in a one-shot manner.
Convergence of overparameterized NNs in FL. Comparison with (Matena and Raffel, 2022). We note that, similar to FedFisher(Diag), FisherMerge (Matena and Raffel, 2022) also proposes using the diagonal Fisher when merging models. However, there are several key differences which we would like to mention. Firstly our problem is more well-defined in the sense that we have an explicit objective in Equation ( 1) that we are trying to solve with Fisher merging. This theoretical formulation is what leads us to propose the additional regularization in FedFisher in Equation ( 7) (and also FedFisher(Diag) subsequently) and guarantees for two-layer networks. Matena and Raffel (2022) on the other hand does not provide any theoretical guarantees. Also note that the procedure to compute the Fisher averaged model is different (gradient descent in our work versus diagonal Fisher inversion in Matena and Raffel ( 2022)). Finally, on a practical side, Matena and Raffel (2022) do not consider the effect of data heterogeneity in their experiments and focus on fusing two models whereas we explicitly account for data heterogeneity in our model training and aggregate up to thirty local models in our experiments (see results in Figure 3).
We have,
Also, arg max
B.2 Proof of Proposition 2.
Proposition 2. Any W which satisfies
Proof.
This implies that G(W ) is a convex function. Therefore any W which satisfies ∇G(W ) = 0 is a minimizer of G(W ). We have,
Setting ∇G(W ) = 0 we get
Lemma 1. Let W (1) , W (2) , . . . be the iterates generated by running the following gradient descent (GD) procedure:
Fi W (t) -Fi Wi with W (0) = M i=1 Wi /M and η S ≤ 1/λ max where λ max is the maximum eigenvalue of M i=1 Fi . Then, lim T →∞ W (T ) = W * .
Recall,
. Therefore, the expression for W * can be simplified as,
Since F is symmetric and PSD (see Equation ( 12)), we have by the spectral decomposition of F ,
Here V ∈ R (d×d) is an orthogonal matrix consisting of the eigenvectors of F , and Σ is a diagonal matrix consisting of all the eigenvalues of F . Σ 1 is a diagonal matrix consisting of only the positive eigenvalues. V 1 consists of the eigenvectors corresponding to the positive eigenvalues while V 2 consists of the eigenvectors corresponding to the zero eigenvalues. Also note that we have V ⊤ 1 V 2 = 0. We first observe that for any W * 1 and W * 2 such that F W * 1 = F W * 2 = b we have,
Next we observe,
The proof for this follows via simple contradiction argument as follows. Suppose
leading to a contradiction.
According to the GD step in Lemma 1 we have,
Therefore,
In the limit T → ∞ and with η S ≤ 1/λ max (Σ 1 ), we have,
Thus,
This completes the proof.
In this subsection, we provide the proof for Theorem 1 in Section 3 of our work. To do so, we first introduce some additional notation and basic results.
Recall the two-layer ReLU NN is modeled as follows,
We can write the output of the neural network alternatively as,
where
Initialization. Recall W 0 is the common initialization point for all the local models before they perform local optimization, i.e., W
Wi /M is the initialization point for the server optimization.
Definition 2. (Matrices A 0 , Ã, H 0 , H) Matrices A 0 and à are defined as follows:
and,
Furthermore we define H 0 = A 0 A ⊤ 0 , H = Ã Ã⊤ ∈ R N ×N . Definition 3. (Vector of true labels y and vector of predicted outputs ỹ)
The vector of true labels y is defined as follows,
Given local models W1 , W2 , . . . , WM , the vector of predicted outputs ỹ is defined as follows,
Definition 4. (Proxy output and vector of proxy outputs at iteration t)
Given local models W1 , W2 , . . . , WM , the proxy output f (W , x ij ) for any W ∈ R d and for any
Let W (0) , W (1) , . . . be the sequence of iterates generated by the global optimization process in FedFisher (Algorithm 1 Lines 12-14). We define f (t) as follows:
Proof.
Firstly observe that,
We have,
Claim 2. (Fisher closed-form expression) Under the squared loss ℓ(z, y) = 1 2 (y -z) 2 , the Fisher at client i simplifies to,
Proof.
Under the squared loss, we have
Thus,
where (a) uses P(y|f
to be a vector such that x(i-1)×n+1:i×n = x and zero everywhere else. Then x⊤
Also define
where H ∞ i is defined in Claim 3 and H ∞ is defined in Definition 1.
Proof. This follows from an application of Claim 4.
Before moving to the theorem proof, we first state some key lemmas and their proofs. Note that the probabilities of all events in this proof are over the random initialization of W 0 .
Lemma 2 is used to bound the local optimization at clients. Note that since the local optimization at each client is independent of all other clients, we can consider them to be M instances of centralized optimization starting from the same initialization. Thus we can use existing results in the centralized setting to give these bounds. and η = O λ0 N 2 , then with probability at least 1 -δ over the random initialization of W 0 , for all clients i ∈ [M ] simultaneously, we have
where wi,r denotes the r-th weight in the local model of the i-th client and w 0,r denotes the r-th weight of the randomly initialized model W 0 .
Firstly note that the local optimization at client i depends on H ∞ i . Here we are using the result in Claim 3 which shows that λ min (H ∞ i ) ≥ λ min (H ∞ ) = λ 0 . Also note that using Arora et al. ( 2019) we can only guarantee that the above event holds at a single client with probability 1 -δ. In order for this event to hold simultaneously with probability 1 -δ for all clients, we use the union bound and set the failure probability at each client to be δ ′ = δ/M . This leads the bound to have a dependence on N instead of n as in the single client case. Also note that substituting the bound in part (1) in part (3), we have the following upper bound on R 0 :
Lemma 3 and Lemma 4 are standard in two-layer NN optimization and are adopted as is.
Lemma 3. (Lemma 3.2 in Du et al. ( 2019)) For a given R, define the following event:
We have Pr(A ijr (R)) ≤ 2R √ 2πκ where the randomness is over the initialization of w 0,r . Lemma 4. (Lemma 3.1 in Du et al. ( 2019
The (k, l)-th entry of H ∈ R N ×N is given by,
where
We have,
where (a) uses
follows from the definition of A ijr in Lemma 3 and definition of R 0 in Lemma 2, (c) uses the result in Lemma 3.
By Markov's inequality, with probability 1 -δ, we have
Thus,
This implies,
using Lemma 4 and substituting m = Ω
in the upper bound on R 0 in ??.
We also have,
and η = O λ0 N 2 , then with probability at least 1 -δ over the random initialization of W 0 , we have
Proof.
Part (1).
We have,
where (a) uses Cauchy-Schwartz and a r I x ⊤ ij wi,r ≥ 0 x ij 2 ≤ 1.
The GD step that the central server performs for FedFisher can be written as,
We have,
Therefore,
Part (3).
From Equation (50) we have,
where in (a) we use the notation A x:y to denote the submatrix of A having row numbers x to y.
Therefore,
With probability 1 -δ over the initialization, we have
where R 0 is defined in Lemma 2 and R 1 is defined in Lemma 6.
We have,
where (a) uses union bound, (b) uses Lemma 3 and ∥w * r -w 0,r ∥ 2 ≤ w * r -w
Thus, using Markov's inequality, with probability at least 1 -δ, we have for any i ∈ [M ] and j ∈ [n],
Setting the failure probability as δ/N in Equation ( 56), we have with probability 1 -δ, for all i ∈ [M ] and j ∈
This implies,
We first state the full theorem statement with the exact dependence of m on (N, λ -1 0 , δ -1 , κ -1 ).
Theorem 1. Under Assumptions 2, 3, 4, for m = Ω N 9 λ 8 0 δ 4 κ 2 , and i.i.d Gaussian initialization weights of W 0 as w 0,r ∼ N (0, κ), and initializing the second layer weights a r = {-1, 1} with probability 1/2 for all r ∈ [m], for step sizes η = O(λ 0 /N 2 ), η S = O(λ 0 /N 2 ) and for a given failure probability δ ∈ (0, 1), the following is true with probability 1 -δ over the random initialization:
Proof.
Note that the conditions on m, η and η S in Lemma 2, Lemma 4, Lemma 5 and Lemma 6 are satisfied by setting
) and hence we can now apply these lemma results for our proof. Furthermore, we can scale down the failure probability in these lemmas by a constant factor to ensure that all the results in the lemmas hold simultaneously with high probability via union bound.
First observe that setting t → ∞ in Lemma 6 part (2), we have,
Now,
T 2 measures how well the local models fit their local data and can be bounded as
using the result from Lemma 2.
We now bound T 1 as follows:
where (a) follows from the definition of S ⊥ ij in Lemma 7. (b) uses the observation that since r ∈ S ⊥ ij we have sign
Now substituting the bounds in Equation ( 62) and Equation ( 63) in Equation ( 61), we have,
which completes the proof.
In this subsection we provide generalization guarantees for the FedFisher global model. To do so, we first introduce some additional assumptions and definitions.
Definition 5.
from D, with probability at least 1 -δ, we have λ min (H ∞ ) ≥ λ 0 > 0.
Assumption 5. Let ξ be the distribution from which the test data points are sampled. ξ is a (λ 0 , δ/3, N )-nondegenerate distribution and the data samples {(x k , y k )} N k=1 ∈ D are i.i.d samples from ξ. Remark 1. Note that in Assumption 5, we are only assuming that the collection of local data across clients is drawn i.i.d from some distribution ξ. This does not imply that the data at any particular client i is drawn from ξ, i.e., D i ∼ ξ or that the data at client i is i.i.d with the data at client j for i ̸ = j. The data in D can be partitioned arbitrarily across clients and we make no assumptions on this. For instance, consider the case with M = 2 clients, Given the test data distribution ξ, for a function f : R d → R, we have the following definitions of population loss Lξ(f ) and empirical loss L(f ):
Definition 6. (Rademacher complexity) Given N samples, x 1 , x 2 , . . . , x N , the Rademacher complexity of a function class F that maps from R d to R is defined as,
where each ϵ i is sampled independently from the Rademacher distribution unif({-1, 1}).
We now state some key lemmas that will be used in the generalization proof.
We begin with the following lemma from Arora et al. ( 2019) which bounds the Rademacher complexity of the class of two-layer neural network functions with bounded distance from initialization.
Lemma 8. (Lemma 5.4 in Arora et al. ( 2019)) Given R > 0, with probability at least 1 -δ over the random initialization of (W 0 , a), simultaneously for every B > 0, the function class
Lemma 9 is standard in generalization theory and used to bound the population loss in terms of the empirical loss and Rademacher complexity of the class of functions to which our estimator belongs. N 2 , then with probability at least 1 -δ over the random initialization, we have
where y i = vec(y i1 , y i2 , . . . , y in ) and H ∞ i is defined in Claim 3.
Note that in order to ensure that part (2) holds simultaneously for all i ∈ [M ] we set the failure probability to be δ/M in Lemma 5.3 in Arora et al. (2019).
The following lemmas are our contribution and used to bound the distance of the FedAvg model from initialization and the FedFisher model from the FedAvg model respectively.
Lemma 11. (Bounding distance of FedAvg model from initialization)
Assuming the conditions in Lemma 10 hold true we have,
(70)
Proof. We have,
Thus we have,
This completes the proof. then with probability at least 1 -δ over the random initialization, we have,
Proof.
We have,
For (a) we use the following argument. Let H = V DV ⊤ be the eigen decomposition of H. We see that
which follows from Lemma 5.
For (b) we use the following argument,
where (c) follows from Lemma 10 part (2), Lemma 11 and setting κ = O λ0δ N . Thus,
We first state the generalization bound and then provide the proof below.
Theorem 2. Fix a failure probability δ ∈ (0, 1). Under Assumptions 2, 3, 5, if we let κ = O λ0δ N , m ≥ κ -2 poly(N, λ -1 0 , δ -1 ), K = Ω 1 ηλ0 log(N/δ) then with probability at least 1 -δ over the random initialization and sampling of training data points, we have comes from bounding the distance of the FedFisher model from FedAvg, i.e., W * -W (0) 2 and can be seen as the additional error incurred in the federated setting by FedFisher. In particular, to bound W * -W (0) 2 we need to bound f (0) -ỹ 2 (see Lemma 11). Clearly, if the data across clients is similar then f (0) -ỹ 2 will be small. However, our current approach does not make any bounded heterogeneity assumptions leading to a pessimistic O √ N λ0
bound. We conjecture that this bound can be improved by explicitly incorporating bounded data heterogeneity assumptions; however we leave the investigation of this as future work.
As a first step, we bound the distance of the W * from W 0 . We have,
where the last line follows from setting κ
Next we aim to bound ∥w * r -w 0,r ∥ 2 for any r ∈ [m]. Recall we already bounded this quantity in Equation ( 63) as follows,
We now use a similar argument as the proof of theorem 5.1 in Arora et al. (2019) to provide the generalization guarantee. Under Assumption 5 we have D ∼ ξ, where ξ is a (λ 0 , δ/3, N ) non-degenerate distribution. This implies with probability 1 -δ/3 we have λ min (H ∞ ) ≥ λ 0 > 0. Assuming this event holds, we note that the following events hold with probability 1 -δ over the random initialization.
i)
which follows from Theorem 1 by setting m = Ω poly(N,
which follows from Equation (77) and Equation ( 78) respectively.
iii
where the last inequality follows from setting m ≥ poly(N, λ -1 0 , δ -1 , κ -1 ), κ = O λδ N and Equation (77). Also note that f W * ,a ∈ F W0,a R,B * k . iv) Using Lemma 9 and a union bound over all k ∈ {0, 1, 2, O(N/δ)} we have,
Note that we can scale down the failure probabilities in each of the above events to ensure that all the above events hold simultaneously with probability 1 -δ, via union bound. Now assuming all the above conditions hold simultaneously, we have, 80) and Equation ( 81))
which completes the proof. Specifically each A l and B l is first compressed to the SVD decomposition corresponding to maintaining the top l v vectors. This ensures that the communication cost of (A l , B l ) is 32(2m l-1 l v + 2m l l v + 2l v ). This SVD decomposition is then further compressed using s q quantization compression to ensure that the communication cost is (32/s q )(2m l-1 l v + 2m l l v + 2l v ). We set s q and l v such that L l=1 (32/s q )(2m l-1 l v + 2m l l v + 2l v ) ≤ 16d. The corresponding s v is then defined as ⌈m/2l v ⌉. Table 5 summarizes the results obtained by different levels of s q and s v for the FashionMNIST and CIFAR datasets respectively. We see that keeping s q = 4 and setting s v accordingly usually gives the best performance and hence we use this setting for all our experiments.
Figure 3 shows the result our experiment where we keep the heterogeneity to be fixed as α = 0.3 and vary the number of clients M = {10, 20, 30}. Note that as the total dataset size is fixed, as we increase M , each client gets assigned fewer data samples. Thus as M increases, local models have a larger tendency to overfit the data they are trained on, making it harder to aggregate such models to achieve a global model with good performance. Nonetheless we see that FedFisher variants, especially FedFisher(K-FAC), continue to outperform baselines even for large M with up to 8% improvement in some cases such as CIFAR-10.
For hyperparameter tuning we assume that the server has access to a dataset of 500 samples, sampled uniformly from the original training set. We describe the hyperparameters tuned for each of the algorithms below.
FedFisher(Diag) and FedFisher(K-FAC). While Algorithm 1 performs the server optimization with GD, this can be replaced with any other suitable GD optimizer. We choose to use the Adam optimizer here. We set η S = 0.01, β 1 = 0.9, β 2 = 0.99 and ϵ = 0.01 for the Adam optimizer and number of steps T = 2000 for all our experiments. We measure the validation performance after every 100 steps and use the model which achieves the best validation performance as the final FedFisher model.
For PFNM we use the version applicable for CNNs whose code is available at https://github.com/IBM/ FedMA. We keep the default values of the hyperparameters σ = 1, σ 0 = 1 and γ = 7 as in the original code in all our experiments.
For DENSE we use the default settings as described in Section 3.1.4 of Zhang et al. (2022) and official code implementation available at https://github.com/zj-jayzhang/DENSE. In particular we set T G = 30, λ 1 = 1, λ 2 = 0.5 and train the server model with SGD optimizer with learning rate η S = 0.01 and momentum factor 0.9. We use the validation data to determine the model which achieves the best validation performance during the server training and use this as the final model.
OTFusion. The official code for OTFusion available at https://github.com/sidak/otfusion contains more than 15 hyperparameters, making it hard to tune each of these parameters. Among these we found that the correction and past-correction hyperparameter affect performance most and hence we focus on tuning these hyperparameters for our experiments. In particular we vary correction ∈ {True, False} and past-correction ∈ {True, False} in our experiments and choose the model which achieves the best validation performance.
RegMean. For RegMean, we keep all parameters the same as the official code available at https://github.com/ bloomberg/dataless-model-merging and only tune α in the range {0.1, 0.9} using validation data as suggested by the authors.
FisherMerge. For FisherMerge, we compute the diagonal Fisher at each client using the entire dataset available at each client with a batch size of 1 and set λ i = p i where p i is the proportion of data available at client i. We also keep the fisher-floor variable to be 10 -6 as in the official code implementation available at https://github.com/mmatena/model_merging/tree/master/model_merging.
Experiments in Table 2, Table 4, Table 5 and Figure 3 were averaged over 5 seeds and run on a NVIDIA TitanX GPU. Results in Figure 2 and Table 3 were averaged over 3 seeds and run on NVIDIA A100 GPU.
Measuring Privacy using Inversion Attacks. As discussed in Section 4, for FedFisher(K-FAC), while the local models can be securely aggregated, the server does need access to individual K-FAC information at clients to perform the global optimization. We argue, however, that this is more privacy-preserving than having access to the local models which is needed for knowledge distillation and neuron matching baselines. We demonstrate this empirically via the following inversion attack on the MNIST dataset LeCun (1998).
We consider a setup where client 1 has all the images corresponding to label 3 along with 1000 randomly sampled images. The remaining data is then split among 9 clients with α = 0.1. The objective at the server is to reconstruct the data at client 1 corresponding to label 3. To do so, we adopt the attack strategy outlined in Zhang et al. ( 2020), a popular model inversion attack using GANS. We assume that the server has access to auxiliary data of size 5000 that is randomly sampled from the data belonging to clients 2 to 10. Note that this ensures that the server does not have any samples corresponding to label 3. We now consider three cases (i) server has access to the local model at client 1 (DENSE, OTFusion, PFNM) (ii) server has access to the aggregate of local models (FedAvg) and (iii) server has access to the aggregate of local model plus the K-FAC computed by client 1 (FedFisher(K-FAC)). To modify the attack to include the K-FAC information, we add an additional loss term that penalizes the difference between the K-FAC on the generated data and the given K-FAC. In particular, we use the squared loss between the elements of the generated K-FAC and given K-FAC, summed over all elements.
Table 6 summarizes the attack accuracy in the different cases. Figure 4 shows the generated images for the first and third case. We see that while adding the K-FAC information slightly improves the attack accuracy compared to case (ii), the overall accuracy is still much lower than case (i) as also evidenced by the generated images. This hints at the difficulty of inverting the K-FAC to generate meaningful client data, especially when we do not have access to the local models themselves. We acknowledge that there could be more sophisticated attacks designed especially for inverting the K-FAC; however we leave a investigation of such methods to future work.
OTFusion, PFNM, DENSE 78.4±3.0 FedAvg 5.0±1.0 FedFisher(K-FAC) 9.6±2.9