Robust and Personalized Federated Learning with Spurious Features: an Adversarial Approach

1 INTRODUCTION . Federated learning ( FL ) is a leading framework for clients to collaboratively train a shared global machine learning ( ML ) model without releasing their local private datasets ( McMahan et al. , 2017 ; Kairouz et al. , 2019 ) . The jointly trained global model could be further fine-tuned on each client ’ s local dataset to produce personalized models ( Fallah et al. , 2020 ; T. Dinh et al. , 2020 ; Li et al. , 2021 ) . While existing theoretical and empirical results highlight how personalized models improve accuracy on local data , few works consider what features the personalized models learn from the local dataset . Our motivating hypothesis is that not all local features are beneficial . For example , consider a gender prediction task using face images , where the ML model learns to predict gender based on hair color because females are more likely to have blond hair ( Sagawa et al. , 2020 ) . In this case , the hair color is called a spurious feature because it only statistically correlates with the gender on biased-aligned samples but not necessarily on the overall population . Thus , the accuracy of a model that relies on spurious features such as hair color is likely to drop significantly on bias-conflicting samples where the spurious correlation does not hold , e.g. , for blond male images ( Sagawa et al. , 2020 ) . This paper calls the accuracy difference of an ML model on the dataset with spurious features and the dataset without spurious features accuracy disparity . More broadly , the accuracy disparity caused by spurious features leads to issues in both fairness ( McNamara et al. , 2019 ; Zhao & Gordon , 2019 ; Agarwal et al. , 2019 ; Chi et al. , 2021 ) , i.e. , racial bias ( Khani & Liang , 2021 ) and robustness , i.e. , accuracy decrease under distribution shift ( Zhao et al. , 2019 ; Koh et al. , 2021 ) . Compared to the global model , because of the local fine-tuning , the personalized models are more vulnerable to spurious features and have a larger accuracy disparity . Empirically , we observe that the typical non-i.i.d . data distributions in FL settings reduce the accuracy disparity of the global model . One potential explanation is that the statistical heterogeneity ( Wang et al. , 2021 ) under spurious features across users is larger than that of non-spurious features . As a result , the aggregation operation in the central server averages out the diverse shifts under the spurious features ( Figure 1 ) so that the global model becomes more robust against spurious features on the benchmark datasets . On the other hand , the local fine-tuning step will result in a personalized model that entangles spurious features and becomes biased . Hence , it remains a key challenge how to debias the personalized models . Various methods ( Wang et al. , 2019 ; Sagawa et al. , 2020 ; Liu et al. , 2021 ) have been developed to disentangle spurious features from ML models , but few can be easily applied to personalized models under the federated learning setting . The main reason lies in their reliance on bias-conflicting samples . The bias-conflicting samples can be rare . For example , in the CelebA dataset , only around 1.7k samples are bias-conflicting out of more than170k samples . If we distribute the CelebA dataset across users , nearly half do not have any bias-conflicting samples . Prior work ( Li & Wang , 2019 ) has tried using a global proxy dataset for training FL models . However , a global proxy dataset may not characterize the local samples well . Furthermore , estimating the accuracy disparity of the personalized model becomes difficult without access to bias-conflicting samples . To approach this problem , we propose a novel method to reduce the accuracy disparity of personalized models . Our proposed method does not rely on the bias-conflicting samples , e.g. , blond male images in the aforementioned task of gender prediction , which are often rare and may not be available for every user . Instead , inspired by prior works ( Tramèr et al. , 2017 ; Liang et al. , 2021 ) on the transferability of adversarial examples , we use the global model as a reference and the adversarial transferability between the global model and the personalized models as a proxy to estimate the accuracy disparity of the personalized models . The intuition is that if two ML models use disjoint subsets of features , the adversarial examples that one ML model generates do not transfer to the other ML model . Based on this intuition , we propose the following hypothesis : If the personalized models entangle spurious features ( thus increasing the accuracy . disparity ) , the adversarial examples generated by the global model ( that uses non-spurious features ) do not transfer to the personalized models . Empirically , we show that the adversarial transferability between the global and personalized models strongly correlates with the accuracy disparity between the global and personalized models ( Section 4 ) , validating our hypothesis . We further theoretically connect the adversarial transferability to the accuracy disparity ( Section 5 ) . Based on the empirical observations and theoretical results , we develop a method that enforces adversarial transferability between the global and the personalized models to reduce the accuracy disparity of personalized models ( Section 6 ) . Our contributions are summarized as follows : • We empirically evaluate the accuracy disparity of the global and personalized models in a federated learning setting with spurious features , highlighting a risk of existing personalization methods . • We design a method to estimate the accuracy disparity of the personalized models , based on the low accuracy disparity global model and the adversarial transferability between the global and personalized models . • We theoretically connect the adversarial transferability and the accuracy disparity of the global and personalized models . • We develop a methods to reduce the accuracy disparity of personalized models by enforcing the adversarial transferability between the global and personalized models . Empirically , we conduct extensive experiments to validate the effectiveness of the proposed methods in mitigating accuracy disparity under the FL setting . Our experiments on MNIST ( Deng , 2012 ) , CelebA ( Liu et al. , 2015 ) , and Coil20 ( Nene et al. , 1996 ) datasets show that the proposed approach reduces the accuracy disparity of personalized models from 15.12 % to 2.15 % , which is closer to that of the global model ( −0.63 % ) . Our method also preserves the benefit of the enhanced average accuracy from fine-tuning , resulting in 3.43 % accuracy improvement on the biased test set and 0.85 % accuracy improvement on the biased-conflicting test set . 2 RELATED WORK . Personalized Federated Learning Fine-tuning is typical for personalization methods . The metalearning-based method first trains a global model and fine-tunes the global model locally ( Fallah et al. , 2020 ) . Other methods using multi-task learning ( Li et al. , 2021 ) or Moreau envelopes ( T. Dinh et al. , 2020 ) have an interpretation as fine-tuning the local model along with training the global model . Fine-tuning is also compatible with clustering-based method ( Ghosh et al. , 2020 ) . Debiasing ML Models A few prior works ( Li & Vasconcelos , 2019 ; Sagawa et al. , 2020 ) utilize group labels , which might require human annotation , to debias ML models . For example , the group distributional robust optimization ( DRO ) method ( Sagawa et al. , 2020 ) aims to optimize the worst-case error rate of ML models across different ( often manually annotated ) groups . Some groups contain bias-conflicting samples while others do not . Residual learning-based methods ( He et al. , 2019 ; Nam et al. , 2020 ; Liu et al. , 2021 ) train a biased ML model and up-weight the residual , which mainly contains bias-conflicting samples that the biased ML model mis-predicts . Chi et al . ( 2021 ) aims to mitigate the accuracy disparity in regression problems via learning the appropriate representations . However , all these methods rely on the explicit access to the bias-conflicting samples , making them difficult to apply on personalized federated learning , where bias-conflicting samples may not be accessible for every client . 3 PRELIMINARIES . Definitions and Notation A data sample is a vector x = [ xr , xs ] , where xr corresponds to the robust and non-spurious features and xs are the spurious features . ds is the dimension of spurious features . Let y be a label , and define ` : Y × Y −→ R to be a λ-smooth , twice differentiable loss function and L ( f , D ) = E ( x , y ) ∼D [ ` ( f ( x ) , y ) ] to be the empirical risk . wg and wp are the weights for the global model fg : X −→ Y and personalized model fp : X −→ Y , respectively . γ is the ratio between the gradient norms of the global and personalized models , ‖∇x ` ( fg ( x ) , y ) ‖‖∇x ` ( fp ( x ) , y ) ‖ . Let supp ( D ) be the support of distribution D. We define a global data distribution Dg , a biased local data distribution Db , and assume a bias-conflicting local data distribution Dbc . We define the “ pseudo-gradient ” as the difference between the updated local model and the global model from the previous round ( sometimes we will use the term “ gradient ” when it is clear from context ) . 〈· , ·〉 denotes an inner product of two vectors and ·_ · denotes a concatenation of two vectors . θ is the angle between ∇x ` ( fg ( x ) , y ) and ∇x ` ( fp ( x ) , y ) . θg is the angle between ∇x ` ( fg ( x ) , y ) and ∇xr ` ( fg ( x ) , y ) _ 0 , which measures the entanglement of the global model to spurious features . Training and Personalization Methods We train the global model using the federated averaging algorithm ( McMahan et al. , 2017 ) , which learns a model f : X −→ Y that minimizes : L ( f , Dg ) =∑N i=1 |Dbi |/|Dg| · Ex , y∼Dbi ` ( f ( x ) , y ) , where N is the number of clients , Dbi is the biased local dataset for client i , and Dg = ∪Ni=1Dbi is the global dataset . When the global model fg converges , fg will be sent to local clients for further fine-tuning by minimizing Ex , y∼Dbi ` ( f ( x ) , y ) . Adversarial Examples and Transferability We can generate an adversarial example xadv given data sample x with label y by solving : xadv = arg max ‖x′−x‖≤ ` ( f ( x′ ) , y ) , ( 1 ) where f is the victim ML model , and is the attack budget . In this example , we consider the L2 attacks , but extensions to general Lp attacks are straightforward . We say an adversarial example to be transferable if it also fools another ML model ( e.g. , a personalized model ) other than the original victim model f ( e.g. , the global model ) .

Summary: The authors have proposed a new FL training strategy to reduce the performance discrepancy between the central model and the client models. The PGD generated adversarial examples are fed into both central model and client models and their outputs are used to minimize the entropy loss. The computation is further simplified using Taylor expansion.