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

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: 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 instance-level feature 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 local feature 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) represent a special type of “features” that encodes domain-speciﬁc characteristics; hence, proper manipulations of the local feature statistics in the federation may beget novel domains, which can potentially alleviate local feature shift and beneﬁt collaborative learning.

Original: Based on this insight, we model each feature statistic probabilistically via a Gaussian distribution, with the mean corresponding to the original statistic and the variance quantifying the augmentation scope, from which novel feature statistics can be drawn to fulﬁll augmentation.
Revised: Based on this insight, we model each feature statistic probabilistically augmented from a Gaussian distribution, with the mean corresponding to the original statistic and the variance quantifying the augmentation scope, from which novel feature statistics can be drawn to fulﬁll augmentation.

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.
Revised: We demonstrate through extensive experiments that F ED FA can signiﬁcantly advance federated learning in diverse scenarios.

Original: The code will be released.
Revised: The code will be released.

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 coming 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 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 (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.

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 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 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).

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.
Revised: To address this, F EDM IX (Yoon et al., 2021) generalizes Mixup (Zhang et al., 2018) into federated learning.

Original: It follows the basic idea of vanilla Mixup, and performs cross-client sample mixup by exchanging averaged data across clients.
Revised: It follows the basic idea of vanilla Mixup, and performs cross-client sample mixup by exchanging 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 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 paper, we introduce a novel feature augmentation technique F ED FA for federated learning, based on a probabilistic modeling of feature statistics ( i .

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

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

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

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 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.

Original: This motivates us to turn the problem of tackling feature shift into tackling feature statistic shift , for which we propose
Revised: This motivates us to turn the problem of tackling feature shift into tackling feature statistic shift , for which we propose

Original: F ED FA to perform feature statistic augmentation within each local client by incorporating more universal feature statistic information characterized by all clients in the federation.
Revised: F ED FA to perform feature statistic augmentation within each local client by incorporating more universal feature statistic information characterized by all clients in the federation.

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

Original: The Gaussian mean corresponds to the original statistic, and variance reﬂects the potential local distribution shift.
Revised: The Gaussian mean corresponds 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, we can sample novel statistics effortlessly from the Gaussian distribution as augmented “features of domain”.

Original: For effective augmentation, we determine a proper variance based on variances of feature statistics within each client, which are regulated by universal variances of feature statistics characterized by all participating clients.
Revised: For effective augmentation, we determine a proper variance based on variances of feature statistics within each client, which are regulated by universal variances of feature statistics 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 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.

Original: F ED FA is a conceptually simple but surprisingly effective method.
Revised: F 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
Revised: It is non-parametric, requires negligible additional computational resources, and can be seamlessly incorporated into arbitrary

Original: CNN architectures.
Revised: 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.
Revised: 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.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Federated Learning.
Revised: Federated Learning.

Original: 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: 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 &

Original: Shmatikov, 2015), i .
Revised: Shmatikov, 2015), i .

Original: e ., train a global model on distributed datasets without disclosing private data information.
Revised: e ., train a global model on distributed datasets without disclosing private data information.

Original: 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.
Revised: 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.

Original: However, F ED A VG is designed for i.i.d.
Revised: However, F ED A VG is designed for i.i.d.

Original: data and suffers in statistical accuracy or even diverge if deployed over non-i.i.d.
Revised: data and suffers in statistical accuracy or even diverge if deployed over non-i.i.d.

Original: client samples.
Revised: 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
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

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).
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).

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.
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.

Original: It adapts the well-known Mixup algorithm (Zhang et al., 2018) from centralized learning into the federated learning scenario.
Revised: 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.
Revised: 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.
Revised: 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.
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.

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).
Revised: 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.
Revised: Data Augmentation.

Original: Data augmentation has a long and rich history in machine learning.
Revised: 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.
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.

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).
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).

Original: In addition, there is a recent trend for label-perturbing augmentation, e . g .,
Revised: 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).
Revised: 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;
Revised: Separate from these input-level augmentation techniques are feature augmentation methods (Verma et al., 2019;

Original: Li et al., 2021;
Revised: Li et al., 2021a;

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

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

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

Original: It is inspired by recent efforts on implicit feature augmentation (Li et al., 2021;
Revised: 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 .
Revised: 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.
Revised: 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  spanning with the variances.
Revised: 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., 2021;
Revised: 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.
Revised: 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 3. 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  m ( w )
Revised: where L ERM m ( w [Equation]

Original: =
Revised: =

Revised: E ( x i

Revised: [Equation]

Original: E ( x i ,y i ) ∼ P m [ (cid:96) i
Revised: y i [Equation] ∼ P m [ (cid:96) [Equation]

Original: ( g ◦ h ( x i ) , y i
Revised: ( g ◦ h ( x [Equation] y i

Original: [Equation]
Revised: [Equation]

Original: w
Revised: w

Original: [Equation] 
Revised: [Equation] [Equation]

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  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 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  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 em ( 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: 3. 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 e m 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 ., Mixup (Zhang et al., 2018), Cutmix (Yun et al., 2019), have been developed.

Original: Simply applying them to local clients, though allowing for performance improvements (see Table 5), 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: 3. 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 Modeling.
Revised: Probabilistic Feature Statistic Augmentation.

Original: Instead of pursuing an explicit modeling of
Revised: Instead of pursuing an explicit modeling of

Original: V km ( ˆ X km | X km ) , our method performs feature augmentation in an implicit manner by augmenting channel-wise feature statistics .
Revised: V km ( ˆ X km | X km ) , our method performs feature augmentation in an implicit manner by augmenting 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 σ km are given as follows:
Revised: e ., mean µ km and standard deviation σ k m are given as follows:

Original: [Equation]
Revised: [Equation]

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

Original: As the abstract of latent features, thesefeature 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 has been recently used for data augmentation in deep recognition (Li et al., 2021).
Revised: They are instrumental to image generation (Huang & Belongie, 2017), and have been recently used for data augmentation in deep 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: Speciﬁcally, we hypothesize that each feature statistic follows 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 . e .,

Original: e ., µ 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 an estimation of 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: 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.
Revised: 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.

Original: 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.
Revised: 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.

Original: This motivates us to estimate the variance of Gaussian from the variance of feature statistics .
Revised: This motivates us to estimate the variance of Gaussian from the variance of feature statistics .

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 particular client.

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 statistic variances are solely computed based on the data in each 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  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: Along with the aggregated model parameters in F ED A VG , these variances are distributed back to each client to inform a global estimation of feature statistic variances.

Original: Note that { Σ 2 µ k , Σ 2 σ k } are shared by all participating clients.

Original: Adaptive Variance Fusion.

Original: Client sharing estimations { Σ 2 µ k , Σ 2 σ 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 { Σ 2 µ k , Σ 2 σ k } , so that each client has a sense of such difference.

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,  the client sharing feature statistic variances of the j th channel at layer h
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: k turn into:
Revised: (Eq. 5), respectively. They are modulated by the t-distribution as follows:

Original: [Equation]
Revised: [Equation]

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

Revised: By applying Eq.

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

Revised: =

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

Original:  In this way, the channels with higher variances in Σ 2 µ k or Σ 2 σ k  will be assigned with even higher importance  than other channels, which enables more extensive augmentation along those directions.
Revised: , γ ( C ) µ k ] ∈ R C and γ σ k = [ γ (1) σ k , . . , γ ( C ) σ k ] ∈ R C as modulated statistic variances of all feature channels at layer h k . 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: Adaptive Variance Fusion.

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.

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 follows:
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]

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 apply the re-parameterization trick (Kingma & Welling, 2013) in our implementation:
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 semantic augmentation.
Revised: 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.

Original: In our implementation, we add a FFA layer after each convolutional stage in AlexNet or U-Net (see Table 8 and Table 12 for detailed network structures).
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) 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).

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

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

Original: E XPERIMENT

Original: We examine the efﬁcacy of F ED FA against relevant algorithms ( e . g ., F ED A VG , F ED P ROX , F ED BN, F ED SAM, F ED M IX ) on three datasets ( i . e ., Ofﬁce-Caltech 10, DomainNet, ProstateMRI) in two tasks, i .
Revised: 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.

Original: e ., image classiﬁcation and medical image segmentation.

Original:   4. E XPERIMENTAL S ETUP
Revised:  E XPERIMENT 4. E XPERIMENTAL S ETUP

Original: Datasets.
Revised: Datasets.

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

Original: e ., Ofﬁce-Caltech 10 (Gong et al., 2012) and DomainNet (Peng et al., 2019) for image classiﬁcation, ProstateMRI (Liu et al., 2020b) for medical image segmentation. Ofﬁce-Caltech 10 has four data sources, three from Ofﬁce- (Saenko et al., 2010) andone from Caltech-256 (Grifﬁn et al., 2007).
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: They are collected from different camera devices or in diverse environments with different background. DomainNet 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. ProstateMRI is a multi-site prostate segmentation dataset consisting of six data sources of T2-weighted MRI from different medical institutions. For all the three datasets, we regard each data source as a client, and thus real-world feature shift exists among clients.
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) and some variants such as 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).
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: We also compare withF ED M IX (Yoon et al., 2021) that generalizes mixup into the federated learning scenarios.
Revised: Moreover, we compare with

Original: Moreover, we compare with a specialized algorithm F ED H ARMO (Jiang et al., 2022)  designed for medical image applications in ProstateMRI.
Revised: F ED H ARMO (Jiang et al., 2022) in ProstateMRI, that is specialized designed for medical imaging.

Original: In addition, to gain more insights into F ED FA, we design two baseline variants: F ED FA-R( ANDOM ) and F ED FA-C( 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 λ is set to 0 .
Revised: e ., ˆΣ µ km = ˆΣ σ km = λ , where λ = 0 .

Original: F ED FA-C performs augmentation based only on client speciﬁcstatistic 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. Following

Original: Two learning models are adopted for the two tasks: i) AlexNet (Krizhevsky et al., 2017) on Ofﬁce-Caltechand DomainNet using the SGD optimizer with learning rate 0 .
Revised: F ED BN (Li et al., 2020b), we adopt AlexNet (Krizhevsky et al., 2017) on Ofﬁce-Caltech 10 and DomainNet, using the SGD optimizer with learning rate 0 .

Original: 01 and batch size 32 and ii) U-Net (Ronneberger et al., 2015) on ProstateMRI using Adam as the optimizer with learning rate 1eand batch size 16 .
Revised: 01 and batch size 32 . Following F EDH ARMO (Jiang et al., 2022), we employ U-Net (Ronneberger et al., 2015) on ProstateMRI using 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: We run each experiment for ﬁve times with different seeds and report the mean performance.
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. writers, and train a LeNet-5 (LeCun et al., 1998) using SGD with batch size 10. The total number of clients is 200 and only 10 clients are sampled per communication round.

Original: For all experiments, we set α and p to 0 .
Revised: We run 200 rounds in total.

Original: 99 and 0 .{+train a ResNet-
 (He et al., 2016). The number of clients is 100 with participation rate 0.1, while the number of communication round is set to 100.+}
Revised: For CIFAR-10, we sample local data based on Dirichlet distribution Dir ( α ) to simulate label distribution heterogeneity. As (Qu et al., 2022; Kim et al., 2022), we set α to 0.3 or 0.6, and {+train a ResNet-
 (He et al., 2016). The number of clients is 100 with participation rate 0.1, while the number of communication round is set to 100.+}

Original: 5 by default. More details are provided in Appendices B-C.
Revised: 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: communi- cation rounds on Ofﬁce-Caltech 10.

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

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

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

Original:  Results on Ofﬁce-Caltech 10 and DomainNet.
Revised: CIFAR-10 in Table 3. 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.This is en- [Table] [Table]
Revised: Ofﬁce-Caltech 10 and 2.2% on DomainNet.

Original: [Table]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure]

Original: [Table]

Original: [Figure] [Table] [Figure] couraging 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.depicts the convergence curves of comparative methods on Ofﬁce-Caltech 10.
Revised: Moreover, Fig. 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. 
Revised: But as the training goes, F ED FA is able to converge to a more optimal solution. [Table] [Table] [Table]

Revised: [Table]

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

Original: Results on ProstateMRI.
Revised: [Figure] in each client [Figure] age

Revised: [Figure]

Revised: FedSAM

Original:  F ED FA shows leading performance with extremely small local datasets.
Revised: [Figure] FedHarmo FedFA Results on ProstateMRI. 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 specially 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 ,/ 6 , 5 / 6 , 1 ) of training samples .
Revised: e ., 1 / 6 , 2 / 6 , 3 / / 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. [Table]

Original: 4. F EDERATED D OMAIN G ENERALIZATION P ERFORMANCE
Revised: Results on CIFAR-10 and EMNIST. 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). As shown in Table 3, in CIFAR10, F ED FA surpasses the second best method, 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: Notably, as the 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.

Revised: scenarios.

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.

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:  Federated learning are dynamic systems, in which novel clients may enter the system after model training, most possibly with test-time distribution shift.
Revised: challenges in FL. 4. F EDERATED D OMAIN G ENERALIZATION P ERFORMANCE 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  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 [Table] 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, F ED FA is supposed to enforce regularization to neural network learning, which could improve generalization capability.
Revised: With a core of data augmentation [Table] F ED FA is supposed to enforce regularization to neural network learning, which could improve generalization capability.

Original: To verify this, we perform experiments for federated domain generalization based on the leaveone-client-out strategy, i .
Revised: To verify this [Table] 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 3.
Revised: The results are presented in Table 4.

Original: As seen, F ED FA achieves leading gen- [Table] eralization 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: [Table]

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 verifythe efﬁcacy of 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 devicedependent augmentation of feature statistics, without explicitly considering any global information as done in F ED FA.
Revised: Both variants involve only device-dependent augmentation of feature statistics, without explicitly considering any global information.

Original: As shown in Table 4, by randomly perturbing feature statistics, F ED FA-R shows no improvements or even suffers performance degradation on 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 client-sharing feature statistic information in federated augmentation.
Revised: As shown in Table 5, by randomly perturbing feature statistics, F ED FA-R shows no improvements or even suffers performance degradation on 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.

Original: [Table]
Revised: [Table]

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

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

Original: e .,
Revised: e .,

Original: Mixup (Zhang et al., 2018), Manifold Mixup (Verma et al., 2019), MixStyle (Zhou et al., 2021) and MoEx (Li et al., 2021).
Revised: Mixup (Zhang et al., 2018), Manifold Mixup (Verma et al., 2019), MixStyle (Zhou et al., 2021) and MoEx (Li et al., 2021a).

Original: The results are presented in Table 5.
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 .
Revised: We show that i) all the four techniques yield non-trivial improvements over F ED A VG , and some of them ( e [Table] g .,

Original: g ., 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, respectively; iii) combining F ED FA with these techniques allows further performance uplifting, verifying the complementary roles of F ED FA to them.
Revised: 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.

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:ˆΣ 2 µ km =
Revised: We design a baseline “Direct Fusion” that directly combines the client-speciﬁc and client-sharing statistic variances as:

Original: (Σ 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 6 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.
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: [Table]
Revised: [Table]

Original: Estimation of Client Speciﬁc Statistic Variances
Revised: Estimation of Client Speciﬁc Statistic Variances

Original: Beyond mini-Batch.
Revised: 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.
Revised: 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.
Revised: 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
Revised: This allows

Original: F ED FA to estimate variances from a larger number of samples, even with a small batch size.
Revised: 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.
Revised: 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%.
Revised: 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 .
Revised: 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.
Revised: 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  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 illustrated in Fig.

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

Original: Notably, the model yields promising performance at α =0 , in which the model only uses limited client information, i .
Revised: Notably, the model yields promising performance at α = 0 , in which the model only uses limited client information, i .

Original: e ., the feature statistics of the last mini-batch in each local epoch, to compute client-sharing statistic variances.
Revised: e ., 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)
Revised: ii)

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

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

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

Original: The model achieves the best performance at p = 0 .
Revised: The model achieves the best performance at p =0 .

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

Original: 1 to 1 . 0 .
Revised: 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. Compared to F ED A VG , it requires 4 (cid:80)

Original: Compared to F ED A VG , it requires only 18 KB and 15.5 KB more GPU memory allocation to store { ¯ µ km , ¯ σ km } Kk =1 and { Σ 2 µ k , Σ 2 σ k } Kk =1 in AlexNet and U-Net, respectively.
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. Here C k is the number of feature channel at each layer k . The costs are in practice very minor, e . g ., 18 KB/15.5 KB for AlexNet/U-Net.

Original: For comparison, F ED M IX requires around twice more GPU memory than F ED A VG .
Revised: For comparison, F ED M IX requires 2 × more GPU memory than F ED A VG .

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

Original: Communication cost.
Revised: Communication cost.

Original: In eachcommunication round, F ED FA incurs additional communication costs since it requires the sending 1) from client to server the momentum feature statistics ¯ µ km∈ R C and ¯ σ km ∈ R C (Eq. 4), as well as 2) from server to client the client sharing feature statistic[-variances Σ 2 µ k ∈ R C and Σ 
-]
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 Σ 2 µ k ∈ R C and Σ 
-]

Original: σ k ∈ R C (Eq. 5). Thus, for each layer, it requires to transmit 4 C × 32 bit or 16 C Byte additional data in  each round.
Revised: Thus, for K layers in total, the extra communication cost for each client is c e =4 (cid:80)

Original: In our implementation, we add one FFA layer after each convolutional stage of the feature extractors in AlexNet and U-Net. Please see Table 8 and Table 12 in Appendix for detailed architectures. This leads to additional communication costs of 18 KB and 15.5 KB for AlexNet and U-Net, respectively.
Revised: Kk =1 C k , where the factor of 4 is for server receiving and sending two statistic values.

Original: However, we note that these additional costs are quite minor in comparison with the costs for exchanging model parameters, which are 99 MB for AlexNet and 59. MB for U-Net.
Revised: We further denote c m as the cost for exchanging model parameters in FedAvg. In general we have c e (cid:28) c m ( e . g ., c e = 18KB vs. c m = 99MB for AlexNet), hence, the extra communication bruden in F ED FA is almost negligible.

Original: More details are provided in Appendix §D.
Revised: More details are provided in Appendix §E.

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 as more universal feature statistic distribution across clients.
Revised: 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 as more universal feature statistic distribution across clients.

Original: Feature augmentation is implicitly achieved by sampling novel feature statistics from the Gaussian distribution.
Revised: Feature augmentation is implicitly achieved by sampling novel feature statistics from the Gaussian distribution.

Original: We evaluate F ED FA across a suite of federated datasets, and the results conﬁrm its superiority.
Revised: We evaluate F ED FA across a suite of federated datasets, and the results conﬁrm its superiority.

Original: A D ETAILED A LGORITHM
Revised: 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: (cid:46) Run the k -th layer of the feature extractor h 6: ˆ X km , ¯ µ km , ¯ σ km =
Revised: X km = h k X k − Run the k -th layer of 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 
Revised: Run loss computation and backward optimization 9: Σ 2 µ k = 1 M (cid:80) Mm =1 (¯ µ km − E [¯ µ k ])

Original: [Equation]

Original: E [¯ µ k ])(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 , ( j ) µ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 , ( j ) σk ) −14: return γ µ k , γ σ k 
Revised: Cc =1 (1+ 1 / Σ2 , ( c ) σk ) −14: return γ µ k , γ σ k B T HEORETICAL I NSIGHTS

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

Revised: We begin with interpreting

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 (§3.1), and show that F ED FA injects federation-aware noises to latent representations (§B.1).

Original: B E XPERIMENTAL D ETAILS FOR I MAGE C LASSIFICATION
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). Without loss of generality, we conduct all analysis for an arbitrary client m ∈ [ M ] . [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Table]

Revised: [Equation]

Revised: [Table]

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

Revised: [Equation]

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

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

Revised: Noise Injection in Neural Networks.

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

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.

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

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

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

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

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).

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.

Revised: Federation-Aware Noise Injection.

Revised: From Eq.

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

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

Revised: Lemma 1.

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 =

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

Revised: [Equation]

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

Revised: [Equation]

Revised: Proof of Lemma 1.

Revised: We can easily prove this by substituting Eq.

Revised: 9 into Eq.

Revised: [Equation]

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 .

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

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

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

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

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).

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

Revised: g ◦ h , where h =

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

Revised: Given a batch of samples X m with labels

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

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

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

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

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.

Revised: For simplicity, we denote

Revised: K = { 1 , . . .

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

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

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

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

Revised: E Z∼K

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

Revised: L ERM m ( cf.

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

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 .

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: where L ERM m is the standard ERM loss, and L REG m is the regularization term:

Revised: [Equation]

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

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

Revised: 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 (

Revised: cid:80) z

Revised: ∈Z J z ( X m )

Revised: e zm .

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

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

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 .

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

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

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

Revised: Proposition 1.

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

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.

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

Revised: [Equation]

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 )

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.

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

Revised: Proof of Proposition 1.

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

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

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 .

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

Revised: [Equation]

Revised: [Equation]

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

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

Revised: [Equation]

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

Revised: [Equation]

Revised: X

Revised: [Equation]

Revised: ∂

Revised: X

Revised: [Equation]

Revised: m e 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: Note that Eq.can be generalized to an arbitrary layer.

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

Revised: [Equation]

Revised: where ˇ e Km =

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

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

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

Revised: ( X m ) completes the proof.

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

Revised: [Equation]

Revised: g

Revised: [Equation]

Revised: h

Revised: [Equation]

Revised: K

Revised: [Equation]

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

Revised: Based on Eq.

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

Revised: [Equation]

Revised: [Table]

Revised: C A DDITIONAL A BLATION S TUDY

Revised: [Table]

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

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.

Revised: The results are shown in Table 9.

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).

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.

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

Revised: D E XPERIMENTAL D ETAILS D. D ATASET

Revised: We conduct extensive experiments on ﬁve datasets:

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).

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

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.

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.

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

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

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

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 .

Revised: e .,

Revised: α = 0 .

Revised: 6 or 0 . 3 .

Revised: [Table]

Revised: [Table]

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.

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

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

Revised: D. 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 8.
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 11.

Original: Training Details.
Revised: Training 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 10, respectively.[Table] [Table]
Revised: The dataset splits of Ofﬁce-Caltech 10 and DomainNet in our experiments are summarized in Tableand Table 13, respectively.

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

Original: [Table]
Revised: [Table]

Original: C E XPERIMENTAL D ETAILS FOR M EDICAL I MAGE S EGMENTATION
Revised: D. 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 12 and Table 13.
Revised: 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.

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 14 provides a detailed performance statistic of different methods on
Revised: Table 17 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] [Table]
Revised: [Table] [Table]

Original: [Table]
Revised: [Table]

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

Original: In eachcommunication round, F ED FA incurs additional communication costsover F ED A VG , since it requires the sending 1) from client to server the averaged feature statistics ¯ µ km∈ R C and ¯ σ km ∈ R C (Eq. 4), as well as 2) from server to client the client-sharing statistic variances Σ 2 µ k∈ R C and Σ 2 σ k ∈
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)

Original: R C (Eq. 5). Thus, for each layer, it requires to transmit 4 C × 32 bit or 16 C Byte additional data in each round.
Revised: Kk =1 C k , where the factor of 4 is for server receiving and sending two statistic values.

Original: As presented in Table 8 and Table 12, we append one FFA layer after each convolutional stage of feature extractors in AlexNet and U-Net.
Revised: 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.

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

Original: [Equation]
Revised: [Equation]

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 .
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 .

Original: 5 MB and 2 × 29 .
Revised: 5 MB and 2 × 29 .

Original: 6 MB for AlexNet and
Revised: 6 MB for AlexNet and

Original: U-Net, respectively.
Revised: U-Net, respectively.

Original: [Table]
Revised: [Table]