Gradient Masked Averaging for Federated Learning
Authors: Irene Tenison
Sai Aravind Sreeramadas
Vaikkunth Mugunthan
Edouard Oyallon
Irina Rish
Eugene Belilovsky
Venue: Accepted by TMLR ()

Abstract: Federated learning (FL) is an emerging paradigm that permits a large number of clients with heterogeneous data to coordinate learning of a unified global model without the need to share data amongst each other. A major challenge in federated learning is the heterogeneity of data across client, which can degrade the performance of standard FL algorithms. Standard FL algorithms involve averaging of model parameters or gradient updates to approximate the global model at the server. However, we argue that in heterogeneous settings, averaging can result in information loss and lead to poor generalization due to the bias induced by dominant client gradients. We hypothesize that to generalize better across non-i.i.d datasets, the algorithms should focus on learning the invariant mechanism that is constant while ignoring spurious mechanisms that differ across clients. Inspired from recent works in Out-of-Distribution generalization, we
propose a gradient masked averaging approach for FL as an alternative to the standard averaging of client updates. This aggregation technique for client updates can be adapted as a drop-in replacement in most existing federated algorithms.  We perform extensive experiments on multiple FL algorithms with in-distribution, real-world, feature-skewed out-of-distribution, and quantity imbalanced datasets and show that it provides consistent improvements, particularly in the case of heterogeneous clients.

Published in Transactions on Machine Learning Research (10/2023)
Gradient Masked Averaging for Federated Learning
Irene Venison*ly, Sai Aravind Sreeramadas Vaikkunth Mugunthan', Edouard Oyallon', Irina Rish2 and Eugene Belilovsky
Massachussetts Institute of Technology, MA, USA, 2 Universite de Montreal, Quebee, Canada, Mila, Quebec Al Institute, Canada, *Sorbonne Universily, Paris, France, $ Concordia Universily, Quebec, Canada itenison @mit.edu
Reviewed on OpenRevicw: https: // openreview. net/ forum? id= REAyrhRYAo
Abstract
Federated learning (FL) is an emerging paradigm that permits a large number clients with heterogencous data to coordinate learning of a unified global model without the need to share data amongst each other. A major challenge in federated learning is the heterogeneity of data across client, which can degrade the performance of standard FT. algorithms. Standard FL algorithms involve averaging of model parameters or gradient updates to approximate the global model at the server. However, we argue that in heterogeneous settings, averaging can result in information loss and lead to poor generalization due to the bias induced by dominant client gradients. We hypothesize that to generalize better across non-i.id datasets, the algorithms should focus on learning the invariant mechanism that is constant while ignoring spurious mechanisms that differ across clients. Inspired from recent works in Out- of-Distribution generalization, we propose gradient masked averaging approach for FI. as an alternative to the standard averaging of client updates. This aggregation technique for client updates can be adapted as a drop-in replacement in most existing federated algorithms. We perform extensive experiments on multiple FL algorithms with in-distribution, real- world, fcature-skewed out-of-distribution, and quantity imbalanced datascts and show that provides consistent improvements, particularly in the case of heterogeneous clients.
1 Introduction
Federated Learning (FL) is a distributed machine learning approach that allows clients with decentralized data to efficiently learn a shared global model without having share their sensitive datasets McMahan et al. (017): : Kairouz et al.] [2021]. This enhances privacy data is neither collected at a central location cloud nor communicated over any channel. Furthermore, Parcollet et al. 2021 and Qiu al. (2021) argue that federated learning has lower carbon footprint than traditional machine learning. A challenge in FL is heterogeneity in the data distributed across clients. The non-i.i.d data distribution degrades the performance of federated learning models Li et al. (2020a]: Wang et al. (2019): ; Zhao et al. (2018). One of the reasons for this is the loss of information on invariances between clients induced by the averaging model parameters or updates. This is further exacerbated by the multiple local steps taken by each client with the aim of reducing communication rounds, which results in 'client drift Karimireddy et al. (2021). Each client after multiple local steps can progress too far towards minimizing their local objective, which may deviate from that of the global objective.
Many existing FL works attempt to tackle this problem through the lens of optimization, proposing constrained gradient optimization based approaches Karimireddy et al. ) Wang et al. (2019: Li et al. (2019) which attempt to maintain a solution close the global optimum across the federated data. Another branch of work has considered approaches based on knowledge distillation Zhu ct al. (2021) address this problem.
*This work was done when Trene Tenison was Universite de Montreal and Mila, Quebec AT Institute.

---

Published in Transactions on Machine Learning Research (10/2023)
My-1 1/ 1.0.
Deo0a
(00000
GOODS
Figure 1: Illustration of Gradient Masked Averaging (GMA). Each client receives a global model and performs multiple iterations to obtain a final update sent to the server. The server uses the sign of each update component to determine an elementwise agreement score. This then used adjust the overall global model update, emphasizing components where client's 's agree on their direction.
On the other hand can consider the challenge of local training and evaluation on different clients with varying distributions as distribution shift problem. One can thus consider leveraging many of the recent. advances in tackling distribution shifts made in areas such as OOD generalization Arjovsky et al. 2020))
OOD generalization often formalized using the notion of domains or environments. Under the formalism of
Arjovsky et al. (2000) an environment corresponds to a data generating distribution that can be related through underlying (potentially unknown) causal variables to set of other environments. Different environments can arise during model training and (esting, while it is typically assumed that all environments (train and test) share some invariant mechanisms. However, they can have spurious mechanisms that differ across environments Parascandolo et al. (2020): Bühlmann (2018). The concept of environments can be related to the FL setting involving decentralized data by considering each client as producing a set of data generated from a different environment. Clients have invariant mechanisms across them that can lead to robust global models, as well as spurious mechanisms associated to their local data. Leveraging the existing strategies in OOD generalization can thus potentially lead more robust global model.
Recently Parascandolo et al. (2020) proposed simple approach to improve generalization in the Out-of- Distribution (OOD) setting. Their approach studied in the centralized setting modifies gradient based optimization by applying masking operation to individual sample gradients. Identifying this is an approach which can create more robust. predictors and tackle the inherent distribution shift problem in FL we propose a new aggregation method called gradient masked averaging (GMA) with the goal of improving generalization across clients and of the global model. GMA operates on client updates from different FL models, determining their component-wee wise agreement, and subsequently applying soft-masking operation to perform the final update. It is illustrated in Figure
GMA can be plugged into any FL algorithm an alternative to naive averaging of model parameters at the server. Intuitively, gradient masking prioritizes gradient components that are aligned with the overall dominant dircction across clients while the inconsistent components of the gradient are given lower importance when making the update. Applying this approach leads to improved performance in a variety of training and evaluation settings across a number of standard datasets used for federated learning. Our performance gains are particularly large in cases where the clients are non-iid. Code for our experiments is include in the supplementary materials and will be made available at the time of publication.
Our Contributions in this work are as follows:
We introduce gradient masked averaging as a drop-in robust alternative to naive averaging of parameters.
We empirically show that applying gradient masking to any FL algorithm consistently improves performance of the algorithm. This improvement was observed on in-distribution evaluations and its various settings, quantity skewed datasets, real world evaluations, and more complex out-of- distribution evaluation datasets.
2

---

Published in Transactions on Machine Learning Research (10/2023)
2 Related Work
Federated Learning In FedAVG McMahan et al. (2017) for each communication round, all randomly selected B fraction of clients (called participating clients) perform E local steps of gradient descent with their local datasets. The model parameters from participating clients are averaged at the server to obtain the global model. It is equivalent to FedSGD McMahan et al. (0177) when E' = and each client performs stochastic gradient descent. Multiple local steps help minimize communication costs, which is major bottleneck in FL. Quantization methods Reisizadeh et al. (020)) and gradient descent acceleration Yuan & Ma 2021 methods have been proposed to reduce communication overhead. Convergence of FedAVG under i.i.d scttings have been analyzed widely (ttchh (2019): Yu al. (2018); : Wang & Joshi (2019). The convergence rate of FedAVG worsens with increasing heterogeneity among client datasets and this has been analyzed by several works Li et al. (2020a]; Wang al. (2019). Multiple variations of FedAVG have been proposed to improve convergence in non-i.i.d data distribution settings, including adding regularization to the client objective Li et al. (020a),, normalized averaging of model parameters Wang et al. (2020b), introducing server momentum Hsu et al. (2019) or adaptive optimizers at. the server Reddi et al. (002)), and introducing control variates Karimiroddy et al. (2021). Generalization in federated learning has been explored in terms
of out-of-sample and participation gaps in Yuan et al. (2021). Francis et al. (2021) explores adaptations of IRM to exploit invariance in federated learning, Gupta et al. (2022) explores the same from a game theoretic
perspective, and Zhhang et al. (2021) employs federated adversarial training for the same. PFNM Yurochkin
ct al.
2019) and MrdMAWangg et al. (2020a) addresses the problems due to permutation variances in the neural networks. Algorithms like PerFedAVG Fallah et. al. (2020). Dilto Li et al. (2021b), FedBABU Oh et al.
(2021)
FedSHIBU
Raj et al. (2022) focus on client personalization.
Out of Distribution Genoralization In traditional machine learning, model is evaluated based on its test performance an unseen dataset drawn iid from the train data distribution. However, this assumption may not be true in real-world datasets and many supervised learning models do not perform well on related but non-i.id test datasets. This problem is often referred to as the out-of-distribution generalization or the closely related domain generalization problem Koyama & Yamaguchi (2021): Ahuja et al. (2020) This
problem has been addressed in several works like Invariant Risk Minimization (TRM) Arjovsky et al. (2020); Risk Extrapolation approaches typically focus introducing penalties that learn invariant reprosentations in sctting with known variations in the data (corresponding to environments). However, this idea cannot be easily ported to
Krueger et al.] (2021), and Gradient Starvation Pezeshki et al. (2021). These
FL. setting as the clients performing the optimization steps would require access to the data of other clients. On the other hand Paarascandooo et al. (2020) proposed a gradient agreement method based on gradient directions to learn features that agree across environments. This was extended by Shahtalebi et al. (021 to include gradient magnitude. In this paper we focus on this class of gradient agreement. methods. Distinct. from the prior work, which considers the case of individual samples and single global updates, we adapt this approach to federated setting, where each client produces an aggregate update based on multiple gradient. iterations.
3 Methods
In the following scction we introduce the notations used, review standard federated aggregation, and then introduce gradient masked averaging for global model approximation.
3.1 Federated Aggregation
Fix D a training set that we assume splitted across N clients in {D") ) local datasets of size sn. The objective of the clients to collectively learn function via supervision, f Rd - that, in our case, corresponds to a neural network model with parameters, w E Rd. For this, we define the global objective function in terms of the local objective, Fa, of each client shown below, assuming that EgnaDe [fn [f](w; =*)] F.(w), where fn
3

---

Published in Transactions on Machine Learning Research (10/2023)
is the clients objective function. In our setting, we are thus interested in minimizing:
min s(w) = min Σ Fw(w),
When using the local data, with cach client having &n samples, we focus on minimizing the empirical objoctive function, given by:
min
Σ fn(w;)")
8n
For this, we consider a standard first order minimization scheme. At each communication round, I, and for a total of T' iterations, the parameters of the global model, Wt, are sent to each of the N participating clients which perform multiple local gradient steps to obtain local update, A7, corresponding to the difference between the clients model after multiple updates and Wr. In most FL algorithms the global model update at tth communication round is then obtained as
w++1 = Wt - - where A
ΣA. IN
(1)
REN
Fg is the global learning rate and A, is the update or 'pseudo-gradient" at the 1 global communication round obtained by aggregating the updates from the participating clients (Sg) In this context, for the local learning rate client m. we obtain a local update during K iterations, via local parameters t and batches 2 given by:
A 711 EESN ST k=1
Σ (n
(2)
In the case of single gradient step at each client, the update A, corresponds to the gradient of the global objective. Each client has different data distribution and thus different loss surface. Paaracandolo et al. (2020) shows that averaging of gradients across environments leads to poor consistency of solutions, and reduced generalization, particularly to unseen environments. Indeed naive averaging of parameters fails to capture the consistencies in the loss landscapes due to the bias that may be induced by dominant features in the environments as explained by Shahtalebi et al. (202I). This is further exacerbated in real world federated settings as there are multiple possible scenarios where some clients dominate over others.
3.2 Gradient Masked Aggregation
Parascandolo el al. 20200) highlights several issues associated with the standard arithmetic mean, which they equate to a logical OR. They propose to use an analog of the logical AND operation resulting in taking geometric mcan of sample gradients. The gradient components that are "inconsistent' in sign across environments are set to 0. Specifically they construct binary vector, m, based on the agreement of gradient components among environments. The components j of the mask, m, is computed as [m-]] = 1ll |V Σee, sign((VL.]j) 2 T||. Here Le is the loss gradient for environment e, and E [0, 1] is hyperparameter.
Direct application of this idea to the FL setting is however challenging. Parascandolo et al. 2020 applies the rule assuming each sample represents an environment, whereas cach client more naturally corresponds to the environment in FL. Furthermore, they show that this can lead to a slower convergence rate in practice as too many components can be masked at each iteration. The same was observed in a federated setting as shown in Appendix. This would be impractical in the federated setting as we would not want to sacrifice convergence speed for generalization. In the federated setting, we propose a variant of this mask that doesn't sacrifice convergence speed while retaining some of the improved generalization properties. Specifically, we propose to use masking at the aggregation stage of standard FI., with mask computed based on cach client update (which arises from multiple local gradient steps). The mask is calculated based on sign agreement among client updates A, and applied on the global model update AT. This masking controls the parameter update based on the agrecment of direction among the gradients across clients or environments. To provide rapid convergence, we apply a soft masking procedure instead of the hard binary mask.
4

---

Published in Transactions on Machine Learning Research (10/2023)
We define agreement score, A, E (0, 1/. given a function of all client updates, and the mask m, is defined element-wise,
Algorithm 1 Gradient Masked FedAVG McMahan
et al. (2017)
Server Executes:
Initialize we Rd randomly for each server epoch, t == 3........ T do
[> = if A, 2 else Ay; A
Σ ASN
sign(A) ))
II
Choose clients at random
(3)
for each client in C do
The global model update is given by m (1Ar)) @ A. This ensures that the updates to the global model are with respect to their agreement across clients. When the agreement across clients is greater than the hyperparam- etcr T, it would be assigned and when the agrecment is lesser than т, the mask value would be equivalent to the agreement. score. This real mask ensures that each parameter updates, but the magnitude adjusted to be proportional the agreement across clients. We observe empirically in Appendix that the fraction of clients whose magnitude gets adjusted as mentioned above is correlated to the heterogeneity in the distribu- tion of data across clients. The method can be applied any base FI. optimizer, in Algorithm we detail its implementation for FedAVG.
wi = Client Update(we-111
A.
Σ
An
(w) w;-1)
LKEC
end for
A, Σ me AP
me = m- (AA
W/ W Ml my C A
end for ClientUpdateew; n):
Iuitialize wa
for each local client iteration, == =,,,,,,,, K do
Sample = D 9k Fr Wx+1 = We - 1lt end for return WK to server
4 Method Analysis
We now analyze the convergence properties of the gra- dient masking, focusing on the case of FedAVG and gradient masked aggregation. To obtain standard convergence result our algorithm requires to lower-bound on the norm of the masked gradients, requiring additional assumptions. We propose a simplification, which allows to obtains similar convergence guarantees as Reddi et al. [2021]: we assume that the binary masks follow independent Bernoulli distributions. Also, note that we propose our proofs the context of FedAVG rather than FedADAM, making the analysis more straightforward. Following Reddi et al. (2021); we make the following standard assumptions.
Assumption 4.1 (Lipschitz gradient). We assume that each client objective has Lipschitz gradient with constant L, meaning that there exists L > 0, Vn S N, Yw, E Rd, (,) - (.)( < L|| [v]] vll Assumption 4.2 (Bounded gradients). We assume that each client has a bounded gradient by G, leading to: 3G 0, Vn < N, Yw E R ( < G.
Assumption 4.3 (Finite variance). We assume global bound on the variance of the gradient estimate of each individual client, meaning that: 3 > 0, Vn S N, Vw Rd, S d, [2(V Fn V fn(w; =))3 S oo,.- Assumption 4.4 (Global variance). We assune global bound on the variance of the gradient estimate of each individual client, meaning that: Vj, 30 0, Vn, Vw, N (VF,
fw))) og.
The next assumption is core and allows to work with random masks m² {0, 1} at step t, assuming every coordinates could be unmasked with a non zero probability.
Assumption 4.5 (Random mask.). We assume that K[m?] 2 a > and sampling is independent. from V Vfn(wzz))
We begin by this first Lemma, which allows to work with weighted agregations of the local gradients:
Lemma 4.6 (Adaptation from Reddi et al, using the notations of Reddi et al. (022I)). If [wt Lhask are the
K local iterations of the n-th machine at time i, we have for any 2n An = X. 2 0:
Σ x w we w
5K 2E Σ
(o + 2K 23)) 30K3 [[]
5

---

Published in Transactions on Machine Learning Research (10/2023)
Proof. As done in Appendix A, Lemma 3 of Reddi et al. (2021)
will prove the result by induction, and
we use the same notations. First, we use the inequality proved by Lemma 3 of Reddi et al. (2021]: I:
E Ellws -wel? (1 2K- + +6K
Σ oi + 6k (V (f(w))))
+ 6K (
Next, we note that by weighted averaging that:
A - wel2 (1
+ +KK
Σ
-- -))
+ 1//T3 Σ of
2K -
+ 6K {.E[I((E,,..))) -X
Now, using our assumption, we get:
m
m
d Σ
o 9.
(4)
n.
j=l
The rest of the proof follows directly as in Reddi al. (2021) as relies purely on Φx - E EI wall
Our proofs will rely on the following Lemma, where the constant C' is specified in Lemma A4.6 given in the Appendix. We emphasize that the logic of this proof and the constants are identical to Reddi ct al. (2021) Lemma 4.7 (Bounded drift from client update, Adapted from Appendix A, Lemma 3 of Reddi et al. (021) 4.4. there cxists C > 0 such that for any time step t and w., A, obtained
Given Assumptions 4.1
from Alg. 1, if m S SLK
E[|IA + Km <
+ LL LE Σ
(o., + 3Ko +
and
||a|A + KnVS(WIII 52 oot +KK + 33
Its proof can be found in the appendix. The next proposition derives an upper bound on the model gradient following results from Reddi et al. (202I), and relates the speed of convergence a.
Proposition 4.8 (Convergence analysis). Given Assumptions 4.1 m O( KT.)), such that one has the following rate over the masked gradients given by the FedAVG algorithm
if fl = O(1) then there crists
vf(ww)) < O
VoeT
a²T
Proof. We follow the approach of Bottou 2010 for obtaining optimal non-convex bounds. Each F, is L-smooth, thus so is f by averaging, thus, by definition using Alg.
L
f(wr+1) S f(w.) + [V w++1 w.) | w
2
L
=(w)) + n, (VS[V 2 A.) + "
m 3 1.11
2
L
S f(()) + ny (Vf( O Ae) +
2 11
A

---

Published in Transactions on Machine Learning Research (10/2023)
Next, observe one side that:
E E.... 1 IIA, + (nvV + +2K [V
On the other side, note that:
No (Vf(w)), my A1) (n(Vf(()),
Vf((w)) -.f((.))-KKmmm (m, KnKKmlll (mmmmmmmm 2 Vf(w))] + T/ (V(w)), my @
(m @ Vf(wi)) m De)
(5)
++)
(6)
Km 3 Vf(ww)) + (m 0 (f(w), Km V (w) + At)
(7)
(8)
Then, taking the expectation over data 2n and application of Cauchy-Schwartz leads to:
Ez., Mg (V m @ A.) <
7 2
K
1[m, @ vf(ww))) +
ng 2m
EE + Kn vf(m))l
2
Now, note that taking the expectation over the mask gives a a|[Vf(w))) and summing for L = O,... T 1, we get:
B m2 o v((w)]l
Combining
T-1 Σ
E f(w))- (((0o)-((w")>)
k Ln n2 (aoaa(K11)) 11mm "1 [|
4 2m
t=
For the sake of simplicity, we introduce ==
= 5K E E (o%.) + +3Ko² 33.) and B = 30K so that using Lemma
5 f(w)) ((000-f(uul f(w" +
(K2 Lif Lniiig -a a(K- |
Ta
- +2Bm ((x))
2m
which also writes:
5 (f + Σ (K2 t2 n1 L2 -aaaaa(K aK-1))n 2 [Ln + 7g (A+B)L'n
Here, let's pick fg = leading to, for 0 <m m 0.001 KL
(where we clearly have suboptimal constants):
T
K
0 < f(wo) - f(w")
+
t=0
Reorganizing the term, this writes fir m < 0.001 a
and o² an upper bound over the sum of the variances:
KL
[[())) f(w)) - f(w")
80Lni0?
cmp + a
Now, we invoke Lemma 17 of Koloskova et al.
2020 which states that there exists of < 0.001 KT such that
there is an absolute C' > 0 satistying:
T
s s
Loz f(w)) - f(ww)
L(f(wo) f(w)))
V (w)))
KT
2²T
7

---

Published in Transactions on Machine Learning Research (10/2023)
and this leads to the final conclusion.
While Assumption 4.5 may not always hold in practice, it allows us to explore the highly non-convex nature of masks given by Eq. 3. This theoretical analysis serves as a sanity check, providing insights that the process of masking is unlikely to significantly impede convergence towards a particular local minimum. However, it is important to nole that this theorem alone does not explain the observed strong performance. Instead, it. indicates that there are no inherent obstacles preventing the algorithm from converging to a local minimum. Therefore, this simplification of the problem sctting serves valuable starting point for further exploration and refinement of the masking approach, allowing for the investigation of more realistic and complex scenarios.
5 Experiments
In this section, show empirically that the proposed GMA tends to outperform standard aggregation (AVG), converging at similar or better than standard aggregation, while enhancing the global model generalization. We observe this for multiple FL algorithms with respect to multiple datasets and data distributions (i.i.d and non-i.i.d).
Implementation We conduct experiments on gradient masked and naive versions of non-adaptive federated optimizers like FedAVG McMahan et al. (2017), FedProx Li al. (2020b) and adaptive federated optimizers like Aed and Fed Reddi al. (2021) across variety of datasets. Our experiment include label distribution skew, feature distribution skew, quantity skew, real-world data distribution, and mixed (label and feature) skew. The code is available at https:/ https://githuub.oo http//gttuu.oo//ariiiiieeeo coo////vi97///1 /FL. Details of the datasets explored and the respective skews induced is summarised in Table] [ The reported performances in Tables 2 and g are average accuracies of 4 independent runs of the model on a test dataset of the corresponding dataset, data skew, and algorithm.
5.1 In-Distribution Evaluation
This the most widely considered setting the FL literature. The global model is evaluated on a test datasel sampled from data at all clients irrespective of the data distribution across clients. This test dataset is a representation of all participating clients. Table 2 shows the test performance of the algorithms and their GMA versions on variety of datasets and data splits. It was observed that with the robust features of GMA, the algorithm is capable of outperforming naive averaging. The difference in improvement is more significant when the data distribution is non-i.i.d. The major reason for this is that gradient masking is capable of focusing on learning the invariances even under increased spuriousness of non-i.i.d data distribution.
Hyperparameters An SGD optimizer with a momentum (p 0.9) and cross-entropy loss was used to train each client before aggregation at the server in all our experiments. The momentum parameters of adaptive federated optimizers are fixed at 31 0.9 and B2 = 0.99 as per Reddi al. (2021). For each of the considerod algorithms we tune the local client. model learning rates, global model learning rates, lau, and number of local epochs (CIFAR and TinylmageNet) to consider the best performances of the algorithms. More details on the hyperparameter tuning is given in Appendix.
Dataset
Skew
# classes
# classes in label skew
sampled # clients
Model
local epochs
# clients
label
10
500, 100,
10, 10,
LoNet.
FMNIST
labe]
10
2
500, 100, 10
10, 10, 10
LeNet
CIFAR10
label, quantity
10
2,
500, 100, 10
10.
10.
ResNet18
3
100
label
100
10
500, 100, 10
10, 10,
ReaNet18
Tiny Imageeee
label real
200
20
500, 100, 10
10, 10, 10
ReaNet18
12
FEMNIST
10
3550, 3550
35, 350
LeNet
FedCMNIST
feature, mixed
10
2
500, 100, 10
10.
10, 10
LeNet
Table 1: Datascts used, heterogencity inducod, and other parameters corresponding to the datasct.
8

---

Published in Transactions on Machine Learning Research (10/2023)
Dataset
Non-IID
Number cliente
FedAVG
FedProx
GMA
FedAdam
AVG GMA
FodYog!
96.24+0.03
97.36-0.02
05.80+0.01
96.05+0.03
97.75++0.05
98.48=0005
38.21+0.06
98.38+0.03
98.610.02
98.75.00.02
97.7610.02
98.641 0.00
98.88110.0
98.9110.03
98.5810001
98.8210.07
MNIST
99,07:00.02
09.0990002
96.91+0.02
09.0840.03
98.51=0.0G
08.0120.00
97.5440.07
98.7410.01
500
97.14.0.02
96.85 86..01
07.27
06.5810.08
98.02.10.01
96.810.02
97.0010.00
96.1810.01
No
100
98, 28.10
98.65.00.02
98.7210.02
98.7910.03
98.9410.022
98.72110.01
98.75.10.00
98.8210.01
10
99.11
99.17=0003
99.1340.02
99.16::0.02
58.7130.02
98.7740.02
98.08:0.01
98.7550.00
500
78.0220.088
78.82:0011
76.2110.12
77.7710.12
78.8010.00
81:28:00.02
78.08160.11
80.5810.14
Yee
100
63.58:10.04
86.27:00.05
85.14:1.122
65.78
65.9110.04
87.18:0001
85.32:00.9
86.4:0.14
FMNIST
$7.43.10.05
88.37_10.07
$8.25.10.14
88.55.10.12
57,2610.077
80.00
37,3910.12
68.610.11
R30070000
43.94--0.12"
83.47-0008
M2.7+-6 122
65444-02288
86.78-0.26
55.60++0.06
86.42-00116
08.03+0
88.43:0.088
67.05+0.02
67.46:10.04
68.6610.18
60.16:0.22
88.6340.06
69.7+0.22
80.35++0.02
89.66=00003
80,92+0.12
90.12:0000
88.66+0.22
88.86#0.18
88.86+00188
89.85+0.21
73.38+0.46
74.44=0.22
72.21:0.22
03.95=00.66
76.41#0.08
77.1240.12
76.1330.008
77.56=0.04
Yes
85.35+0.41
85.62-0.08
84.11+0.36
65.37+0.22
84.43+0.08
85.52+0.08
84.62+0.11
85,95+0.12
CIEAR
84.0610.11
85,21110,26
$1.6510441
85.84110338
$4.6310.0
85.7710.03
84.2111.04
85:3210.04
10
500
76.81
76.90=00.11
74.43:0.08
16.19:20.08
77.02:0122
71.1720.12
17.87100006
78.11:0006
No
100
80.42:0.114
80.07+0.16
79.84=0.09
79.6=0.12
84.18+0.08
86.59=0007
84.13+0.11
86.27=0.09
10
96.47.10.04
87.3810.03
87.04_0.066
86.8910.14
87.0610.12
87.28:10.16
87,06.10.11
87.1710.14
500
12.68.1..55
44.0110.61
38.45.10.04
00.74.10.33
47.5710.12
50.37_0042
66.27.10.62
48.8910.46
Yes
100
51.6510055
53.83:10:42
49.8110.44
50.76:10.21
53.410082
56.68100.66
54.6410.43
56.8810.38
CIFAR 100
80.88:10
52.03::0.18
51.1240.36
51.76::0.36
54.6610.28
50.93:10.20
54.55:10.22
56.73:10.28
500
16.15.10.28
56.110.1
38.55110:2
16.36_0.18
"T.1220.122
44.6510:08
No
86.88+0.06
57.19-0.03
58.31+0.12
64.13+0..08
51.33+-0111
51.65++0.08
51.19++0006
51.41+0.12
36.00-00.2
37.21:0.08
58.16+0.04
58.21:00044
37.1190U8
67.77:10.12
55.98:00.66
66.91:0008
20.24+0.04
21.42:00.08
21.18+0.08
22.36+00.06
23.65+0.12
25.92:0.06
23.81+0.13
25.85+0.18
Tay
Yes
09.00
26.94=0.08
25.64=0.03
25.11#0.11
25.67+0.17
27.41220.13
26.3240.06
28.83#0.32
Image
10
25. 17
26.17=0.08
26.32+0.04
26.71+0.04
27.41+0008
28.11-0.12
27.12+0.08
27.88+0.12
Nec
500
24.1210.18
24.87-0.12
24.73=00.15
20.56110.22
21.1710.12
23.73:10.14
25.80110.1
25:32-0.11
No
100
26.71:0012
27.00:0.16
26.59=0.15
26.56:40.35
26:9940.277
29.07:00.19
27.86:+0.17
29.32-0.12
30.66=0.02
30,45+0.03
20:04+0.04
29.83+0.04
28.98+0.04
32.02=0.08
27.77/+0.06
30.78+0.05
Table 2: In-Distribution evaluations in both iid and Non-iid settings on MNIST, FMNIST, CIFAR 10, CIFAR 100, and TinyImageNel distributed across 500, 100, and 10 clients. In all cases, 10 clients are randomly sample fr participation in each round. Average best test performance of the algorithms and their GMA versions are reported below. In most scttings and FL algorithms GMA provides substantial advantages. Improvements are more substantial Non-iid settings. The best result in each setting across all FL algorithuss shown in red (if GMA) and blue (if AVG). Except in one case GMA version of the algorithm provides the overall best performance.
5.2 Real-World Evaluation
In the practical federated setting, the data across clients is heterogeneous and the clients which deploy the global model (including test clients, non-participating clients, and new clients in the federated network) can have data distribution different from that at any train clients. This can be simulated by using a realistic federated data characterised by feature distribution skew unique to each client including the test clients. Specifically, we use train data from the same domain distributed across clients such that each client has data corresponding to user (or set of users) unique to the client. The test data consists of data from the same domain as the train dataset distributed across clients but from one user (or set of users) not included in the set of train clients. The test clients do not participate in any training rounds. That is, it is cit-of-dditributiin to the train data. similar feature skew is provided by Federated EMNIST [Caldas et al. (2019) where the data points have a user identifier. The test performance of the algorithms and their gradient masked alternatives are given in Table 3. We observe that gradient masking outperforms naive averaging. In the next section we consider more complex feature skewed out-of-distribution to further evaluate GMA.
9

---

Published in Transactions on Machine Learning Research (10/2023)
Mixed Non-1ID
# clients, sampled
FodAVG
FedProx AVG GMA
FedAdam AVG GMA
FodYogi AVG
Dataset
AVG
GMA
Real-World Evaluation PEAINIST
3000 30
04.2
05.57+0.05
94,75+0000
96.01=0.01
90 28+ 0.83
93.02 0.02
91 22 0.02
92.79H 0.02
Dut-of-Distribution Evaluation
500, 10
81/480022
86.22
0.22
83 15:11771
86.25:10.28
81:22:0012
87.16:0.337
84.06110.66
0..00510.44
Yos
100, 10
85.0211 0:28
88:321 0.32
83.310.33
86.92.1 0.28
81.191 0.12
80.261: 0.38
63.3210.05 18
81.791 0.22
Fed- GMNNITI'
10, 10
86.77+ 0:43
89.17-
0.34
86.84 0.38
89.33+ 0.34
85.75+
0.53
89.28- 0.44
86.674 0.53
89.76+ 0.37
300,
60:0620:46
000.12:0.26
87.13-0.444
87.02-00088
87.1820.00
81.7840008
88.44:00008
No
100,
89.88.1 0.32
91:371 0.16
90.62:10:16
90.831 0.22
87.3820.12
68.2110.16
90.27.00008
91.10:10:00
10, 10
89.37+ 0:83
900
80.61+0.88
90.22-0.78
88.39-0.86
80.4-0.84
88.78-0.83
80.88-0.76
Quautity Skew CIFARIO
100, 10
79.36:00:2
80.61=0.24
50. 12
80.88
79.24=0.38
82.96200.41
79.81:50.12
81.02:0.48
0.5)
10, 10
81.7810.44
86.28 0.38
83. 11:00.5 81.42:10.46
85.92-0.16
86.58:10.26
84.98.10.15
00.17:10.29
Table 3: Real-world evaluation on FEMNIST, feature skew and mixed (feature and label) skew on FedCMNIST and FedRT and Quantity skew on CIFAR 10. Total number of clients and number of clients sampled for participation in each communication round is as given in the corresponding row. Average best test performance of the algorithms and their GMA versions are reported below. The best average result among AVG and GMA having atleast 1% higher than the other algorithm is shown in bold.
CIPAAaD medavo & PolAND UMA Festaic SOpe
TingtnagsMeg nodAnn PedAMo 6PA Fedivic Tedtin 2he
Heal World Podeau & redAro OMA festis Feools CAA.
PSSCHMOT PedR/o & redlvo SMA
-aadin OP Fudoiic
033
34
105
founs
CIARLOO togi& Yogi CHA
Tieplmggehst Yogl & Yogi CMA post
Deal World Yogl & hogi
Ped Yagg 4 CMA regi
Yogo
C2
0
Figure 2: Test accuracy VS. communication rounds of non-adaptive (FedAVG) and adaptive (FedYogi) federated algorithms and their GMA versions on in-distribution CIFAR 100 and TinylmageNet, real-world FEMNIST, and mixed skew FedCMNIST datasets.
5.3
cut-of-Distribuutee Evaluation
A more complex OOD test can be implemented to better understand the performance of gradient masking in federated learning. For this we induce unique spurious mechanisms or features in the clients (including test) besides the class label based heterogeneity. The global model would be tested on a dataset having a spurious mechanism that was not present in any of the train clients, while the spurious mechanism at each train client is unique to itself. This ensures that all clients (including test clients) are out-of-distributiin to each other. We combine this with label skew that is similar to the non-i.id heterogeneity in McMahan et al. (2017). FedCMNIST Francis et al. (202)), federated multiclass version of CMNIST Arjovsky et al. (2020) with multiple color-label correlations was used for this. The invariant mechanism here is the digit. There also exists a spurious mechanism marked by a color given to the numbers. The color is digit specific to induce correlation to the label. The performance of various algorithms and their gradient masked averaging counterparts on these datascts is given in Table 3 It is to be noted that almost all cases gradient masking outperforms naive averaging. Figure 3 a. shows train and test curves of the algorithms and their gradient masked alternatives on non-i.i.d distribution across 10 clients with full participation on FedCMNIST. It is to
10

---

Published in Transactions on Machine Learning Research (10/2023)
a) Acc. v8. Comm Round b) Test Acc. Local Epochs
c) Test Acc.
# Clients
d) Acc. V8. 7
Figure 3: Ablations: a) Train and test accuracies of FedAVG on FedCMNIST with 10 clients and full participation. b) Test accuracies of FedAVG and GMA on 10 client FMNIST with full participation is recorded against increasing number of local epochs. c) Test accuracies against number of clients in the federated network while all clients participate training in each round (full participation). d) Test accuracy against (0, 1} on label skewed CIFAR10 distributed across 10 clients with full participation.
be noted that although GMA train accuracies are less than equal to that of naive averaging, the GMA test accuracies are higher. This indicates that gradient masked versions generalize better than the naive versions.
5.4 Quantity Skew
In a practical federated setting the quantity of data available for update at cach client varies drastically. Depending on connectivily, processing power, user behaviour, and various other factors, the number of data samples generated at each client can differ from zero or one data point to an extremely large number. To simulate this quantity imbalance, we have used a Dirichlet distribution based quantity skew with 0.5 as Li et al. (2021a) on CIFAR-10 across 100 and 10 clients with 10 clients participating in each communication round. The Dirichlet sampling is independent of the labels features of the dataset. The test data contains samples from all classes. Table 3 shows that GMA outperforms AVG under quantity imbalances.
5.5 Ablations and Analysis
Ablations
Using the non-iid distributed FMNIST data study how the performance is affected as the number of clients grows while all clients participate in training in each round (N 10, 50, 100, 250) (shown in Figure c.) and the number of local epochs increases (E = 1, 5, 10, 20) (shown in Figure b.).
In the former, the data is distributed across 10 clients with varying number of local epochs and in the latter, the data, is distributed across a varying number clients with local epochs =1. In both cases all clients partici- pale in training each communication round. It can observed from various datasets and settings that gradi- ent masking increasingly outperforms naive averaging in complex datascts and sampled settings. This validates the enhanced invulnerability of gradient masking to the bias that could be induced by one or more clients in the network.
Table 4: Test accuracies of paticipating and non- participating clients of single FedAVG round with and without GMA on FMNIST distributed across 10 cicnts. Participating clients of round those that are randomly sampled to contribute to train- ing in that round and the remaining are the non- participating clients.
Participating
PontPaatiiipating GMA
Trial
AVG
GMA
AVG
Convex Objective To further understand the per- formance of gradient masking in the convex setting, experiment with MNIST and FedAVG on both i.id and non-iid data distributions and it was observed that GMA outperforms naive averaging in the non-iid setting. A logistic regression model with SGD with momentum optimizer was used at the clients for these experiments.
84.34
85.67
81.63
85.39
2
87.22
89.44
72.38
73.04
89.34
92.11
71.21
77.18
MMan:
86.97
89.07
75.07
78.53
"Relative"
Improv.:
2.4%
4.6%
When data distribution was i.i.d, GMA was converging an average (over last 10 communcattin rounds) of 92.5% test accuracy while naive averaging obtains to
11

---

Published in Transactions on Machine Learning Research (10/2023)
92.4%. Furthermore, when the data distribution across clients was non-iid,, the enhancement in performance was more significant with gradient masking. While naive averaging converging to 87.0% test and 92.0% train, while GMA reached 88.5% test and 92.2% train. Further demonstrating GMA can generalized better in the non-iid casc.
Effect on non-participating clients GMA tends to improve performance gains in the relevant setting with lower client participation. We can gain insight into this by considering an individual FL round where the non-participating clients data are out-of-distribution with respect to the obtained model, as their data is not. used to update the global model in that round. Motivated by this observation of a single FedAVG round with and without GMA the FMNIST dataset using 3 local epochs 10 clients with 5 clients participating in each round. The results are shown in Table 5 where we report the accuracy over participating and non-participating clients. We observe that GMA improves performance in both cases but the relative performance change on non-participating clients is higher (4.6%) versus 2.4%). This suggests the performance gain lower client participation is stronger due to improvement on non-participating client. performance.
further study the behavior
6 Conclusion
We proposed a new aggregation scheme applicable to wide variety of foderated learning algorithms. The proposed method, gradient masking enhances generalization performance of the global model in FL by focusing learning the invariances across clients. The simple masking outperforms their naive averaging versions across variety of algorithms and datasets. Our theoretical analysis shows the convergence of the proposed masking algorithm and the stability of the proposed mask. Future directions include exploration of masks incorporating magnitude and other methods to better capture the invariances, thus leading to better generalization at the global model.
References
Kartik Ahuja, Karthikeyan Shanmugam, Kush R. Varshney, and Amit Dhurandhar. Invariant risk minimization games, 2020.
Martin Arjovsky, Léon Bottou, Ishaan Gulrajani, and David Lopez-Paz. Invariant risk minimization, 2020.
Léon Bottou. Largo-scale machine learning with stochastic gradient descent. In Proceedings of COMP- 201 pp. 177 186. Springer, 2010.
Peter Bühlmann. Invariance, causality and robustness, 2018.
Sebastian Caldas, Sai Meher Karthik Duddu, Peter Wi, Tian Li, Jakub Konečný, H. Brendan McMahan, Virginia Smith, and Ameet Talwalkar. Leaf: A benchmark for federated settings, 2019.
Alireza Fallah, Aryan Mokhtari, and Asuman Ozdaglar. Personalized federated learning: A meta-learning approach, 2020.
Sreya Francis, Irene Tenison, and Irina Rish. Towards causal federated learning for enhanced robustness and privacy, 2021.
Sharut Gupla, Kartik Ahuja, Mohammad Havaei, Niladri Chatterjee, and Yoshua Bengio. Fl gamess: A federated learning framework for distribution shifts, 2022. URL https: org/ abs////........///////./.2 11101
Tzu-Ming Harry Hsu, Hang Qi, and Matthew Brown. Measuring the effects of non-identical data distribution for federated visual classification, 2019.
Peter Kairouz, H. Brendan McMahan, Brendan Avent, Aurélien Bellet, Mehdi Bennis, Arjun Nitin Bhagoji, Kallista Bonawitz, Zachary Charles, Graham Cormode, Rachel Cummings, Rafacl L. D'Oliveira, Hubert Eichner, Salim El Rouayheb, David Evans, Josh Gardner, Zachary Garrett, Adrià Gascón, Badih Ghazi, Phillip B. Gibbons, Marco Gruteser, Zaid Harchaoui, Chaoyang He, Lie He, Zhouyuan Huo, Ben Hutchinson,
12

---

Published in Transactions on Machine Learning Research (10/2023)
Justin Hsu, Martin Jaggi, Tara Javidi, Gauri Joshi, Mikhail Khodak, Jakub Konečný, Aleksandra Korolova, Farinaz Koushanfar, Sanmi Koyejo, Tancrède Lepoint, Yang Liu, Prateek Mittal, Mehryar Mohri, Richard Nock, Ayfer Özgür, Rasımus Pagh, Mariana Raykova, Hang Qi, Daniel Ramage, Ramesh Raskar, Dawn Song. Weikang Song, Scbastian U. Stich, Ziteng Sun, Ananda Theertha Suresh, Florian Tramèr, Pranceth Vepakomma,, Jianyu Wang, Li Xiong, Zheng Xu, Qiang Yang, Felix X. Yu, Han Yu, and Sen Zhao. Advances and open problems in federated learning, 2021.
Sai Praneeth Karimireddy, Satyen Kale, Mchryar Mohri, Sashank J. Reddi, Sebastian U. Stich, and Ananda Theertha Suresh. Scaffold: Stochastic controlled averaging for federated learning, 2021.
Anastasia Koloskova, Nicolas Loizou, Sadra Borciri, Martin Jaggi, and Sebastian Stich. A unified theory of decentralized sgd with changing topology and local updates. In International Conference on Machine Learning, pp. 1381-5394 PMLR, 2020.
Masanori Koyama and Shoichiro Yamaguchi. When is invariance useful in an out-of-distribution generalization problem 2, 2021.
David Krueger, Ethan Caballero, Joern-Henrik Jacobsen, Amy Zhang Jonathan Binas, Dinghuai Zhang, Remi Le Priol, and Aaron Courville. Out-of-distributio generalization via risk extrapolation (rex), 2021.
Qinbin Li, Yiqun Diao, Quan Chen, and Bingsheng He. Federated learning on non-iid data silos: An experimental study. arXie preprint SKk::2102.02019 2021a.
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smithy. Feddane: A federated newton-type method. In 2019 53rd Asilomar Conference on Signals, Systems, and Computers, pp.
IEEE, 2019.
Tian Li, Amit Kumar Sahu, Ameet Talwalkar, and Virginia Smith. Federated learning: Challenges, methods, and future directions. IEEE Signal Processing Magazine, 37(3):50 60, May 2020a. ISSN 1558-0792. doi: 10.1109 http://ppp.....9111/
URL lttp:: doi. org/10. 1109/MSP. 2020. 2975749
Tian Li, Anit Kumar Sahu, Manzil Zaheer, Maziar Sanjabi, Ameet Talwalkar, and Virginia Smith. Federated optimization in heterogeneous networks, 2020b.
Tian Li, Shengyuan Hu, Ahmad Beirami, and Virginia Smith. Ditto: Fair and robust federated learning through personalization, 2021b.
Brendan McMahan, Eider Moore, Daniel Ramage, Seth Hampson, and Blaise Aguera y Arcas. Communication- efficient. learning of deep networks from decentralized data. In Artificial intelligence and statislics, pp- 1273 1282. PMLR, 2017.
Jachoon Oh, Sangmook Kim, and Se-Young Yun. Fedbabu: Towards enhanced representation for federated image classification, 2021.
Giambattista Parascandolo, Alexander Neitz, Antonio Orvieto, Luigi Gresele, and Bernhard Schölkopf. Learning explanations that are hard to vary, 2020.
Titouan Parcollet, Xinchi Qiu, Daniel Beutel, Taner Topal, Akhil Mathur, and Nicholas D. Lane. Can federated learning save the planet?, 2021.
Mohammad Pezeshki, Sékou- Oumar Kaba, Yoshua Bengio, Aaron Courville, Doina Precup, and Guillaume Lajoie. Gradient starvation: A learning proclivity in neural networks. In A. Beygelzimer, Y. Dauphin, P. Liang, and J. Wortman Vaughan (eds.), Advances in Neural Information Processing Systems, 2021. URL https://openrevieu hwt.a//eeee.eeoooooe//d//oooo...
Xinchi Qiu, Titouan Parcollet, Javier FernandezMarquos, Podro Porto Buarque de Gusmao, Daniel J. Beutel, Taner Topal, Akhil Mathur, and Nicholas D. Lane. A first look into the carbon footprint of foderated learning, 2021.
13

---

Published in Transactions on Machine Learning Research (10/2023)
Athul Sreemathy Raj, Irene Tenison, Kacem Khaled, Felipe Gohring de Magalhães, and Gabriela Nico- lescu. FedSHIBU: Federated similarity-baset head independent body update. In Workshop on Federated Learning: Recent Advances and New Challenges (in Conjunction with NeurIPS 2022), 2022. URL
net/forumidd=02y 2Fge2_222
Sashank Reddi, Zachary Charles, Manzil Zaheer, Zachary Garrett, Keith Rush, Jakub Konečný, Sanjiv Kumar, and H. Brendan McMahan. Adaptive federated optimization, 2021.
Amirhosscin Reisizadeh, Aryan Mokhtari, Hamed Hassani, Ali Jadbabaie, and Ramtin Pedarsani. Fedpaq: A commmication-efficion foderated learning method with periodic averaging and quantization, 2020.
Soroosh Shahtalebi, Jean- MhanCChistoppe Gagnon-Audet, Touraj Laleh, Mojtaba Faramarzi, Kartik Ahuja, and Trina Rish. Sand-mask: An enhanced gradient masking strategy for the discovery of invariances in domain generalization, 2021.
Sebastian U. Stich. Local SGD converges fast and communicates little. In International Conference Learning Representations, 2019. URL https::/openreviice htt.://www.....ooooooooooaaaaaaa
Shruti Tople, Amit Sharma, and Aditya Nori. Alleviating privacy attacks via causal learning, 2020.
Hongyi Wang, Mikhail Yurochkin, Yuekai Sun, Dimitris Papailiopoulos, and Yasaman Khazaeni. Federated learning with matched averaging, 2020a.
Jianyu Wang and Gauri Joshi. Cooperative sgd: A unified framework for the design and analysis of communication-efieee sgd algorithms, 2019.
Jianyu Wang, Qinghua Liu, Hao Liang, Gauri Joshi, and H. Vincent Poor. Tackling the objective inconsistency problem in heterogencous foderatod optimization, 2020b.
Shiqiang Wang, Tiffany Tuor, Thcodoros Salonidis, Kin K. Leung, Christian Makaya, Ting He, and Kevin Chan. Adaptive federated learning in resource constrained edge computing systems, 2019.
Yuxin Wu and Kaiming He. Group normalization, 2018.
Jing Xu, Sen Wang, Liwei Wang, and Andrew Chi-Chih Yao. Fedem: Foderated learning with client-level momentum, 2021.
Hao Yu, Sen Yang, and Shenghuo Zhu. Parallel restarted sgd with faster convergence and less communication: Demystifying why model averaging works for deep learning, 2018.
Honglin Yuan and Tengyu Ma. Federated accelerated stochastic gradient descent, 2021.
Honglin Yuan, Warren Morningstar, Lin Ning, and Karan Singhal. What do we mean by generalization in federated learning?, 2021.
Mikhail Yurochkin, Mayank Agarwal, Soumya Ghosh, Kristjan Greenewald, Nghia Hoang, and Yasaman Khazaeni. Probabilistic federated neural matching, 2019. URL sttpp://openreviier.ne///ooo nww.:frmuidi SygHGnRqk7/
Liling Zhang, Xinyu Lei, Yichun Shi, Hongyu Huang, and Chao Chen. Federated learning with domain generalization, 2021. URL https://arxiv: arg/abb/2211.1099999
Yue Zhao, Meng Li, Liangzhen Lai, Naveen Suda, Damon Civin, and Vikas Chandra. Federated learning with non-iid data, 2018.
Zhuangdi Zhu, Junyuan Hong, and Jiay Zhou. Data-free knowledge distillation for heterogeneous federated learning. In International Conference on Machine Learning, pp. 12878-12222 12889. PMLR, 2021.
14

---

Published in Transactions on Machine Learning Research (10/2023)
A Appendix
A. Drifting Bound of k
Below the proof of Lemma 4.7
Proof. All the core idea this proof is to observe that for the batch of data =
N
| Knvs(uu))
=n
Sn
(V
V (w:)))]
(9)
k=1 nm
S nの K
8n
(V (w -
(10)
N
8n
K
N
8n
[V =k - (w)e1))
(11)
k=1 n=
m=1
8n
(12)
We also observe that:
K
N
E[A + Kn
8n
(VFn(W())) Fl.(w)))
A=
==
2n=1
8n
Now, we obtain that by taking the expectation and L-smoothness that:
K
N
E||l: + Kn
Sn
[[ (ur(2)))1 +
Σ
(13)
N
k=1
R=1
m=1
8n
J=1
K
Sn
(L2 (L21| wzl12+
2
(14)
m=
8n
and we can use Lemma 4.6 and we obtain the conclusion (using 1 S K) and similar reasoning for the other part of the lemma.
B. Algorithms
This section gives the pscudocode of the the proposed gradient masking on FodAVG, FodADAM, and FodYogi. GMA can be plugged into most federated learning algorithms easily. Algorithm shows how to extend FedAVG to also include Fed and Fed and highlights how the masking interplays with this.
15

---

Published in Transactions on Machine Learning Research (10/2023)
Algorithm Gradient Masked FedAVG McMahan et al. (0117) FedADAM and
Fed Yog
Reddi et al.
(2021)
Server Executes:
Initialize wio
for each server epoch, f = 1,2,3,.. do
Choose clients at random
for each client in C, do
W ( (w. wi-1)
A E1er
end for
A, EnEc A
my
m ( (nE))
2t = -112--1 + (1-B1)A,
1/ Ut-1 (1 32) (2iig) Af)
1/2 821--1 1 + (1 A2)
21 A
Vi++t"
W/ Wr-1 + Mlg ml O A.
end for ClientUpdato(w, n):
Initialize Wig = w
for each local client iteration, k 0,,,,,,, K do Sample D
Ok V (www;;))
Wk+1 1n 9k
end for return WK to server
Additional Motivation
C.1 Agreement and Homogeneity
Figure 4: fraction of gradients having agreement tau (0.4) across clients vs. homogeneity heterogeneity od data distribution across clients.
We analyze the effect of data heterogeneity on the fraction gradients that were less that 7 = 0.4. The gradients were adjusted make update proportional their agreement across clients. was observed that with decreasing heterogeneity or increasing homogeneity, this fraction was reduced as shown in Figure 4. That is, with increasing homogeneity, the fraction of clients that were proportionately reduced decreased.
16

---

Published in Transactions on Machine Learning Research (10/2023)
Alternatively, with increasing homogeneity, the number of gradients having agreement >> tau increased. This indicates the relevance of the gradient agreement across clients in federated learning. The heterogeneity induced using Dirichlet distribution with a = 0.1 representing heterogeneous case and a = 100 representing homogencous case. The experiments were based on label skewed MNIST.
D.2 GMA vs. Binary Mask
According to Parascandolo et al. 2020), though the binary AND-Mask they proposed improves generalization performance, is slower in convergence. The same was observed in a federated setting. The binary mask converges to a higher test accuracy than FedAVG and approximately equal the proposed G.M.A However, the number of communication rounds required achieve it was higher when the data was non-iid as shown in Figure In federated setting, communication bottleneck is severe and improved convergence speed cannot be sacrificed for minor improvement in generalization performance.
Test Accuracy Piak
FadA//
TedAVG GPA
0.6
0.2
Step
100
150
Figure 5: Comparison on FodAVG, proposed FedGMA, and the binary mask proposed by Parascandolo et al.
(2200. The experiments were conducted on MNIST with label distribution skew.
D.3 Effect on non-participating clients 2 clients per communication round
Table 5: Test accuracies of paticipating and con-participating clients of single FedAVG round with and without GMA on FMNIST distributed across 10 cients. Participating clients of round are those that are randomly sampled to contribute to training in that round and the remaining clients the non-participating clients.
Participating AVG
Non- Non-Paaticipating
Trial
GMA
AVG
GMA
89.73
87.72
71.61
74.34
2
90.70
85.57
71.01
76,56
88.5
86.73
72.09
76.43
Mean:
89.64
86.67
71.57
75.77
Relative
Improv.:
-3.3%
5.8%
In this section extend effect non-participating clients subsection of section 5.5 of the paper. We look at the effect of GMA on non-participating clients when only 20% of the clients participate in each communication round. In the lable above, we consider 10 client with only 2 randomly chosen clients participaling in each round. We observe that the relative performance improvement on non- participating clients is significant with GMA. We also note that performance of GMA on participating clients is lower than that of AVG. This is bocause GMA would update only if both clients agree on the sign. Due this most parameters do not update in every round.
17

---

Published in Transactions on Machine Learning Research (10/2023)
D. Hyperparameters
D.1 Effect of Tau
E 10 is hyperparameter introduced to threshold agreement across clients in Equation 2. It marks the minimum agreement required to consider the gradient for aggregation. When the gradients of all parameters will have agreement score grealer than equal to This makes all gradients consistent across clients. The agreement score would be over-written by 1 and the equation becomes equal to that of naive foderated aggregation where all gradients are averagod with the same importance. This equivalent to an underfit condition. The opposite overfit criterion can happen with a high value. In this case, no gradient. would be considered dominant. and all parameters updates would be diminished corresponding to their agreement score. The agreement score can be when the data distribution across clients is idcal i.id distribution where all gradicnts across clients would be along the same direction. But in practical scenario, such data distributions are rare in federated setting. When agreement = 1, w* = w= 7g C A = w - ny equivalent to naive federated ag- gregation. This implies that the naive fedrated aggregation is case of the proposed gradient masked averaging,
or negligible,
Throughout the experiments we have maintained 5 0.4. This value was observed to give the best performance upon searching on {0, 1.0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9). T 0.4 implies that the gradient being considered have 40% excess or total of 60% of the client gradients along the dominant direction. From our experiments on CIFAR10 distributed non-iid across 10 clients with full- hull-participation, it was observed that when 7 is low, the model underfits. The accuracy was best at T = 0.4 and on further increasing it was overfitting. This is visible from the test accuracy vs. T plot in Figure 3 d. and test loss vs. plot in Figure 3d.
D.2 Effect of Client Momentum
In contrast to the experiments in Reddi et al. (2021). use an SGD optimizer with momentum (p= 0.9) at each client for all our experiments. This was primarily because of the increase in test accuracy observed during our experiments on FedAVG with and without momentum. However, the enhancements due to gradient masking was independent of the momentum induced at the client optimizer. In both cases (with and without momentum), gradient masking was outperforming naive averaging in most of the algorithms and datasots. Table p shows the performance of the algorithms and their GMA versions on non-i.i.d distributod FMNIST using an LeNet model. The client optimizers used in our experiments is a naive SGD optimizer with momentum parameter and it does not involve the correction parameter introduced in Xu et al. 2021
D.3 Effect of Group Norm
The test accuracies reported in paper corresponding to CIFAR-10 used ResNet18 model with batch normalization layers replaced by group normalization & & He 2018 similar to the experiments in Reddi et al. (002I). Our initial experiments involved batch normalization as in the original ResNet and was obscrved that the replacement of batch norm with group norm improved the test accuracics. Table 7 shows the comparison of the algorithms and their GMA versions on CIFAR-10 using ResNet model having batch normalivation and group normalization layers. It is to be noted that irrespective of the normalization layer used, gradient masking outperforming naive averaging across all algorithms and data distributions. This further validates the capabilities of the proposed GMA.
D.4 Grid Search Range
For all our experiments a search for client learning rate (πτ) and global learning (7) rate across the grid specified below was conducted and the best performing learning rates of each algorithm was used for experiments that reported the test performances in the paper. The grid was fixed the same for all datasets and algorithms for easy experimentation.
18

---

Published in Transactions on Machine Learning Research (10/2023)
Table Performance of the algorithms and their GMA versions with and without momentum(p) on non-i.id distributed FMNIST using an LeNet model on 10 clients with pullpparticiaation. Momentum improves performance of the algorithms. Irrespective of momentum, GMA outperforms AVG.
FedAVG (p=0) 0)
FedAVG (p=0.9)
Dataset (Model)
AVG
GMA
AVG GMA
MNIST (LeNet)
IDD Non-IID
99.01
98.96
99.1
99.16
98.43
98.55
98.87
98.9
FMNIST (LeNet)
IID Non-IID
88.61
88.49
89. 14
90.52
86.95
87.8
88.1
88.38
FEMNIST (LeNet)
IID Non-IID
98.8
98.92
99.7
99.68
92.17
94.61
94.2
96.04
CIFAR-10
IID Non-IID
85.8
86.31
87.3
87.61
(ResNet)
81.1
82.28
83.25
83.95
Table 7: Average in-distribution test performance(%) over the last 10 communication rounds of FedAVG, FedProx, SCAFFOLD, FodAdam, FedYogi and their GMA versions on i.idd and non-i.id distributions of CIFAR-10 on ResNet18 models using batch normalization and group normalization. The best result among AVG and GMA versions of each algorithm is shown in bold.
Dataset (Model)
FedAVG AVG GMA
FedProx AVG
FedADAM AVG GMA
Yogi AVG
GMA
GMA
ResNet BatchNorm
IID Non-IID
87.11
87.42
87.2
87.52
77.32
80.65
78.78
80.55
77.3
79.9
78.4
80.2
69.82
74.05
67.29
71.51
ResNet GroupNorm
IID
87.3
87.61
87.18
87.5
86.9
87.7
87.53
87.78
Non-IID
83.25
83.66
83.87
84.4
83.53
84.84
83.17
84.55
m E (10 10-2,,,11,,.00 2,,.00,022 2,10-
1/ E 11,1,.2,
E 0..1, 0.3, 0.5.0.........500..,... 0.6 0.7,0,.
The values for global learning rate and client learning rale was observed to vary across algorithms and datasets. However, it was observed that the same combination was giving the best performance for GMA and AVG versions of the same algorithm for the same dataset. This enables better comparison of naive averaging and gradient masked versions of algorithms better. For 7, it was observed that 0.4 was giving the best result acriss algorithms and datasets. More details about hyperparameter given in Appendix D.1.
19

---

Published in Transactions on Machine Learning Research (10/2023)
D.5 Increased Clients and Local Epochs
Table and Table g shows average accuracies corresponding to Figure 3 b and These report the average test performance over last 10 communication rounds convergence of FedAVG and its GMA version across increasing number of selected clients and number of local epochs at each client per communication round.
Table 8: Average test performance values of non-iidd FMNIST on varying number of clients
Number of Selected Clients
10
50
100
250
AVG
88.1
82.42
74.39
71.11
FedAVG
GMA
88.38
83.11
76.11
72.87
Table Average test performance values of non-i.dd FMNIST with varying number of local client epochs per communication round
Number of Local Epochs
3
5
10
20
AVG
89.91
85.71
84.13
80.27
FedAVG
GMA
91.73
87.29
87.11
86.66
D.6 Performance for same learning rates
Table 10 shows the in-distribution test performance of naive and gradient masked versions of FedAVG on non-iid FMNIST dataset. Performance corresponding to the entire range of hyperparameters mentioned in Appendix C.4 has been given here. It can be noted that across learning rates where the model learns, gradient masking outperforms naive averaging. This suggest the gains are highly robust to hyperparameter choices.
Table 10: Performance of FedAVG across a range of global learning rate (ng) and client learning rate (m) on non-iid FMNIST distributed across 10 clients with full participation. can be observed that GMA outperforms AVG in most of the cases where the algorithms learn and converge.
57.72
65.11
72.56
GMA
0.1
56.03
57.51
73.9
78.66
$3.39
88.70
85.22
87.22
0.1
0.1
AVG
GMA
72.09
86.49
87.70
88.31
0.1
AVG
1.0
79.54
87.0
88.14
88.4
0.1
GMA
77.13
87.04
88.85
0.001
0.05
1.0
20

---

Published in Transactions on Machine Learning Research (10/2023)
E. Convex Objective
For the experiments on convex objective we employed logistic regression model on MNIST with FedAVG for global model approximation at the server. The model was created using a single linear layer in a perceptron model without any non-lincar activation functions. cross-cntropy loss was used and the model was updates by an SGD optimizer with momentum at each client. Each client undergoes update step before aggregation approximate the global model. Experiments included i.i.d and
(gradient masked or naive avergaing) non-ii.d data distribution of data across clients with in-distribution test for global model evaluation. The number of clients selected at each round was set to 10 and the models were run until convergence. The various hyperparameters involved in the experiments are global model learning rate of 1.0, client model learning rate of 0.01, and t 0.4.
F. Membership Inerence Attack
Total Dataset 60K Samples
GMA
Target Model Training Set 50K Samples
Target Model Test 10K Samples
Text Loss
0.72
7K
7K
3K
Attacker Model Training Set 14K Samples
Attacker Model Training 6K Samples
0.8
1000
2000
3000
5000
b)
Figure 6: (a) Data split and creation for attacker model (b) Test loss vs. epochs of the logistic regression attacker model.
Most machine learning models tends to overfit on their training data and such models are susceptible to membership inference attacks that can accurately prodict whether a data sample was present in the training set of the model given the model output logits Tople et al. (2020). This major privacy breach and it can simulated by using a black-box adversarial attacker model.
The attacker model we have employed is binary logistic regression model with binary cross-entropy loss. The input to this attacker model is the logits of the converged gradient masked averaging model and naive averaging model. The attacker model is supposed to identify whether the input of the global model corresponding to the logit given was present in the global model's training set or not. For our experiments, we used CIFAR-10 and ResNet models. Firstly, the GMA and AVG models were trained and tested. For each model, the logits corresponding to the train and test set data and their labels (whether train data or test data) were stored. The data is split as shown in Figure a Tople et al. (2020) and the attacker model is trained for 5000 rounds. The accuracy is reported in the paper and loss shown in Figure 7 b A lower attack accuracy of gradient masked implics that GMA has better immunity to membership inference attacks than naive averaging global models.
This experiment based on Tople et al. (2020) which suggests that algorithms focusing on lcarning the causal mechanisms provide stronger privacy guarantees in certain cases, for example they can be more robust to membership inference attacks and model inversion attacks. Based experiments observe that the
21

---

Published in Transactions on Machine Learning Research (10/2023)
attack accuracy with respect to the GMA model is 55% while that of naive averaging model is 57%. This suggests that gradient masking can potentially enhance robustness to membership inference attacks.
G. Mask Stability
Below we establish concentration bound, to determine the likelihood of applying the masking operation to coordinate
Proposition A.1 (Mask stability). Denote 8,6" the corresponding respectivcly to coordinate of A, A Purthermore consider the r.v. for each coordinale of A, where A & A. Assurn that 8 is o-sub Gaussian, that the ön are mutually independent and write μ²" = E[0"]. If N cara((nlp"" > 0}) > T, then,
with probabilily 1 1 (u nnn ) we oblain == 6.
Proof. We show a lower bound on 8 With probability 1 e 2,, we get 10th -0"-""" inf [0= 4 By considering the intersection of those events, it implies that with probability at least 1 0z, 5th 1tt 2 2/a² 0. Consequently, the mask is equal to and 5t = 8".
1²1 < Let's thus pick
22


Reviews:
[{"review_id": "8cN6r11soB", "reviewer": "Reviewer_Q7tq", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n * The paper suggests a novel idea of using a \"soft\" gradient masking. \n * A binary masking which can be thought of as a predecessor to this method does underperform against the proposed aggregation scheme which shows the need for softer masks.\n * The theory stated in the paper is simple to understand and for the most part, beautifully put.\n * From a practical point of view, the proposed method outperforms the AVG aggeration method which can be considered a naive and widely used aggeration method. \n * The authors have conducted extensive experimentation highlighting and showing the pros and cons of their proposed method.\n\nWeaknesses:\n  * The authors propose the use of soft masks as opposed to using binary masks. However, their theoretical guarantees are done with respect to binary masks, via the use of Bernoulli variables. This is understandable since analyzing the convergence associated with softer masks is hard to obtain.\n  * There are some typos across the paper; see the following section.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper provides a new aggregation scheme that can be used by various federated learning algorithms. The proposed method hinges upon soft masking of the local gradients prior to updating the global model. The authors then conduct extensive experimentation showing the efficacy of their approach and its impact in various settings. ", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1) At Algorithm 1, please replace $n_s$ with $s_n$.\n2) At Assumption 4.3, please make that the expectation sign is concerning a valid norm; currently, the norm is not written correctly.\n3) What do the lines with light colors mean in Figure 5 in Section D.2 (in the appendix)?", "broader_impact_concerns": "There are no broader impact concerns.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "eUTEu7sgYx", "reviewer": "Reviewer_54rX", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths\n1. The proposed method can be applied directly to many existing federated learning algorithm, by serving as a drop-in robust alternative to naive averaging of parameters. \n2. The proposed GMA is a simple mechanism and is easy to be adopted. Using FedADAM as an example, the authors have shown that under the assumption that binary mask follows independent Bernoulli distribution, the modified FedADAM with GMA still converges. \n\nWeaknesses\nThe authors claimed that the proposed GMA, when applied to existing FL algorithms, can generalize better to out-of-distribution data without sacrificing the convergence rate. However,  there is no rigorous support for such a statement. Theoretically, the authors only provide the convergence analysis for a special case: FedAvg with GMA under the assumption that binary mask follows independent Bernoulli. It is unclear whether GMA can really help generalization to out-of-distribution data. Empirically, the paper only applies GMA to 4 existing FL algorithms: FedAvg, FedProx, FedAdam, FedYogi, but the improvement over naive averaging mechanism doesn’t seem to be significant. To fully evaluate GMA, many other baselines should be considered. \n\n", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes an aggregation mechanism, Gradient Masked Averaging (GMA),  for federated learning. Compared to conventional FL that simply averages over local parameters, the proposed method is more robust to out-of-distribution data and can be used as a component for many existing federated learning algorithms. Under GMA, each client after receiving global model performs multiple local updates, the sign of each update component is then used to compute element-wise agreement score. A soft-masking operation is applied subsequently to determine the final global update.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. The authors mentioned that due to highly-correlated masked gradients, the convergence analysis of FedAvg with GMA is hard to obtain. Can you elaborate more on this? Compared to conventional FedAvg, how is the convergence rate affected when applying GMA? Under heterogeneous clients, isn’t the assumption that binary mask follows independent Bernoulli distribution easily violated? \n2. How do we know theoretically GMA improevs generalization to out-of-distribution data?  ", "broader_impact_concerns": "N/A", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "TIoKi1por3", "reviewer": "Reviewer_ffYZ", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "**Strength**\n* The method of masking is well-motivated and novel in the context of FL. Specifically, the idea of using masking to encourage consensus and avoid learning from spurious features has been explored in OOD generalization literature. This paper treats each client as a separate environment whereas each sample is treated as an individual environment in OOD generalization. Furthermore, establishing a high-level connection between OOD generation and federated learning is also an interesting aspect as both fields share some common underlying challenges such as spurious correlations in the form of heterogeneity.  \n\n* The experiments are extensive. The paper provides experiments on many different FL settings including common simulated IID and Non-IID splits and real-world datasets. Moreover, the paper also provides results on OOD experiments where OOD data have different spurious mechanisms. It is intriguing to see that masking can improve both federated performance and also OOD performance in a federated setting. \n\n**Weakness**\n* The theory is not specifically enlightening. The main convergence result (proposition 4.8) relies on an assumption of a randomly sampled mask instead of the real mechanism proposed. While the theoretical difficulty of analyzing the original algorithm can be understood, this nevertheless weakens the contribution.  \n* The writing of the theory section is also confusing. There is little discussion on theoretical results and their implications. For example, it is not clear why Proposition 4.9 (mask stability) is important and what it indicates about the masking method. \n\n**Minor**\n* At the beginning of 3.1, the paper wrote \"$f_n$ is an estimate of the corresponding gradient\". Should $f_n$ here denote the empirical loss/objective function given sample data? \n* In equation 1, $\\eta_g\\Delta_t^n$ should be  $\\eta_g\\Delta_t$.\n* At the beginning of 3.2, the notation $m_\\tau(\\delta)$ is not clarified. ", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes gradient maks averaging (GMA) for Federated Learning (FL). Specifically, the method computes an element-wise soft mask, which has values between 0 and 1, and applies it to the aggregated gradients from all clients before applying the update to the global model.  The mask is calculated based on the principle of agreement among clients, inspired by related works in out-of-distribution (OOD) generalization. The paper applies the method to several existing FL algorithms and demonstrates better performance than the original algorithms on various FL settings including artificial IID and Non-IID datasets, real-world FL datasets, and OOD datasets with spurious features. ", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "* The authors should provide more clear writing of the theory component and discuss the motivations and implications of theoretical results. It is ok to have assumptions but still, it is not clear what the convergence results indicate under the current assumptions. Also, proposition 4.9 is very confusing in its current state. (**critical**)\n\n* Please fix the typos mentioned in the previous section. (**improving presentation**)", "broader_impact_concerns": "The reviewer does not see potential ethical implications of this submission. ", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Main comments by the reviewers:
- the idea is "well-motivated" [ffYZ], "novel" [ffYZ, D7tq] and simple and "easy to be adopted" [54rX]. It "can be applied to many existing federated learning algorithms" [54rX].
- the theory is clear and "beautifully put" according to [D7tq] but on the other hand, the "writing of the theory section is also confusing" [ffYZ].
- the theory is limited: "the authors only provide the convergence analysis for a special case" [54rX], that is, i.i.d. Bernoulli masks. Thus, they do not completely explain why the method works. "While the theoretical difficulty of analyzing the original algorithm can be understood, this nevertheless weakens the contribution" [ffYZ]. This is acknowledged by the authors. They highlight the difficulty of studying the general case.
- the experiments are extensive [ffYZ] and honest, as they show "the pros and cons of the proposed method" [D7tq]. I will note a claim by [54rX]: "the improvement over naive averaging mechanism doesn’t seem to be significant", that seems to be in contradiction with most of the experimental results (if we interpret 'significant' as 'statistically significant') and with the claims of the other reviewers: " the proposed method outperforms the AVG aggeration method" [D7tq], "it is intriguing to see that masking can improve both federated performance and also OOD performance in a federated setting" [ffYZ].

The few typos pointed out by [ffYZ, D7tq] were fixed by the authors. They also slightly re-organized the theoretical part to address the comment on clarity by [ffYZ].

As a conclusion, the revised version of the paper proposes a new method that is easy to use in practice, well supported by experiments. They proposed a preliminary theory that, despite its clarity, does not explain why the method works in its general form. For [D7tq], the practical advantages are enough to support fully publication. [ffYZ] has more reservation because of the theory. Still, this reviewer is also leaning towards acceptance. I assume that he/she considers that the experimental results are enough to support the publication of the paper. On the other hand, [54rX] is leaning towards rejection. However, he/she did not reply to the authors rebuttal, and did not provide any explanation motivating this decision. I can only infer that for this reviewer, the experimental results do not outweigh the lack of theory. As both reviewers and myself disagree with [54rX]'s evaluation of the empirical result, I feel safe enough to recommend to accept this paper (I also note that, would [54rX] had strong objections to the acceptance of the paper, he/she would have stated it clearly anyway). Overall, it's a promising method well supported by a strong experimental analysis. A full theory seems to be quite challenging, but we can hope this problem will receive attention in the coming months.
Decision: Accept (poster)