InfoBERT: Improving Robustness of Language Models from An Information Theoretic Perspective

1 INTRODUCTION . Self-supervised representation learning pre-trains good feature extractors from massive unlabeled data , which show promising transferability to various downstream tasks . Recent success includes large-scale pre-trained language models ( e.g. , BERT , RoBERTa , and GPT-3 ( Devlin et al. , 2019 ; Liu et al. , 2019 ; Brown et al. , 2020 ) ) , which have advanced state of the art over a wide range of NLP tasks such as NLI and QA , even surpassing human performance . Specifically , in the computer vision domain , many studies have shown that self-supervised representation learning is essentially solving the problem of maximizing the mutual information ( MI ) I ( X ; T ) between the input X and the representation T ( van den Oord et al. , 2018 ; Belghazi et al. , 2018 ; Hjelm et al. , 2019 ; Chen et al. , 2020 ) . Since MI is computationally intractable in high-dimensional feature space , many MI estimators ( Belghazi et al. , 2018 ) have been proposed to serve as lower bounds ( Barber & Agakov , 2003 ; van den Oord et al. , 2018 ) or upper bounds ( Cheng et al. , 2020 ) of MI . Recently , Kong et al . point out that the MI maximization principle of representation learning can be applied to not only computer vision but also NLP domain , and propose a unified view that recent pre-trained language models are maximizing a lower bound of MI among different segments of a word sequence . On the other hand , deep neural networks are known to be prone to adversarial examples ( Goodfellow et al. , 2015 ; Papernot et al. , 2016 ; Eykholt et al. , 2017 ; Moosavi-Dezfooli et al. , 2016 ) , i.e. , the outputs of neural networks can be arbitrarily wrong when human-imperceptible adversarial perturbations are added to the inputs . Textual adversarial attacks typically perform word-level substitution ( Ebrahimi et al. , 2018 ; Alzantot et al. , 2018 ; Ren et al. , 2019 ) or sentence-level paraphrasing ( Iyyer et al. , 2018 ; Zhang et al. , 2019 ) to achieve semantic/utility preservation that seems innocuous to human , while fools NLP models . Recent studies ( Jin et al. , 2020 ; Zang et al. , 2020 ; Nie et al. , 2020 ; Wang et al. , 2020 ) further show that even large-scale pre-trained language models ( LM ) such as ∗Work was done during Boxin Wang ’ s Summer internship in Microsoft Dynamics 365 AI Research . BERT are vulnerable to adversarial attacks , which raises the challenge of building robust real-world LM applications against unknown adversarial attacks . We investigate the robustness of language models from an information theoretic perspective , and propose a novel learning framework InfoBERT , which focuses on improving the robustness of language representations by fine-tuning both local features ( word-level representation ) and global features ( sentence-level representation ) for robustness purpose . InfoBERT considers two MI-based regularizers : ( i ) the Information Bottleneck regularizer manages to extract approximate minimal sufficient statistics for downstream tasks , while removing excessive and noisy information that may incur adversarial attacks ; ( ii ) the Anchored Feature regularizer carefully selects useful local stable features that are invulnerable to adversarial attacks , and maximizes the mutual information between local stable features and global features to improve the robustness of the global representation . In this paper , we provide a detailed theoretical analysis to explicate the effect of InfoBERT for robustness improvement , along with extensive empirical adversarial evaluation to validate the theory . Our contributions are summarized as follows . ( i ) We propose a novel learning framework InfoBERT from the information theory perspective , aiming to effectively improve the robustness of language models . ( ii ) We provide a principled theoretical analysis on model robustness , and propose two MIbased regularizers to refine the local and global features , which can be applied to both standard and adversarial training for different NLP tasks . ( iii ) Comprehensive experimental results demonstrate that InfoBERT can substantially improve robust accuracy by a large margin without sacrificing the benign accuracy , yielding the state-of-the-art performance across multiple adversarial datasets on NLI and QA tasks . 2 RELATED WORK . Textual Adversarial Attacks/Defenses Most existing textual adversarial attacks focus on wordlevel adversarial manipulation . Ebrahimi et al . ( 2018 ) is the first to propose a whitebox gradientbased attack to search for adversarial word/character substitution . Following work ( Alzantot et al. , 2018 ; Ren et al. , 2019 ; Zang et al. , 2020 ; Jin et al. , 2020 ) further constrains the perturbation search space and adopts Part-of-Speech checking to make NLP adversarial examples look natural to human . To defend against textual adversarial attacks , existing work can be classified into three categories : ( i ) Adversarial Training is a practical method to defend against adversarial examples . Existing work either uses PGD-based attacks to generate adversarial examples in the embedding space of NLP as data augmentation ( Zhu et al. , 2020a ) , or regularizes the standard objective using virtual adversarial training ( Jiang et al. , 2020 ; Liu et al. , 2020 ; Gan et al. , 2020 ) . However , one drawback is that the threat model is often unknown , which renders adversarial training less effective when facing unseen attacks . ( ii ) Interval Bound Propagation ( IBP ) ( Dvijotham et al. , 2018 ) is proposed as a new technique to consider the worst-case perturbation theoretically . Recent work ( Huang et al. , 2019 ; Jia et al. , 2019 ) has applied IBP in the NLP domain to certify the robustness of models . However , IBPbased methods rely on strong assumptions of model architecture and are difficult to adapt to recent transformer-based language models . ( iii ) Randomized Smoothing ( Cohen et al. , 2019 ) provides a tight robustness guarantee in ` 2 norm by smoothing the classifier with Gaussian noise . Ye et al . ( 2020 ) adapts the idea to the NLP domain , and replace the Gaussian noise with synonym words to certify the robustness as long as adversarial word substitution falls into predefined synonym sets . However , to guarantee the completeness of the synonym set is challenging . Representation Learning MI maximization principle has been adopted by many studies on selfsupervised representation learning ( van den Oord et al. , 2018 ; Belghazi et al. , 2018 ; Hjelm et al. , 2019 ; Chen et al. , 2020 ) . Specifically , InfoNCE ( van den Oord et al. , 2018 ) is used as the lower bound of MI , forming the problem as contrastive learning ( Saunshi et al. , 2019 ; Yu et al. , 2020 ) . However , Tian et al . ( 2020 ) suggests that the InfoMax ( Linsker , 1988 ) principle may introduce excessive and noisy information , which could be adversarial . To generate robust representation , Zhu et al . ( 2020b ) formalizes the problem from a mutual-information perspective , which essentially performs adversarial training for worst-case perturbation , while mainly considers the continuous space in computer vision . In contrast , InfoBERT originates from an information-theoretic perspective and is compatible with both standard and adversarial training for discrete input space of language models . 3 INFOBERT . Before diving into details , we first discuss the textual adversarial examples we consider in this paper . We mainly focus on the dominant word-level attack as the main threat model , since it achieves higher attack success and is less noticeable to human readers than other attacks . Due to the discrete nature of text input space , it is difficult to measure adversarial distortion on token level . Instead , because most word-level adversarial attacks ( Li et al. , 2019 ; Jin et al. , 2020 ) constrain word perturbations via the bounded magnitude in the semantic embedding space , by adapting from Jacobsen et al . ( 2019 ) , we define the adversarial text examples with distortions constrained in the embedding space . Definition 3.1 . ( -bounded Textual Adversarial Examples ) . Given a sentence x = [ x1 ; x2 ; ... ; xn ] , where xi is the word at the i-th position , the -bounded adversarial sentence x′ = [ x′1 ; x ′ 2 ; ... ; x ′ n ] for a classifier F satisfies : ( 1 ) F ( x ) = o ( x ) = o ( x′ ) but F ( x′ ) 6= o ( x′ ) , where o ( · ) is the oracle ( e.g. , human decision-maker ) ; ( 2 ) ||ti − t′i||2 ≤ for i = 1 , 2 , ... , n , where ≥ 0 and ti is the word embedding of xi . 3.1 INFORMATION BOTTLENECK AS A REGULARIZER . In this section , we first discuss the general IB implementation , and then explain how IB formulation is adapted to InfoBERT as a regularizer along with theoretical analysis to support why IB regularizer can help improve the robustness of language models . The IB principle formulates the goal of deep learning as an information-theoretic trade-off between representation compression and predictive power ( Tishby & Zaslavsky , 2015 ) . Given the input source X , a deep neural net learns the internal representation T of some intermediate layer and maximizes the MI between T and label Y , so that T subject to a constraint on its complexity contains sufficient information to infer the target label Y . Finding an optimal representation T can be formulated as the maximization of the Lagrangian LIB = I ( Y ; T ) − βI ( X ; T ) , ( 1 ) where β > 0 is a hyper-parameter to control the tradeoff , and I ( Y ; T ) is defined as : I ( Y ; T ) = ∫ p ( y , t ) log p ( y , t ) p ( y ) p ( t ) dy dt . ( 2 ) Since Eq . ( 2 ) is intractable , we instead use the lower bound from Barber & Agakov ( 2003 ) : I ( Y ; T ) ≥ ∫ p ( y , t ) log qψ ( y | t ) dy dt , ( 3 ) where qψ ( y|t ) is the variational approximation learned by a neural network parameterized by ψ for the true distribution p ( y|t ) . This indicates that maximizing the lower bound of the first term of IB I ( Y ; T ) is equivalent to minimizing the task cross-entropy loss ` task = H ( Y | T ) . To derive a tractable lower bound of IB , we here use an upper bound ( Cheng et al. , 2020 ) of I ( X ; T ) I ( X ; T ) ≤ ∫ p ( x , t ) log ( p ( t | x ) ) dx dt − ∫ p ( x ) p ( t ) log ( p ( t | x ) ) dx dt . ( 4 ) By combining Eq . ( 3 ) and ( 4 ) , we can maximize the tractable lower bound L̂IB of IB in practice by : L̂IB = 1 N N∑ i=1 [ log qψ ( y ( i ) | t ( i ) ) ] − β N N∑ i=1 [ log ( p ( t ( i ) | x ( i ) ) ) − 1 N N∑ j=1 log ( p ( t ( j ) | x ( i ) ) ) ] ( 5 ) with data samples { x ( i ) , y ( i ) } Ni=1 , where qψ can represent any classification model ( e.g. , BERT ) , and p ( t | x ) can be viewed as the feature extractor fθ : X → T , where X and T are the support of the input source X and extracted feature T , respectively . The above is a general implementation of IB objective function . In InfoBERT , we consider T as the features consisting of the local word-level features after the BERT embedding layer fθ . The following BERT self-attentive layers along with the linear classification head serve as qψ ( y|t ) that predicts the target Y given representation T . Formally , given random variablesX = [ X1 ; X2 ; ... ; Xn ] representing input sentences withXi ( word token at i-th index ) , let T = [ T1 ; ... ; Tn ] = fθ ( [ X1 ; X2 ; ... ; Xn ] ) = [ fθ ( X1 ) ; fθ ( X2 ) ; ... ; fθ ( Xn ) ] denote the random variables representing the features generated from input X via the BERT embedding layer fθ , where Ti ∈ Rd is the high-dimensional word-level local feature for word Xi . Due to the high dimensionality d of each word feature ( e.g. , 1024 for BERT-large ) , when the sentence length n increases , the dimensionality of features T becomes too large to compute I ( X ; T ) in practice . Thus , we propose to maximize a localized formulation of IB LLIB defined as : LLIB : = I ( Y ; T ) − nβ n∑ i=1 I ( Xi ; Ti ) . ( 6 ) Theorem 3.1 . ( Lower Bound of LIB ) Given a sequence of random variables X = [ X1 ; X2 ; ... ; Xn ] and a deterministic feature extractor fθ , let T = [ T1 ; ... ; Tn ] = [ fθ ( X1 ) ; fθ ( X2 ) ; ... ; fθ ( Xn ) ] . Then the localized formulation of IB LLIB is a lower bound of LIB ( Eq . ( 1 ) ) , i.e. , I ( Y ; T ) − βI ( X ; T ) ≥ I ( Y ; T ) − nβ n∑ i=1 I ( Xi ; Ti ) . ( 7 ) Theorem 3.1 indicates that we can maximize the localized formulation of LLIB as a lower bound of IB LIB when I ( X ; T ) is difficult to compute . In Eq . ( 6 ) , if we regard the first term ( I ( Y ; T ) ) as a task-related objective , the second term ( −nβ ∑n i=1 I ( Xi ; Ti ) ) can be considered as a regularization term to constrain the complexity of representation T , thus named as Information Bottleneck regularizer . Next , we give a theoretical analysis for the adversarial robustness of IB and demonstrate why localized IB objective function can help improve the robustness to adversarial attacks . Following Definition 3.1 , let T = [ T1 ; T2 ; ... ; Tn ] and T ′ = [ T ′1 ; T ′ 2 ; ... ; T ′ n ] denote the features for the benign sentence X and adversarial sentence X ′ . The distributions of X and X ′ are denoted by probability p ( x ) and q ( x ) with the support X and X ′ , respectively . We assume that the feature representation T has finite support denoted by T considering the finite vocabulary size in NLP . Theorem 3.2 . ( Adversarial Robustness Bound ) For random variables X = [ X1 ; X2 ; ... ; Xn ] and X ′ = [ X ′1 ; X ′ 2 ; ... ; X ′ n ] , let T = [ T1 ; T2 ; ... ; Tn ] = [ fθ ( X1 ) ; fθ ( X2 ) ; ... ; fθ ( Xn ) ] and T ′ = [ T ′1 ; T ′ 2 ; ... ; T ′ n ] = [ fθ ( X ′ 1 ) ; fθ ( X ′ 2 ) ; ... ; fθ ( X ′ n ) ] with finite support T , where fθ is a deterministic feature extractor . The performance gap between benign and adversarial data |I ( Y ; T ) − I ( Y ; T ′ ) | is bounded above by |I ( Y ; T ) − I ( Y ; T ′ ) | ≤ B0 +B1 n∑ i=1 √ |T | ( I ( Xi ; Ti ) ) 1/2 +B2 n∑ i=1 |T |3/4 ( I ( Xi ; Ti ) ) 1/4 +B3 n∑ i=1 √ |T | ( I ( X ′i ; T ′i ) ) 1/2 +B4 n∑ i=1 |T |3/4 ( I ( X ′i ; T ′i ) ) 1/4 , ( 8 ) where B0 , B1 , B2 , B3 and B4 are constants depending on the sequence length n , and p ( x ) . The sketch of the proof is to express the difference of |I ( Y ; T ) − I ( Y ′ ; T ) | in terms of I ( Xi ; Ti ) . Specifically , Eq . ( 25 ) factorizes the difference into two summands . The first summand , the conditional entropy |H ( T | Y ) − H ( T ′ | Y ) | , can be bound by Eq . ( 42 ) in terms of MI between benign/adversarial input and representation I ( Xi ; Ti ) and I ( X ′i ; T ′ i ) . The second summand |H ( T ) − H ( T ′ ) | has a constant upper bound ( Eq . ( 85 ) ) , since language models have bounded vocabulary size and embedding space , and thus have bounded entropy . The intuition of Theorem 3.2 is to bound the adversarial performance drop |I ( Y ; T ) − I ( Y ; T ′ ) | by I ( Xi ; Ti ) . As explained in Eq . ( 3 ) , I ( Y ; T ) and I ( Y ; T ′ ) can be regarded as the model performance on benign and adversarial data . Thus , the LHS of the bound represents such a performance gap . The adversarial robustness bound of Theorem 3.2 indicates that the performance gap becomes closer when I ( Xi ; Ti ) and I ( X ′i ; T ′ i ) decrease . Note that our IB regularizer in the objective function Eq . ( 6 ) achieves the same goal of minimizing I ( Xi ; Ti ) while learning the most efficient information features , or approximate minimal sufficient statistics , for downstream tasks . Theorem 3.2 also suggests that combining adversarial training with our IB regularizer can further minimize I ( X ′i ; T ′ i ) , leading to better robustness , which is verified in §4 .

Summary: This work (InfoBERT) proposes additional objectives for transformer finetuning to obtain models more robust to adversarial inputs. The authors first propose a mutual information based information bottleneck objective, next the authors propose an adversarial loss inspired method for identifying robust features and a subsequent objective to emphasize the mutual information between global representations and these robust features. The experiments demonstrate that InfoBERT consistently outperforms other adversarial training approaches on a variety of adversarial evaluations.