SlowGuess/NAACL_2021_OCR huggingface dataset:
Posterior Differential Regularization with  $f$ -divergence for Improving Model Robustness
Hao Cheng $^{1}$ , Xiaodong Liu $^{1}$ , Lis Pereira $^{2}$ , Yaoliang Yu $^{3}$ , Jianfeng Gao $^{1}$
<sup>1</sup>Microsoft Research <sup>2</sup>Ochanomizu University <sup>3</sup>University of Waterloo & Vector Institute
{chehao,xiaodl,jfgao}@microsoft.com
{kanashiro.pereira}@ocha.ac.jp
{yaoliang.yu}@uwaterloo.ca
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 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. $^{1}$
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-of-domain generalization to datasets with domain shift or adversarial scenarios (Nie et al., 2019; Hsieh et al., 2019). For example, large-scale pretrained neural language models (Devlin et al., 2019; Liu et al., 2019) have better generalization, but experience performance reduction caused by domain shift (Hendrycks et al., 2020). Textual entailment models trained on MNLI (Williams et al., 2018) have picked up superficial cues focusing on either the
Figure 1: BERT-based question answering (QA) and natural language inference (NLI) model robustness towards domain shift and adversarial attack. The numbers in percentage indicate the corresponding relative performance drop between in-domain and out-of-domain test sets.
presence of certain keywords (Gururangan et al., 2018) or whether similar words are mentioned in the sentence pairs (McCoy et al., 2019). Jia and Liang (2017) 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 BERT-based (Devlin et al., 2019) 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 (Blitzer et al., 2007; Daumé III, 2007) 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 (Gururangan et al., 2018; Jia and Liang, 2017; McCoy et al., 2019), unintended correlations between input and output learned by statistical models, ignites the develop
$^{1}$ Code is available at https://github.com/hao-cheng/f-divergence.
Proceedings of the 2021 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, pages 1078-1089 June 6-11, 2021. ©2021 Association for Computational Linguistics
ment of model debiasing techniques (Clark et al., 2019; He et al., 2019). These methods leverage discovered data biases to improve model generalization over adversarial datasets (Jia and Liang, 2017; McCoy et al., 2019) 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 (Madry et al., 2018; Goodfellow et al., 2014) and NLP (Zhu et al., 2020; Jiang et al., 2020; Cheng et al., 2019; Wang et al., 2019), 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) (Miyato et al., 2018) and Jacobian Regularization (JR) (Sokolic et al., 2017). In addition, we propose to generalize VAT and random perturbation training (RPT) (Miyato et al., 2018) 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 (Devlin et al., 2019; Liu et al., 2019) and mismatched evaluations (Hendrycks et al., 2020), we systematically study the regularization framework using BERT (Devlin et al., 2019) on a diverse set of tasks in terms of both in-domain and out-of-domain generalization. Specifically, we use representative datasets (Socher et al., 2013; Williams et al., 2018; Rajpurkar et al., 2016, 2018) 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 (Tsatsaronis et al., 2015; Maas et al., 2011) and challenge adversarial datasets (Jia and Liang, 2017; McCoy et al., 2019) 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.
2 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_{\Theta}(\mathbf{x}):\mathbb{R}^n\to \mathbb{R}^m$  to denote the posterior function which is a neural network parameterized by  $\Theta$  that maps the input  $\mathbf{x}\in \mathbb{R}^n$  to the output probability space with  $m$  discrete classes.
Both adversarial learning (Goodfellow et al., 2014) and the posterior difference regularization aim at making the model more robust. Adversarial learning focuses on minimizing the following objective
$$
\min  _ {\Theta} \max  _ {\| \epsilon \| \leq c} L \left(f _ {\Theta} (\mathbf {x} + \epsilon), y\right), \tag {1}
$$
where  $L$  is the cross-entropy loss,  $\epsilon \in \mathbb{R}^n$  is a random vector bounded in a norm by  $c$ , a small positive constant and  $y$  is the target label. Instead, the posterior differential regularization directly promotes the model local smoothness, e.g. stabilizing the model posterior distribution towards small input perturbations. Typically, it is in the form of
$$
\min  _ {\Theta} L \left(f _ {\Theta} (\mathbf {x}), \mathbf {y}\right) + \alpha R \left(f _ {\Theta} (\mathbf {x}), f _ {\Theta} (\hat {\mathbf {x}})\right), \tag {2}
$$
where  $\hat{\mathbf{x}} = \mathbf{x} + \epsilon$ ,  $R$  is a regularization term penalizing the model instability, and  $\alpha$  is a hyperparameter for balancing the classification loss and the regularization term. As we can see, posterior differential regularization is a task-agnostic method which makes it applicable to semi-supervised, self-supervised and unsupervised learning. For simplicity, we will use  $f$  and  $R$  to denote  $f_{\Theta}$  and  $R(f_{\Theta}(\mathbf{x}), f_{\Theta}(\hat{\mathbf{x}}))$ , respectively.
Jacobian Regularization: A recent regularization approach to stabilize the model is Jacobian regularization (Sokolic et al., 2017; Li et al., 2016). Specifically, using the input-output Jacobian matrix,  $J = J_{f}(\mathbf{x}) \in \mathbb{R}^{m \times n}$ , we can get the first-order Taylor approximation
$$
f (\hat {\mathbf {x}}) = f (\mathbf {x} + \epsilon) = f (\mathbf {x}) + J _ {f} (\mathbf {x}) \epsilon . \tag {3}
$$
In order to reduce the overall model sensitivity to the input perturbation, Sokolic et al. (2017) propose to directly regularize the Frobenius norm of the input-output Jacobian matrix so that
$$
\begin{array}{r c l} \| f (\hat {\mathbf {x}}) - f (\mathbf {x}) \| _ {2} ^ {2} & = & \| J \epsilon \| _ {2} ^ {2} = \epsilon^ {T} J ^ {T} J \epsilon \\ & \leq & \| \epsilon \| _ {2} ^ {2} \| J \| _ {s p} ^ {2} \leq \| \epsilon \| _ {2} ^ {2} \| J \| _ {F} ^ {2}, \end{array}
$$
where  $\| \cdot \|_2$  is the  $L2$  norm,  $\| \cdot \|_{sp}$  is the spectral norm, and  $\| J \|_F^2 = \operatorname{tr}(J^T J)$  is the Frobenius norm of the Jacobian matrix with tr as the trace operator. In other words, by letting  $R = \| J \|_F^2$ , the  $L2$  difference between clean and noisy inputs is thus being effectively regularized.
Virtual Adversarial Training: Motivated by the adversarial learning objective used in (Goodfellow
et al., 2014), Miyato et al. (2018) introduce a regularized objective to enhance the model robustness towards small input perturbations
$$
\min  \left[ L (\mathbf {y}, \hat {\mathbf {y}}) + \alpha \underbrace {\max  _ {\| \epsilon \| \leq c} \mathrm {K L} (\hat {\mathbf {y}} , f (\mathbf {x} + \epsilon))} _ {R} \right], \tag {4}
$$
where KL is the well-known Kullback-Leibler divergence, and  $\hat{\mathbf{y}} = f(\mathbf{x})$ . Based on the above definition, VAT essentially regularizes the KL-based worst-case posterior difference between the clean and noisy input using an inner loop to search for the most adversarial direction.
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
$$
\| J \| = \sup  _ {\| \nu \| = 1} \| J \nu \|, \tag {5}
$$
where the matrix norm on the left side is induced by the corresponding vector norm on the right side. It is easy to show that
$$
\begin{array}{l} c ^ {2} \| J \| _ {s p} ^ {2} \approx \sup  _ {\| \epsilon \| _ {2} = c} \| J \epsilon \| _ {2} ^ {2} \\ \leq \sup  _ {\| \epsilon \| _ {2} = c} \| f (\hat {\mathbf {x}}) - f (\mathbf {x}) \| _ {1} ^ {2} \\ \leq 2 \sup  _ {\| \epsilon \| _ {2} = c} \mathrm {K L} (f (\mathbf {x}), f (\hat {\mathbf {x}})), \\ \end{array}
$$
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 (Abbas et al., 2016), 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  $f$ -divergence: 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 (Csiszár and Shields, 2004). 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.
$$
D _ {g} \left(f (\hat {\mathbf {x}}), f (\mathbf {x})\right) = \sum_ {i} f _ {i} (\mathbf {x}) g \left(\frac {f _ {i} (\hat {\mathbf {x}})}{f _ {i} (\mathbf {x})}\right), \tag {6}
$$
where the generator function  $g:\mathbb{R}_{+}\mapsto \mathbb{R}$  is a convex and lower-semicontinuous function satisfying  $g(1) = 0$ $\hat{\mathbf{x}} = \mathbf{x} + \epsilon$  and  $f_{i}$  indicates the  $i$ -th element of vector  $f$ . Different choices of  $g$  lead to several popular divergences, e.g. KL, squared Hellinger and Jensen-Shannon divergence. Based on this, it is easy to show that the corresponding second order approximation is
$$
D _ {g} (f (\hat {\mathbf {x}}), f (\mathbf {x})) \approx \frac {g ^ {\prime \prime} (1)}{2} \epsilon^ {T} J ^ {T} \operatorname {d i a g} \left(\frac {1}{f}\right) J \epsilon , \tag {7}
$$
where  $J$  is the input-output Jacobian of  $f$ , and  $\mathrm{diag}\left(\frac{1}{f}\right)$  is a diagonal matrix with elements equal to  $\frac{1}{f}$  (See Appendix A for full derivation).
Compared with the Frobenius norm of Jacobian matrix  $\| J\| _F^2$ , Equation 7 can be seen as a weighted version of JR where each row is rescaled by the model confidence  $f_{i}$  for the corresponding class. In other words, it is close to JR for more confident classes, whereas for uncertain classes it allows less Jacobian variance. Additionally, although  $g''(1)$  is a constant once the generator function is selected, various  $f$  -divergences can lead to different approximations which might result in task-dependent benefits. Therefore, different from KL-based VAT or its sampling alternative without the inner search for the most adversarial direction as proposed in (Miyato et al., 2018), we generalize the posterior differential regularization with the family of  $f$  -divergences and show that they all provide an approximation to a variant of JR which adapts the regularization based on model confidence.
Given its superior performance over a wide range of NLP tasks, we focus on exploring different train
ing techniques using BERT (Devlin et al., 2019). 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}, t = 1, \dots, T$ . The tokenized input sequence is then transformed into embeddings,  $\mathbf{x}_1, \ldots, \mathbf{x}_T \in \mathbb{R}^n$ , 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 element-wise summation. The embedding layer is used as the input to multiple transformer layers (Vaswani et al., 2017) to generate the contextual representations,  $\mathbf{h}_1, \ldots, \mathbf{h}_T \in \mathbb{R}^d$ , which are the hidden states of the last layer of the BERT model. For all regularizations, we sample noise vectors  $\epsilon_1, \ldots, \epsilon_T$  from  $\mathcal{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.  $\mathbf{x}_1 + c\epsilon_1, \ldots, \mathbf{x}_T + c\epsilon_T$ . Here, we fix  $c = 1\mathrm{e} - 3$  in this paper.
Sentence-level Classifier: Following the standard setup of BERT-based textual entailment model (Devlin et al., 2019), 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  $\mathbf{h}_{[CLS]}$  produced by BERT, corresponding to the [CLS] token. Then, the probability of class  $k$  is
$$
P (k | w _ {1}, \dots , w _ {T}) \propto W _ {k} ^ {C} \mathbf {h} _ {[ C L S ]}, \tag {8}
$$
where  $W^{C} \in \mathbb{R}^{m \times d}$  is the learnable parameter, the subscript  $k$  indicates the  $k$ -th row of the matrix, and the bias term is left out for simplicity. For standard BERT training, the log-likelihood based on Equation 8 is used. For regularized models, the regularization term is added to stabilize the class probability change with regard to the input noise.
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  $\mathbf{h}_t \in \mathbb{R}^d$  is used as the contextualized token embedding, where  $d$  is the vector dimension.
The span-begin score is computed as  $s_b(i) = \mathbf{w}_b^T\mathbf{h}_i$  using a weight vector  $\mathbf{w}_b\in \mathbb{R}^d$ . The probability for a start position  $i$  is
$$
P _ {b} (i) = \frac {\exp \left(s _ {b} (i)\right)}{Z _ {b}}, \tag {9}
$$
where  $Z_{b}$  is the normalizing factor computed by normalizing over  $\mathcal{I}$  (the set of all possible positions in the passage), i.e.  $Z_{b} = \sum_{i\in \mathcal{I}}\exp (s_{b}(i))$ . The span-end score  $s_e(j)$ , the probability  $P_{e}(j)$  for an end position  $j$ , and the normalizing factor  $Z_{e}$  are defined in the same way. The probability of an answer span  $(i,j)$  is
$$
P (i, j) = P _ {b} (i) P _ {e} (j) = \frac {\exp (s _ {b} (i) + s _ {e} (j))}{Z _ {b} Z _ {e}}.
$$
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.
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 (Miyato et al., 2018), 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  $f$ -divergences, i.e. KL, squared Hellinger (SHL) and Jensen-Shannon divergence (JSD).
Dataset: All the datasets used in this paper are summarized in Table 1. We consider three tasks, i.e. QA, textual entailment, and sentiment analysis, where the last two are sentence classification
tasks. Following the literature, we report the exact match (EM) and F1 scores for QA datasets and classification accuracy for textual entailment and sentiment analysis. For model training, we use MNLI (Williams et al., 2018) and SST-2 (Socher et al., 2013) and SQuAD v1.1/v2.0 (Rajpurkar et al., 2016, 2018), respectively. The corresponding development set is used for evaluating the indomain generalization.
To evaluate the out-of-domain generalization with domain shift, we use the BioAQS dataset (Tsatsaronis et al., 2015) from MRQA (Fisch et al., 2019) and the IMDB dataset (Maas et al., 2011). 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 (Socher et al., 2013), IMDB includes lengthy movie reviews from non-experts (Maas et al., 2011). 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 (Jia and Liang, 2017) and HANS (McCoy et al., 2019) 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 (Rajpurkar et al., 2016) 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 (Jia and Liang, 2017). The HANS dataset is built using three heuristics to ensure that the hypothesis sentence only contains words from the premise sentence (McCoy et al., 2019). 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 (Devlin et al., 2019). We select the learning rate from  $\{3\mathrm{e} - 5,4\mathrm{e} - 5\}$  for QA models and  $\{2\mathrm{e} - 5,3\mathrm{e} - 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
Task	Training	Metrics	Evaluation
Domain Shift	Adversarial Attack
Question Answering	SQuAD v1.1	F1/Exact Match (EM)	N/A	Adversarial SQuAD
Question Answering	SQuAD v2.0	F1/Exact Match (EM)	BioASQ	N/A
Textual Entailment	MNLI	Accuracy (Acc)	N/A	HANS
Sentiment Analysis	SST-2	Accuracy (Acc)	IMDB	N/A