RVFR: Robust Vertical Federated Learning via Feature Subspace Recovery

1 INTRODUCTION . Federated Learning ( FL ) involves distributed training where a central server coordinates with multiple agents to collaboratively train a machine learning model , and the agents keep their own training data , e.g. , due to privacy concerns . The theory and practice of FL has progressed rapidly in recent years ( Konečnỳ et al. , 2016 ; McMahan et al. , 2017 ; Bonawitz et al. , 2017 ; Dayan et al. , 2021 ) . Two broad categories ( Yang et al. , 2019 ) of FL are Horizontal FL ( HFL ) and Vertical FL ( VFL ) . In HFL , each agent has a different training set , while in VFL ( Liu et al. , 2019 ; Chen et al. , 2020 ; Liu et al. , 2020 ) , different agents hold different parts of features for the same set of training data . For example , in VFL for credit score prediction , Agent 1 may have the banking data of a user and Agent 2 may have the shopping history of the same user , while the server holds corresponding labels . Given the distributed training , FL raises new security concerns as the server has less control of the training data and agent devices . A variety of attacks have been demonstrated on HFL and on VFL . Several studies have focused on making HFL more robust , wherein each agent sends the model updates based on local data to the server and the server aggregates these updates . Robust aggregation protocols have been proposed as defenses in HFL against malicious attacks ( Yin et al. , 2018 ; Guerraoui et al. , 2018 ; Blanchard et al. , 2017 ; Fu et al. , 2019 ; Pillutla et al. , 2019 ; Fung et al. , 2020 ; Chen et al. , 2017 ; Xie et al. , 2019b ) . However , it is challenging to defend against malicious attacks in VFL , as there is no clear redundancy among the agents . In fact , there are few studies on robustness exploration for VFL . In this paper , we propose a robust VFL training and inference framework via feature subspace recovery ( RVFR ) , which is able to defend against many types of attacks during both training and inference ( see the attack taxonomy in Section 4.1 ) . In particular , during training , we propose to train each agent ’ s feature extractor separately based on feedback from the server ( Quarantine training stage ) ; the server then performs robust feature subspace recovery given the embedded features provided by different agents ( Robust feature subspace recovery stage ) ; and the server further purifies the features based on the assumption that the fraction of malicious agents is relatively small ( Feature purifying stage ) . Finally , the server trains its global model based on the purified features ( Server training stage ) . An overview of the training process is illustrated in Figure 1 . During the inference phase , the server first purifies the embedded features provided by the agents and then feeds it to the trained global model for prediction . Building upon our framework , we aim to answer the following questions : Is it possible to train a clean model with theoretical guarantees in the presence of poisoned features ? Can we make predictions based on incomplete ( corrupted ) features ? If yes , how many malicious agents/instances and missing features can we tolerate ? Under the RVFR framework , in the robust feature subspace recovery stage , we propose a variant of the Robust AutoEncoder ( Zhou & Paffenroth , 2017 ) to recover the feature subspace . While the existing Robust AutoEncoder has no theoretical justification , we provide theoretical support for it by proving that when the underlying feature subspace is linear , it can exactly recover the linear feature subspace in the presence of corruption/poisoning . In the feature purifying stage ( during both training and inference ) , we propose a novel AutoEncoder-based robust decomposition method that decomposes the potentially corrupted feature vectors into two parts : one lies on the learnt feature subspace ( can be non-linear ) ; and the other contains a block-sparse structure . We theoretically show that under certain conditions , the proposed method can recover the underlying feature vector exactly despite corruption from malicious agents . Technical Contributions : We take a first step towards providing robust VFL against diverse adversarial attacks during both training and inference . Our main contributions include the following : • We propose a novel robust training procedure to defend against backdoor attacks in VFL , and prove that under certain conditions it can exactly recover the underlying uncorrupted features . To the best of our knowledge , this is the first defense for VFL against training attacks with theoretical guarantees . • We provide the first theoretical justification for the Robust AutoEncoder , which may be of independent interest . • We propose a robust VFL inference procedure to defend against a variety of targeted or untargeted attacks ( e.g. , adversarial and missing feature attacks ) . To our knowledge , this is the first defense for VFL against inference-phase attacks with theoretical guarantees . • We conduct extensive experiments on NUS-WIDE and CIFAR-10 datasets , and show that RVFR is significantly more robust than baselines against diverse types of attacks . 2 RELATED WORK . In this section , we briefly review prior work on backdoor attacks and defenses in different types of federated learning frameworks . Backdoor attack and defense in Horizontal Federated Learning . Many recent works , e.g. , Bhagoji et al . ( 2019 ) ; Bagdasaryan et al . ( 2020 ) ; Wang et al . ( 2020 ) ; Xie et al . ( 2019a ) , have demonstrated the vulnerability of HFL to backdoor attacks . For instance , Bhagoji et al . ; Bagdasaryan et al . show that training poisoned local models and submit malicious model updates to the server can mislead the global model effectively . While Xie et al . exploits the distributed nature of HFL and propose a distributed backdoor attack . To defend against backdoor attacks in HFL , several robust federated protocols are proposed to mitigate the attacks during training empirically . For example , Sun et al . shows that clipping the norm of model updates and adding Gaussian noise can mitigate backdoor attacks . Andreina et al . incorporates an additional validation phase to each round of FL to detect backdoor . To provide certain robustness guarantees , some robust aggregation methods are proposed ( Yin et al. , 2018 ; Guerraoui et al. , 2018 ; Blanchard et al. , 2017 ; Fu et al. , 2019 ; Pillutla et al. , 2019 ; Fung et al. , 2020 ; Chen et al. , 2017 ; Xie et al. , 2019b ) ; while such robustness guarantees are not for backdoor attacks directly . Recently , Cao et al . proposes Ensemble FL to defend against backdoor attacks with certifiable robustness . However , the proposed majority voting strategy requires training hundreds of FL models . A recent work exploits model clipping and smoothing in HFL , aiming to provide certified robustness to backdoor attacks ( Xie et al. , 2021 ) . However , none of these works can provide robustness guarantees on the learned feature subspace of the server model , which is the focus of our work . Backdoor attack and defense in Vertical Federated Learning . Backdoor attack against VFL is challenging since in the typical setting ( Chen et al. , 2020 ) the agent does not have the label information . Recent studies assume that the malicious agent knows at least one training instance belonging to the target class ( Liu et al. , 2020 ) . Under this assumption , they propose the gradient-replacement method to perform backdoor attack . Some defense strategies are proposed against such attacks in VFL , e.g. , sparsify the intermediate gradient before sending to the server ( Liu et al. , 2020 ) . However , there is no work providing any robustness certification against backdoor attacks in VFL to our best knowledge . 3 PRELIMINARIES . Vertical Federated Learning . We first describe the basic setup of a typical VFL framework ( Chen et al. , 2020 ; Liu et al. , 2020 ) . There are M agents and a server collaboratively training a machine learning model based on a set of n training data D , { x { j } 1 , ... , x { j } M , y { j } } nj=1 . The server ( e.g. , a third party ) holds the label y { j } , while agent i holds partial feature x { j } i of instance x { j } , [ x { j } 1 ; ... ; x { j } M ] . Due to privacy concerns , the raw agent data x { j } i are not shared with other agents and the server . Instead , each agent i learns a local embedding gi parameterized by θi that maps the original feature vector x { j } i to feature vector h { j } i , gi ( x { j } i ; θi ) . The dimension of h { j } i is usually smaller than x { j } i . The server only has access to the embedded features from the agents . The learning objective of VFL is to minimize the following loss : L ( D ; θ0 , θ1 , ... , θM ) = 1 n n∑ j=1 ` ( ŷ { j } , y { j } ) + M∑ i=1 γ ( θi ) , ( 1 ) where ŷ { j } = fθ0 ( h { j } 1 , ... , h { j } M ) , fθ0 ( g1 ( x { j } 1 ; θ1 ) , ... , gM ( x { j } M ; θM ) ) , and fθ0 is the global model of server . ` is the loss function and γ is the regularizer that confines the model complexity . Robust Subspace Recovery . Robust subspace recovery aims to recover the underlying subspace of data despite that some data points are arbitrarily corrupted ( a.k.a . outliers ) . Mathematically , one observes a data matrix H , where each column corresponds to a data point . We can decompose H as L + E , where L is the underlying uncorrupted data matrix , and the matrix E models the outlier corruptions . Since the fraction of the outliers is usually small , most columns of E are zero , i.e. , E is column-sparse . Assume the underlying uncorrupted data points lie on a low-dimensional linear subspace , one of the notable Robust PCA works ( Xu et al. , 2012 ) showed that it is possible to recover the column-space of L exactly by solving the following convex optimization : min L , E ‖L‖∗ + λ‖ET ‖2,1 s.t.H = L + E. ( 2 ) Later we will see that in our proposed RVFR framework , H = [ h { 1 } , ... , h { n } ] , where the j-th column of H corresponds to the embedded features for the j-th instance . In the backdoor attack , a few columns of H are poisoned . In general the low-dimensional manifold that the embedded features L lies on is not linear , and we use a Robust AutoEncoder instead of Robust PCA to capture this non-linearity . We will show that Robust PCA can be viewed as a linearized version of the proposed Robust AutoEncoder , and we will provide theoretical justification for it . Formally , the Robust PCA result from Xu et al . ( 2012 ) is as follows : Theorem 1 . ( Exact subspace recovery ( Xu et al. , 2012 ) ) Suppose we observe H = L∗ + E∗ , where L∗ has rank r and incoherence parameter µ ( see Definition 1 in Appendix A.2 ) . Suppose further that E∗ is supported on at most βn columns . Any output to Eq . 2 recovers the column space of L∗ exactly , as long as the fraction of corrupted columns , β , satisfies β1−β ≤ c1 µr , where c1 = 9/121 . This can be achieved by setting the parameter λ in Eq . 2 to be 3 7 √ βn ( in fact it holds for any λ in a certain range ) . Inspired by the Robust PCA properties , we propose the Robust AutoEncoder in this work , which can capture the more general non-linear manifold . So far there is no theoretical guarantees for Robust AutoEncoders . The proposed Robust AutoEncoder is slightly different from an existing one ( Zhou & Paffenroth , 2017 ) , and we theoretically justify the proposed one by showing that when the underlying feature subspace is linear , after transformation based on the proposed Robust AutoEncoder , it exactly recovers the linear feature subspace as Robust PCA . Note that our established connection between the proposed Robust AutoEncoder and Robust PCA is non-trivial , especially due to the proposed ` 2,1-norm regularizer on the latent layer . Note that the exact recovery of the feature subspace is not enough for our goal of robust VFL , and we need to recover the underlying L∗ . We further propose a robust decomposition technique to recover L∗ exactly by utilizing our learnt feature subspace based on the assumption that the fraction of malicious agents is small . More details can be found in Theorem 3 .

Summary: The paper proposes a robust vertical federated learning framework called RVFR against backdoor attacks. In the framework, the agents train the local feature extractors and send the embedded features to the server. Then, the server trains a robust auto-encoder to recover and purify the features. Last, the purified features are used to train the global model. The paper provides theoretical analyses of the feature recovery under different threat models. The experiments show that RVFR is more effective when defending the backdoor attacks compared with the other baselines.