summaries: This paper studies the robustness of vertical federated learning. They use intuitions from robust feature subspace recovery for low rank matrices to design a defense framework for vertical federated learning. They empirically show the effectiveness of their framework for some settings. They also try to provide theoretical evidence for the robustness of their scheme, but I do not find their theoretical studies illuminating.

summaries: This paper studies robust vertical federated learning framework against backdoor attacks. The authors utilize robust autoencoders with purified training to remove the backdoor features and obtain clean global model. The authors also proved that, with some assumptions, the exact uncorrupted features can be recovered.

summaries: The paper tackles the problem of adversarial attacks against federated learning. The main proposal is RVFR, an original method that allows to protect against backdoor attacks targeting the specific “vertical federated learning” setting. The proposal is provided with a theoretical analysis of its effectiveness, and is then evaluated on two well-known datasets, showing better results than prior works in adversarial scenarios.  Overall, the presentation of the paper should be improved. The quality of the English text is appropriate. Figures and Tables should be improved. The topic addressed by the manuscript is relevant and in-line with ICLR. The references must be expanded with security-related works. The contribution is potentially significant.  STRENGTHS: + Relevant and trendy subject + Evaluation on two datasets + The theoretical analysis is appreciable  WEAKNESSES - Very poor Introduction, Abstract and Background (Related Work) - Unclear threat model - No Tradeoff

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

abstractText: Vertical Federated Learning (VFL) is a distributed learning paradigm that allows multiple agents to jointly train a global model when each agent holds a different subset of features for the same sample(s). VFL is known to be vulnerable to backdoor attacks, where data from malicious agents are manipulated during training, and vulnerable to inference-phase attacks, where malicious agents manipulate the test data. However, unlike the standard horizontal federated learning, improving the robustness of VFL remains challenging, as there is no clear redundancy among the agents in VFL. To this end, we propose RVFR, a novel robust VFL training and inference framework. Under certain conditions, RVFR can recover the underlying uncorrupted features with provable guarantees and thus sanitizes the model against a vast range of backdoor attacks. Further, RVFR also defends against inferencephase adversarial and missing feature attacks. We conduct extensive experiments on NUS-WIDE and CIFAR-10 datasets and show that RVFR outperforms different baselines in terms of robustness against diverse types of attacks.

authors:
Sebastien Andreina
Giorgia Azzurra Marson
Helen Möllering
Ghassan Karame

title: Baffle: Backdoor detection via feedback-based federated learning

venue: arXiv preprint arXiv:2011.02167,

authors:
Eugene Bagdasaryan
Andreas Veit
Yiqing Hua
Deborah Estrin
Vitaly Shmatikov

title: How to backdoor federated learning

venue: In International Conference on Artificial Intelligence and Statistics,

authors:
Arjun Nitin Bhagoji
Supriyo Chakraborty
Prateek Mittal
Seraphin Calo

title: Analyzing federated learning through an adversarial lens

venue: In International Conference on Machine Learning,

authors:
Peva Blanchard
El Mahdi El Mhamdi
Rachid Guerraoui
Julien Stainer

title: Machine learning with adversaries: Byzantine tolerant gradient descent

venue: In Proceedings of the 31st International Conference on Neural Information Processing Systems,

authors:
Keith Bonawitz
Vladimir Ivanov
Ben Kreuter
Antonio Marcedone
H Brendan McMahan
Sarvar Patel
Daniel Ramage
Aaron Segal
Karn Seth

title: Practical secure aggregation for privacypreserving machine learning

venue: In proceedings of the 2017 ACM SIGSAC Conference on Computer and Communications Security,

authors:
Xiaoyu Cao
Jinyuan Jia
Neil Zhenqiang Gong

title: Provably secure federated learning against malicious clients

venue: In Proceedings of the AAAI Conference on Artificial Intelligence,

authors:
Tianyi Chen
Xiao Jin
Yuejiao Sun
Wotao Yin

title: Vafl: a method of vertical asynchronous federated learning

venue: arXiv preprint arXiv:2007.06081,

authors:
Yudong Chen
Lili Su
Jiaming Xu

title: Distributed statistical machine learning in adversarial settings: Byzantine gradient descent

venue: Proceedings of the ACM on Measurement and Analysis of Computing Systems,

authors:
Ittai Dayan
Holger R Roth
Aoxiao Zhong
Ahmed Harouni
Amilcare Gentili
Anas Z Abidin
Andrew Liu
Anthony Beardsworth Costa
Bradford J Wood
Chien-Sung Tsai

title: Federated learning for predicting clinical outcomes in patients with covid-19

venue: Nature medicine,

authors:
Shuhao Fu
Chulin Xie
Bo Li
Qifeng Chen

title: Attack-resistant federated learning with residualbased reweighting

venue: arXiv preprint arXiv:1912.11464,

authors:
Clement Fung
Chris JM Yoon
Ivan Beschastnikh

title: The limitations of federated learning in sybil settings

venue: In 23rd International Symposium on Research in Attacks, Intrusions and Defenses ({RAID} 2020),

authors:
Ian J Goodfellow
Jonathon Shlens
Christian Szegedy

title: Explaining and harnessing adversarial examples

venue: arXiv preprint arXiv:1412.6572,

authors:
Rachid Guerraoui
Sébastien Rouault

title: The hidden vulnerability of distributed learning in byzantium

authors:
Jakub Konečnỳ
H Brendan McMahan
Daniel Ramage
Peter Richtárik

title: Federated optimization: Distributed machine learning for on-device intelligence

venue: arXiv preprint arXiv:1610.02527,

authors:
Yang Liu
Yan Kang
Xinwei Zhang
Liping Li
Yong Cheng
Tianjian Chen
Mingyi Hong
Qiang Yang

title: A communication efficient collaborative learning framework for distributed features

authors:
Yang Liu
Zhiqian Yi
Tianjian Chen

title: Backdoor attacks and defenses in feature-partitioned collaborative learning

authors:
Brendan McMahan
Eider Moore
Daniel Ramage
Seth Hampson
Blaise Aguera y Arcas

title: Communication-efficient learning of deep networks from decentralized data

venue: In Artificial intelligence and statistics,

authors:
Krishna Pillutla
Sham M Kakade
Zaid Harchaoui

title: Robust aggregation for federated learning

venue: arXiv preprint arXiv:1912.13445,

authors:
Ziteng Sun
Peter Kairouz
Ananda Theertha Suresh
H Brendan McMahan

title: Can you really backdoor federated learning

venue: arXiv preprint arXiv:1911.07963,

authors:
Hongyi Wang
Kartik Sreenivasan
Shashank Rajput
Harit Vishwakarma
Saurabh Agarwal
Jy-yong Sohn
Kangwook Lee
Dimitris Papailiopoulos

title: Attack of the tails: Yes, you really can backdoor federated learning

authors:
Chulin Xie
Keli Huang
Pin-Yu Chen
Bo Li

title: Dba: Distributed backdoor attacks against federated learning

venue: In International Conference on Learning Representations,

authors:
Chulin Xie
Minghao Chen
Pin-Yu Chen
Bo Li

title: Crfl: Certifiably robust federated learning against backdoor attacks

venue: arXiv preprint arXiv:2106.08283,

authors:
Cong Xie
Sanmi Koyejo
Indranil Gupta

title: Zeno: Distributed stochastic gradient descent with suspicion-based fault-tolerance

authors:
Huan Xu
Constantine Caramanis
Sujay Sanghavi

title: Robust pca via outlier pursuit

venue: IEEE Transactions on Information Theory,

authors:
Qiang Yang
Yang Liu
Tianjian Chen
Yongxin Tong

title: Federated machine learning: Concept and applications

venue: ACM Transactions on Intelligent Systems and Technology (TIST),

authors:
Dong Yin
Yudong Chen
Ramchandran Kannan
Peter Bartlett

title: Byzantine-robust distributed learning: Towards optimal statistical rates

authors:
Chong Zhou
Randy C Paffenroth

title: Anomaly detection with robust deep autoencoders

venue: In Proceedings of the 23rd ACM SIGKDD international conference on knowledge discovery and data mining,

heading: 1 INTRODUCTION

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

heading: 2 RELATED WORK

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

heading: 3 PRELIMINARIES

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

heading: 4 ROBUST VFL VIA FEATURE SUBSPACE RECOVERY (RVFR)

text: We first describe the threat models during VFL training and inference, and then present the training and inference procedures of the proposed RVFR framework.

heading: 4.1 THREAT MODEL

text:
There are M agents which hold different parts of the feature of the same set of instances, and the label y is held by the server. The attacker knows the feature information held by every agent and the label information on the server, but it can only choose and control αM agents to perform malicious attacks. During training, the goal of the malicious agents is to perform backdoor attacks, i.e., to achieve high accuracy on both the original main task and the targeted backdoor tasks.
Training phase threat model (backdoor attack): There are α fraction of agents that are malicious, and total β fraction of instances with backdoor trigger. With the adversarial target yA, the backdoor attacks are conducted by optimizing the following objective for the instances with backdoor trigger:
min hB
`(fθB ({hB ,hbenign}), y A),where θB = arg min
θ∈Θ ∑ j `(fθ(h {j}), y{j}), (3)
here hB represents overall embedded features provided by the malicious agents for the instance with backdoor trigger, its magnitude can be arbitrarily large unless regularized by the server; while hbenign denotes the overall embedded features provided by benign agents. If j-th instance has backdoor trigger, h{j} contains hB .
During inference, the malicious agents can perform adversarial attacks as well. Assume the total number of agents is M , and the fraction of malicious agents is α. We use hbenign to denote the embedded features provided by the (1−α)M benign agents (these agents are indexed by Ωbenign and |Ωbenign| = (1− α)M ), and we use hadv to denote the embedded features provided by malicious agents (indexed by Ωadv). We categorize the inference phase attacks into the following three types:
Inference phase threat model A (adversarial attack): The malicious agents aim to send adversarial features hadv to mislead the prediction, whose magnitude can be arbitrarily large unless regularized by the server:
min hadv
`(fθ0({hadv,hbenign}), y A), (4)
Inference phase threat model B (missing-feature attack): On the other hand, the malicious agents also have the option to NOT provide any features to the server. Such missing feature attack can make the malicious agents less detectable and still influence server’s prediction. The presence of missing feature attack can be viewed as that the server only observes partial blocks indexed by Ω of the overall embedded features h, i.e., hΩ. We use Ωc to denote the index set of agents that perform missing feature attacks. The attacker can compromise arbitrary αM agents.
min Ωc
`(fθ0(hΩ), y A), s.t. |Ωc| ≤ αM (5)
Inference phase threat model C (combined attack): The attackers can also perform a combined attack of the adversarial and missing feature attacks:
min Ωc, hadv
`(fθ0({hadv,hbenign}Ω), y A), s.t. |Ωc|+ |Ωadv| ≤ αM (6)
Besides above mentioned targeted attacks, the attacker can also perform untargeted attacks by maximizing the distance between the prediction and the ground truth, which should be easier and here we focus on defending these strong targeted attacks.

heading: 4.2 TRAINING PROCEDURES OF RVFR

text:
The proposed training procedure has four stages: quarantine training, robust feature subspace recovery, feature purifying, and server training.
Stage 1 (Quarantined training): In VFL, the server trains the combined ML model, and each agent trains its local feature extractor. A strong adversary could interfere with the server training and further propagate errors to benign agents. So we propose that the server first connects to each agent separately to train the feature extractors gi parameterized by θi, i = 1, ...,M . The training can be viewed as a special case of (Chen et al., 2020, Algorithm 1) with only one agent.
In particular, agent i sends h{j}i , gi(x {j} i ; θi) to the server. The server uses it to update/train a temporary global model (whose input dimension equals that of h{j}i ) with learning rate η0 and batch size B0. Then the server calculates ∂`
∂h {j} i
and sends it back to agent i. The benign agent i
updates its local feature extractor parameterized by θi using the combined gradient ∂` ∂h {j} i
∂h {j} i
∂θi with
learning rate ηi and batch size Bi. After a number of iterations between the server and agent i, the server records the final embedded feature h{j}i for every instance x {j} i , j = 1, ..., n. Finally, for each instance x{j}, the server holds {h{j}1 ,h {j} 2 , ...,h {j} M } from all the M agents, and concatenates them
into a long column vector h{j}. Let H = [ h{1}, ...,h{n} ] , the j-th column of H corresponds to the concatenated embeded features for the j-th instance.
Stage 2 (Robust feature subspace recovery): Since the fraction β of the backdoored training instances is small, this provides us the redundancy across instances for robustifying the server’s model. In this stage, the server disconnects from all the agents, and uses H to train a Robust AutoEncoder by minimizing the following objective:
min L,E,φ,ψ
‖Eψ(L)‖2,1 + λ‖ET ‖2,1 s.t.H = L + E, L = Dφ(Eψ(L)) (7)
where L models the underlying feature subspace, and E models the outlier corruptions due to backdoor attacks. In Eq. 7, we only enforce E to be column-sparse. Dφ is the decoder parameterized by φ, and Eψ is the encoder parameterized by ψ. Note that Eq. 7 is equivalent to the following:
min L,φ,ψ
‖Eψ(L)‖2,1 + λ‖(H −L)T ‖2,1 s.t. L = Dφ(Eψ(L)) (8)
In practice, we further relax the constraint L = Dφ(Eψ(L)) and solve the following instead:
min L,φ,ψ
‖Eψ(L)‖2,1 + λ‖(H −L)T ‖2,1 + η‖L−Dφ(Eψ(L))‖2F (9)
To optimize this, one can perform alternating optimization for L and parameters of φ and ψ.
Stage 3 (Feature purifying based on the recovered feature subspace): Recovering the feature subspace is not enough to achieve robust VFL, as we need to recover the original features to train the server. Fortunately, there is another important prior information that we can leverage: the fraction α of the malicious agents is small. For each training instance h{j}, we use the learned AutoEncoder to decompose it as h{j} = Dφ(Eψ(l{j})) + e{j}, where e{j} models the block-sparse corruptions (each block corresponds to an agent). Ideally we want to solve the following:
min l,e M∑ i=1 1{ei 6= 0} s.t. h = Dφ(Eψ(l)) + e (10)
which is equivalent to: min
l M∑ i=1 1{[h−Dφ(Eψ(l))]i 6= 0} (11)
However, the above block-sparsity minimization objective is NP-Complete, so we relax it to the `2,1-norm:
min l M∑ i=1 ‖[h−Dφ(Eψ(l))]i‖2 (12)
Our theoretical analysis in next section shows that under certain conditions, solving the above relaxed objective is able to recover the underlying true features.
Stage 4 (Server training): The server trains its own global model parameterized by θ0 based on the above purified features Dφ(Eψ(l)), with learning rate η0 and batch size B0.

heading: 4.3 INFERENCE PROCEDURES OF RVFR

text:
We propose an all-in-one solution to defend against both training phase and inference phase malicious attacks defined in Section 4.1, and thus during inference we will also perform feature purifying using the recovered feature subspace and the sparse nature of the malicious agents.
Let Ω be the index set of the observed agent blocks. Assume we have successfully trained the AutoEncoder, such that we have l = Dφ(Eψ(l)), where l is any uncorrupted embedded feature vector. We want to use the AutoEncoder to decompose the observed hΩ as hΩ = lΩ + eΩ with l = Dφ(Eψ(l)), where eΩ models the block-sparse corruptions (each block corresponds to an agent) on the observed set Ω. Note that if all the blocks (agents) are observed, Ω becomes the whole set. Since eΩ is block-sparse, similar to the feature purifying during training, we naturally aim to solve the following:
min l ∑ i∈Ω 1{[h−Dφ(Eψ(l))]i 6= 0} (13)
However, the above block-sparsity minimization objective is NP-Complete, we instead relax it to the `2,1-norm:
min l ∑ i∈Ω ‖[h−Dφ(Eψ(l))]i‖2 (14)
Our theoretical analysis in Theorems 3 & 4 show that under certain conditions, solving the above relaxed objective is able to recover the underlying true features, despite the partial observations. After obtained l̂, we feed Dφ(Eψ(l̂)) to server’s trained machine learning model for the prediction.

heading: 5 THEORETICAL ANALYSIS OF RVFR

text:
We first analyze the proposed Robust AutoEncoder, which is the backbone of the proposed defense framework. Then, we prove that the proposed feature purifying method can recover the underlying features exactly under certain conditions, which means that the server’s global model can be both trained and tested on correct features, and thus is robust against diverse types of attacks.
The Robust AutoEncoder can be viewed as the generalization of Robust PCA, where the lowdimensional linear subspace is extended to the more general low-dimensional manifold. It is challenging to prove the exact recovery of the low-dimensional manifold in the presence of outlier corruptions. However, we next show that when the underlying feature subspace is linear, the Robust AutoEncoder can exactly recover the underlying subspace.
Theorem 2. (Exact subspace recovery) Assume l∗{i} ∈ Rd, i = 1, ..., n lie on a low-dimension subspace, i.e., rank(L∗) = r d. For the AutoEncoder, assume linear activation functions and no bias terms, while setting the dimension of the latent layer of AutoEncoder as r = rank(L∗) and restricting the weight matrices to be orthonormal. Then under the conditions in Theorem 1, the global optimal solution of Eq. 7 is equivalent to the solution of Eq. 2, and the corresponding weight matrix of the encoder is guaranteed to recover the subspace of L∗ exactly.
The proof is omitted to Appendix A.3. The key is to show that the objective value of Eq. 7 is greater than or equal to the objective value of Eq. 2, with equality achieved when the weight matrix of the encoder equals the top-r left singular vectors of L. This is non-trivial especially due to the proposed `2,1-norm regularizer on the latent layer of the Robust AutoEncoder.
Remark 1. The conditions required in Theorem 2 provide some useful insights. First, the total fraction β of poisoned instances needs to be small. Second, the dimension of the underlying feature subspace needs to be small. If the dimension of the feature subspace is too large, it is difficult to distinguish corrupted features from the uncorrupted ones. On the other hand, low dimensionality implies that there is enough redundancy among the features among agents, which is preferred since redundancy provides more robustness. Third, as discussed in Appendix A.2 regarding the incoherence condition, preferring small incoherence parameters implies that we do not want any dimension of the feature subspace to be defined by very few data points. Otherwise, it is risky if those instances are poisoned, as it becomes impossible to recover that dimension.
Note that exact recovery of the subspace does not mean exact recovery of L. For the corrupted/poisoned instances, usually it is impossible to restore them as the corruptions can be arbitrary. Fortunately, once we have recovered the underlying feature subspace, it is possible to recover the original features exactly by utilizing the sparseness of the malicious agents. This holds for both training and inference phases. The following theorem shows that we can recover the underlying embedded features exactly against the inference phase threat models A & C.
Theorem 3. (Exact feature recovery under threat model A & C) Let Wend be the weight matrix of the last layer of the AutoEncoder. Assume there is no bias term nor non-linear activation function in the last layer. Assume the trained AutoEncoder captures the underlying feature subspace (i.e., Dφ(Eψ(l)) = l for uncorrupted feature vector l). Let Ω be the index set of the observed blocks (agents). If ∀v ∈ Range(Wend)\0, for any partition {SΩ, S̄Ω} of Ω with |SΩ| = g > |Ω|/2, it holds that ∑ i∈SΩ ‖vi‖2 > ∑ i∈S̄Ω ‖vi‖2, then for any hΩ = l ∗ Ω + e ∗ Ω with Dφ(Eψ(l∗)) = l∗ and∑
i∈Ω 1{e∗i = 0} ≥ g, Dφ(Eψ(l∗)) is the unique solution of Eq. 14 and Eq. 13.
The proof can be found in Appendix A.4, where we use contradiction to show that there does not exist any global optimal solution of Eq. 14 or Eq. 13 that is different from Dφ(Eψ(l∗)). Remark 2. (Exact feature recovery via feature purifying during training) Theorem 3 directly justifies the proposed feature purifying during training. Under certain conditions, we can recover each column of L exactly during training.
Next, we analyze the exact feature recovery property under inference phase threat model B (missing feature attack), where there are no adversarial features.
Theorem 4. (Exact feature recovery under threat model B) Let Wend be the weight matrix of the last layer of the AutoEncoder. Assume there is no bias term nor non-linear activation function in the last layer. Assume the trained AutoEncoder captures the underlying feature subspace (i.e., Dφ(Eψ(l)) = l for uncorrupted feature vector l). Let Ω be the index set of the observed blocks (agents). If ∀v ∈ Range(Wend)\0, it holds that vΩ 6= 0, then given hΩ = l∗Ω, Dφ(Eψ(l∗)) is the unique solution of Eq. 14 and Eq. 13.
The proof is deferred to Appendix A.5, where we use contradiction to show that there does not exist any global optimal solution of Eq. 14 or Eq. 13 that is different from Dφ(Eψ(l∗)). The condition vΩ 6= 0 for any v ∈ Range(Wend)\0 required by Theorem 4 is much weaker than that required by Theorem 3, which is reasonable as there are no corruptions on the observed blocks in Theorem 4. This required condition also provides very useful insights on how many agents we need to observe in order to make a prediction. It can be understood through the following toy example. Suppose the prediction task is credit approval. Agent 1 holds the banking data, Agent 2 holds the shopping history, while Agent 3 holds age and gender features. Intuitively, if we only observe the features from Agent 3, it is very hard to make the prediction, since there are definitely Person A and B of the same age and gender, but with quite different other features. Let lA be Person A’s overall features, lB be Person B’s overall features. In this toy example, Ω corresponds to Agent 3 only. We have lA 6= lB but lAΩ = lBΩ , or say (lA − lB)Ω = 0. Note that both lA and lB are in the range space of Wend, so (lA − lB) ∈ Range(Wend)\0. The smaller the set Ω implies the higher chance that (lA − lB)Ω = 0. In general, there is no exact value for the minimal fraction of agents we need to observe in order to make a prediction. It highly depends on whether the observed agents cover enough information to distinguish different classes in the classification task. For example, if there is enough redundancy across the agents, say Agent 3 holds age, gender, banking data, and shopping history, then observing Agent 3 could be sufficient.
In summary, Theorem 2 shows that under certain natural conditions (e.g., the fraction of poisoning instances β is small), we can recover the underlying feature subspace exactly. Further, Theorems 3 & 4 show that under certain natural conditions (e.g., the fraction of malicious agents α is small and the AutoEncoder we have learnt captures the underlying feature subspace), we can recover the underlying embedded features exactly during both training and inference.

heading: 6 EMPIRICAL EVALUATION

text: We first describe experimental setup in Section 6.1, and present experimental results in Section 6.2. Additional implementation details and results can be found in Appendix A.1.

heading: 6.1 EXPERIMENTAL SETUP

text:
Dataset and Models. We study the classification task on two datasets: NUS-WIDE and CIFAR-10. Following (Liu et al., 2020), which proposed the backdoor attack against VFL, we first use NUSWIDE dataset to evaluate our defense. In NUS-WIDE, each sample has 634 image features, 1000 text features, and 5 different labels. We consider two VFL settings (1 server and M agents): M = 2 and M = 4. We denote G1 as Agent 1, G2 as Agent 2, etc. for simplicity. Following (Liu et al., 2020), when M = 2, G1 holds image features and G2 holds text features. When M = 4, G1 (G2) holds 225 (409) image features while G3 (G4) holds 500 (500) text features. In all settings, only the server holds the labels. In CIFAR-10 dataset, each sample have 32x32x3 image features. We consider the VFL settings for CIFAR-10 with M = 2 and M = 3. When M = 2, each agent has 16x32x3 features, i.e., half image. When M = 3, G1, G2, G3 have 9x32x3, 10x32x3, 13x32x3 features respectively. Following the standard setup, we use SGD to update the local models and server model for E epochs with learning rate 0.01 and batch size Bs. The detailed dataset and model descriptions as well as training parameter setups are summarized in Table 1 in Appendix.
Attack Setup. We consider the backdoor attack, its combination with the missing feature attack, as well as the combination of adversarial evasion attack (Goodfellow et al., 2014) and missing feature attack: (1) The backdoor attack is proposed by (Liu et al., 2020) and the detailed description of the attack method can be found in Appendix. A.1. In NUS-WIDE, the attacker’s backdoor attack goal is to change the predicted label of the instance with a backdoor trigger to be a specific target class (i.e., ‘buildings’), where the backdoor trigger is ‘the last text feature equals 1’. There are 152 backdoored training samples. In CIFAR-10, the backdoor targeted label is ‘truck’ and the trigger pattern is a white rectangle with size 4x6x3 in the lower right corner. The number of backdoored training samples is 100 when M = 3, and 600 when M = 2. (2) In the missing feature attack, the malicious agent does not send its local embedded feature to the server during inference phase. The server uses a zero vector as its local embedded feature. We combine this attack with other attacks as this attack alone is not powerful enough. (3) In the evasion attack, the attacker use the information from server to locally generate the partial adversarial input features (i.e., adversarial examples) based on the Fast Gradient Signed Method (FGSM) attack method (Goodfellow et al., 2014). The detailed setup of this attack method against VFL is presented in Appendix. A.1. The evasion attack (or its combination with missing feature attack) is only conducted during the VFL inference phase (VFL training procedure is attack-free).
Baseline methods. Besides the Vanilla VFL framework without defense (Liu et al., 2020, Algorithm 1), we compare our proposed defense method with two state-of-the-art VFL defense techniques, e.g., Gradient Sparsification method that sparsifies the intermediate gradients sent from the server to agents, and Differential Privacy (DP) that adds noise to such exchanged gradients (Liu et al., 2020). Laplacian Noise is used as a realization of the DP method. The noise variance is 0.05 for NUS-WIDE and 0.0001 for CIFAR-10 to preserve reasonable clean accuracy of the model. In Gradient Sparsification, the hyperparameter controlling the sparsity is 0.999 for NUS-WIDE and 0.95 for CIFAR-10. As for our proposed RVFR, in Stage 1 (i.e., quarantine training for local models) and Stage 4 (i.e., training for server model), we use the same training parameters as in Table 1. The robust autoencoder is a four-layer fully connected neural network. The details about the training of robust autoencoder, the alternative optimization, and the purified training of the server in Stage 2, 3, 4 and details of the inference procedures can be founded in Appendix. A.1. All of our experiments are run three times.

heading: 6.2 EMPIRICAL RESULTS

text:
Evaluation Metrics. For the backdoor attack (BKD) and its combination with missing feature attack (Miss), the evaluation metrics for defense methods are the Backdoor Accuracy (the lower the better)
on the backdoored test data and the main task Clean Accuracy (the higher the better) on the clean test dataset. There are 102 backdoored testing samples in NUS-WIDE and 10000 in CIFAR-10. For the adversarial evasion attack (Adv) and missing feature attack, the evaluation metric is the main task Accuracy on the adversarial examples. 10000 adversarial examples are generated for both datasets.
Figure 2 & 3 show the performance of the proposed RVFR and baseline methods on NUS-WIDE and CIFAR-10. When defending against backdoor attack and its combination with missing feature attack, the proposed RVFR method exhibits significantly lower Backdoor Accuracy and similar Clean Accuracy on both datasets, compared to other methods. Note that under the backdoor attack, in some cases, the Clean Accuracy of all defense methods is lower than the vanilla VFL, which exhibits a trade-off between robustness and accuracy. As shown in the right column of Figure 2 & 3, when defending against the combination of adversarial evasion attack and missing feature attack, the proposed RVFR method achieves consistently higher accuracy than the baseline methods on both datasets, which demonstrates its robustness against such attacks.

heading: 7 CONCLUSIONS

text: In this work, we proposed a novel robust feature subspace recovery based VFL framework to defend against backdoor attacks during training, and a variety of attacks during inference, both with theoretical guarantees. An important byproduct of our analysis is the first theoretical justification for the Robust AutoEncoder, which may be of independent interest. We further validate the robustness of our proposed framework through extensive experiments on NUS-WIDE and CIFAR-10 datasets.

heading: ETHICS STATEMENT

text: We propose a novel defense algorithm for vertical federated learning, and provide analysis from theoretical and empirical perspectives. All the datasets and packages we use are open-sourced. We do not have ethical concerns in our paper.

heading: REPRODUCIBILITY STATEMENT

text: The source code of the proposed RVFR method is available as the supplemental material, which also includes the detailed data processing steps. The complete proof of the theorems can be found in the Appendix.

heading: A APPENDIX

heading: A.1 ADDITIONAL EXPERIMENTAL DETAILS

text:
Simulation Environment. We simulate the VFL setup (1 server and multiple agents) on a Linux machine with Intel® Xeon® Gold 6132 CPUs and 8 NVidia® 1080Ti GPUs. Our code is based on Tensorflow 2.0.
VFL Backdoor Attacks. For the j-th instance with backdoor trigger, the malicious agent i does not update its local feature extractor (parameterized by θi) like benign agents via using the combined
gradient ∂` ∂h {j} i
∂h {j} i
∂θi . Instead, it replaces the intermediate gradient received from the server ∂`
∂h {j} i
by
γ ∂` ∂h {target} i and uses the manipulated gradient γ ∂` ∂h {target} i
∂h {j} i
∂θi to update its local feature extractor,
where ∂` ∂h {target} i is the intermediate gradient received from the server for any instance belonging to the target class, and γ is the gradient amplify ratio. In our experiments, we use γ = 10 for NUS-WIDE and γ = 1 for CIFAR-10. Moreover, since the label owned by the server is clean, the attacker would submit zero vector to the server to prevent the server mapping the poisoned embedding to the true label.
Evasion Attacks. Given the j-th test sample, the attacker adds small perturbations xpertest {j} i generated from FGSM to the original data xtest {j} i to obtain the adversarial example x adv test {j} i = xtest {j} i + x per test {j} i where controlling the degree of attack. For NUS-WIDE, we set = 0.01 when M = 2 and = 0.1 when M = 4; for CIFAR-10, we set = 0.01.
Training and Inference Details of RVFR. In Stage 2, we first train the AutoEncoder for Erae epochs, and alternatively update AutoEncoder and the estimated uncorrupted embedded features L for Ealt epochs using SGD with an initial learning rate 0.01 which exponentially decays by a factor of 0.99 every 10000 steps. During the alternating optimization, λ and η in the objective function are set to be 0.1. In NUS-WIDE, Erae = 100 when M = 2 and Erae = 200 when M = 4; in CIFAR-10, Erae = 100. We use Ealt = 100 for all settings. In Stage 3 & 4, for each batch of training data, we first purifies the embedded features according to Eq. 12 (using the estimated L from Stage 2 as initialization) for Epur train = 2 epochs, then use the purified embedded feature to train server’s global ML model. We perform such procedure for Eserver = 100 epochs using the same SGD optimizer as in Stage 2. In the inference phase, the server purifies the embedded features according to Eq. 14 (use the embedded features as initialization) for Epur test = 2 epochs, and feeds the purified feature to the trained server model for prediction.
Clean Accuracy under Non-adversarial Setting In this section, we report the clean accuracy of different VFL methods under non-adversarial setting in Table 2.

heading: A.2 DEFINITION OF INCOHERENCE PARAMETER

text:
Definition 1. (Incoherence parameter (Xu et al., 2012)) A matrix L ∈ Rd×n with thin SVD L = UΣV T and (1 − β)n whose columns are non-zero, is said to be column-incoherent with parameter µ if
max i ‖V Tei‖22 ≤
µr
(1− β)n , (15)
where {ei} are the coordinate unit vectors, r is the rank of matrix L.
A small incoherence parameter implies the right singular vectors of L are not ‘spiky’. As mentioned in Xu et al. (2012), if the left hand side of Eq. 15 is as big as 1, it essentially means that one of the directions of the column space which we wish to recover, is defined by only a single observation. Given the regime of a constant fraction of arbitrarily chosen and arbitrarily corrupted points, such a setting is highly undesirable.

heading: A.3 PROOF OF THEOREM 2

text:
Proof. First note that Wencoder ∈ Rr×d since we limit the dimension of the latent layer of the AutoEncoder to be r = rank(L∗). The problem we consider is:
min L,E,Wencoder∈Rr×d
‖WencoderL‖2,1 + λ‖ET ‖2,1, s.t.H = L + E,L = Dφ(Eψ(L)) (16)
Since we have the constraint L = Dφ(Eψ(L)), the rank of any feasible L should be no more than r, and L must lie within the row-space of Wencoder. So we can write Wencoder = RUT where R ∈ Rr×r, and U are the top-r left singular vectors of L. As the rows of Wencoder are orthornormal, i.e., WencoderW Tencoder = I , so RU
TURT = I and therefore RRT = I (so this square matrix R is unitary). Further, since setting Wdecoder = W Tencoder will always meet the constraint L = Dφ(Eψ(L)), the main problem we need to consider becomes:
min {L|rank(L)≤r},E,
Wencoder=RU T ,R is unitary
‖WencoderL‖2,1 + λ‖ET ‖2,1, s.t.H = L + E, (17)
Since ‖L‖∗ = ‖RUTL‖∗ = ‖WencoderL‖∗ = ‖I(WencoderL)‖∗ = ‖ ∑r i=1 ei(WencoderL)i,:‖∗ ≤ ∑r i=1 ‖ei(WencoderL)i,:‖ = ∑r i=1 ‖(WencoderL)i,:‖2 , ‖WencoderL‖2,1. The equality is achieved when R = I , i.e., Wencoder = UT , because ‖WencoderL‖2,1 = ‖UTL‖2,1 = ‖UTUΣV T ‖2,1 = ‖ΣV T ‖2,1 = ∑r i=1 σi = ‖L‖∗.
So
min {L|rank(L)≤r},E,
Wencoder=RU T ,R is unitary
‖WencoderL‖2,1 + λ‖ET ‖2,1 ≥ min {L|rank(L)≤r},E ‖L‖∗ + λ‖ET ‖2,1,
(18) with equality achieved when Wencoder equals top-r left singular vectors of L.
Then we only need to consider the right-hand-side of Eq. 18, i.e.,
min {L|rank(L)≤r},E
‖L‖∗ + λ‖ET ‖2,1 s.t.H = L + E (19)
Recall that under the conditions of Theorem 1, the solution L̂ of Eq. 2 has exactly the same column space as L∗ (L̂ may not and not necessary to be equal to L∗), so rank(L̂) = r. Then the solution L̂ of Eq. 2 must also be the global optimal solution of Eq. 19. Finally, as the row-space of Wencoder equals the column-space of L̂, it recovers the underlying subspace of L∗ exactly.

heading: A.4 PROOF OF THEOREM 3

text:
Proof. A) Suppose Dφ(Eψ(l′)) is the global optimal solution of Eq. 14 that is different from Dφ(Eψ(l∗)). Let Dφ(Eψ(l′)) = Dφ(Eψ(l∗))− v. So we have v ∈ Range(Wend)\0.
∑ i∈Ω ‖[h−Dφ(Eψ(l′))]i‖2 (20)
= ∑ i∈Ω ‖[h−Dφ(Eψ(l∗)) + v]i‖2 (21)
= ∑ i∈Ω ‖[h−Dφ(Eψ(l∗))]i + vi‖2 (22)
= ∑ i∈SΩ ‖[h−Dφ(Eψ(l∗))]i + vi‖2 + ∑ i∈S̄Ω ‖[h−Dφ(Eψ(l∗))]i + vi‖2 (23)
= ∑ i∈SΩ ‖vi‖2 + ∑ i∈S̄Ω ‖[h−Dφ(Eψ(l∗))]i + vi‖2 (24)
≥ ∑ i∈SΩ ‖vi‖2 + ∑ i∈S̄Ω ‖[h−Dφ(Eψ(l∗))]i‖2 − ∑ i∈S̄Ω ‖vi‖2 (25)
> ∑ i∈S̄Ω ‖[h−Dφ(Eψ(l∗))]i‖2 (26)
= ∑ i∈Ω ‖[h−Dφ(Eψ(l∗))]i‖2 (27)
where SΩ is the index set of size g such that e∗i = 0,∀i ∈ SΩ. And the last inequality follows from the assumed range space property since v ∈ Range(Wend)\0. The above contradicts the assumption that Dφ(Eψ(l′)) is the global optimal solution that is different from Dφ(Eψ(l∗)). B) First, note that hΩ = Dφ(Eψ(l∗))Ω+e∗Ω and ∑ i∈Ω 1{e∗i 6= 0} ≤ |Ω|−g. SupposeDφ(Eψ(l′)) is the global optimal solution of Eq. 13 that is different fromDφ(Eψ(l∗)). Let hΩ = Dφ(Eψ(l′))Ω +e′Ω, then we have ∑
i∈Ω 1{e′i 6= 0} ≤ ∑ i∈Ω 1{e∗i 6= 0} ≤ |Ω| − g (28)
and Dφ(Eψ(l′))Ω + e′Ω = Dφ(Eψ(l∗))Ω + e∗Ω. From Eq. 28 we know that
∑ i∈Ω 1{[Dφ(Eψ(l∗))−Dφ(Eψ(l′))]i = 0} ≥|Ω| − ∑ i∈Ω 1{e∗i 6= 0} − ∑ i∈Ω 1{e′i 6= 0} (29)
≥|Ω| − (|Ω| − g)− (|Ω| − g) = 2g − |Ω| (30)
Let Dφ(Eψ(l′)) = Dφ(Eψ(l∗)) − v, so we have v ∈ Range(Wend)\0 and ∑ i∈Ω 1{vi = 0} ≥ 2g − |Ω|. Now split Ω into 3 disjoint sets {Ω0,Ω1,Ω2}, where Ω0 is any subset of Ω with size 2g − |Ω| such that vΩ0 = 0, and |Ω1| = |Ω2| = |Ω| − g. Since |Ω0 ∪ Ω1| = g, by our assumption, we have ∑ i∈Ω0∪Ω1 ‖vi‖2 > ∑ i∈Ω2 ‖vi‖2. Since |Ω0 ∪ Ω2| = g, by our assumption, we have∑
i∈Ω0∪Ω2 ‖vi‖2 > ∑ i∈Ω1 ‖vi‖2. However, this leads to a contradiction since ∑ i∈Ω0 ‖vi‖2 =
0.

heading: A.5 PROOF OF THEOREM 4

text:
Proof. SupposeDφ(Eψ(l′)) is the global optimal solution of Eq. 14 that is different fromDφ(Eψ(l∗)). Let Dφ(Eψ(l′)) = Dφ(Eψ(l∗))− v. So we have v ∈ Range(Wend)\0.
∑ i∈Ω ‖[h−Dφ(Eψ(l′))]i‖2 (31)
= ∑ i∈Ω ‖[h−Dφ(Eψ(l∗)) + v]i‖2 (32)
= ∑ i∈Ω ‖vi‖2 (33)
>0 (34) = ∑ i∈Ω ‖[h−Dφ(Eψ(l∗))]i‖2, (35)
where the inequality is due to the condition that vΩ 6= 0 for any v ∈ Range(Wend)\0. The above contradicts the assumption that Dφ(Eψ(l′)) is the global optimal solution of Eq. 14 that is different from Dφ(Eψ(l∗)). Similarly, suppose Dφ(Eψ(l′)) is the global optimal solution of Eq. 13 that is different from Dφ(Eψ(l∗)). Let Dφ(Eψ(l′)) = Dφ(Eψ(l∗))− v. So we have v ∈ Range(Wend)\0.
∑ i∈Ω 1{[h−Dφ(Eψ(l′))]i 6= 0} (36)
= ∑ i∈Ω 1{[h−Dφ(Eψ(l∗)) + v]i 6= 0} (37)
= ∑ i∈Ω 1{vi 6= 0} (38)
>0 (39) = ∑ i∈Ω 1{[h−Dφ(Eψ(l∗))]i 6= 0} (40)
The above contradicts the assumption that Dφ(Eψ(l′)) is the global optimal solution that is different from Dφ(Eψ(l∗)).

title: RVFR: ROBUST VERTICAL FEDERATED LEARNING


summaries: This paper studies pre-training of Siamese Transformer encoders for the zero-shot dense retrieval problem. The author consider contrastive learning loss function to optimize the encoders and investigate variants of ICT for data augmentations. While not much novelty in the techincal sides, the author indeed present solid empirical study on zero-shot retrieval benchmark as well as extension to the few-shot evaluation.

summaries: This paper proposes to use cropping strategy to train a neural model for IR. It is compared to BM25 and several other neural models, including BERT and the ones using inverse cloze task. The goal of cropping is to train a model that fits better the IR tasks. It is hypothesized that this task is more similar to IR than other pre-training strategies. The method is tested on a set of benchmarks of zero-shot retrieval and MSMARCO with training data. It is shown than the proposed method fine-tuned on MSMARCO can outperform the compared methods, including BM25. When no fine-tuning is used, the method is almost at par with BM25.

summaries: This paper presents a contrastive learning approach for unsupervised text retrieval. As the approach is unsupervised, similar to the Inverse Cloze Task (ICT) approach, the authors propose to do independent cropping of a paragraph to generate a pair of "pseudo-question" and "pseudo-document", which is used to train the model. The model is trained using contrastive learning following the MoCo algorithm. Experiments are performed on the recent BEIR benchmark for zero-shot and few-shot retrieval evaluation. The proposed approach often performs better than the BM25 algorithm on zero-shot retrieval on the BEIR benchmark. There are also some ablations presented to analyze different aspects of model training such as the effect of number of negative examples, training data, data augmentations etc.

summaries: This paper proposed a contrastive learning-based method for zero-shot dense IR. Based on the contrastive learning framework, 3 positive pair construction methods and 2 negative pair construction methods are used and compared. Evaluations are performed on the BEIR benchmark, with the best average result achieved compared to a list of IR (both sparse and dense) baselines in the zero-shot setting.

abstractText: Information retrieval is an important component in natural language processing, for knowledge intensive tasks such as question answering and fact checking. Recently, information retrieval has seen the emergence of dense retrievers, based on neural networks, as an alternative to classical sparse methods based on term-frequency. These models have obtained state-of-the-art results on datasets and tasks where large training sets are available. However, they do not transfer well to new domains or applications with no training data, and are often outperformed by term-frequency methods such as BM25 which are not supervised. Thus, a natural question is whether it is possible to train dense retrievers without supervision. In this work, we explore the limits of contrastive learning as a way to train unsupervised dense retrievers, and show that it leads to strong retrieval performance. More precisely, we show on the BEIR benchmark that our model outperforms BM25 on 8 out of 14 datasets. Furthermore, when a few thousands examples are available, we show that fine-tuning our model on these leads to strong improvements compared to BM25. Finally, when used as pre-training before fine-tuning on the MS MARCO dataset, our technique obtains state-of-the-art results on the BEIR benchmark.

authors:
Payal Bajaj
Daniel Campos
Nick Craswell
Li Deng
Jianfeng Gao
Xiaodong Liu
Rangan Majumder
Andrew McNamara
Bhaskar Mitra
Tri Nguyen

title: Ms marco: A human generated machine reading comprehension dataset

venue: arXiv preprint arXiv:1611.09268,

authors:
Adam Berger
Rich Caruana
David Cohn
Dayne Freitag
Vibhu Mittal

title: Bridging the lexical chasm: statistical approaches to answer-finding

venue: In Proceedings of the 23rd annual international ACM SIGIR conference on Research and development in information retrieval,

authors:
Mathilde Caron
Ishan Misra
Julien Mairal
Priya Goyal
Piotr Bojanowski
Armand Joulin

title: Unsupervised learning of visual features by contrasting cluster assignments

venue: arXiv preprint arXiv:2006.09882,

authors:
Mathilde Caron
Hugo Touvron
Ishan Misra
Hervé Jégou
Julien Mairal
Piotr Bojanowski
Armand Joulin

title: Emerging properties in self-supervised vision transformers

venue: arXiv preprint arXiv:2104.14294,

authors:
Wei-Cheng Chang
Felix X Yu
Yin-Wen Chang
Yiming Yang
Sanjiv Kumar

title: Pre-training tasks for embedding-based large-scale retrieval

venue: arXiv preprint arXiv:2002.03932,

authors:
Danqi Chen
Adam Fisch
Jason Weston
Antoine Bordes

title: Reading Wikipedia to answer open-domain questions

venue: In Proc. ACL,

authors:
Ting Chen
Simon Kornblith
Mohammad Norouzi
Geoffrey Hinton

title: A simple framework for contrastive learning of visual representations

venue: In International conference on machine learning,

authors:
Scott Deerwester
Susan T Dumais
George W Furnas
Thomas K Landauer
Richard Harshman

title: Indexing by latent semantic analysis

venue: Journal of the American society for information science,

authors:
Jacob Devlin
Ming-Wei Chang
Kenton Lee
Kristina Toutanova

title: BERT: Pre-training of deep bidirectional transformers for language understanding

venue: In Proc. NAACL,

authors:
Hongchao Fang
Sicheng Wang
Meng Zhou
Jiayuan Ding
Pengtao Xie

title: Cert: Contrastive self-supervised learning for language understanding

venue: arXiv preprint arXiv:2005.12766,

authors:
Luyu Gao
Jamie Callan

title: Unsupervised corpus aware language model pre-training for dense passage retrieval

venue: arXiv preprint arXiv:2108.05540,

authors:
Daniel Gillick
Alessandro Presta
Gaurav Singh Tomar

title: End-to-end retrieval in continuous space

venue: arXiv preprint arXiv:1811.08008,

authors:
John M Giorgi
Osvald Nitski
Gary D Bader
Bo Wang

title: Declutr: Deep contrastive learning for unsupervised textual representations

venue: arXiv preprint arXiv:2006.03659,

authors:
Kelvin Guu
Kenton Lee
Zora Tung
Panupong Pasupat
Ming-Wei Chang

title: Realm: Retrievalaugmented language model pre-training

venue: arXiv preprint arXiv:2002.08909,

authors:
Kaiming He
Haoqi Fan
Yuxin Wu
Saining Xie
Ross Girshick

title: Momentum contrast for unsupervised visual representation learning

venue: In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition,

authors:
Sebastian Hofstätter
Sheng-Chieh Lin
Jheng-Hong Yang
Jimmy Lin
Allan Hanbury

title: Efficiently teaching an effective dense retriever with balanced topic aware sampling

venue: arXiv preprint arXiv:2104.06967,

authors:
Po-Sen Huang
Xiaodong He
Jianfeng Gao
Li Deng
Alex Acero
Larry Heck

title: Learning deep structured semantic models for web search using clickthrough data

venue: In Proceedings of the 22nd ACM international conference on Information & Knowledge Management,

authors:
Samuel Humeau
Kurt Shuster
Marie-Anne Lachaux
Jason Weston

title: Poly-encoders: Transformer architectures and pre-training strategies for fast and accurate multi-sentence scoring

authors:
Gautier Izacard
Edouard Grave

title: Leveraging passage retrieval with generative models for open domain question answering

venue: arXiv preprint arXiv:2007.01282,

title: Distilling knowledge from reader to retriever for question answering

authors:
Gautier Izacard
Fabio Petroni
Lucas Hosseini
Nicola De Cao
Sebastian Riedel
Edouard Grave

title: A memory efficient baseline for open domain question answering

venue: arXiv preprint arXiv:2012.15156,

authors:
Yacine Jernite
Samuel R Bowman
David Sontag

title: Discourse-based objectives for fast unsupervised sentence representation learning

venue: arXiv preprint arXiv:1705.00557,

authors:
Jeff Johnson
Matthijs Douze
Hervé Jégou

title: Billion-scale similarity search with gpus

venue: IEEE Transactions on Big Data,

title: A statistical interpretation of term specificity and its application in retrieval

venue: Journal of documentation,

authors:
Mandar Joshi
Eunsol Choi
Daniel S. Weld
Luke Zettlemoyer

title: Triviaqa: A large scale distantly supervised challenge dataset for reading comprehension

authors:
Vladimir Karpukhin
Barlas Oğuz
Sewon Min
Ledell Wu
Sergey Edunov
Danqi Chen
Wen-tau Yih

title: Dense passage retrieval for open-domain question answering

venue: arXiv preprint arXiv:2004.04906,

authors:
Omar Khattab
Christopher Potts
Matei Zaharia

title: Relevance-guided supervision for openqa with colbert

authors:
Tom Kwiatkowski
Jennimaria Palomaki
Olivia Redfield
Michael Collins
Ankur Parikh
Chris Alberti
Danielle Epstein
Illia Polosukhin
Jacob Devlin
Kenton Lee

title: Natural questions: a benchmark for question answering research

venue: Transactions of the Association for Computational Linguistics,

authors:
Jungmin Kwon
Jeongseop Kim
Hyunseo Park
In Kwon Choi

title: ASAM: adaptive sharpnessaware minimization for scale-invariant learning of deep neural networks

venue: arXiv preprint arXiv:2102.11600,

authors:
Zhenzhong Lan
Mingda Chen
Sebastian Goodman
Kevin Gimpel
Piyush Sharma
Radu Soricut

title: Albert: A lite bert for self-supervised learning of language representations

authors:
Kenton Lee
Ming-Wei Chang
Kristina Toutanova

title: Latent retrieval for weakly supervised open domain question answering

authors:
Mike Lewis
Marjan Ghazvininejad
Gargi Ghosh
Armen Aghajanyan
Sida Wang
Luke Zettlemoyer

title: Pre-training via paraphrasing

venue: arXiv preprint arXiv:2006.15020,

authors:
Jimmy Lin
Rodrigo Nogueira
Andrew Yates

title: Pretrained transformers for text ranking: Bert and beyond

venue: arXiv preprint arXiv:2010.06467,

authors:
Ilya Loshchilov
Frank Hutter

title: Decoupled weight decay regularization, 2019

authors:
Yi Luan
Jacob Eisenstein
Kristina Toutanova
Michael Collins

title: Sparse, dense, and attentional representations for text retrieval

venue: arXiv preprint arXiv:2005.00181,

authors:
Xueguang Ma
Kai Sun
Ronak Pradeep
Jimmy Lin

title: A replication study of dense passage retriever

authors:
Christopher D Manning
Hinrich Schütze
Prabhakar Raghavan

title: Introduction to information retrieval

venue: Cambridge university press,

authors:
Yu Meng
Chenyan Xiong
Payal Bajaj
Saurabh Tiwary
Paul Bennett
Jiawei Han
Xia Song

title: Coco-lm: Correcting and contrasting text sequences for language model pretraining

venue: arXiv preprint arXiv:2102.08473,

authors:
Tomas Mikolov
Ilya Sutskever
Kai Chen
Greg S Corrado
Jeff Dean

title: Distributed representations of words and phrases and their compositionality

venue: In Advances in neural information processing systems,

authors:
Bhaskar Mitra
Nick Craswell

title: An introduction to neural information retrieval

venue: Foundations and Trends® in Information Retrieval,

authors:
Tri Nguyen
Mir Rosenberg
Xia Song
Jianfeng Gao
Saurabh Tiwary
Rangan Majumder
Li Deng

title: Ms marco: A human generated machine reading comprehension

authors:
Rodrigo Nogueira
Kyunghyun Cho

title: Passage re-ranking with bert

venue: arXiv preprint arXiv:1901.04085,

authors:
Hamid Palangi
Li Deng
Yelong Shen
Jianfeng Gao
Xiaodong He
Jianshu Chen
Xinying Song
Rabab Ward

title: Deep sentence embedding using long short-term memory networks: Analysis and application to information retrieval

venue: IEEE/ACM Transactions on Audio, Speech, and Language Processing,

authors:
Colin Raffel
Noam Shazeer
Adam Roberts
Katherine Lee
Sharan Narang
Michael Matena
Yanqi Zhou
Wei Li
Peter J Liu

title: Exploring the limits of transfer learning with a unified text-to-text transformer

venue: arXiv preprint arXiv:1910.10683,

authors:
Nils Reimers
Iryna Gurevych

title: Sentence-bert: Sentence embeddings using siamese bert-networks

venue: arXiv preprint arXiv:1908.10084,

authors:
Stephen Robertson
Hugo Zaragoza

title: The probabilistic relevance framework: BM25 and beyond

venue: Now Publishers Inc,

authors:
Stephen E Robertson
Steve Walker
Susan Jones
Micheline M Hancock-Beaulieu
Mike Gatford

title: Okapi at TREC-3

venue: NIST Special Publication Sp,

authors:
Devendra Singh Sachan
Mostofa Patwary
Mohammad Shoeybi
Neel Kant
Wei Ping
William L Hamilton
Bryan Catanzaro

title: End-to-end training of neural retrievers for open-domain question answering

venue: arXiv preprint arXiv:2101.00408,

authors:
Devendra Singh Sachan
Siva Reddy
William Hamilton
Chris Dyer
Dani Yogatama

title: End-to-end training of multi-document reader and retriever for open-domain question answering, 2021b

authors:
Yelong Shen
Xiaodong He
Jianfeng Gao
Li Deng
Grégoire Mesnil

title: Learning semantic representations using convolutional neural networks for web search

venue: In Proceedings of the 23rd international conference on world wide web,

authors:
Nandan Thakur
Nils Reimers
Andreas Rücklé
Abhishek Srivastava
Iryna Gurevych

title: Beir: A heterogenous benchmark for zero-shot evaluation of information retrieval models

venue: arXiv preprint arXiv:2104.08663,

authors:
James Thorne
Andreas Vlachos
Christos Christodoulopoulos
Arpit Mittal

title: Fever: a large-scale dataset for fact extraction and verification

venue: arXiv preprint arXiv:1803.05355,

authors:
Guillaume Wenzek
Marie-Anne Lachaux
Alexis Conneau
Vishrav Chaudhary
Francisco Guzmán
Armand Joulin
Edouard Grave

title: CCNet: Extracting high quality monolingual datasets from web crawl data

venue: In Proceedings of the 12th Language Resources and Evaluation Conference,

authors:
Zhirong Wu
Yuanjun Xiong
Stella X Yu
Dahua Lin

title: Unsupervised feature learning via nonparametric instance discrimination

venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,

authors:
Zhuofeng Wu
Sinong Wang
Jiatao Gu
Madian Khabsa
Fei Sun
Hao Ma

title: Clear: Contrastive learning for sentence representation

venue: arXiv preprint arXiv:2012.15466,

authors:
Lee Xiong
Chenyan Xiong
Ye Li
Kwok-Fung Tang
Jialin Liu
Paul Bennett
Junaid Ahmed
Arnold Overwijk

title: Approximate nearest neighbor negative contrastive learning for dense text retrieval

authors:
Sohee Yang
Minjoon Seo

title: Is retriever merely an approximator of reader

venue: arXiv preprint arXiv:2010.10999,

authors:
Tiancheng Zhao
Xiaopeng Lu
Kyusong Lee

title: Sparta: Efficient open-domain question answering via sparse transformer matching retrieval

venue: arXiv preprint arXiv:2009.13013,

heading: 1 INTRODUCTION

text:
Document retrieval is the task of finding relevant documents in a large collection to answer specific queries. This is an important task by itself and a core component of many natural language processing (NLP) problems, such as open domain question answering (Chen et al., 2017) or fact checking (Thorne et al., 2018). Traditionally, retrieval systems, or retrievers, leverage lexical similarities to match queries and documents, using, for instance, TF-IDF or BM25 weighting (Robertson & Zaragoza, 2009). These approaches, based on exact matches between tokens of the queries and documents, suffer from the lexical gap and do not generalize well (Berger et al., 2000). By contrast, approaches based on neural networks allow learning beyond lexical similarities, resulting in state-ofthe-art performance on question answering benchmarks, such as MS-MARCO (Nguyen et al., 2016) or NaturalQuestion (Kwiatkowski et al., 2019).
As training neural networks requires large amount of data, their strong retrieval results were possible for domains and applications where large training datasets are available. In the case of retrieval, creating these datasets requires manually matching queries to the relevant documents in the collection. This is hardly possible when the collection contains millions or billions of element, resulting in many scenarios where only a few in-domain examples, if any, are available. A potential solution is to train a dense retriever on a large retrieval dataset such as MS MARCO, and then apply it to new domains, a setting refered to as zero-shot. Unfortunately, in this setting dense retrievers are often outperformed by classical methods based on term-frequency, which do not require supervision (Thakur et al., 2021).
Thus, a natural alternative to transfer learning is unsupervised learning, which raises the following question: is it possible to train dense retrievers without supervision, and match the performance of BM25? Training dense retrievers without supervision can be achieved by using a pretext task that approximates retrieval. Given a document, one can generate a synthetic query and then train the network to retrieve the original document, among many others, given the query. The inverse Cloze task (ICT), proposed by Lee et al. (2019) to pre-train retrievers, uses a given sentence as query and predict the context surrounding it. While showing promising results as pre-training (Chang et al., 2020; Sachan et al., 2021a), this approach still lags behind BM25 when used as a zero-shot retriever.
ICT is strongly related to contrastive learning (Wu et al., 2018), which has been widely applied in computer vision (Chen et al., 2020; He et al., 2020). In particular, computer vision models trained with the latest contrastive learning methods lead to features well suited to retrieval (Caron et al., 2021). In this work, we thus propose to revisit how well contrastive learning performs to train dense retrievers without supervision. We want to investigate how much the recent developments from computer vision, such as MoCo, can improve dense retrievers. Here, our goal is not to develop new techniques, but to determine how far we can go by pushing existing methods to their limits.
In this paper, we thus make the following contributions. First, we show that contrastive learning can lead to strong unsupervised retrievers: our model achieves Recall@100 results competitive with BM25 on most of the BEIR benchmark. Second, in a few-shot setting, we show that our model benefits from few training examples, and obtains better results than transfering models from large datasets such as MS MARCO. Third, when used as a pre-training method before fine-tuning on MS MARCO, our technique leads to state-of-the-art retrieval performance on the BEIR benchmark. Finally, we perform ablations to motivate our design choices, and show that cropping works better than the inverse Cloze task.

heading: 2 RELATED WORK

text:
In this section, we briefly review relevant work in information retrieval, and application of machine learning to this problem. This is not an exhaustive review, and we refer the reader to Manning et al. (2008), Mitra et al. (2018) and Lin et al. (2020) for a more complete introduction to the field.
Term-frequency based information retrieval. Historically, in information retrieval, documents and queries are represented as sparse vectors where each element of the vectors corresponds to a term of the vocabulary. Different weighing schemes have been proposed, to determine how important a particular term is to a document in a large dataset. One of the most used weighing scheme is known as TF-IDF, and is based on inverse document frequency, or term specificity (Jones, 1972). BM25, which is still widely used today, extends TF-IDF (Robertson et al., 1995). A well known limitation of these approaches is that they rely on exact match to retrieve documents. This lead to the introduction of latent semantic analysis (Deerwester et al., 1990), in which documents are represented as low dimensional dense vectors.
Neural network based information retrieval. Following the successful application of deep learning methods to natural language processing, neural networks techniques were introduced for information retrieval. Huang et al. (2013) proposed a deep bag-of-words model, in which representations of queries and documents are computed independently. A relevance score is then obtained by taking the dot product between representations, and the model is trained end-to-end on click data from a search engine. This method was later refined by replacing the bag-of-words model by convolutional neural networks (Shen et al., 2014) or recurrent neural network (Palangi et al., 2016). A limitation of bi-encoders is that queries and documents are represented by a single vector, preventing the model to capture fine-grained interactions between terms. Nogueira & Cho (2019) thus introduced a cross-encoder model, based on a pre-trained BERT model (Devlin et al., 2019), which jointly encodes queries and documents. The application of a strong pre-trained model, as well as the cross-encoder architecture, lead to important improvement on the MS-MARCO benchmark (Bajaj et al., 2016).
The methods described in the previous paragraph were applied to re-rank documents, which were retrieved with a traditional IR system such as BM25. Gillick et al. (2018) first studied whether continuous retrievers, based on bi-encoder neural models, could be viable alternative to re-ranking. In the context of question answering, Karpukhin et al. (2020) introduced a dense passage retriever (DPR) based on the bi-encoder architecture. This model is initialized with a BERT network, and trained discriminatively using pairs of queries and relevant documents, with hard negatives from BM25. Xiong et al. (2020) further extended this work by mining hard negatives with the model itself during optimization, and trained on the MS-MARCO dataset. Once a collection of documents, such as Wikipedia articles, is encoded, retrieval is performed with a fast k-nearest neighbors library such as FAISS Johnson et al. (2019). To alleviate the limitations of bi-encoders, Humeau et al. (2019) introduces the poly-encoder architecture, where documents are encoded by multiple vectors. Similarly, Khattab et al. (2020) proposes the ColBERT model, which keeps a vector representation for each term of the queries and documents. To make the retrieval tractable, the term-level function
is approximated to first retrieve an initial set of candidates, which are then re-ranked with the true score. In the context of question answering, knowledge distillation has been used to train retrievers, either using the attention scores of the reader of the downstream task as synthetic labels (Izacard & Grave, 2021), or the relevance score from a cross encoder (Yang & Seo, 2020). Luan et al. (2020) compares, theoretically and empirically, the performance of sparse and dense retrievers, including bi-, cross- and poly-encoders. Dense retrievers, such as DPR, can lead to indices that weigh close to 100 GB when encoding document collections such as Wikipedia. Izacard et al. (2020) shows how to compress such indices, with limited impact on performance, making them more practical to use.
Self-supervised learning for NLP. Following the success of word2vec (Mikolov et al., 2013), many self-supervised techniques have been proposed to learn representation of text. Here, we briefly review the ones that are most related to our approach: sentence level models and contrastive techniques. Jernite et al. (2017) introduced different objective function to learn sentence representations, including next sentence prediction and binary order prediction. These objectives were later used in pre-trained models based on transformers, such as BERT (Devlin et al., 2019) and AlBERT (Lan et al., 2019). In the context of retrieval, Lee et al. (2019) introduced the inverse cloze task (ICT), whose purpose is to predict the context surrounding a span of text. Guu et al. (2020) integrated a bi-encoder retriever model in a BERT pre-training scheme. The retrieved documents are used as additional context in the BERT task, and the whole system is trained end-to-end in an unsupervised way. Similarly, Lewis et al. (2020) proposed to jointly learn a retriever and a generative seq2seq model, using self-supervised training. Chang et al. (2020) compares different pre-training tasks for retrieval, including the inverse cloze task. Constrastive learning was introduced in computer vision by Wu et al. (2018), followed by several modifications to improve the training (He et al., 2020; Chen et al., 2020; Caron et al., 2020). In the context of natural language processing, Fang et al. (2020) proposed to apply MoCo where positive pairs of sentences are obtained using back-translation. Different works augmented the masked language modeling objective with a contrastive loss (Giorgi et al., 2020; Wu et al., 2020; Meng et al., 2021). Finally, SBERT (Reimers & Gurevych, 2019) uses a siamese network similar to contrastive learning to learn a BERT-like model that is adapted to matching sentence embeddings. Their formulation is similar to our work but requires aligned pairs of sentences to form positive pairs while we propose to use data augmentation to leverage large unaligned textual corpora. Concurrent to this work, Gao & Callan (2021) have also shown the potential of contrastive learning for information retrieval; building on the same observation that both tasks share a similar structure.

heading: 3 METHOD

text:
In this section, we describe how to train a dense retriever with no supervision. We review the model architecture and then describe contrastive learning - a key component of its training.
The objective of a retriever is to find relevant documents in a large collection for a given query. Thus, the retriever takes as input the set of documents and the query and outputs a relevance score for each document. A standard approach is to encode each query–document pair with a neural network (Nogueira & Cho, 2019). This procedure requires re-encoding every document for any new query and hence does not scale to large collections of documents. Instead, we propose to use a bi-encoder architecture, where documents and queries are encoded independently (Huang et al., 2013; Karpukhin et al., 2020). One can compute the relevance score by taking the dot product (or cosine similarity) between the document’s representation and the representation of the query. More precisely, given a pair of query q and document d, we encode each of them independently using the same model, fθ, parametrized by θ. The relevance score s(q, d) is then the dot product of the resulting representations:
s(q, d) = 〈fθ(q), fθ(d)〉.
In practice, we use a transformer network for fθ to embed both queries and documents. The representation fθ(q) (resp. fθ(d)) for a query (resp. document) is obtained by averaging the hidden representations of the last layer. Following previous work on dense retrieval with neural networks, we use the BERT base uncased architecture and refer the reader to Devlin et al. (2019) for more details.

heading: 3.1 UNSUPERVISED TRAINING ON UNALIGNED DOCUMENTS

text: In this section, we describe our unsupervised training pipeline. We briefly review the loss function traditionally used in contrastive learning. We then discuss obtaining positive pairs from a single text document, a critical ingredient for this training paradigm.

heading: 3.1.1 CONTRASTIVE LEARNING

text:
Contrastive learning is an approach that relies on the fact that every document is, in some way, unique. This signal is the only information available in the absence of manual supervision. The resulting algorithm learns by discriminating between documents, using a contrastive loss (Wu et al., 2018). This loss compares either positive (from the same document) or negative (from different documents) pairs of document representations. More formally, given a positive pair of representations (q, k+) and a set of negative pairs (q, ki)i=0..K , the contrastive InfoNCE loss is defined as:
L(q, k+) = exp(s(q, k+)/τ)∑K i=0 exp(s(q, ki)/τ) ,
where τ is a temperature parameter. This loss encourages the relevance score of similar examples to be high and that of dissimilar examples to be low. Another interpretation of this loss function is the following: given the query representation q, the goal is to recover, or retrieve, the representation k+ corresponding to the positive document, among all the negatives ki. In the following, we refer to the left-hand side representations in the score s as queries and the right-hand side representations as keys.

heading: 3.1.2 BUILDING POSITIVE PAIRS FROM A SINGLE DOCUMENT

text:
A crucial element of contrastive learning is how to build positive pairs from a single input. In computer vision, this step relies on applying two independent data augmentations to the same image, resulting in two “views” that form a positive pair (Wu et al., 2018; Chen et al., 2020). While we consider similar independent text transformation, we also explore dependent transformations designed to reduce the correlation between views.
Inverse Cloze Task is a data augmentation that generates two mutually exclusive views of a document, introduced in the context of retrieval by Lee et al. (2019). The idea is to take a span of tokens to form one view, and its complementary to form the other view. More precisely, given a sequence of text (w1, ..., wn), ICT samples a span a, b, where 1 ≤ a < b ≤ n, uses the tokens of the span as the query and the rest of the tokens as the document (or key). In the original implementation by Lee et al. (2019) the span corresponds to a sentence, and is kept in the document 10% of the time to encourage lexical matching.
Independent cropping is a common independent data augmentation used for images where views are generated independently by cropping the input. In the context of text, cropping is equivalent to sampling a span of tokens. This strategy thus samples independently two spans from a document to form a positive pair. As opposed to the inverse Cloze task, in cropping both views of the example correspond to contiguous subsequence of the original data. A second difference between cropping and ICT is the fact that the task is symmetric: both the queries and documents follow the same distribution. Independent cropping can also lead to overlap between the two views of the data, hence encouraging the network to learn exact matches between the query and document, in a way that is similar to lexical matching methods like BM25. In practice, we can either fix the length of the span for the query and the key, or sample them.
Additional data augmentation. Finally, we also consider additional data augmentations such as random word deletion, random word replacement or random word masking. We use these perturbation in addition to random cropping.

heading: 3.1.3 BUILDING LARGE SET OF NEGATIVE PAIRS

text:
An important aspect of contrastive learning is to maintain a large set of negative pairs. Most standard frameworks differ from how the negatives are handled, and we briefly describe two of them, SimCLR and MoCo, that we use in this work.
Negative pairs within a batch. A first solution, called SimCLR (Chen et al., 2020), is to generate the negative pairs by using the other examples from the same batch: each example in a batch is transformed twice to generate positive pairs, and we generate negative pairs by using the views from the other examples in the batch. In that case, the gradient is back-propagated through the representations of both the queries and the keys. A downside of this approach is that it requires extremely large batch sizes to work well. This technique was also considered when using ICT to pre-train retrievers by Lee et al. (2019).
Negative pairs across batches. An alternative approach is to store representations from previous batches in a queue and use them as negative examples in the loss (Wu et al., 2018). This allows for smaller batch size but slightly changes the loss by making it asymmetric between “queries” (one of the view generated from the elements of the current batch), and “keys” (the other generated views, as well as the elements stored in the queue). Gradient is only backpropagated through the “queries”, and the representation of the “keys” are considered as fixed. In practice, the features stored in the queue from previous batches comes form previous iterations of the network. This leads to a drop of performance when the network rapidly changes during training. Instead, He et al. (2020) proposed to generate representations of keys from a second network that is updated more slowly. This approach, called MoCo, considers two networks: one for the keys, parametrized by θk, and one of the query, parametrized by θq. The parameters of the query network are updated with backpropagation and stochastic gradient descent, like in SimCLR, while the parameters of the key network, or Momentum encoder, is updated from the parameters of the query network by using a exponential moving average:
θk ← mθk + (1−m)θq, where m is the momentum parameter that takes its value in [0, 1].

heading: 4 EXPERIMENTS

text: In this section, we empirically evaluate contrastive learning as a way to train dense retrievers. First, we compare our best models to the state-of-the-art on competitive retrieval benchmarks for both zero-shot and few-shot retrieval. Then, we provide an ablation study to motivate the technical choices leading to our best retriever, shedding some light on the important ingredients of our approach. We give technical details about our models in Appendix A.

heading: 4.1 BEIR BENCHMARK

text: The BEIR benchmark, introduced by Thakur et al. (2021), contains 18 retrieval datasets with a focus on diversity. Each dataset is made of a set of queries, the corresponding relevant documents and a large collection of documents to retrieve from. These datasets correspond to nine different retrieval tasks, such as fact checking, question answering, entity retrieval or citation prediction. Moreover, the benchmark covers a large number of domains, such as Wikipedia, news articles, scientific publications or social media posts. Finally, there is also diversity in terms of documents and queries length, with queries ranging from 3 to 190 words and documents from 11 to 630 words. Most datasets from the BEIR benchmark do not contain a training set, and the focus of the benchmark is zero-shot retrieval. However, most machine learning based retrievers are still trained on supervised data, such as the large scale retrieval dataset MS-MARCO (Bajaj et al., 2016). Following standard practice, we report two metrics on this benchmark, the nDGC@10 and the Recall@100. These two metrics are complementary and both important: nDCG focuses on the ranking of the top 10 retrieved documents, and is good at evaluating documents returned to human, for example in a search engine. On the other hand, the Recall@100 is relevant to evaluate retrievers that are included in marchine learning systems, such as question answering. Indeed, such models can process hundreds of documents, and ignore the ranking of these documents (Izacard & Grave, 2020).

heading: 4.2 BASELINES

text:
First, we compare our approach to the BM25 ranking function, which does not require any training data. However, the relevance of individual terms depends on the inverse document frequency, which is computed over each document collection independently. In that sense, the ranking function used in the BM25 baseline is adapted to each dataset on which retrieval is performed. We also compare
to machine learning based retrievers, which can be classified in three categories: sparse, dense and late-interaction. For sparse methods, we compare to docT5query (Nogueira et al., 2019), which is a document expansion technique based on the T5 model (Raffel et al., 2019). The second sparse model we use as baseline is SPARTA (Zhao et al., 2020), which computes sparse representations of documents using a BERT model. More precisely, the ranking function of SPARTA is based on interactions between terms of the document and the query. Because the representations of the query terms are non-contextualized, the scores can be pre-computed when indexing the database. For dense methods, we use DPR (Karpukhin et al., 2020) and ANCE (Xiong et al., 2020) as baselines, which are bi-encoder models trained on supervised data such as NaturalQuestions or MS-MARCO. We also include the TAS-B technique (Hofstätter et al., 2021), which performs distillation from a cross-encoder to a bi-encoder model. Finally, we compare to GenQ, which uses a generative query model to create synthetic query-document pairs, and is applied to each dataset. Thus, a different GenQ model is used for each dataset of the BEIR benchmark. For late-interaction methods, we use ColBERT (Khattab et al., 2020), which computes pairwise scores between contextualized representations of the query and the document. To make this approach scalable, this scoring function is only applied on documents obtained by applying an approximation of the full score. Interestingly, it should be noted that BM25 is a very competitive baseline in the zero-shot setting, outperforming all dense retrievers on average over the BEIR benchmark. Only the re-ranking models and the query expansion (docT5query), which is also based on BM25, obtain better results than BM25.

heading: 4.3 RESULTS

text:
In Table 1, we report the nDCG@10 on the BEIR benchmark for different retrievers. The corresponding Recall@100 performances can be found in Table 7 of appendix. We individually report results on each dataset as well as the average (“Avg.”) over 14 datasets of the BEIR Benchmark (excluding 3 for license reasons). We omit retrievers that use a re-ranker during inference since it is orthogonal to our focus and can be applied to any method.
We observe that contrastive pre-training leads to strong performance among retrievers: our model obtains state-of-the art performance when finetuned on the MS-MARCO dataset. It should be noted that our finetuning procedure on MS-MARCO is simpler than for other dense retrievers: we use a simple strategy for negative mining and do not use distillation. Our model would probably also benefits from improvements proposed by these retrievers, but this is beyond the scope of this paper.
finetuning SciFact NFCorpus FiQA
Dataset size 809 2,590 5,500
BM25 - 66.5 32.5 23.6
BERT - 75.2 29.9 26.1 Ours - 84.0 33.9 31.5
BERT MS-MARCO 80.9 33.2 30.9 Ours MS-MARCO 84.1 34.9 34.4
We also illustrate the benefit of our dense retriever compared to BM25 by evaluating it in a few shot setting, where we have access to a small number of training queries and relevant documents to finetune locally the retriever. This setting is common in practice, and lexical based methods, like BM25, cannot benefit leverage the small training sets to adapt its weights. In Table 2, we report the nDCG@10 results on the three datasets from the BEIR benchmark associated with the smallest training sets. These training sets range from 809 to almost 5,500 queries. We observe that on these small datasets, our pre-training leads to significantly better results than BERT pre-training only, even when BERT model is finetuned on MS-MARCO. Our pre-trained model also outperforms BM25, showing the advantage of dense retriever over lexical methods in the low-shot setting.
Finally, we compare the performance of fully unsupervised models, i.e., without finetuning on MS-MARCO or other annotated data. In Table 8 we report the retrieval performance of our model on two standard question answering benchmarks: NaturalQuestions (Kwiatkowski et al., 2019) and TriviaQA (Joshi et al., 2017). In Figure 1 we report the performance of unsupervised models on the BEIR benchmark. Interestingly, we observe that, in this setting, our pre-training can lead to a model with competitive retrieval performance compared to BM25. Our approach outperforms previously proposed unsupervised dense retrievers. On question answering benchmarks it also outperforms a strong BM25 baseline (Ma et al., 2021) by 3% for the Recall@100, and is on par with BM25 on TriviaQA for the Recall@100. On datasets of the BEIRbenchmark, we observe that our unsupervised retriever is competitive with BM25 on every datasets, but TREC-COVID and Tóuche-2020. These results show the potential of contrastive learning to develop fully unsupervised dense retrievers.

heading: 5 ABLATION STUDIES

text:
In this section, we investigate the influence of different design choices on our method. In these ablations, all the models are pre-trained on Wikipedia for 200k gradient steps, with a batch size of
2,048 (on 32 GPUs). Each fine-tuning on MS-MARCO takes 20k gradient steps with a batch size of 512 (on 8 GPUs), using AdamW and no hard negatives.
MoCo vs. SimCLR. First, we compare the two contrastive pre-training methods: MoCo and SimCLR. As in SimCLR, the number of negative examples is equal to the batch size, we train models with a batch size of 4,096 and restrict the queue in MoCo to the same number of elements. This experiment measures the effect of using of momentum encoder for the keys instead of the same network as for the queries. Using a momentum also prevents from backpropagating the gradient through the keys. We report results, with fine-tuning on MS-MARCO in Table 5. We observe that the difference of performance between the two methods is small, especially after fine-tuning on MS-MARCO. This validates our choice of MoCo as our contrastive learning framework, since it scales to a larger number of negative examples without the need to increase the batch size.
Number of negative examples. Next, we study the influence of the number of negatives used in the contrastive loss, by varying the queue size of the MoCo algorithm. We consider values ranging from 2,048 to 131,072, and report results in Figure 2. We see that on average over the BEIR benchmark, increasing the number of negatives leads to better retrieval performance, especially in the unsupervised setting. However, we note that this effect is not equally strong for all datasets.
Data augmentations. Third, we compare different ways to generate pairs of positive examples from a single document or chunk of text. In particular, we compare random cropping, which leads to pairs with overlap, and the inverse cloze task, which was previously considered to pre-train retrievers. Interestingly, as shown in Table 4, the random cropping strategy strongly outperforms the inverse cloze task in our setting. We argue that overlap between the two views of the data is beneficial, as it allows the network to learn lexical matching, similarly to BM25. We also investigate
whether additional data perturbations, such as random word deletion or replacement, are beneficial for retrieval. After fine-tuning on MS-MARCO, it seems that such data perturbations slightly degrade the performance of dense retrievers.
Training data. Finally, we study the impact of the pre-training data on the performance of our retriever, by training on Wikipedia, CCNet or a mix of both sources of data. We report results in Table 5, and observe that there is no clear winner between the two data sources. Unsurprisingly, pre-training on Wikipedia leads to better performance on datasets for which documents also come from Wikipedia, such as NaturalQuestions or FEVER. On the other hand, on datasets from different domains than Wikipedia such as FiQA or ArguAna, training on the more diverse CCNet data leads to better results. To get the best of both worlds, we consider to strategies to mix the two data sources. In the “50/50%” strategy, examples are sampled uniformly across domain, meaning that half the batches are from Wikipedia and the other half from CCNet. In the “uniform” strategy, examples are sampled uniformly over the union of the dataset. Since CCNet is significantly larger than Wikipedia, this means that most of the batches are from CCNet.
Impact of fine-tuning on MS-MARCO. To isolate the impact of our pre-training method as opposed to the impact of our fine-tuning on MS-MARCO, we propose to apply the same fine-tuning to the BERT base uncased model. We report results in Table 6, and observe that when applied to BERT, our fine-tuning leads to results that are lower than the state-of-the-art. Hence, we believe that most of the improvement compared to the state-of-the-art retrievers can be attributed to our contrastive pre-training strategy.
DISCUSSION
In this work, we propose to explore the limits of contrastive pre-training to learn dense text retrievers. We use the MoCo technique, which allows us to train with a large number of negative examples. We make several interesting observations: first, while the inverse Cloze task has been introduced in the context of retrieval, we show that independent random cropping leads to better results. Second, we show that neural networks trained without supervision using contrastive learning exhibits good (albeit not state-of-the-art) retrieval performance. Based on these observations, we trained a model that obtains state-of-the-art performance on the BEIR benchmark when fine-tuned on MS-MARCO.

heading: A TECHNICAL DETAILS

text:
Contrastive pre-training. For the model with fine-tuning on MS-MARCO (respectively, without fine-tuning), we use the MoCo algorithm He et al. (2020) with a queue of size 65536 (resp. 131,072), a momentum value of 0.999 (resp. 0.9995) and a temperature of 0.05. We use the random cropping data augmentation, with documents of 256 tokens and span sizes sampled between 10% and 50% (resp. 5% and 50%) of the document length. We also apply word deletion with a probability of 20%. We optimize the model with the AdamW (Loshchilov & Hutter, 2019) optimizer, with learning rate of 10−4 (resp. 5 · 10−5), batch size of 2,048 and 500,000 (resp. 200,000) steps. We initialize the network with the publicly available BERT base uncased model. For training data, we use a data from Wikipedia (resp. a mix between Wikipedia and CC-net data (Wenzek et al., 2020), where half the batches are sampled from each source).
Fine-tuning on MS-MARCO. When fine-tuning our model on MS-MARCO, we use the ASAM optimizer (Kwon et al., 2021), with a learning rate of 10−5 and a batch size of 1024 with a temperature of 0.05, also used during pretraining. We train an initial model with random negative examples for 20000 steps, mine hard negatives with this first model, and re-train a second model with those. Each query is associated with a gold document and a negative document, which is a random document in the first phase and a hard negative 10% of the time in the second phase. All documents from a batch are used as negatives for each query.

title: TOWARDS UNSUPERVISED DENSE INFORMATION RE-