Original: F ED FA: F EDERATED F EATURE A UGMENTATION
Revised: F ED FA: F EDERATED F EATURE A UGMENTATION

Original: Anonymous authors

Original: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: Federated learning is a distributed paradigm that allows multiple parties to collaboratively train deep models without exchanging the raw data.
Revised: Federated learning is a distributed paradigm that allows multiple parties to collaboratively train deep models without exchanging the raw data.

Original: However, the data distribution among clients is naturally non-i.i.d., which leads to severe degradation of the learnt model.
Revised: However, the data distribution among clients is naturally non-i.i.d., which leads to severe degradation of the learnt model.

Original: The primary goal of this paper is to develop a robust federated learning algorithm to address feature shift in clients’ samples, which can be caused by various factors, e .
Revised: The primary goal of this paper is to develop a robust federated learning algorithm to address feature shift in clients’ samples, which can be caused by various factors, e .

Original: g ., acquisition differences in medical imaging.
Revised: g ., acquisition differences in medical imaging.

Original: To reach this goal, we propose F ED FA to tackle federated learning from a distinct perspective of federated feature augmentation.
Revised: To reach this goal, we propose F ED FA to tackle federated learning from a distinct perspective of federated feature augmentation.

Original: F ED FA is based on a major insight that instance-level feature statistics ( i .
Revised: F ED FA is based on a major insight that each client’s data distribution can be characterized by statistics ( i .

Original: e ., mean and standard deviation)represent a special type of “features” that encodes domain-speciﬁc characteristics; hence, proper manipulations of the localfeature statistics  in the federation may beget novel domains, which can potentially alleviate local feature shift and beneﬁt collaborative learning.
Revised: e ., mean and standard deviation) of latent features; and it is likely to manipulate these local statistics globally , i . e ., based on information in the entire federation, to let clients have a better sense of the underlying distribution and therefore alleviate local data bias.

Original: Based on this insight, we model each  feature statistic probabilistically augmented from a Gaussian distribution, with the mean corresponding to the original statistic andthe variance quantifying the augmentation scope, from which novel feature statistics can be drawn to fulﬁll augmentation.
Revised: Based on this insight, we propose to augment each local feature statistic probabilistically based on a normal distribution, whose mean is the original statistic and variance quantiﬁes the augmentation scope.

Original: Key to our approach is the determination of a meaningful Gaussian variance, which is accomplished by taking into account not only biased data of each individual client, but also underlying feature statistics characterized by all participating clients.
Revised: Key to our approach is the determination of a meaningful Gaussian variance, which is accomplished by taking into account not only biased data of each individual client, but also underlying feature statistics characterized by all participating clients.

Original: We demonstrate through extensive experiments that F ED FA can signiﬁcantly advance federated learning in diverse scenarios. The code will be released.
Revised: We offer both theoretical and empirical justiﬁcations to verify the effectiveness of F ED FA. Our code is available at https://github.com/tfzhou/FedFA .

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: Federated learning (FL) (Kone ˇ cn ` y et al., 2016) is an emerging collaborative training framework that enables training on decentralized data residing devices like mobile phones.
Revised: Federated learning (FL) (Kone ˇ cn ` y et al., 2016) is an emerging collaborative training framework that enables training on decentralized data residing devices like mobile phones.

Original: It comes with the promise of training centralized models using local data points such that the privacy of participating devices is preserved, and has attracted signiﬁcant attention in critical ﬁelds like healthcare or ﬁnance.
Revised: It comes with the promise of training centralized models using local data points such that the privacy of participating devices is preserved, and has attracted signiﬁcant attention in critical ﬁelds like healthcare or ﬁnance.

Original: Since data come from different users, it is inevitable that the data of each user have a different underlying distribution, incurring large heterogeneity (non-iid-ness) among users’ data.
Revised: Since data come from different users, it is inevitable that the data of each user have a different underlying distribution, incurring large heterogeneity (non-iid-ness) among users’ data.

Original: In this work, we focus on feature shift (Li et al., 2020b), which is common in many real-world cases, like medical data acquired from different  devices or natural image collected in diverse environments.
Revised: In this work, we focus on feature shift (Li et al., 2020b), which is common in many real-world cases, like medical data acquired from different medical devices or natural image collected in diverse environments.

Original: While the problem of feature shift has been studied in classical centralized learning tasks like domain generalization, little is understood how to tackle it in federated learning.
Revised: While the problem of feature shift has been studied in classical centralized learning tasks like domain generalization, little is understood how to tackle it in federated learning.

Original: (Li et al., 2020b; Reisizadeh et al., 2020; Jiang et al., 2022; Liu et al., 2020a) are rare exceptions.
Revised: (Li et al., 2020b; Reisizadeh et al., 2020; Jiang et al., 2022; Liu et al., 2020a) are rare exceptions.

Original: F ED R OBUST (Reisizadeh et al., 2020) and F ED BN (Li et al., 2020b) solve the problem through client-dependent learning by either ﬁtting the shift with a client-speciﬁc afﬁne distribution or learning unique BN parameters for each client.
Revised: F ED R OBUST (Reisizadeh et al., 2020) and F ED BN (Li et al., 2020b) solve the problem through client-dependent learning by either ﬁtting the shift with a client-speciﬁc afﬁne distribution or learning unique BN parameters for each client.

Original: However, these algorithms may still suffer signiﬁcant local dataset bias.
Revised: However, these algorithms may still suffer signiﬁcant local dataset bias.

Original: Other works (Liu et al., 2020a; Jiang et al.,  2022) learn robust models by adopting Sharpness Aware Minimization (SAM) (Foret et al., 2021) as the local optimizer, which, however, doubles the computational cost compared to SGD or Adam.
Revised: Other works (Qu et al., 2022; Jiang et al., 2022; Caldarola et al., 2022) learn robust models by adopting Sharpness Aware Minimization (SAM) (Foret et al., 2021) as the local optimizer, which, however, doubles the computational cost compared to SGD or Adam.

Original: In addition to model optimization, F ED H ARMO (Jiang et al., 2022) has investigated  image normalization techniques to mitigate feature shift in medical domains.
Revised: In addition to model optimization, F ED H ARMO (Jiang et al., 2022) has investigated specialized image normalization techniques to mitigate feature shift in medical domains.

Original: Despite the progress, there leaves an alternative space – data augmentation – largely unexplored in federated learning scenarios, even that it has been extensively investigated in centralized setting to impose regularization and improve generalizibility (Wang et al., 2019; Zhou et al., 2021; Zhang et al., 2018).
Revised: Despite the progress, there leaves an alternative space – data augmentation – largely unexplored in federated learning, even though it has been extensively studied in centralized setting to impose regularization and improve generalizibility (Zhou et al., 2021; Zhang et al., 2018).

Original: While seemingly straightforward, it is non-trivial to perform effective data augmentation in federated learning because users have no direct access to external data of other users.
Revised: While seemingly straightforward, it is non-trivial to perform effective data augmentation in federated learning because users have no direct access to external data of other users.

Original: Simply applying conventional augmentation techniques to each client is sub-optimal since without injecting global information, augmented samples will most likely still suffer local dataset bias.
Revised: Simply applying conventional augmentation techniques to each client is sub-optimal since without injecting global information, augmented samples will most likely still suffer local dataset bias.

Original: To address this, F EDM IX (Yoon et al., 2021) generalizes Mixup (Zhang et al., 2018) into federated learning. It follows the basic idea of vanilla Mixup, and performs cross-client sample mixup by exchanging averaged data across clients.
Revised: To address this, F EDM IX (Yoon et al., 2021) generalizes M IXUP (Zhang et al., 2018) into federated learning, by mixing averaged data across clients.

Original: The method performs augmentation in the input level, which is naturally weak to create  meaningful semantic transformations, e .
Revised: The method performs augmentation in the input level, which is naturally weak to create complicated and meaningful semantic transformations, e .

Original: g ., make-bespectacled .
Revised: g ., make-bespectacled .

Original: Moreover, allowing exchange of averaged data will suffer certain levels of privacy issues.
Revised: Moreover, allowing exchange of averaged data will suffer certain levels of privacy issues.

Original: In this paper, we introduce a novel feature augmentation technique F ED FA for federated learning, based on a probabilistic modeling of feature statistics ( i .
Revised: In this work, we introduce a novel federation-aware augmentation technique, called FedFA, into federated learning.

Original: e ., mean and standard deviation) over all participating clients.

Original: F ED FA is based on the insight thatfeature statistics  can convey essential domain-aware characteristics (Huang & Belongie, 2017;
Revised: F ED FA is based on the insight that statistics of latent features can capture essential domain-aware characteristics (Huang & Belongie, 2017;

Original: Zhou et al., 2021; Li et al., 2021a; 2022), and can be treated as “features of domain”.
Revised: Zhou et al., 2021; Li et al., 2022a;b; 2021a), thus can be treated as “features of participating client”.

Original: Accordingly, we argue that  no matter the shift of each local data distribution from the true distribution, or local distribution differences among clients, even test-time distribution shift, can be interpreted as the shift of feature statistics.
Revised: Accordingly, we argue that the problem of feature shift in FL, no matter the shift of each local data distribution from the underlying distribution, or local distribution differences among clients, even test-time distribution shift, can be interpreted as the shift of feature statistics.

Original: This motivates us to turn the problem of tackling feature shift into tackling feature statistic shift, for which we propose F ED FA to perform feature statistic augmentation within each local client by incorporatingmore universalfeature statisticinformation characterized by all clients in the federation. 
Revised: This motivates us to directly addressing local feature statistic shift by incorporating universal statistic characterized by all participants in the federation. F ED FA instantiates the idea by online augmenting feature statistics of each sample during local model training, so as to make the model robust to certain changes of “features of participating client”.

Original: In particular, we model the augmentation of each feature statistic in a probabilistic manner via a multi-variate Gaussian distribution.
Revised: Concretely, we model the augmentation procedure in a probabilistic manner via a multivariate Gaussian distribution.

Original: The Gaussian mean corresponds to the original statistic, and variance reﬂects the potential local distribution shift.
Revised: The Gaussian mean is ﬁxed to the original statistic, and variance reﬂects the potential local distribution shift.

Original: In this manner,we can sample novel statistics  effortlessly  from the Gaussian distribution as augmented “features of domain”.
Revised: In this manner, novel statistics can be effortlessly synthesized by drawing samples from the Gaussian distribution.

Original: For effective augmentation, we determine a proper variance based on  variances of feature statistics within each client, which are regulated by universal variancesof feature statistics characterized by all participating clients.
Revised: For effective augmentation, we determine a reasonable variance based on not only variances of feature statistics within each client, but also universal variances characterized by all participating clients.

Original: The augmentation in F ED FA allows the local model to be trained over samples drawn from more diverse feature distributions, facilitating local distribution shift alleviation and client-invariant representation learning, eventually contributing to a better global model.
Revised: The augmentation in F ED FA allows each local model to be trained over samples drawn from more diverse feature distributions, facilitating local distribution shift alleviation and client-invariant representation learning, eventually contributing to a better global model.

Original: F ED FA is a conceptually simple but surprisingly effective method.
Revised: ED FA is a conceptually simple but surprisingly effective method.

Original: It is non-parametric, requires negligible additional computational resources, and can be seamlessly incorporated into arbitrary CNN architectures.
Revised: It is non-parametric, requires negligible additional computation and communication costs, and can be seamlessly incorporated into arbitrary CNN architectures.

Original: We provide extensive empirical evaluations on competitive image classiﬁcation and medical image segmentation benchmarks, and the results demonstrate that: i) F ED FA works favorably with extremely small local datasets; ii) it shows remarkable generalization performance to unseen test clients outside of the federation; iii) it outperforms traditional data augmentation techniques by solid margins, and can complement them quite well in the federated learning setup.  R ELATED W ORK Federated Learning. Recent years have witnessed tremendous progress in federated learning (Kone ˇ cn ` y et al., 2016), which opens the door for privacy-preserving deep learning (Shokri &
Revised: We propose both theoretical and empirical insights.

Original: Shmatikov, 2015), i .

Original: e ., train a global model on distributed datasets without disclosing private data information. F ED A VG (McMahan et al., 2017) is a milestone; it trains localmodels independently in multiple clients and then averages the resulting model updates via a central server once in a while. However, F ED A VG is designed for i.i.d. data and suffers in statistical accuracy or even diverge if deployed over non-i.i.d.
Revised: Theoretically, we show that F ED FA implicitly introduces regularization to local model learning by regularizing the gradients of latent representations, weighted by variances of feature statistics estimated from the entire federation.

Original: client samples.

Original: To address this issue, numerous efforts have been devoted to handling heterogeneous federated environments, by, for example, adding a dynamic regularizer to local objectives in F ED P ROX (Li et al., 2020a) and

Original: F ED D YN (Acar et al., 2021), correcting client drift through variance reduction in SCAFFOLD (Karimireddy et al., 2020), adaptive server optimization in F ED O PT (Reddi et al., 2021), local batch normalization in F ED BN (Li et al., 2020b), or training a perturbed loss in F ED SAM (Qu et al., 2022) and F ED R OBUST (Reisizadeh et al., 2020).

Original: F ED M IX (Yoon et al., 2021), as far as we know, is the only existing method that solves federated learning based on data augmentation.

Original: It adapts the well-known Mixup algorithm (Zhang et al., 2018) from centralized learning into the federated learning scenario.

Original: Nevertheless, F ED M IX requires exchanging local data (or averaged version) across clients for data interpolation, thereby suffering privacy issues.

Original: In addition, F ED M IX operates on the input level, while our approach focuses on latent feature statistic augmentation.

Original: Since deeper representations tend to disentangle the underlying factors of variation better (Bengio et al., 2013), traversing along latent space will potentially make our method encounter more realistic samples.

Original: This is supported by the fact that F ED FA achieves consistent performance improvements over F ED M IX in diverse scenarios (see §4).

Original: Data Augmentation.

Original: Data augmentation has a long and rich history in machine learning.

Original: Early studies (Sch¨olkopf et al., 1996; Kukaˇcka et al., 2017) focus on label-preserving transformations to employ regularization via data, alleviating overﬁtting and improving generalization.

Original: For image data, some techniques, like random horizontal ﬂipping and cropping are commonly used for training of advanced neural networks (He et al., 2016).

Original: In addition, there is a recent trend for label-perturbing augmentation, e . g .,

Original: Mixup (Zhang et al., 2018) or Cutmix (Yun et al., 2019).

Original: Separate from these input-level augmentation techniques are feature augmentation methods (Verma et al., 2019;

Original: Li et al., 2021a; Zhou et al., 2021) that make augmentation in latent feature space. These various data augmentation techniques have shown great successes to learn domain-invariant models in thecentralized setup. Our method is an instance of label-preserving feature augmentation, designed for federated learning.
Revised: Empirically, we demonstrate that F ED FA (1) works favorably with extremely small local datasets; (2) shows remarkable generalization performance to unseen test clients outside of the federation; (3) outperforms traditional data augmentation techniques by solid margins, and can complement them quite well in the federated learning setup.

Original: It is inspired by recent efforts on implicit feature augmentation (Li et al., 2021a;

Original: Zhou et al., 2021) that synthesize samples of novel domains by swapping or mixing instance-level feature statistics .

Original: In these works, feature statistics are treated as ‘features’, which capture essential domainspeciﬁc characteristics.

Original: In F ED FA, instead of linearly combining feature statistics, we estimate appropriate variances of feature statistics from a federated perspective, and draw novel statistics probablistically from a distribution centered on old statistics, while spanning with the variances.

Original: F ED FA avoids statistics mixing of instances from different clients, as done in (Li et al., 2021a;

Original: Zhou et al., 2021) and recent federated method F ED M IX , thus can better preserve data privacy.

Original:  O UR A PPROACH 3. P RELIMINARY : F EDERATED L EARNING
Revised:  O UR A PPROACH 2. P RELIMINARY : F EDERATED L EARNING

Original: We assume a standard federated learning setup with a server that can transmit and receive messages from M client devices.
Revised: We assume a standard federated learning setup with a server that can transmit and receive messages from M client devices.

Original: Each client m ∈
Revised: Each client m ∈

Original: [ M ] has access to N m training instances (cid:8) ( x i , y i ) (cid:9)
Revised: [ M ] has access to N m training instances (cid:8) ( x i , y i ) (cid:9)

Original: Nm i =1 in the form of image x i ∈X and corresponding labels
Revised: Nm i =1 in the form of image x i ∈X and corresponding labels

Original: y
Revised: y

Original: i ∈Y that are drawn i.i.d.
Revised: i ∈Y that are drawn i.i.d.

Original: from a device-indexed joint distribution, i .
Revised: from a device-indexed joint distribution, i .

Original: e ., ( x i , y i ) ∼ P m ( x, y ) .
Revised: e ., ( x i , y i ) ∼ P m ( x, y ) .

Original: The goal of standard federated learning is to learn a deep neural network: f ( w g , w h ) (cid:44) g
Revised: The goal of standard federated learning is to learn a deep neural network: f ( w g , w h ) (cid:44) g

Original: ( w g ) ◦
Revised: ( w g ) ◦

Original: h ( w h ) , where h : X → Z is a feature extractor with K convolutional stages:
Revised: h ( w h ) , where h : X → Z is a feature extractor with K convolutional stages:

Original: h = h K ◦ h K − 1 ◦ · · · ◦ h 1 , and g : Z → Y is a classiﬁer.
Revised: h = h K ◦ h K − 1 ◦ · · · ◦ h 1 , and g : Z → Y is a classiﬁer.

Original: To learn network parameters w = { w g , w h } , the empirical risk minimization (ERM) is widely used:
Revised: To learn network parameters w = { w g , w h } , the empirical risk minimization (ERM) is widely used:

Original: [Equation]
Revised: [Equation]

Original: where L ERM m ( w [Equation] =
Revised: where L ERM m ( w ) =

Original: E ( x i [Equation] y i 
Revised: E ( x i ,y i )

Original: [Equation]
Revised: [Equation]

Original: ∼ P m [ (cid:96)
Revised: P m [ (cid:96)

Original: [Equation]
Revised: [Equation]

Original: ( g ◦ h ( x

Original: [Equation]
Revised: [Equation]

Original: y i
Revised: h ( x

Original: [Equation]
Revised: [Equation]

Original:  w [Equation] [Equation]
Revised: ) , y i ; w )] .

Original: Here the global objective L ERM is decomposable as a sum of device-level empirical loss objectives ( i .
Revised: Here the global objective L ERM is decomposable as a sum of device-level empirical loss objectives ( i .

Original: e ., {L ERM m } m ).
Revised: e ., {L ERM m } m ).

Original: Each L  m is computed based on a per-data loss function (cid:96) i .
Revised: Each L ERM m is computed based on a per-data loss function (cid:96) i .

Original: Due to the separation of clients’ data, L ERM ( w ) cannot be solved directly.
Revised: Due to the separation of clients’ data, L ERM ( w ) cannot be solved directly.

Original: F ED A VG (McMahan et al., 2017) is a leading algorithm to address this.
Revised: F ED A VG (McMahan et al., 2017) is a leading algorithm to address this.

Original: It starts with client training of all the clients in parallel, with each client optimizing L ERM m independently.
Revised: It starts with client training of all the clients in parallel, with each client optimizing L ERM m independently.

Original: After local client training, F ED A VG performs model aggregation to average all client models into a updated global model, which will be distributed back to the clients for the next round of client training.
Revised: After local client training, F ED A VG performs model aggregation to average all client models into a updated global model, which will be distributed back to the clients for the next round of client training.

Original: Here the client training objective in F ED A VG is equivalent to empirically approximating the local distribution P m by a ﬁnite N m number of examples, i .
Revised: Here the client training objective in F ED A VG is equivalent to empirically approximating the local distribution P m by a ﬁnite N m number of examples, i .

Original: e .,
Revised: e .,

Original: P em ( x, y )
Revised: P e m ( x, y )

Original: = 1 / N m (cid:80) N m i =

Revised: = 1 / N m (cid:80) N m i =

Original: δ
Revised: δ

Original: ( x = x i , y = y i ) , where δ ( x = x i , y = y i ) is a Dirac mass centered at ( x i , y i ) .
Revised: ( x = x i , y = y i ) , where δ ( x = x i , y = y i ) is a Dirac mass centered at ( x i , y i ) .

Original: 3. M OTIVATION
Revised: 2. M OTIVATION

Original: While the ERM-based formulation has achieved great success, it is straightforward to see that the solution would strongly depend on how each approximated local distribution P
Revised: While the ERM-based formulation has achieved great success, it is straightforward to see that the solution would strongly depend on how each approximated local distribution P

Original: em mimics the underlying universal distribution P .
Revised: em mimics the underlying universal distribution P .

Original: In real-world federated learning setup however, in all but trivial cases each P em exhibits a unique distribution shift from P , which causes not only inconsistency between local and global empirical losses (Acar et al., 2021;
Revised: In real-world federated learning setup however, in all but trivial cases each P em exhibits a unique distribution shift from P , which causes not only inconsistency between local and global empirical losses (Acar et al., 2021;

Original: Wang et al., 2020), but also generalization issues (Yuan et al., 2022).
Revised: Wang et al., 2020), but also generalization issues (Yuan et al., 2022).

Original: In this work, we circumvent this issue by ﬁtting each local dataset a richer distribution (instead of the delta distribution) in the vicinal region of each sample ( x i , y i ) so as to estimate a more informed risk.
Revised: In this work, we circumvent this issue by ﬁtting each local dataset a richer distribution (instead of the delta distribution) in the vicinal region of each sample ( x i , y i ) so as to estimate a more informed risk.

Original: This is precisely the principle behind vicinal risk minimization (VRM) (Chapelle et al., 2000).
Revised: This is precisely the principle behind vicinal risk minimization (VRM) (Chapelle et al., 2000).

Original: Particularly, for data point ( x i , y i ) , a vicinity distribution V m (ˆ x i , ˆ y i | x i , y i ) is deﬁned, from which novel virtual samples can be generated to enlarge the support of the local data distribution.
Revised: Particularly, for data point ( x i , y i ) , a vicinity distribution V m (ˆ x i , ˆ y i | x i , y i ) is deﬁned, from which novel virtual samples can be generated to enlarge the support of the local data distribution.

Original: In this way, we obtain an improved approximation of P m as
Revised: In this way, we obtain an improved approximation of P m as

Original: P vm = 1 / N m (cid:80) N m i
Revised: P vm = 1 / N m (cid:80) N m i

Original: =1 V m (ˆ x i , ˆ y i | x i , y i ) .
Revised: =1 V m (ˆ x i , ˆ y i | x i , y i ) .

Original: In centralized learning scenarios, various successful instances of V m , e . g ., Mixup (Zhang et al., 2018), Cutmix (Yun et al., 2019), have been developed.
Revised: In centralized learning scenarios, various successful instances of V m , e . g . , M IXUP (Zhang et al., 2018), C UTMIX (Yun et al., 2019), have been developed.

Original: Simply applying them to local clients, though allowing for performance improvements (see Table 6), is suboptimal since, without injecting any global information, P vm only provides a better approximation to the local distribution P m , rather than the true distribution P .
Revised: Simply applying them to local clients, though allowing for performance improvements (see Table 6), is suboptimal since, without injecting any global information, P vm only provides a better approximation to the local distribution P m , rather than the true distribution P .

Original: We solve this by introducing a dedicated method F ED FA to estimate more reasonable V m in federated learning.
Revised: We solve this by introducing a dedicated method F ED FA to estimate more reasonable V m in federated learning.

Original: 3. F ED FA: F EDERATED F EATURE A UGMENTATION
Revised: 2. F ED FA: F EDERATED F EATURE A UGMENTATION

Original: F ED FA belongs to the family of label-preserving feature augmentation (Xie et al., 2020).
Revised: F ED FA belongs to the family of label-preserving feature augmentation (Xie et al., 2020).

Original: During training, it estimates a vicinity distribution V km at each layer h k to augment hidden features in client m .
Revised: During training, it estimates a vicinity distribution V km at each layer h k to augment hidden features in client m .

Original: Considering X km ∈ R B × C × H × W as the intermediate feature representation of B mini-batch images, with spatial size ( H × W ) and channel number ( C ), and Y km as corresponding label.
Revised: Considering X km ∈ R B × C × H × W as the intermediate feature representation of B mini-batch images, with spatial size ( H × W ) and channel number ( C ), and Y km as corresponding label.

Original: V km is label-preserving in the sense that V km ( ˆ X km ,
Revised: V km is label-preserving in the sense that V km ( ˆ X km ,

Original: ˆ Y m | X km , Y m ) (cid:44) V km ( ˆ X km | X km ) δ ( ˆ Y m =
Revised: ˆ Y m | X km , Y m ) (cid:44) V km ( ˆ X km | X km ) δ ( ˆ Y m =

Original: Y m ) , i .
Revised: Y m ) , i .

Original: e ., it only transforms the latent feature X km to ˆ X km , but preserves the original label Y km .
Revised: e ., it only transforms the latent feature X km to ˆ X km , but preserves the original label Y km .

Original: Probabilistic Feature Statistic Augmentation.
Revised: 2.3. F EDERATED F EATURE A UGMENTATION FROM A P ROBABILISTIC V IEW

Original: Instead of pursuing an explicit modelingof V km ( ˆ X km | X km ) , our method performs  feature augmentationin an implicit manner by augmenting channel-wise feature statistics.
Revised: Instead of explicitly modeling V km ( ˆ X km | X km ) , our method performs implicit feature augmentation by manipulating channel-wise feature statistics.

Original: Speciﬁcally, for X km , its channel-wise statistics, i .
Revised: Speciﬁcally, for X km , its channel-wise statistics, i .

Original: e ., mean µ km and standard deviation σ k m are given as follows:
Revised: e ., mean µ km and standard deviation σ km are given as follows:

Original: [Equation]
Revised: [Equation]

Original: where X k, ( h,w ) m ∈ R B × C represents features at spatial location ( h, w ) .
Revised: where X k, ( h,w ) m ∈ R B × C represents features at spatial location ( h, w ) .

Original: As the abstract of latent features, these statistics carry domain-speciﬁc information ( e .
Revised: As the abstract of latent features, these statistics carry domain-speciﬁc information ( e .

Original: g ., style).
Revised: g ., style).

Original: They are instrumental to image generation (Huang & Belongie, 2017), and have been recently used for data augmentation in deep recognition (Li et al., 2021a).
Revised: They are instrumental to image generation (Huang & Belongie, 2017), and have been recently used for data augmentation in image recognition (Li et al., 2021a).

Original: In heterogeneous federated learning scenarios, the feature statistics among local clients will be inconsistent, and exhibit uncertain feature statistic shifts from the statistics of the true distribution.
Revised: In heterogeneous federated learning scenarios, the feature statistics among local clients will be inconsistent, and exhibit uncertain feature statistic shifts from the statistics of the true distribution.

Original: Our method explicitly captures such shift via probabilistic modeling.
Revised: Our method explicitly captures such shift via probabilistic modeling.

Original: Concretely, instead of representing each feature X km with deterministic statistics { µ km , σ km } , we hypothesize that the feature is conditioned on probabilistic statistics { ˆ µ km , ˆ σ km } , which are sampled around the original statistics based on a multi-variate Gaussian distribution, i .
Revised: Concretely, instead of representing each feature X km with deterministic statistics { µ km , σ km } , we hypothesize that the feature is conditioned on probabilistic statistics { ˆ µ km , ˆ σ km } , which are sampled around the original statistics based on a multi-variate Gaussian distribution, i .

Original: e .,
Revised: e .,

Original: ˆ µ km ∼N ( µ km , ˆΣ 2 µ km ) and ˆ σ km ∼N ( σ km , ˆΣ 2 σ km ) , where each Gaussian’s center corresponds to the original statistic, and the variance is expected to capture the potential feature statistic shift from the true distribution.
Revised: ˆ µ km ∼N ( µ km , ˆΣ 2 µ km ) and ˆ σ km ∼N ( σ km , ˆΣ 2 σ km ) , where each Gaussian’s center corresponds to the original statistic, and the variance is expected to capture the potential feature statistic shift from the true distribution.

Original: However, determining an appropriate variant range is challenging since each client has access to only the data itself, and has no sense of its statistic bias. Recent research (Wang et al., 2019; Upchurch et al., 2017) shows that deep feature space contain many semantic directions, and feature variances provide a reasonable measurement to potential meaningful semantic changes along the directions. This motivates us to estimate the variance of Gaussian from the variance of feature statistics .
Revised: Our core goal is thus to estimate proper variances ˆΣ 2 µ km / ˆΣ 2 σ km for reasonable and informative augmentation.

Original: Client Speciﬁc Statistic Variances.
Revised: Client-speciﬁc Statistic Variances.

Original: In client-side , we compute client-speciﬁc variances of feature statistics based on the information within each mini-batch:
Revised: In client-side , we compute client-speciﬁc variances of feature statistics based on the information within each mini-batch:

Original: [Equation]
Revised: [Equation]

Original: where Σ 2 µ km and Σ 2 σ km denote the variance of feature mean µ km and standard deviation σ km that are speciﬁc to each client.
Revised: where Σ 2 µ km and Σ 2 σ km denote the variance of feature mean µ km and standard deviation σ km that are speciﬁc to each client.

Original: Each value in Σ 2 µ km or Σ 2 σ km is the variance of feature statistics in a particular channel, and its magnitude manifests how the channel will change potentially in the particular client.
Revised: Each value in Σ 2 µ km or Σ 2 σ km is the variance of feature statistics in a particular channel, and its magnitude manifests how the channel will change potentially in the feature space.

Original: Client Sharing Statistic Variances.
Revised: Client-sharing Statistic Variances.

Original: The client speciﬁc statistic variances are solely computed based on the data in each  client, and thus likely biased due to local dataset bias.
Revised: The client-speciﬁc variances are solely computed based on the data in each individual client, and thus likely biased due to local dataset bias.

Original: To solve this, we further estimate client sharing feature statistic variances taking information of all clients into account.
Revised: To solve this, we further estimate client-sharing feature statistic variances taking information of all clients into account.

Original: Particularly, we maintain a momentum version of feature statistics for each client, which are online estimated during training:
Revised: Particularly, we maintain a momentum version of feature statistics for each client, which are online estimated during training:

Original: [Equation]
Revised: [Equation]

Original: where ¯ µ km and ¯ σ km are the momentum updated feature statistics of layer h k in client m , and they are initialized as C -dimensional all-zero and all-one vectors, respectively.
Revised: where ¯ µ km and ¯ σ km are the momentum updated feature statistics of layer h k in client m , and they are initialized as C -dimensional all-zero and all-one vectors, respectively.

Original: α is a momentum coefﬁcient.
Revised: α is a momentum coefﬁcient.

Original: We set a same α for both updating, and found no beneﬁt to set it differently.
Revised: We set a same α for both updating, and found no beneﬁt to set it differently.

Original: In each communication, these accumulated local feature statistics are sent to the server along with model parameters.
Revised: In each communication, these accumulated local feature statistics are sent to the server along with model parameters.

Original: Let ¯ µ k
Revised: Let ¯ µ k

Original: =
Revised: =

Original: [¯ µ k 1 , · · · , ¯ µ kM ] ∈ R M × C and ¯ σ k 
Revised: [¯ µ k 1 , . . , ¯ µ kM ] ∈ R M × C and ¯ σ k =

Original: = [¯ σ k 1 , · · · , ¯ σ kM ] ∈ R M × C denote collections of accumulated feature statistics of all clients, the client sharing statistic variances are determined in server-side by:
Revised: [¯ σ k 1 , . . , ¯ σ kM ] ∈ R M × C denote collections of accumulated feature statistics of all clients, the client sharing statistic variances are determined in server-side by:

Original: [Equation]
Revised: [Equation]

Original: In addition, it is intuitive that some channels are more potentially to change than others, and it will be favorable to highlight these channels to enable a sufﬁcient and reasonable exploration of the space of feature statistics.
Revised: In addition, it is intuitive that some channels are more potentially to change than others, and it will be favorable to highlight these channels to enable a sufﬁcient and reasonable exploration of the space of feature statistics.

Original: To this end, we further modulate client sharing estimations with a
Revised: To this end, we further modulate client sharing estimations with a

Original: Student’s t-distribution (Student, 1908; Van der Maaten & Hinton, 2008) with one degree of freedom to convert the variances to probabilities.
Revised: Student’s t-distribution (Student, 1908; Van der Maaten & Hinton, 2008) with one degree of freedom to convert the variances to probabilities.

Original: The t-distribution has heavier tails than other alternatives such as Gaussian distribution, allowing to highlight the channels with larger statistic variance, at the same time, avoiding overly penalizing the others.
Revised: The t-distribution has heavier tails than other alternatives such as Gaussian distribution, allowing to highlight the channels with larger statistic variance, at the same time, avoiding overly penalizing the others.

Original: Formally, denote Σ 2 , ( j ) µ k and Σ 2 , ( j ) σ k as the shared variances of the j th channel in Σ 2 µ k and Σ 2 σ k
Revised: Formally, denote Σ 2 , ( j ) µ k and Σ 2 , ( j ) σ k as the shared variances of the j th channel in Σ 2 µ k and Σ 2 σ k

Original: (Eq. 5), respectively.
Revised: (Eq. 5), respectively.

Original: They are modulated by the t-distribution as follows:
Revised: They are modulated by the t-distribution as follows:

Original: [Equation]
Revised: [Equation]

Original: where γ ( j ) µ k and γ ( j ) σ k refer to the modulated variances of the j -th channel.
Revised: where γ ( j ) µ k and γ ( j ) σ k refer to the modulated variances of the j -th channel.

Original: By applying Eq.
Revised: By applying Eq.

Original: 6 to each channel separately, we obtain γ µ k
Revised: 6 to each channel separately, we obtain γ µ k

Original: =
Revised: =

Original: [ γ (1) µ k , . .
Revised: [ γ (1) µ k , . .

Original: , γ ( C ) µ k ] ∈ R C and γ σ k =
Revised: , γ ( C ) µ k ] ∈ R C and γ σ k =

Original: [ γ (1) σ k , . .
Revised: [ γ (1) σ k , . .

Original: , γ ( C ) σ k ] ∈ R C as modulated statistic variances of all feature channels at layer h k .
Revised: , γ ( C ) σ k ] ∈ R C as modulated statistic variances of all feature channels at layer h k .

Original: In this way, the channels with large values in Σ 2 µ k (or Σ 2 σ k ) will be assigned with much higher importance in γ µ k (or γ σ k ) than other channels, allowing for more extensive augmentation along those directions.
Revised: In this way, the channels with large values in Σ 2 µ k (or Σ 2 σ k ) will be assigned with much higher importance in γ µ k (or γ σ k ) than other channels, allowing for more extensive augmentation along those directions.

Original: Adaptive Variance Fusion.
Revised: Adaptive Variance Fusion.

Original: The modulated client sharing estimations { γ µ k , γ σ k } provide a quantiﬁcation of distribution difference among clients, and larger values imply potentials of more signiﬁcant changes of corresponding channels in the true feature statistic space.
Revised: The modulated client sharing estimations { γ µ k , γ σ k } provide a quantiﬁcation of distribution difference among clients, and larger values imply potentials of more signiﬁcant changes of corresponding channels in the true feature statistic space.

Original: Therefore, for each client, we weight the client speciﬁc statistic variances { Σ 2 µ km , Σ 2 σ km } by { γ µ k , γ σ k } , so that each client has a sense of such difference.
Revised: Therefore, for each client, we weight the client speciﬁc statistic variances { Σ 2 µ km , Σ 2 σ km } by { γ µ k , γ σ k } , so that each client has a sense of such difference.

Original: To avoid overly modiﬁcation of client speciﬁc statistic variances, we add a residual connection for fusion, yielding an estimation of Gaussian ranges as:
Revised: To avoid overly modiﬁcation of client speciﬁc statistic variances, we add a residual connection for fusion, yielding an estimation of Gaussian ranges as:

Original: [Equation]
Revised: [Equation]

Original: where (cid:12) denotes the Hadamard product.
Revised: where (cid:12) denotes the Hadamard product.

Original: Implementation of Feature Augmentation.
Revised: Implementation of Feature Augmentation.

Original: After establishing the Gaussian distribution, we synthesize novel feature ˆ X km in the vicinity of X km as follows:
Revised: After establishing the Gaussian distribution, we synthesize novel feature ˆ X km in the vicinity of X km as follows:

Original: [Equation]
Revised: [Equation]

Original: Here X km is ﬁrst normalized with its original statistics by ( X km − µ km ) / σ km , and further scaled with novel statistics { ˆ µ km , ˆ σ km } that are randomly sampled from corresponding Gaussian distribution.
Revised: Here X km is ﬁrst normalized with its original statistics by ( X km − µ km ) / σ km , and further scaled with novel statistics { ˆ µ km , ˆ σ km } that are randomly sampled from corresponding Gaussian distribution.

Original: To make the sampling differentiable, we use the re-parameterization trick (Kingma & Welling, 2013):
Revised: To make the sampling differentiable, we use the re-parameterization trick (Kingma & Welling, 2013):

Original: [Equation]
Revised: [Equation]

Original: where (cid:15) µ ∼N (0 , 1) and (cid:15) σ ∼N (0 , 1) follow the normal Gaussian distribution.
Revised: where (cid:15) µ ∼N (0 , 1) and (cid:15) σ ∼N (0 , 1) follow the normal Gaussian distribution.

Original: The proposed federated feature augmentation (FFA) operation in Eq. 8 is a plug-and-play layer, which can be inserted at arbitrary positions of the feature extractor h to facilitate latent feature augmentation.
Revised: The proposed federated feature augmentation (FFA) operation in Eq. 8 is a plug-and-play layer, i . e ., it can be inserted at arbitrary layers in the feature extractor h .

Original: In our implementation, we add a FFA layer after each convolutional stage of the networks.
Revised: In our implementation, we add a FFA layer after each convolutional stage of the networks.

Original: During training, we follow the stochastic learning strategy (Verma et al., 2019; Zhou et al., 2021) to activate each FFA layer with a probability of p , which allows for more diverse augmentation from iteration to iteration (based on the activated FFA layers).
Revised: During training, we follow the stochastic learning strategy (Verma et al., 2019; Zhou et al., 2021; Li et al., 2022b) to activate each FFA layer with a probability of p . This allows for more diverse augmentation from iteration to iteration (based on the activated FFA layers).

Original: At test time, no augmentation is performed.
Revised: At test time, no augmentation is applied.

Original: In Appendix A, we provide detailed descriptions of F ED FA in Algorithm 1 and the
Revised: In Appendix A, we provide detailed descriptions of F ED FA in Algorithm 1 and FFA in Algorithm 2.

Revised:  T HEORETICAL I NSIGHTS

Revised: In this section, we provide mathematical analysis to gain deeper insights into F ED FA.

Revised: To begin with, we show that F ED FA is a noise injection process (Bishop, 1995; Camuto et al., 2020;

Revised: Lim et al., 2022) that injects federation-aware noises to latent features.

Original: FFA layer  in Algorithm 2.
Revised: Lemma 1. Consider client m ∈ [ M ] , for a batch-wise latent feature X km at layer k , its augmentation in F ED FA (cf. Eq.

Revised: 8) follows a noising process ˆ X km =

Revised: X km + e km , with the noise e km taking the form:

Revised: [Equation]

Revised: where (cid:15) µ ∼N (0 , 1) , (cid:15) σ ∼N (0 , 1) , ¯ X km =

Revised: ( X km − µ km ) / σ km .

Revised: Based on Lemma 1, we can identify the federation-aware implicit regularization effects of F ED FA.

Revised: Theorem 1.

Revised: In F ED FA , the loss function L F ED FA m of client m can be expressed as:

Revised: [Equation]

Revised: [Equation]

Revised: where L ERM m is the standard ERM loss, and L REG m is the regularization term:

Revised: [Equation]

Revised: [Equation]

Original: In addition, to gain more theoretical insights to F ED FA, we provide mathematical analysis in Appendix B and show that F ED FA exhibits a natural form of federationaware implicit regularization to local client training.
Revised: where J z denotes the Jacobian of layer z (see Proposition 1 in Appendix for its explicit expression). Theorem 1 implies that, F ED FA implicitly introduces regularization to local client learning by regularizing the gradients of latent representations ( i . e ., (cid:53)

Revised: h 1: K ( X m ) (cid:96) ( g ( h 1: K ( X m )) , Y m ) (cid:62) ), weighted by federation-aware noises in Lemma 1, i .

Revised: e ., (cid:80)

Revised: z ∈Z J z

Revised: ( X m ) e zm .

Original:  E XPERIMENT 4.E XPERIMENTAL S ETUP
Revised:  E MPIRICAL R ESULT 4. S ETUP

Original: Datasets.
Revised: Datasets.

Original: We conduct extensive experiments on ﬁve datasets: Ofﬁce-Caltech 10 (Gong et al., 2012),
Revised: We conduct extensive experiments on ﬁve datasets: Ofﬁce-Caltech 10 (Gong et al., 2012),

Original: DomainNet(Peng et al., 2019) and ProstateMRI(Liu et al., 2020b) for validation of F ED FA in terms of feature-shift non-IID, as well as larger-scale datasets CIFAR-10 (Krizhevsky & Hinton, 2009) and
Revised: DomainNet(Peng et al., 2019) and ProstateMRI(Liu et al., 2020b) for validation of F ED FA in terms of feature-shift non-IID, as well as larger-scale datasets CIFAR-10 (Krizhevsky & Hinton, 2009) and

Original: EMNIST (Cohen et al., 2017) for cases of label distribution and data size heterogeneity, respectively.
Revised: EMNIST (Cohen et al., 2017) for cases of label distribution and data size heterogeneity, respectively.

Original: Baselines.
Revised: Baselines.

Original: For comprehensive evaluation, we compare F ED FA against several state-of-the-art federated learning techniques, including F ED A VG (McMahan et al., 2017), F ED A VG M (Hsu et al., 2019), F ED P ROX (Li et al., 2020a), F ED SAM (Qu et al., 2022), F ED BN (Li et al., 2020b), F ED R OBUST (Reisizadeh et al., 2020), and F ED M IX (Yoon et al., 2021).
Revised: For comprehensive evaluation, we compare F ED FA against several state-of-the-art federated learning techniques, including F ED A VG (McMahan et al., 2017), F ED A VG M (Hsu et al., 2019), F ED P ROX (Li et al., 2020a), F ED SAM (Qu et al., 2022), F ED BN (Li et al., 2020b), F ED R OBUST (Reisizadeh et al., 2020), and F ED M IX (Yoon et al., 2021).

Original: Moreover, we compare with
Revised: Moreover, we compare with

Original: F ED H ARMO (Jiang et al., 2022) in ProstateMRI, that is specialized designed for medical imaging.
Revised: F ED H ARMO (Jiang et al., 2022) in ProstateMRI, that is specialized designed for medical imaging.

Original: To gain more insights into F ED FA, we develop two baselines: F ED FA-R( ANDOM ) and F ED FAC( LIENT ).
Revised: To gain more insights into F ED FA, we develop two baselines: F ED FA-R( ANDOM ) and F ED FAC( LIENT ).

Original: F ED FA-R randomly perturbs feature statistics based on Gaussian distribution with a same standard deviation for all channels, i .
Revised: F ED FA-R randomly perturbs feature statistics based on Gaussian distribution with a same standard deviation for all channels, i .

Original: e ., ˆΣ µ km = ˆΣ σ km = λ , where λ = 0 .
Revised: e ., ˆΣ µ km = ˆΣ σ km = λ , where λ = 0 .

Original: F ED FA-C performs augmentation based only on client speciﬁc variances, i .
Revised: F ED FA-C performs augmentation based only on client speciﬁc variances, i .

Original: e .
Revised: e .

Original: , Eq.
Revised: , Eq.

Original: 7 turns into ˆΣ 2 µ km =
Revised: 7 turns into ˆ Σ 2 µ km =

Original: Σ 2 µ km , ˆΣ 2 σ km =Σ 2 σ km .
Revised: Σ 2 µ km , ˆ Σ 2 σ km =Σ 2 σ km .

Original: Metrics.
Revised: Metrics.

Original: As conventions, we use top-1 accuracy for image classiﬁcation and Dice coefﬁcient for medical image segmentation, respectively.
Revised: As conventions, we use top-1 accuracy for image classiﬁcation and Dice coefﬁcient for medical image segmentation, respectively.

Original: We report the performance only for the global model.
Revised: We report the performance only for the global model.

Original: Implementation Details.
Revised: Implementation Details.

Original: We use PyTorch to implement F ED FA and other baselines.
Revised: We use PyTorch to implement F ED FA and other baselines.

Original: Following
Revised: Following

Original: F ED BN (Li et al., 2020b), we adopt AlexNet (Krizhevsky et al., 2017) on Ofﬁce-Caltech 10 and 
Revised: F ED BN (Li et al., 2020b), we adopt AlexNet (Krizhevsky et al., 2017) on Ofﬁce-Caltech 10 and [Table] [Table] [Table]

Revised: [Table]

Revised: [Table]

Original: DomainNet, using the SGD optimizer with learning rate 0 .
Revised: [Table] [Table] [Table] [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: F ED FA

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: DomainNet, using the SGD optimizer with learning rate 0 .

Original: 01 and batch size 32 .
Revised: 01 and batch size 32 .

Original: Following F EDH ARMO (Jiang et al., 2022), we employ U-Net
Revised: Following F EDH ARMO (Jiang et al., 2022), we employ U-Net

Original: (Ronneberger et al., 2015) on ProstateMRI using
Revised: (Ronneberger et al., 2015) on ProstateMRI using

Original: Adam as the optimizer with learning rate 1e-4 and batch size 16 .
Revised: Adam as the optimizer with learning rate 1e-4 and batch size 16 .

Original: The communication rounds are 400 for Ofﬁce-Caltech 10 and DomainNet, and 500 for ProstateMRI, with the number of local update epoch setting to 1 in all cases.
Revised: The communication rounds are 400 for Ofﬁce-Caltech 10 and DomainNet, and 500 for ProstateMRI, with the number of local update epoch setting to 1 in all cases.

Original: For EMNIST, we strictly follow F ED M IX (Yoon et al., 2021) to introduce data size heterogeneity by partitioning data w.r.t.
Revised: For EMNIST, we strictly follow F ED M IX (Yoon et al., 2021) to introduce data size heterogeneity by partitioning data w.r.t.

Original: writers, and train a LeNet-5 (LeCun et al., 1998) using SGD with batch size 10.
Revised: writers, and train a LeNet-5 (LeCun et al., 1998) using SGD with batch size 10.

Original: The total number of clients is 200 and only 10 clients are sampled per communication round.
Revised: The total number of clients is 200 and only 10 clients are sampled per communication round.

Original: We run 200 rounds in total.
Revised: We run 200 rounds in total.

Original: For CIFAR-10, we sample local data based on Dirichlet distribution Dir ( α ) to simulate label distribution heterogeneity.
Revised: For CIFAR-10, we sample local data based on Dirichlet distribution Dir ( α ) to simulate label distribution heterogeneity.

Original: As (Qu et al., 2022; Kim et al., 2022), we set α to 0.3 or 0.6, and train a ResNet-

Revised: As (Qu et al., 2022; Kim et al., 2022), we set α to 0.3 or 0.6, and train a ResNet-

Original: (He et al., 2016).
Revised: (He et al., 2016).

Original: The number of clients is 100 with participation rate 0.1, while the number of communication round is set to 100.
Revised: The number of clients is 100 with participation rate 0.1, while the number of communication round is set to 100.

Original: Due to the limit of space, we present full details about the datasets and implementations in §D.

Original: 4. M AIN R ESULTS
Revised: 4. M AIN R ESULTS

Original: [Figure]
Revised: [Figure]

Original: We ﬁrst present the overall results on the ﬁve benchmarks, i .
Revised: We ﬁrst present the overall results on the ﬁve benchmarks, i .

Original: e ., Ofﬁce-Caltech 10 and DomainNet in Table 1 and
Revised: e ., Ofﬁce-Caltech 10 and DomainNet in Table 1 and

Original: Fig.
Revised: Fig.

Original: 1, ProstateMRI in Table 2 and Fig. 2, EMNIST and
Revised: 1, ProstateMRI in Table 2 and Fig. 2, EMNIST and

Original: CIFAR-10 in Table 3.
Revised: CIFAR-10 in Table 3.

Original: Results on Ofﬁce-Caltech 10 and DomainNet.
Revised: Results on Ofﬁce-Caltech 10 and DomainNet.

Original: F ED FA yields solid improvements over competing methods for image classiﬁcation.
Revised: F ED FA yields solid improvements over competing methods for image classiﬁcation.

Original: As presented in Table 1, F ED FA leads to consistent performance gains over the competitors across the benchmarks.
Revised: As presented in Table 1, F ED FA leads to consistent performance gains over the competitors across the benchmarks.

Original: The improvements over F ED A VG can be as large as 4.6% and 3.7% on OfﬁceCaltech 10 and DomainNet, respectively.
Revised: The improvements over F ED A VG can be as large as 4.6% and 3.7% on OfﬁceCaltech 10 and DomainNet, respectively.

Original: Moreover, in comparison to prior data augmentation-based algorithm
Revised: Moreover, in comparison to prior data augmentation-based algorithm

Original: F ED M IX , F ED FA also brings solid gains, i .
Revised: F ED M IX , F ED FA also brings solid gains, i .

Original: e ., 3.3% on
Revised: e ., 3.3% on

Original: Ofﬁce-Caltech 10 and 2.2% on DomainNet.
Revised: Ofﬁce-Caltech 10 and 2.2% on DomainNet.

Original: This is encouraging since our approach in nature better preserves privacy than F ED M IX , which requires exchanging averaged data across clients.
Revised: This is encouraging since our approach in nature better preserves privacy than F ED M IX , which requires exchanging averaged data across clients.

Original: Moreover, Fig.
Revised: Moreover, Fig.

Original: 1 depicts the convergence curves of comparative methods on Ofﬁce-Caltech 10.
Revised: 1 depicts the convergence curves of comparative methods on Ofﬁce-Caltech 10.

Original: At the early training stage ( 0 ∼ 200 rounds), F ED FA shows similar training efﬁciency as other baselines like F ED P ROX and F ED M IX .
Revised: At the early training stage ( 0 ∼ 200 rounds), F ED FA shows similar training efﬁciency as other baselines like F ED P ROX and F ED M IX .

Original: But as the training goes, F ED FA is able to converge to a more optimal solution.[Table] [Table] [Table]
Revised: But as the training goes, F ED FA is able to converge to a more optimal solution.

Original: [Table]

Original: 1/ 2/ 3/ 4/ 5/ 1 Fraction of the training

Original: [Figure]

Original: in each client

Original: [Figure]

Original: age

Original: [Figure]

Original: FedSAM

Original: [Figure]

Original: FedHarmo FedFA

Original: Results on ProstateMRI.
Revised: Results on ProstateMRI.

Original: F ED FA shows leading performance with extremely small local datasets.
Revised: F ED FA shows leading performance with extremely small local datasets.

Original: In some practical scenarios like healthcare, the size of local dataset can be very small, which poses a challenge for federated learning.
Revised: In some practical scenarios like healthcare, the size of local dataset can be very small, which poses a challenge for federated learning.

Original: To examine the performance of federated learning algorithms in this scenario, we build mini -ProstateMRI by randomly sampling only 1 / 6 of all training samples in each client for training.
Revised: To examine the performance of federated learning algorithms in this scenario, we build mini -ProstateMRI by randomly sampling only 1 / 6 of all training samples in each client for training.

Original: Results are summarized in Table 2.
Revised: Results are summarized in Table 2.

Original: F ED FA outperforms F ED A VG by signiﬁcant margins ( i .
Revised: F ED FA outperforms F ED A VG by signiﬁcant margins ( i .

Original: e ., 3.0% ) and it even performs better than F ED H ARMO , which is speciﬁcally designed for medical scenarios.
Revised: e ., 3.0% ) and it even performs better than F ED H ARMO , which is speciﬁcally designed for medical scenarios.

Original: In addition, Fig. 2 shows how the performance of methods varies with respect to the size of local dataset.
Revised: In addition, Fig. 2 shows how the performance of methods varies with respect to the size of local dataset.

Original: We train methods with different fractions ( i .
Revised: We train methods with different fractions ( i .

Original: e ., 1 / 6 , 2 / 6 , 3 / / 6 , 5 / 6 , 1 ) of training samples .
Revised: e ., 1 / 6 , 2 / 6 , 3 / 6 ,/ 6 , 5 / 6 , 1 ) of training samples .

Original: F ED FA shows promising performance in all cases.
Revised: F ED FA shows promising performance in all cases.

Original: [Table]
Revised: [Table]

Revised: [Table]

Original: Results on CIFAR-10 and EMNIST.
Revised: Results on CIFAR-10 and EMNIST.

Original: In addition to feature-shift non-i.i.d., F ED FA shows consistent improvements in label distribution heterogeneity (CIFAR-10) and data size heterogeneity (EMNIST).
Revised: In addition to feature-shift non-i.i.d., F ED FA shows consistent improvements in label distribution heterogeneity (CIFAR-10) and data size heterogeneity (EMNIST).

Original: As shown in Table 3, in CIFAR10, F ED FA surpasses the second best method,
Revised: As shown in Table 3, in CIFAR10, F ED FA surpasses the second best method,

Original: F ED M IX , by 0.8% and 1.2% with respect to two non-i.i.d levels Dir(0.6) and Dir(0.3), respectively.
Revised: F ED M IX , by 0.8% and 1.2% with respect to two non-i.i.d levels Dir(0.6) and Dir(0.3), respectively.

Original: Notably, as the non-i.i.d.
Revised: Notably, as the non-i.i.d.

Original: level increasing from Dir(0.6) to Dir(0.3), F ED FA tends to yield a larger gap of performance gain, showing a strong capability in handling severe non-i.i.d.
Revised: level increasing from Dir(0.6) to Dir(0.3), F ED FA tends to yield a larger gap of performance gain, showing a strong capability in handling severe non-i.i.d.

Original: scenarios.
Revised: scenarios.

Original: In addition, in EMNIST, F ED FA outperforms F ED R OBUST by 1.1% and F ED M IX by 1.2% , respectively.
Revised: In addition, in EMNIST, F ED FA outperforms F ED R OBUST by 1.1% and F ED M IX by 1.2% , respectively.

Original: These results reveal that though designed for feature shift non-i.i.d., F ED FA’s very nature of data augmentation makes it a fundamental technique to various non-i.i.d.
Revised: These results reveal that though designed for feature shift non-i.i.d., F ED FA’s very nature of data augmentation makes it a fundamental technique to various non-i.i.d.

Original: challenges in FL.
Revised: challenges in FL.

Revised: [Table]

Revised: [Table]

Original: 4. F EDERATED D OMAIN G ENERALIZATION P ERFORMANCE
Revised: F EDERATED D OMAIN G ENERALIZATION P ERFORMANCE

Original: Federated learning are dynamic systems, in which novel clients may enter the system after model training, most possibly with test-time distribution shift.
Revised: Federated learning are dynamic systems, in which novel clients may enter the system after model training, most possibly with test-time distribution shift.

Original: However, most prior federated learning[Table] algorithms focus only on improving model performance on the participating clients, while neglecting model generalizability to unseen non-participating clients.
Revised: However, most prior federated learning algorithms focus only on improving model performance on the participating clients, while neglecting model generalizability to unseen non-participating clients.

Original: Distribution shift often occurs during deployment, thus it is essential to evaluate the generalizability of federated learning algorithms.
Revised: Distribution shift often occurs during deployment, thus it is essential to evaluate the generalizability of federated learning algorithms.

Original: With a core of data augmentation
Revised: With a core of data augmentation

Original: [Table]
Revised: [Table]

Original: F ED FA is supposed to enforce regularization to neural network learning, which could improve generalization capability.
Revised: F ED FA is supposed to enforce regularization to neural network learning, which could improve generalization capability.

Original: To verify this [Table] we perform experiments for federated domain generalization based on the leaveone-client-out strategy, i .
Revised: To verify this, we perform experiments for federated domain generalization based on the leaveone-client-out strategy, i .

Original: e ., training on M − 1 distributed clients and testing on the held-out unparticipating client.
Revised: e ., training on M − 1 distributed clients and testing on the held-out unparticipating client.

Original: The results are presented in Table 4.
Revised: The results are presented in Table 4.

Original: As seen, F ED FA achieves leading generalization performance on most unseen clients.
Revised: As seen, F ED FA achieves leading generalization performance on most unseen clients.

Original: For example, it yields consistent improvements as compared to F ED M IX , i .
Revised: For example, it yields consistent improvements as compared to F ED M IX , i .

Original: e ., 1.7% on Ofﬁce-Caltech 10, 1.6% on DomainNet, and 1.3% on
Revised: e ., 1.7% on Ofﬁce-Caltech 10, 1.6% on DomainNet, and 1.3% on

Original: ProstateMRI, in terms of average performance.
Revised: ProstateMRI, in terms of average performance.

Original: Despite the improved performance, we ﬁnd by comparing to the results reported in Table 1 and Table 2 that, current federated learning algorithms still encounter signiﬁcant participation gap (Yuan et al., 2022), i .
Revised: Despite the improved performance, we ﬁnd by comparing to the results reported in Table 1 and Table 2 that, current federated learning algorithms still encounter signiﬁcant participation gap (Yuan et al., 2022), i .

Original: e ., the performance difference between participating and non-participating clients, which is a critical issue that should be tackled in future.
Revised: e ., the performance difference between participating and non-participating clients, which is a critical issue that should be tackled in future.

Original: 4. D IAGNOSTIC E XPERIMENT
Revised: 4. D IAGNOSTIC E XPERIMENT

Original: We conduct a set of ablative experiments to enable a deep understanding of F ED FA.
Revised: We conduct a set of ablative experiments to enable a deep understanding of F ED FA.

Original: [Table]
Revised: [Table]

Original: F ED FAvs.
Revised: F ED FAvs.

Original: F ED FA-C and F ED FA-R.
Revised: F ED FA-C and F ED FA-R.

Original: We ﬁrst verify F ED FA against the two baseline variants mentioned in §4.1.
Revised: We ﬁrst verify F ED FA against the two baseline variants mentioned in §4.1.

Original: Both variants involve only device-dependent augmentation of feature statistics, without explicitly considering any global information.
Revised: Both variants involve only device-dependent augmentation of feature statistics, without explicitly considering any global information

Original:  As shown in Table 5, by randomly perturbing feature statistics, F ED FA-R shows no improvements or even suffers performance degradation on
Revised: [Table] As shown in Table [Table] by randomly perturbing feature statistics, F ED FA-R shows no improvements or even suffers performance degradation on

Original: DomainNet and ProstateMRI against F ED A VG ; F ED FA-C yields promising performance gains by taking into account client-speciﬁc feature statistic variances; by comparing F ED FA and F ED FA-C, we conﬁrm the signiﬁcance of universal feature statistic information in federated augmentation. 
Revised: DomainNet and ProstateMRI against F ED A VG ; F ED FA-C yields promising performance gains by taking into account client-speciﬁc feature statistic variances; by comparing F ED FA and F ED FA-C, we conﬁrm the signiﬁcance of universal feature statistic information in federated augmentation. [Table] F ED FA + M O E X 83.

Original: [Table]
Revised: [Table]

Original: F ED FA vs. traditional augmentation methods.
Revised: F ED FA vs. traditional augmentation methods.

Original: Next, we compare F ED FA with four conventional data/feature augmentation techniques, i .
Revised: We compare F ED FA with four conventional data/feature augmentation techniques, i .

Original: e ., Mixup (Zhang et al., 2018), Manifold Mixup (Verma et al., 2019), MixStyle (Zhou et al., 2021) and MoEx (Li et al., 2021a).
Revised: e ., M IXUP (Zhang et al., 2018), M ANIFOLD M IXUP (Verma et al., 2019), M IX S TYLE (Zhou et al., 2021) and M O E X (Li et al., 2021a).

Original: The results are presented in Table 6.
Revised: The results are presented in Table 6.

Original: We show that i) all the four techniques yield non-trivial improvements over F ED A VG , and some of them ( e [Table] g .,
Revised: We show that i) all the four techniques yield non-trivial improvements over F ED A VG , and some of them ( e .

Original: MoEx) even outperform well-designed federated learning algorithms (as compared to Tables 1-2); by accounting for global feature statistics, F ED FA surpasses all of them, yielding 2.9% / 1.9% / 1.0% improvements over the second-best results on Ofﬁce/DomainNet/ProstateMRI [Table] respectively; iii) combining F ED FA with these techniques allows further performance uplifting, verifying the complementary roles of F ED FA to them.
Revised: g ., M O E X ) even outperform welldesigned federated learning algorithms (as compared to Tables 1-2); by accounting for global feature statistics, F ED FA surpasses all of them, yielding 2.9% / 1.9% / 1.0% improvements over the secondbest results on Ofﬁce/DomainNet/ProstateMRI, respectively; iii) combining F ED FA with these techniques allows further performance uplifting, verifying the complementary roles of F ED FA to them.

Original: [Table]
Revised: [Table]

Original: Adaptive Variance Fusion.
Revised: Adaptive Variance Fusion.

Original: Next, we examine the effect of adaptive variance fusion in Eqs.
Revised: Next, we examine the effect of adaptive variance fusion in Eqs.

Original: We design a baseline “Direct Fusion” that directly combines the client-speciﬁc and client-sharing statistic variances as:
Revised: We design a baseline “Direct Fusion” that directly combines the client-speciﬁc and client-sharing statistic variances as:

Original: ˆΣ 2 µ km =(Σ 2 µ k +1)Σ 2 µ km , ˆΣ 2 σ km =(Σ 2 σ k +1)Σ 2 σ km .
Revised: ˆΣ 2 µ km =(Σ 2 µ k +1)Σ 2 µ km , ˆΣ 2 σ km =(Σ 2 σ k +1)Σ 2 σ km .

Original: We ﬁnd from Table 7 that the baseline encounters severe performance degradation across all three benchmarks.
Revised: We ﬁnd from Table 7 that the baseline encounters severe performance degradation across all three benchmarks.

Original: A possible reason is that the two types of variances are mostly mis-matched, and the simple fusion strategy may cause signiﬁcant changes of client-speciﬁc statistic variances, which would be harmful for local model learning.[Table] Estimation of Client Speciﬁc Statistic Variances
Revised: A possible reason is that the two types of variances are mostly mis-matched, and the simple fusion strategy may cause signiﬁcant changes of client-speciﬁc statistic variances, which would be harmful for local model learning.

Original: Beyond mini-Batch.

Original: The client speciﬁc statistic variances are computed based on feature statistics of a mini-batch (Eq. 3), however, this may be inaccurate in cases of small batch size.

Original: To deal with such situations, we investigate an extension of F ED FA by equipping it an external memory bank (He et al., 2020; Wu et al., 2018) that stores the feature statistics in past mini-batches.

Original: This allows

Original: F ED FA to estimate variances from a larger number of samples, even with a small batch size.

Original: Tablesummarizes the results of F ED FA with respect to different combinations of batch size and memory size on Ofﬁce-Caltech 10.

Original: The default model with ( bs = 32 , ms = 0) achieves a score of 83.1%.

Original: It is not surprising to see the improvements after increasing the batch size to 64 or 96 .

Original: But it is indeed appealing that F ED FA delivers comparable performance by using the external memory bank (ms = 32 or 64 ) to mimic the behavior of large batch sizees.

Original: [Figure]
Revised: [Figure]

Original: Hyper-parameter analysis.
Revised: Hyper-parameter analysis.

Original: F ED FA includes only two hyperparameters, i .
Revised: F ED FA includes only two hyperparameters, i .

Original: e ., momentum coefﬁcient α in Eq.
Revised: e ., momentum coefﬁcient α in Eq.

Original: 4 and stochastic learning probability p to apply feature statistic augmentation during training.
Revised: 4 and stochastic learning probability p to apply feature statistic augmentation during training.

Original: As illustrated in Fig.
Revised: As shown in Fig.

Original: 3, i) the model is overall robust to α .
Revised: 3, (1) the model is overall robust to α .

Original: Notably, the model yields promising performance at α = 0 , in which the model only useslimited client information, i . e ., the feature statistics of the last mini-batch in each local epoch, to compute client-sharing statistic variances.
Revised: Notably, it yields promising performance at α = 0 , in which the model only uses the feature statistics of the last mini-batch in each local epoch to compute client-sharing statistic variances.

Original: This result reveals that F ED FA is insensitive to errors of client sharing statistic variances.
Revised: This result reveals that F ED FA is insensitive to errors of client-sharing statistic variances.

Original: ii)

Original: For the probability p , we see that F ED FA signiﬁcantly promotes the baseline ( i .
Revised: For the probability p , we see that F ED FA signiﬁcantly improves the baseline ( i .

Original: e ., p = 0 ), even with avery small probability ( e .
Revised: e ., p = 0 ), even with a small probability ( e .

Original: g ., p = 0 . 1 ).
Revised: g ., p =0 . 1 ).

Original: Themodel achieves the best performance  at p =0 .
Revised: The best performance is reached at p =0 .

Original: 5 , and its performance is relatively stable when p varies from 0 .

Original: 1 to 1 . 0 .

Original: 4. C OMPLEXITY A NALYSIS
Revised: 4. C OMPLEXITY A NALYSIS

Original: Computation and memory costs.
Revised: Computation and memory costs.

Original: F ED FA involves only several basic matrix operations, thus incurring negligible extra computation cost.
Revised: F ED FA involves only several basic matrix operations, thus incurring negligible extra computation cost.

Original: Compared to F ED A VG , it requires 4 (cid:80)
Revised: Compared to F ED A VG , it requires 4 (cid:80)

Original: Kk =1 C k more GPU memory allocation to store four statistic values ( ¯ µ km , ¯ σ km , γ µ k , γ σ k ) at each of the K FFA layers.
Revised: Kk =1 C k more GPU memory allocation to store four statistic values ( ¯ µ km , ¯ σ km , γ µ k , γ σ k ) at each of the K FFA layers.

Original: Here C k is the number of feature channel at each layer k .
Revised: Here C k is the number of feature channel at each layer k .

Original: The costs are in practice very minor, e .
Revised: The costs are in practice very minor, e .

Original: g ., 18 KB/15.5 KB for AlexNet/U-Net.
Revised: g ., 18 KB/15.5 KB for AlexNet/U-Net.

Original: For comparison, F ED M IX requires 2 × more GPU memory than
Revised: For comparison, F ED M IX requires 2 × more GPU memory than

Original: F ED A VG .
Revised: F ED A VG .

Original: The low computation/memory costs make F ED FA favorable for edge devices.
Revised: The low computation/memory costs make F ED FA favorable for edge devices.

Original: Communication cost.
Revised: Communication cost.

Original: In each round, F ED FA incurs additional communication costs since it requires the sending 1) from client to server the momentum feature statistics ¯ µ km and ¯ σ km , as well as 2) from server to client the client sharing feature statistic 
Revised: In each round, F ED FA incurs additional communication costs since it requires the sending 1) from client to server the momentum feature statistics ¯ µ km and ¯ σ km , as well as 2) from server to client the client sharing feature statistic variances γ µ k and γ σ k at each layer k .

Original: Thus, for K layers in total, the extra communication cost for each client is c e =4 (cid:80)
Revised: Thus, for K layers in total, the extra communication cost for each client is c e =4 (cid:80)

Original: Kk =1 C k , where the factor of 4 is for server receiving and sending two statistic values.
Revised: Kk =1 C k , where the factor of 4 is for server receiving and sending two statistic values.

Original: We further denote c m as the cost for exchanging model parameters in FedAvg.
Revised: We further denote c m as the cost for exchanging model parameters in FedAvg.

Original: In general we have c e (cid:28) c m ( e . g .,
Revised: In general we have c e (cid:28) c m ( e . g .,

Original: c e = 18KB vs. c m = 99MB for AlexNet), hence, the extra communication bruden in F ED FA is almost negligible.  
Revised: c e = 18KB vs. c m = 99MB for AlexNet), hence, the extra communication bruden in F ED FA is almost negligible.  R ELATED W ORK Federated Learning.

Revised: Recent years have witnessed tremendous progress in federated learning (Koneˇcn`y et al., 2016), which opens the door for privacy-preserving deep learning (Shokri &

Revised: Shmatikov, 2015), i .

Original: More details are provided in Appendix §E.
Revised: e ., train a global model on distributed datasets without disclosing private data information. F ED A VG (McMahan et al., 2017) is a milestone; it trains local models independently in multiple clients and then averages the resulting model updates via a central server once in a while. However, F ED A VG is designed for i.i.d. data and suffers in statistical accuracy or even diverge if deployed over non-i.i.d.

Revised: client samples.

Revised: To address this issue, numerous efforts have been devoted to handling heterogeneous federated environments, by, for example, adding a dynamic regularizer to local objectives in F ED P ROX (Li et al., 2020a) and

Revised: F ED D YN (Acar et al., 2021), correcting client drift through variance reduction in SCAFFOLD (Karimireddy et al., 2020), adaptive server optimization in F ED O PT (Reddi et al., 2021), local batch normalization in F ED BN (Li et al., 2020b), or training a perturbed loss in F ED SAM (Qu et al., 2022) and F ED R OBUST (Reisizadeh et al., 2020).

Revised: F ED M IX (Yoon et al., 2021), as far as we know, is the only existing method that solves federated learning based on data augmentation.

Revised: It adapts the well-known M IXUP algorithm

Revised: (Zhang et al., 2018) from centralized learning into the federated learning scenario.

Revised: Nevertheless, F ED M IX requires exchanging local data (or averaged version) across clients for data interpolation, thereby suffering privacy issues.

Revised: In addition, F ED M IX operates on the input level, while our approach focuses on latent feature statistic augmentation.

Revised: Since deeper representations tend to disentangle the underlying factors of variation better (Bengio et al., 2013), traversing along latent space will potentially make our method encounter more realistic samples.

Revised: This is supported by the fact that F ED FA achieves consistent performance improvements over F ED M IX in diverse scenarios.

Revised: Data Augmentation.

Revised: Data augmentation has a long and rich history in machine learning.

Revised: Early studies (Sch ¨ olkopf et al., 1996; Kuka ˇ cka et al., 2017) focus on label-preserving transformations to employ regularization via data, alleviating overﬁtting and improving generalization.

Revised: For image data, some techniques, like random horizontal ﬂipping and cropping are commonly used for training of advanced neural networks (He et al., 2016).

Revised: In addition, there is a recent trend for label-perturbing augmentation, e .

Revised: g .

Revised: , M IXUP (Zhang et al., 2018) or C UTMIX (Yun et al., 2019).

Revised: Separate from these input-level augmentation techniques are feature augmentation methods (Verma et al., 2019;

Revised: Li et al., 2021a; 2022b;

Revised: Zhou et al., 2021) that make augmentation in latent feature space.

Revised: These various data augmentation techniques have shown great successes to learn domain-invariant models in the centralized setup.

Revised: Our method is an instance of label-preserving feature augmentation, designed for federated learning.

Revised: It is inspired by recent efforts on implicit feature augmentation (Li et al., 2021a; 2022b;

Revised: Zhou et al., 2021) that synthesize samples of novel domains by manipulating instance-level feature statistics.

Revised: In these works, feature statistics are treated as ‘features’, which capture essential domain-speciﬁc characteristics.

Revised: In F ED FA, we estimate appropriate variances of feature statistics from a federated perspective, and draw novel statistics probablistically from a distribution centered on old statistics, while spanning with the variances.

Revised: F ED FA avoids statistics mixing of instances from different clients, as done in F ED M IX (Yoon et al., 2021), thus can better preserve data privacy.

Original:  C ONCLUSION
Revised:  C ONCLUSION

Original: This work solves federated learning from a unique perspective of feature augmentation, yielding a new algorithm F ED FA that shows strong performance across various federated learning scenarios.
Revised: This work solves federated learning from a unique perspective of feature augmentation, yielding a new algorithm F ED FA that shows strong performance across various federated learning scenarios.

Original: F ED FA is based on a probabilistic (precisely Gaussian) modeling of feature statistics, where Gaussian variances are estimated in a federated manner, based on both local feature statistic distribution within each client, as well asmore universal feature statistic distribution across clients. Feature augmentation is implicitly achieved by sampling novel feature statistics from the Gaussian distribution.
Revised: F ED FA is based on a Gaussian modeling of feature statistic augmentation, where Gaussian variances are estimated in a federated manner, based on both local feature statistic distribution within each client, as well as universal feature statistic distribution across clients. We identify the implicit federation-aware regularization effects of F ED FA through theoretical analysis, and conﬁrm its empirical superiority across a suite of benchmarks in federated learning.  R EPRODUCIBILITY

Original: We evaluate F ED FA across a suite of federated datasets, and  the results conﬁrm its superiority.
Revised: Throughout the paper we have provided details facilitating reproduction of our empirical results. All our experiments are ran with a single GPU (we used NVIDIA GeForce RTX 2080 Ti with a 11G memory), thus can be reproduced by researchers with computational constraints as well. The source code has been made publicly available in https://github.com/tfzhou/FedFA. For the theoretical results, all assumptions, proofs and relevant discussions are provided in the Appendix.

Original:  A D ETAILED A LGORITHM
Revised:  A CKNOWLEDGMENTS

Revised: We would like to thank the anonymous referees for their valuable comments for improving the paper. This study was partly supported by Varian Research, Switzerland Grant. A D ETAILED A LGORITHM

Original: In Algorithm 1, we illustrate the detailed training procedure of our F ED FA.
Revised: In Algorithm 1, we illustrate the detailed training procedure of our F ED FA.

Original: It is consistent with algorithms such as F ED A VG (McMahan et al., 2017).
Revised: It is consistent with algorithms such as F ED A VG (McMahan et al., 2017).

Original: In each communication round, the client performs local model training of the feature extractor h ( w h ) and classiﬁer g ( w g ) .
Revised: In each communication round, the client performs local model training of the feature extractor h ( w h ) and classiﬁer g ( w g ) .

Original: We append a
Revised: We append a

Original: FFA layer (Algorithm 2) after each convolutional stage h k .
Revised: FFA layer (Algorithm 2) after each convolutional stage h k .

Original: Each client additionally maintains a pair of momentum feature statistics { ¯ µ m , ¯ σ m } , which is updated in a momentum manner during training.
Revised: Each client additionally maintains a pair of momentum feature statistics { ¯ µ m , ¯ σ m } , which is updated in a momentum manner during training.

Original: The parameters from local training ( i .
Revised: The parameters from local training ( i .

Original: e ., w = { w h , w g } ), which are omitted in Algorithm 1, along with the momentum feature statistics are sent to server for model aggregation and computation of client-sharing statistic variances, which will be distributed back to clients for the next round of local training.
Revised: e ., w = { w h , w g } ), which are omitted in Algorithm 1, along with the momentum feature statistics are sent to server for model aggregation and computation of client-sharing statistic variances, which will be distributed back to clients for the next round of local training.

Original: X km = h k X k − Run the k -th layer of the feature extractor h 6: ˆ X km , ¯ µ km , ¯ σ km =
Revised: the feature extractor h 6: ˆ X km , ¯ µ km , ¯ σ km =

Original: FFA ( X km , ¯ µ km , ¯ σ km ) (cid:46) FFA layer in Algorithm7:
Revised: FFA ( X km , ¯ µ km , ¯ σ km ) (cid:46) FFA layer in Algorithm7:

Original: Y = g
Revised: Y = g

Original: ( ˆ X Km )
Revised: ( ˆ X Km )

Original: (cid:46) Run classifer g to get predictions 8:
Revised: (cid:46) Run classifer g to get predictions 8:

Original: Run loss computation and backward optimization 9: Σ 2 µ k = 1 M (cid:80) Mm =1 (¯ µ km − E [¯ µ k ])
Revised: Run loss computation and backward optimization 9: Σ 2 µ k = 1 M (cid:80) Mm =1 (¯ µ km − E [¯ µ k ])

Original: (cid:46) Compute client sharing statistic variance (Eq. 5) 10: Σ 2 σ k = 1 M (cid:80) Mm =1 (¯ σ km − E [¯ σ k ])11: for each channel j ∈ [ C ] do (cid:46) Compute adaptive fusion coefﬁcients (Eq. 6) 12: γ ( j ) µ k =
Revised: (cid:46) Compute client sharing statistic variance (Eq. 5) 10: Σ 2 σ k = 1 M (cid:80) Mm =1 (¯ σ km − E [¯ σ k ])11: for each channel j ∈ [ C ] do (cid:46) Compute adaptive fusion coefﬁcients (Eq. 6) 12: γ ( j ) µ k =

Original: C (1+ 1 / Σ2 , ( j ) µk ) −(cid:80)
Revised: C (1+ 1 / Σ2 , ( j ) µk ) −(cid:80)

Original: Cc =1 (1+ 1 / Σ2 , ( c ) µk ) −13: γ ( j ) σ k =
Revised: Cc =1 (1+ 1 / Σ2 , ( c ) µk ) −13: γ ( j ) σ k =

Original: C (1+ 1 / Σ2 , ( j ) σk ) −(cid:80)
Revised: C (1+ 1 / Σ2 , ( j ) σk ) −(cid:80)

Original: Cc =1 (1+ 1 / Σ2 , ( c ) σk ) −14: return γ µ k , γ σ k
Revised: Cc =1 (1+ 1 / Σ2 , ( c ) σk ) −14: return γ µ k , γ σ k

Original: B T HEORETICAL I NSIGHTS
Revised: B T HEORETICAL I NSIGHTS

Original: In this section, we provide mathematical analysis to understand F ED FA.
Revised: In this section, we provide mathematical analysis to understand F ED FA.

Original: We begin with interpreting
Revised: We begin with interpreting

Original: F ED FA as a noise injection process (Bishop, 1995; Camuto et al., 2020; Lim et al., 2022), which is a case of VRM (§3.1), and show that F ED FA injects federation-aware noises to latent representations (§B.1).
Revised: F ED FA as a noise injection process (Bishop, 1995; Camuto et al., 2020; Lim et al., 2022), which is a case of VRM (§2.1), and show that F ED FA injects federation-aware noises to latent representations (§B.1).

Original: Next, we demonstrate that, induced by federation-aware noise injection, F ED FA exhibits a natural form of federation-aware implicit regularization to local client training (§B.2).
Revised: Next, we demonstrate that, induced by federation-aware noise injection, F ED FA exhibits a natural form of federation-aware implicit regularization to local client training (§B.2).

Original: Without loss of generality, we conduct all analysis for an arbitrary client m ∈
Revised: Without loss of generality, we conduct all analysis for an arbitrary client m ∈

Original: [ M ] .
Revised: [ M ] .

Original: [Table]
Revised: [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: Momentum updating feature statistics (Eq. 5) 12:
Revised: Momentum updating feature statistics (Eq. 5) 12:

Original: [Equation]
Revised: [Equation]

Original: 13: return ˆ X km , ¯ µ km , ¯ σ km
Revised: 13: return ˆ X km , ¯ µ km , ¯ σ km

Original: B. U NDERSTANDING F ED FA AS F EDERATION -A WARE N OISE I NJECTION
Revised: B. U NDERSTANDING F ED FA AS F EDERATION -A WARE N OISE I NJECTION

Original: Noise Injection in Neural Networks.
Revised: Noise Injection in Neural Networks.

Original: Let x be a training sample and x k its latent representation at the k -th layer, with no noise injections.
Revised: Let x be a training sample and x k its latent representation at the k -th layer, with no noise injections.

Original: The x k can be noised under a process ˆ x k = x k + e k , where e k is an addition noise drawn from a probability distribution, and x k is the noised representation.
Revised: The x k can be noised under a process ˆ x k = x k + e k , where e k is an addition noise drawn from a probability distribution, and x k is the noised representation.

Original: A popular choice of e k is isotropic Gaussian noise (Camuto et al., 2020), i .
Revised: A popular choice of e k is isotropic Gaussian noise (Camuto et al., 2020), i .

Original: e ., e k ∼N (0 , σ 2 I ) , where
Revised: e ., e k ∼N (0 , σ 2 I ) , where

Original: I is an identity matrix and σ is a scalar, controlling the amplitude of e k .
Revised: I is an identity matrix and σ is a scalar, controlling the amplitude of e k .

Original: To avoid over-perturbation that may cause model collapse, σ is typically set as a small value.
Revised: To avoid over-perturbation that may cause model collapse, σ is typically set as a small value.

Original: Despite its simplicity, the strategy is conﬁrmed as a highly effective regularized for tackling domain generalization (Li et al., 2021b) and adversarial samples (Lecuyer et al., 2019; Cohen et al., 2019).
Revised: Despite its simplicity, the strategy is conﬁrmed as a highly effective regularized for tackling domain generalization (Li et al., 2021b) and adversarial samples (Lecuyer et al., 2019; Cohen et al., 2019).

Original: However, as shown in Table 5, its performance (see F ED FA-R) is only marginally better or sometimes worse than F ED A VG in FL.
Revised: However, as shown in Table 5, its performance (see F ED FA-R) is only marginally better or sometimes worse than F ED A VG in FL.

Original: Federation-Aware Noise Injection.
Revised: Federation-Aware Noise Injection.

Original: From Eq.
Revised: From Eq.

Original: 9, we can clearly see that the feature statistic augmentation in our approach follows the noise injection process above.
Revised: 9, we can clearly see that the feature statistic augmentation in our approach follows the noise injection process above.

Original: Next we show that this eventually results in features perturbed under a federation-aware noising process.
Revised: Next we show that this eventually results in features perturbed under a federation-aware noising process.

Original: Lemma 1.
Revised: Lemma 1.

Original: Consider client m , for a batch-wise latent feature X km at the k -th layer, its augmentation in F ED FA follows a noising process ˆ X km =
Revised: Consider client m , for a batch-wise latent feature X km at the k -th layer, its augmentation in F ED FA follows a noising process ˆ X km =

Original: X km + e km , with the noise e km taking the form:
Revised: X km + e km , with the noise e km taking the form:

Original: [Equation]
Revised: [Equation]

Original: (10) where (cid:15) µ ∼N (0 , 1) , (cid:15)
Revised: (14) where (cid:15) µ ∼N (0 , 1) , (cid:15)

Original: [Equation]
Revised: [Equation]

Original: Proof of Lemma 1.
Revised: Proof of Lemma 1.

Original: We can easily prove this by substituting Eq.
Revised: We can easily prove this by substituting Eq.

Original: 9 into Eq.
Revised: 9 into Eq.

Original: [Equation]
Revised: [Equation]

Original: As compared to the Gaussian noise injections (Camuto et al., 2020; Cohen et al., 2019; Lecuyer et al., 2019), the noise term e km in F ED FA shows several desirable properties: it is 1) data-dependent, adaptively determined based on the normalized input feature ¯ X km ; 2) channel-independent, allowing for more extensive exploration along different directions in the feature space; 3) most importantly federation-aware, i .
Revised: As compared to the Gaussian noise injections (Camuto et al., 2020; Cohen et al., 2019; Lecuyer et al., 2019), the noise term e km in F ED FA shows several desirable properties: it is 1) data-dependent, adaptively determined based on the normalized input feature ¯ X km ; 2) channel-independent, allowing for more extensive exploration along different directions in the feature space; 3) most importantly federation-aware, i .

Original: e ., its strength is controlled by statistic variances ˆΣ µ km and ˆΣ σ km ( cf.
Revised: e ., its strength is controlled by statistic variances ˆΣ µ km and ˆΣ σ km ( cf.

Original: Eq. 7), which are known carrying universal statistic information of all participating clients.
Revised: Eq. 7), which are known carrying universal statistic information of all participating clients.

Original: B. F EDERATION -A WARE I MPLICIT R EGULARIZATION IN F ED FA
Revised: B. F EDERATION -A WARE I MPLICIT R EGULARIZATION IN F ED FA

Original: Next we show that with noise injections, F ED FA imposes federation-aware implicit regularization to local client training.
Revised: Next we show that with noise injections, F ED FA imposes federation-aware implicit regularization to local client training.

Original: By this, we mean regularization imposed implicitly by the stochastic learning strategy, without explicit modiﬁcation of the loss, and the regularization effect is affected by the federation-aware noise (Lemma 1).
Revised: By this, we mean regularization imposed implicitly by the stochastic learning strategy, without explicit modiﬁcation of the loss, and the regularization effect is affected by the federation-aware noise (Lemma 1).

Original: Recall the deep neural network f deﬁned in §3.1: f (cid:44)
Revised: Recall the deep neural network f deﬁned in §2.1: f (cid:44)

Original: g ◦ h , where h =
Revised: g ◦ h , where h =

Original: h K ◦ h K − 1 ◦ · · · ◦ his a K -layer CNN feature extractor and g is a classiﬁer.
Revised: h K ◦ h K − 1 ◦ · · · ◦ his a K -layer CNN feature extractor and g is a classiﬁer.

Original: Given a batch of samples X m with labels
Revised: Given a batch of samples X m with labels

Original: Y m , its latent representation at the k -th layer is computed as X km =
Revised: Y m , its latent representation at the k -th layer is computed as X km =

Original: h k ◦ h k − 1 ◦ · · · ◦ h 1 ( X m ) , or we write it in a simpler form, X km
Revised: h k ◦ h k − 1 ◦ · · · ◦ h 1 ( X m ) , or we write it in a simpler form, X km

Original: = h 1: k ( X m ) .
Revised: = h 1: k ( X m ) .

Original: Note that we only add noises to layers in h , but not to g .
Revised: Note that we only add noises to layers in h , but not to g .

Original: Concretely, in each mini-batch training, F ED FA follows a stochastic optimization strategy to randomly select a subset of layers from { h k } Kk =1 and add noises to them.
Revised: Concretely, in each mini-batch training, F ED FA follows a stochastic optimization strategy to randomly select a subset of layers from { h k } Kk =1 and add noises to them.

Original: For simplicity, we denote
Revised: For simplicity, we denote

Original: K = { 1 , . . .
Revised: K = { 1 , . . .

Original: , K } as the index of all layers in h , Z ⊆K as the subset of layer indexes that are selected,
Revised: , K } as the index of all layers in h , Z ⊆K as the subset of layer indexes that are selected,

Original: E = { e zm } ∀ z ∈Z as the corresponding set of noises.
Revised: E = { e zm } ∀ z ∈Z as the corresponding set of noises.

Original: Then, the loss function L F ED FA m of client m in
Revised: Then, the loss function L F ED FA m of client m in

Original: F ED FA can be equivalently written as L F ED FA m =
Revised: F ED FA can be equivalently written as L F ED FA m =

Original: E Z∼K
Revised: E Z∼K

Original: L Z m , where L Z m is a standard loss function
Revised: L Z m , where L Z m is a standard loss function

Original: L ERM m ( cf.
Revised: L ERM m ( cf.

Original: Eq. 1) imposed by adding noises to layers in Z .
Revised: Eq. 1) imposed by adding noises to layers in Z .

Original: In the remainder, we relate the loss function L F ED FA m to the original ERM loss L ERM m as well as a regularization term conditioned on E .
Revised: In the remainder, we relate the loss function L F ED FA m to the original ERM loss L ERM m as well as a regularization term conditioned on E .

Original: Theorem 1.
Revised: Theorem 1.

Original: In F ED FA , the loss function L F ED FA m of client m can be expressed as:
Revised: In F ED FA , the loss function L F ED FA m of client m can be expressed as:

Original: [Equation]
Revised: [Equation]

Original: where L ERM m is the standard ERM loss, and L REG m is the regularization term:
Revised: where L ERM m is the standard ERM loss, and L REG m is the regularization term:

Original: [Equation]
Revised: [Equation]

Original: where J z denotes the Jacobian of layer z (deﬁned in Proposition 1).
Revised: where J z denotes the Jacobian of layer z (deﬁned in Proposition 1).

Original: Theorem 1 implies that, F ED FA implicitly introduces regularization to local client learning by regularizing the gradients of latent representations ( i .
Revised: Theorem 1 implies that, F ED FA implicitly introduces regularization to local client learning by regularizing the gradients of latent representations ( i .

Original: e ., (cid:53)
Revised: e ., (cid:53)

Original: h 1: K ( X m ) (cid:96) ( g ( h 1: K ( X m )) , Y m ) (cid:62) ), weighted by federation-aware noises (
Revised: h 1: K ( X m ) (cid:96) ( g ( h 1: K ( X m )) , Y m ) (cid:62) ), weighted by federation-aware noises (

Original: cid:80) z
Revised: cid:80) z

Original: ∈Z J z ( X m )
Revised: ∈Z J z ( X m )

Original: e zm .
Revised: e zm .

Original: In the remainder of this section, we prove Theorem 1.
Revised: In the remainder of this section, we prove Theorem 1.

Original: For the sake of analysis, we ﬁrst marginalize the effect of the noising process.
Revised: For the sake of analysis, we ﬁrst marginalize the effect of the noising process.

Original: We do so by deﬁning an accumulated noise ˆ e Km in the ﬁnal layer K , which originates from the forward propagation of all noises in E .
Revised: We do so by deﬁning an accumulated noise ˆ e Km in the ﬁnal layer K , which originates from the forward propagation of all noises in E .

Original: We compute the accumulated noise based on (Camuto et al., 2020) that examines
Revised: We compute the accumulated noise based on (Camuto et al., 2020) that examines

Original: Gaussian noise injection into every latent layer in a neural network.
Revised: Gaussian noise injection into every latent layer in a neural network.

Original: Formally, the accumulated noise on the ﬁnal convolutional layer K can be expressed as follows:
Revised: Formally, the accumulated noise on the ﬁnal convolutional layer K can be expressed as follows:

Original: Proposition 1.
Revised: Proposition 1.

Original: Consider a K -layer neural network, in which a random noise e zm is added to the activation of each layer z ∈ Z .
Revised: Consider a K -layer neural network, in which a random noise e zm is added to the activation of each layer z ∈ Z .

Original: Assuming the Hessians, of the form (cid:53) 2 h 1: k ( X m ) | h 1: n ( X m ) where k, n are the indexes over layers, are ﬁnite.
Revised: Assuming the Hessians, of the form (cid:53) 2 h 1: k ( X m ) | h 1: n ( X m ) where k, n are the indexes over layers, are ﬁnite.

Original: Then, the accumulation noise ˆ e Km is approximated as:
Revised: Then, the accumulation noise ˆ e Km is approximated as:

Original: [Equation]
Revised: [Equation]

Original: where J z ∈ R C K × C z indicates the Jacobian of layer z , i.e., J z ( X ) i,j = ∂h 1: K ( X m )
Revised: where J z ∈ R C K × C z indicates the Jacobian of layer z , i.e., J z ( X ) i,j = ∂h 1: K ( X m )

Original: i ∂h 1: z ( X m ) j , where C K and C z denote the number of neurons in layer K and z , respectively.
Revised: i ∂h 1: z ( X m ) j , where C K and C z denote the number of neurons in layer K and z , respectively.

Original: O ( β ) represents higher order terms in E that tend to be zero in the limit of small noises.
Revised: O ( β ) represents higher order terms in E that tend to be zero in the limit of small noises.

Original: Proof of Proposition 1.
Revised: Proof of Proposition 1.

Original: Starting with layer 1 as the ﬁrst convolution layer, the accumulated noise on layer K can be approximated through recursion.
Revised: Starting with layer 1 as the ﬁrst convolution layer, the accumulated noise on layer K can be approximated through recursion.

Original: If K = 1 , the accumulated noise is equal to ˆ e Km = e 1 m .
Revised: If K = 1 , the accumulated noise is equal to ˆ e Km = e 1 m .

Original: For K =2 , we apply Taylor’s theorem on h 2 ( X 1 m + e 1 m ) around the output feature X 1 m at h 1 .
Revised: For K =2 , we apply Taylor’s theorem on h 2 ( X 1 m + e 1 m ) around the output feature X 1 m at h 1 .

Original: If we assume that all values in Hessian of h 2 ( X 1 m ) is ﬁnite, the following approximation holds:[Equation]
Revised: If we assume that all values in Hessian of h 2 ( X 1 m ) is ﬁnite, the following approximation holds:

Original: [Equation]
Revised: [Equation]

Original: where O ( κ 1 ) denotes asymptotically dominated higher order terms given the small noise.
Revised: where O ( κ 1 ) denotes asymptotically dominated higher order terms given the small noise.

Original: In this special case of K = 2 , we obtain the accumulated noise as
Revised: In this special case of K = 2 , we obtain the accumulated noise as

Original: [Equation]
Revised: [Equation]

Original: The noise consists of two components: (cid:16) ∂h
Revised: The noise consists of two components: (cid:16) ∂h

Original: [Equation]
Revised: [Equation]

Original: X [Equation] ∂ X
Revised: e

Original: [Equation]
Revised: [Equation]

Original: me 1 m + O ( κ 1 ) (cid:17) is the noise propagated from h 1 , while ˇ e 2 m = e 2 m is the noise added to h 2 if the layer is activated; otherwise, ˇ e 2 m =0 .
Revised: m + O ( κ 1 ) (cid:17) is the noise propagated from h 1 , while ˇ e 2 m = e 2 m is the noise added to h 2 if the layer is activated; otherwise, ˇ e 2 m =0 .

Original: Note that Eq.can be generalized to an arbitrary layer.
Revised: Note that Eq.can be generalized to an arbitrary layer.

Original: Repeating this process for each layer z ∈ Z , and assuming that all Hessians of the form (cid:53) 2 h k ( X m ) | h n ( X m ) , ∀ k < n are ﬁnite, we obtain the accumulated noise for layer K as
Revised: Repeating this process for each layer z ∈ Z , and assuming that all Hessians of the form (cid:53) 2 h k ( X m ) | h n ( X m ) , ∀ k < n are ﬁnite, we obtain the accumulated noise for layer K as

Original: [Equation]
Revised: [Equation]

Original: where ˇ e Km =
Revised: where ˇ e Km =

Original: e Km is the noise added to h K if layer K is activated; otherwise, ˇ e Km =0 .
Revised: e Km is the noise added to h K if layer K is activated; otherwise, ˇ e Km =0 .

Original: Denoting ∂h 1: K ( X m )
Revised: Denoting ∂h 1: K ( X m )

Original: ∂h 1: z ( X m ) as the Jacobian J z
Revised: ∂h 1: z ( X m ) as the Jacobian J z

Original: ( X m ) completes the proof.
Revised: ( X m ) completes the proof.

Original: Based on Proposition 1, we provide a linear approximation of the loss (cid:96) for samples ( X m , Y m ) as 
Revised: Based on Proposition 1, we provide a linear approximation of the loss (cid:96) for samples ( X m , Y m ) as (cid:96) ( g

Original: [Equation]
Revised: [Equation]

Original: g
Revised: h

Original: [Equation]
Revised: [Equation]

Original: h
Revised: K

Original: [Equation]
Revised: [Equation]

Original: K
Revised: , Y

Original: [Equation]
Revised: [Equation]

Original: in which the higher order terms in Proposition 1 are neglected.
Revised: in which the higher order terms in Proposition 1 are neglected.

Original: Based on Eq.
Revised: Based on Eq.

Original: 19, we further approximate the local training objective L F ED FA m in client m and derive the regularization term as follows:
Revised: 23, we further approximate the local training objective L F ED FA m in client m and derive the regularization term as follows:

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: C A DDITIONAL A BLATION S TUDY
Revised: C A DDITIONAL A BLATION S TUDY

Original: [Table]
Revised: [Table]

Original: In this section, we study the sensitivity of FedFA to the set of eligible layers to apply FFA.
Revised: In this section, we study the sensitivity of FedFA to the set of eligible layers to apply FFA.

Original: For notation, we use { 1 } to represent that FFA is applied to the 1st convolutional stage; { 1 , 2 } to represent that FFA is applied to both the 1st and 2nd convolutional stages; and so forth.
Revised: For notation, we use { 1 } to represent that FFA is applied to the 1st convolutional stage; { 1 , 2 } to represent that FFA is applied to both the 1st and 2nd convolutional stages; and so forth.

Original: The results are shown in Table 9.
Revised: The results are shown in Table 8.

Original: We observe that i) our default design (using ﬁve layers) always shows the best performance on the three datasets (Ofﬁce, DomainNet and ProstateMRI).
Revised: We observe that i) our default design (using ﬁve layers) always shows the best performance on the three datasets (Ofﬁce, DomainNet and ProstateMRI).

Original: We conjecture that this is due to its potential to beget more comprehensive augmentation; ii) applying FFA to only one particular layer brings minor gains against FedAvg; but iii) by adding more layers, the performance tends to improve.
Revised: We conjecture that this is due to its potential to beget more comprehensive augmentation; ii) applying FFA to only one particular layer brings minor gains against FedAvg; but iii) by adding more layers, the performance tends to improve.

Original: This implies that our approach beneﬁts from inherent complementarity of features in different network layers.
Revised: This implies that our approach beneﬁts from inherent complementarity of features in different network layers.

Original: D E XPERIMENTAL D ETAILS D. D ATASET
Revised: D A NALYSIS OF A DDITIONAL C OMMUNICATION C OST IN F ED FA

Revised: In each round, F ED FA incurs additional communication costs since it requires the sending 1) from client to server the momentum feature statistics ¯ µ km and ¯ σ km , as well as 2) from server to client the client sharing feature statistic variances γ µ k and γ σ k at each layer k . Thus, for K layers in total, the extra communication cost for each client is c e = 4 (cid:80) Kk =1 C k , where the factor of 4 is for server receiving and sending two statistic values.

Revised: As presented in Table 10 and Table 14, we append one

Revised: FFA layer after each convolutional stage of feature extractors in AlexNet and U-Net.

Revised: Hence, the total additional communication costs for AlexNet and U-Net are:

Revised: AlexNet: 4 × (64 + 192 + 384 + 256 + 256) / 1024 × 4 = 18 KB , U-Net: 4 × (32 + 64 + 128 +

Original: {+256 + 512) / 
+}
Revised: {+256 + 512) / 
+}

Revised: × 4 = 15 .

Revised: 5 KB .

Revised: [Equation]

Revised: However, it should be noted that these additional costs are minor in comparison with the cost required for exchanging model parameters, which are 2 × 49 .

Revised: 5 MB and 2 × 29 .

Revised: 6 MB for AlexNet and

Revised: U-Net, respectively.

Revised: E E XPERIMENTAL D ETAILS E. D ATASET

Original: We conduct extensive experiments on ﬁve datasets:
Revised: We conduct extensive experiments on ﬁve datasets:

Original: Ofﬁce-Caltech 10 (Gong et al., 2012) has four data sources, three from Ofﬁce-31 (Saenko et al., 2010) and one from Caltech-256 (Grifﬁn et al., 2007).
Revised: Ofﬁce-Caltech 10 (Gong et al., 2012) has four data sources, three from Ofﬁce-31 (Saenko et al., 2010) and one from Caltech-256 (Grifﬁn et al., 2007).

Original: They are collected from different camera devices or in diverse environments with different background.
Revised: They are collected from different camera devices or in diverse environments with different background.

Original: DomainNet (Peng et al., 2019) contains images from six domains (clipart, infograph, painting, real, and sketch), which are collected by searching a category name along with a domain name in different search engines.
Revised: DomainNet (Peng et al., 2019) contains images from six domains (clipart, infograph, painting, real, and sketch), which are collected by searching a category name along with a domain name in different search engines.

Original: ProstateMRI (Liu et al., 2020b) is a multi-site prostate segmentation dataset consisting of six data sources of T2-weighted MRI from different medical institutions.
Revised: ProstateMRI (Liu et al., 2020b) is a multi-site prostate segmentation dataset consisting of six data sources of T2-weighted MRI from different medical institutions.

Original: For all the three datasets, we regard each data source as a client, and thus real-world feature shift exists among clients.
Revised: For all the three datasets, we regard each data source as a client, and thus real-world feature shift exists among clients.

Original: CIFAR-10 (Krizhevsky & Hinton, 2009) is a popular natural image classiﬁcation dataset for federated learning.
Revised: CIFAR-10 (Krizhevsky & Hinton, 2009) is a popular natural image classiﬁcation dataset for federated learning.

Original: It contains 50,000 training and 10,000 test images.
Revised: It contains 50,000 training and 10,000 test images.

Original: We introduce label distribution heterogeneity for the dataset, by sampling local data based on the Dirichlet distribution Dir ( α ) , and consider two different concentration parameters, i .
Revised: We introduce label distribution heterogeneity for the dataset, by sampling local data based on the Dirichlet distribution Dir ( α ) , and consider two different concentration parameters, i .

Original: e .,
Revised: e .,

Original: α = 0 .
Revised: α = 0 .

Original: 6 or 0 . 3 .
Revised: 6 or 0 . 3 .

Original: [Table] [Table] EMNIST (Cohen et al., 2017) is an image classiﬁcation dataset with 62 classes, including allcapital and small letter of alphabet as well as numbers.
Revised: EMNIST (Cohen et al., 2017) is an image classiﬁcation dataset with 62 classes, including allcapital and small letter of alphabet as well as numbers.

Original: We follow the setup in F ED M IX to simulate data size heterogeneity.
Revised: We follow the setup in F ED M IX to simulate data size heterogeneity.

Original: In Table 10, we summarize the conﬁguration of our experiments for each of the datasets.
Revised: In Table 9, we summarize the conﬁguration of our experiments for each of the datasets.

Original: D. E XPERIMENTAL D ETAILS FOR I MAGE C LASSIFICATION
Revised: E. E XPERIMENTAL D ETAILS FOR I MAGE C LASSIFICATION

Original: Network Architecture.
Revised: Network Architecture.

Original: For the image classiﬁcation tasks on Ofﬁce-Caltech 10 (Gong et al., 2012) and DomainNet (Peng et al., 2019), we use an adapted AlexNet (Krizhevsky et al., 2017), with the detailed network architecture shown in Table 11.
Revised: For the image classiﬁcation tasks on Ofﬁce-Caltech 10 (Gong et al., 2012) and DomainNet (Peng et al., 2019), we use an adapted AlexNet (Krizhevsky et al., 2017), with the detailed network architecture shown in Table 10.

Original: Training Details.
Revised: Details.

Original: For each training image in Ofﬁce-Caltech10 and DomainNet, we reshape its size into 256 × 256 .
Revised: For each training image in Ofﬁce-Caltech10 and DomainNet, we reshape its size into 256 × 256 .

Original: We train AlexNet with the SGD optimizier with a learning rate of 0 .
Revised: We train AlexNet with the SGD optimizier with a learning rate of 0 .

Original: 01 , a mini-batch size of 32 , using the standard cross-entropy loss.
Revised: 01 , a mini-batch size of 32 , using the standard cross-entropy loss.

Original: The total number of communication round is set to 400 , with one local epoch per round by default.
Revised: The total number of communication round is set to 400 , with one local epoch per round by default.

Original: Two basic data augmentation techniques are also applied for training, i .
Revised: Two basic data augmentation techniques are also applied for training, i .

Original: e ., random horizontal ﬂipping and random rotation with degree in [ − 30 , 30] .
Revised: e ., random horizontal ﬂipping and random rotation with degree in [ − 30 , 30] .

Original: The dataset splits of Ofﬁce-Caltech 10 and DomainNet in our experiments are summarized in Tableand Table 13, respectively. 
Revised: The dataset splits of Ofﬁce-Caltech 10 and DomainNet in our experiments are summarized in Tableand Table 12, respectively. [Table] [Table] [Table]

Original: [Table] [Table]
Revised: [Table] [Table]

Original: D. E XPERIMENTAL D ETAILS FOR M EDICAL I MAGE S EGMENTATION
Revised: E. E XPERIMENTAL D ETAILS FOR M EDICAL I MAGE S EGMENTATION

Original: Network Architecture.
Revised: Network Architecture.

Original: For medical image segmentation on ProstateMRI (Liu et al., 2020b), we use a vanilla U-Net architecture, as presented in Table 15 and Table 16.
Revised: For medical image segmentation on ProstateMRI (Liu et al., 2020b), we use a vanilla U-Net architecture, as presented in Table 14 and Table 15.

Original: Training Details.
Revised: Training Details.

Original: Following F ED H ARMO , we use a combination of standard cross-entropy and
Revised: Following F ED H ARMO , we use a combination of standard cross-entropy and

Original: Dice loss to train the network, using the Adam optimizer with learning rate 1e-4, batch size 16, and weight decay 1e-4.
Revised: Dice loss to train the network, using the Adam optimizer with learning rate 1e-4, batch size 16, and weight decay 1e-4.

Original: No any data augmentation techniques are applied.
Revised: No any data augmentation techniques are applied.

Original: The dataset splits of
Revised: The dataset splits of

Original: ProstateMRI used in our experiments are summarized in Table
Revised: ProstateMRI used in our experiments are summarized in Table

Original: Additional Results.
Revised: Additional Results.

Original: Table 17 provides a detailed performance statistic of different methods on
Revised: Table 16 provides a detailed performance statistic of different methods on

Original: ProstateMRI, w.r
Revised: ProstateMRI, w.r

Original: [Table]
Revised: [Table]

Original: t. different fractions ( 1 / 6 , 2 / 6 , 3 / 6 , 4 / 6 , 5 / 6 , 1 ) of training samples used in each client.
Revised: t. different fractions ( 1 / 6 , 2 / 6 , 3 / 6 , 4 / 6 , 5 / 6 , 1 ) of training samples used in each client.

Original: The table corresponds to the plot in Fig.
Revised: The table corresponds to the plot in Fig.

Original: [Table]

Original: [Table]

Original: [Table]
Revised: [Table]

Original: In each round, F ED FA incurs additional communication costs since it requires the sending 1) from client to server the momentum feature statistics ¯ µ km and ¯ σ km , as well as 2) from server to client the client sharing feature statistic variances γ µ k and γ σ k at each layer k . Thus, for K layers in total, the extra communication cost for each client is c e = 4 (cid:80) Kk =1 C k , where the factor of 4 is for server receiving and sending two statistic values.
Revised: [Table]

Original: As presented in Table 11 and Table 15, we append one FFA layer after each convolutional stage of feature extractors in AlexNet and U-Net. Hence, the total additional communication costs for AlexNet and U-Net are: [Equation]
Revised: [Table]

Original: However, it should be noted that these additional costs are minor in comparison with the cost required for exchanging model parameters, which are 2 × 49 . 5 MB and 2 × 29 . 6 MB for AlexNet and U-Net, respectively.
Revised: [Table]

Original: [Table]
Revised: [Table]