Posterior Differential Regularization with f-divergence for Improving Model Robustness (2021)

Abstract: We address the problem of enhancing model robustness through regularization. Specifically, we focus on methods that regularize the model posterior difference between clean and noisy inputs. Theoretically, we provide a connection of two recent methods, Jacobian Regularization and Virtual Adversarial Training, under this framework. Additionally, we generalize the posterior differential regularization to the family of f-divergences and characterize the overall framework in terms of the Jacobian matrix. Empirically, we compare those regularizations and standard BERT training on a diverse set of tasks to provide a comprehensive profile of their effect on model generalization. For both fully supervised and semi-supervised settings, we show that regularizing the posterior difference with f-divergence can result in well-improved model robustness. In particular, with a proper f-divergence, a BERT-base model can achieve comparable generalization as its BERT-large counterpart for in-domain, adversarial and domain shift scenarios, indicating the great potential of the proposed framework for enhancing NLP model robustness.

## introduction although recent neural network based models have achieved great success in a wide range of natural language processing (nlp) tasks, these models may still suffer catastrophic degradation in out-ofdomain generalization to datasets with domain shift or adversarial scenarios @xcite . for example, large-scale pretrained neural language models @xcite have better generalization, but experience performance reduction caused by domain shift @xcite . textual entailment models trained on mnli @xcite have picked up superficial cues focusing on either the 1 code is available at https://github.com/ hao-cheng/f-divergence . presence of certain keywords @xcite or whether similar words are mentioned in the sentence pairs @xcite . @xcite have also shown that squad models are very easily distracted by irrelevant sentences that contain many question words regardless of context, despite of their human-performance on in-domain data. as shown in figure 1 , three bertbased @xcite models perform well for in-domain evaluation data, but transfer poorly to out-of-domain datasets with domain shift or adversarial attack, i.e. more than 25% relative performance reduction. achieving good generalizations towards datasets with domain shift has been a long-standing goal of domain adaptation. various methods @xcite have been developed for training models to learn effectively from both in-domain (source) and out-of-domain (target) datasets. additionally, the recent discovery on the prevalence of data biases @xcite @xcite , unintended correlations between input and output learned by statistical models, ignites the develop-ment of model debiasing techniques @xcite . these methods leverage discovered data biases to improve model generalization over adversarial datasets @xcite designed to fool naively trained models. instead of relying on knowledge of the target dataset, we focus on task-agnostic training techniques for enhancing model robustness with access to only in-domain data. motivated by recent success of adversarial training in computer vision @xcite and nlp @xcite @xcite , we investigate a regularization framework, which directly regularizes the model posterior difference for clean and noisy inputs, as a means to enhance the model robustness. here, we first provide a theoretical connection of two recent methods under this framework, i.e. virtual adversarial training (vat) @xcite and jacobian regularization (jr) @xcite . in addition, we propose to generalize vat and random perturbation training (rpt) @xcite with a family of probability distribution metrics, f -divergences, and characterize their connection with jr. given that large-scale pretrained neural language models have demonstrated their superior generalization for downstream nlp tasks under both matched @xcite and mismatched evaluations @xcite , we systematically study the regularization framework using bert @xcite on a diverse set of tasks in terms of both in-domain and out-ofdomain generalization. specifically, we use representative datasets @xcite @xcite ) from sentiment analysis, textual entailment and question answering (qa) for in-domain training and evaluation. in order to assess the resulting model generalization over domain shift and adversarial attack, we then consider out-of-domain datasets @xcite and challenge adversarial datasets @xcite in a zero-shot learning fashion. our experiments show that regularizing the posterior difference for clean and noise inputs is very effective in improving model generalization under both supervised and semi-supervised learning settings. based on our theoretical analysis, both vat and rpt variants, unlike jr, incorporate model confidence for adaptive regularization, which leads to consistently better empirical robustness over bert with the standard fine-tuning. furthermore, we find that different f -divergences lead to different generalization behaviors for in-domain, domain shift and adversarial settings. in our study, vat with symmetric divergence achieve better generalization for in-domain and domain shift cases, while vat with asymmetric divergence achieve more robustness toward adversarial attack. more importantly, we show that a bert-base model trained with a proper f -divergence can perform comparably to its corresponding bert-large counterpart. it is also worth noting that vat with symmetric divergence lead to improved data efficiency, i.e. achieving comparable in-domain performance as fully-supervised models with only 50% labelled data. this further illustrates the great potential of the proposed general regularization framework for the semi-supervised setting. our main contributions are summarized as follows: 1) we generalize the posterior differential regularization framework to the family of f -divergences and provide additional divergence functions with different characteristics for regularization; 2) based on our framework, we analyze the family of regularization methods and show the theoretical connection of recently proposed methods, jr, vat and rpt, in terms of their regularization effect on the input-output jacobian matrix; 3) we provide a comprehensive profile of different regularization methods over a diverse set of nlp tasks and experimental insight into which f -divergence is more suitable for improving nlp model robustness under both supervised and semi-supervised settings. ## posterior differential regularization in this section, we first introduce the regularization framework that penalizes the difference of posterior between clean and noisy inputs. based on this framework, we set up the basic notions of two recent methods, i.e. vat and jr, and show their theoretical connection. finally, we generalize the posterior differential regularization with any function from the family of f -divergences and characterize their local smoothness promoting in terms of jacobian matrix. in the following, we use f θ (x) : r @xmath0. both adversarial learning @xcite and the posterior difference regularization aim at making the model more robust. adversarial learning focuses on minimizing the following objective in order to reduce the overall model sensitivity to the input perturbation, @xcite propose to directly regularize the frobenius norm of the input-output jacobian matrix so that although sharing with jr the same spirit of encouraging the local smoothness of the model, jr and vat are not fully theoretically connected. in what follows, we use a simple approach to draw the theoretical connection between these two methods. connection between vat and jr: here, we show that vat and jr can be directly related through the definition of induced matrix norm. specifically, the matrix norm of the jacobian matrix is and the last inequality is attained based on pinsker's inequality. therefore, the vat regularization provides an upper bound for the spectral norm of the jacobian matrix. although a similar attempt to relate jr and vat has been first explored in @xcite , we provide a simple and comprehensive connection. specifically, both vat and jr regularize the upper bound of the spectral norm of the jacobian matrix. posterior differential regularization with fdivergence: although both vat and jr have been successful in improving model robustness, they are both special cases of regularizing the model posterior difference between the clean and noisy inputs. one natural question is whether we can use other probability distribution metrics for regularization and characterize them in terms of jacobian matrix. in the following, we extend the posterior difference regularization with the family of f -divergences @xcite . furthermore, we show that posterior differential regularization with all f -divergences results in an adaptive variant of jr which incorporates model confidence. first, let's define the f -divergence for measuring the posterior difference between the clean and noisy inputs, i.e. ## model given its superior performance over a wide range of nlp tasks, we focus on exploring different train-ing techniques using bert @xcite . we first describe the bert representations used for all tasks considered in this paper. then, two variants of task-specific bert-based models are introduced: 1) the sentence-level classifier for textual entailment and sentiment analysis, and 2) the extractive qa model. specifically, we focus on different ways of encoding input text and building task-specific layers using bert representations. bert representation: for all tasks considered in this work, an input text sequence is divided into subword units w t , @xmath1, . . . , t . the tokenized input sequence is then transformed into embeddings, x 1 , . . . , x @xmath2, through a token encoder, which combines a token embedding, a (token) position embedding and a segment embedding (i.e., which text span the token belongs to) by elementwise summation. the embedding layer is used as the input to multiple transformer layers @xcite to generate the contextual representations, h 1 , . . . , h @xmath3, which are the hidden states of the last layer of the bert model. for all regularizations, we sample noise vectors 1 , . . . , t from n (0, i), and normalize each vector into l2 unit vector. the noise input is then constructed by adding the normalized noise vector to the token embeddings, i.e. x 1 + c 1 , . . . , x t + c t . here, we fix @xmath4. sentence-level classifier: following the standard setup of bert-based textual entailment model @xcite , a pair of premise and hypothesis is converted into an input sequence in the form of "[cls]premise [sep]hypothesis [sep]". here, [cls]is a special token indicating the start of the whole sequence and [sep]is another special token for separating the two sentences. for sentiment analysis, a single sentence is converted to the form of "[cls]sentence [sep]". for both classification tasks, the task-specific layer only takes the first hidden vector h [cls] produced by bert, corresponding to the [cls]token. then, the probability of class k is extractive qa model: for extractive qa, the probability space outcomes consist of token positions of answer spans. given a pair of question q and a passage p in the form of "[cls]question [sep]passage [sep]", the bert encoder produces contextualized representations for all tokens in the input. specifically, for each token position t in p, the final hidden vector h @xmath5, where d is the vector dimension. the span-begin score is computed as s b (i) = w t b h i using a weight vector w @xmath6. the probability for a start position i is maximizing the log-likelihood of the above equation is equivalent to maximizing the log probabilities for the correct start and end position, respectively. for regularized models, given it is computationally expensive to enumerate all possible spans, we apply two separate regularization terms for the start and end position probabilities, respectively. ## experiments in this section, we apply the regularization methods discussed so far to bert and evaluate their performance on the model robustness. specifically, we consider two types of posterior regularization with f -divergences. in addition to a vat-like regularization with an inner search for the most adversarial direction, following @xcite , we also evaluate the random perturbation training (rpt) with the family of f -divergences which only uses randomly sampled noise for regularization. in this work, we focus on three representative fdivergences, i.e. kl, squared hellinger (shl) and jensen-shannon divergence (jsd). dataset: all the datasets used in this paper are summarized in to evaluate the out-of-domain generalization with domain shift, we use the bioaqs dataset @xcite from mrqa @xcite and the imdb dataset @xcite . unlike squad which is based on wikipedia, bioaqs is a biomedical qa dataset constructed on pubmed articles. compared with sst-2 containing pithy export reviews @xcite , imdb includes lengthy movie reviews from nonexperts @xcite . we directly apply the qa model trained on squad v2.0 and the sentiment classifier trained on sss-2 to bioaqs and imdb, respectively. to evaluate the model robustness towards adversarial attack, we use two challenging adversarial datasets, i.e. adversarial squad @xcite and hans @xcite for evaluating qa model trained on squad v1.1 and the textual entailment model trained on mnli, respectively. the adversarial squad is constructed based on squad v1.1 @xcite by adding distracting sentences that have high overlap with the question and contain plausible answer candidates. naively trained models tend to exploit the word overlap with the given question and thus are fooled by those distracting sentences @xcite . the hans dataset is built using three heuristics to ensure that the hypothesis sentence only contains words from the premise sentence @xcite . similarly, standard training results in models failing catastrophically, even for bert. implementation: we follow the default setting used for fine-tuning the uncased bert base model @xcite . we select the learning rate from {3e-5, 4e-5} for qa models and {2e-5, 3e-5} for classification models. for both tasks, we tune the number of training epochs in {2, 3, 4, 5}. in addition, we search regularization weight in {0.001, 0.01, 0.1} for jr, and {1, 4, 10} for vat and rpt. we use the in-domain dev set for validation and select the best model based on f1 generalization. consistent with findings in @xcite , by searching for the most adversarial perturbation direction, vat variants achieve the largest boost for in-domain generalization. moreover, we find that both rpt and vat with shl and jsd provides additional improvement over their corresponding counterpart with kl which suggests the benefit of using alternative f -divergences for posterior difference regularization. lastly, by selecting the proper divergence, the performance gap between the bert-base and bert-large model is dramatically narrowed which indicates the advantage of applying posterior difference regularization with f -divergences on top of powerful text representations. ## related work with the goal of developing more robust nlp models, a line of recent work has been devoted to identifying various kinds of superficial patterns learned by high-performance models over many popular datasets @xcite @xcite . the prevalence of data biases over popular datasets poses a real challenge of accurately estimating the model capacity for practical applications, because a closed dataset evaluation usually inflates the model performance. this concern about dataset biases has led researchers to develop new diagnostic datasets @xcite @xcite and training techniques @xcite to overcome those discovered biases. recent debiasing methods @xcite require learning multiple models to access known data biases for the target dataset. moreover, they achieve more robust outof-domain generalization at the price of in-domain performance degradation. in contrast, we focus on the task-agnostic robust learning framework for enhancing model robustness and empirically show that regularization approaches under this framework can result in superior in-domain and out-ofdomain generalization. training with noise has been a very popular approach for enhancing model robustness. the dropout is a widely-used approach in deep learning to improve model generalization @xcite . for adversarial learning methods, the main theme is reducing the model sensitivity toward small input perturbations @xcite , which has been recently applied to both fine-turning @xcite @xcite and pre-training @xcite . however, models trained with adversarial learning are found to have at-odd generalization @xcite . our work studies learning methods with the goal of regularizing the model posterior difference of clean and noisy inputs. we show that compared with the standard bert training, the proposed posterior differential regularization with f -divergence lead to better nlp model robustness. ## conclusion in this paper, we investigate methods regularizing the posterior difference between the clean and noisy inputs for improving model generalization for both in-domain and out-of-domain settings. specifically, we present theoretical analyses of three methods under this framework, i.e. vat, jr, and rpt. we further extend both vat and prt to the family of f -divergences and theoretically characterize them in terms of jacobian matrix. we also demonstrate their effectiveness in enhancing model robustness over a diverse set of nlp tasks under both fullysupervised and semi-supervised scenarios. for future work, it is interesting to explore posterior differential regularization methods for weaklysupervised learning, such as relation extraction and qa with distant supervision.