SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
R = \frac {\sigma}{2} \left[ \Phi^ {- 1} \left(\mathbb {E} _ {\delta} [ h (x + \delta) = y ]\right) - \Phi^ {- 1} \left(\max  _ {y _ {i} \in \mathcal {Y}, y _ {i} \neq y} \mathbb {E} _ {\delta} [ h (x + \delta) = y _ {i} ]\right) \right], \tag {3}
$$
where  $\Phi^{-1}$  is the inverse of the standard Gaussian CDF, and  $\hat{R}$  is calculated from Eq. (1). Due to the space limit, the proof and certification algorithm are shown in Sec. A.7 in the Appendix.
3.4 Training Loss
To provide better supervision for training the base model, we propose a decoupled regularization (DR) learning paradigm that directly conducts modular regularization to aid the CE loss. We decouple the regularization on the whole model into one on the feature extraction module and the other on the classifier module. Thus, this pathway includes two branches. Such a regularization decomposition requires each module to be robust under the perturbation, which contributes to enhancing the robustness of the base model with more flexibility, leading to better training and certification results.
Feature Extraction Regularization. For the feature extraction module, its responsibility is to output a deep representation of a given text. The stability of representation, i.e., the degree of change in the representation space given a small change of input, depicts the robustness of the feature extraction module. To regularize it, we calculate the  $l_{2}$  norm between perturbed representation  $z^{\prime}$  and original representation  $z$
Figure 3: DR loss contains a pathway (MR term) for providing modular regularization for CE loss.
$$
\ell_ {f r} = \left| \left| z ^ {\prime} - z \right| \right|, \tag {4}
$$
where  $||\cdot ||$  is the  $l_{2}$  norm.
We choose  $l_{2}$  norm over cosine similarity because it is directly related to the prediction margin defined in Sec. 2.2 in calculation. Thus, it can limit the influence of perturbation on downgrading the certification performance. Intuitively, it functions as a regularization term that asks the representations obtained from the feature extraction module given the original sample and perturbed sample to be similar to each other, which is helpful for obtaining a similar logit for  $z'$  given the same classifier module. Hence, it further enhances the representational stability and the model robustness.
Classifier Robustness Regularization. Since the representational stability regularization in Eq. (4) only concerns the feature extraction module, another direction for strengthening the base model robustness is to design a feedback signal in the classifier level that is complementary to Eq. (4). We know that the classifier module is the one that outputs the logits. Recall that in Sec. 2.2 we discuss that the certification is highly dependent on the prediction margin output by the smoothed model, which is the difference between the logits of the ground truth class and the runner-up class. By connecting these two observations together, in the context of certified robust training, we can interpret the robustness of the classifier module as being able to generate a positive prediction margin despite the perturbation added to the original input  $z$ .
Consequently, we introduce a novel regularization term for the classifier module based on its ability to output a positive prediction margin against a small perturbation. Since this regularizer is confined to the classifier, we introduce a small amount of Gaussian noise  $\epsilon$  here to model the perturbation and sample a point  $z + \epsilon$  from the neighbor of  $z$ . Note that  $\epsilon$  only serves for modeling small
perturbation in training. It will not be added in certification and is not related to the Gaussian noise  $\delta$  used for certification in Type II scenario. That is, it starts with adding Gaussian noise  $\epsilon \in \mathbf{R}^d$  to  $z$ , where  $\epsilon \sim \mathcal{N}(0,\nu^2 I)$  is drawn from the zero-mean isotropic Gaussian distribution with  $\nu$  as the standard deviation hyperparameter. Denote the prediction margin given a vector  $k \in \mathbf{R}^d$  as  $\mathcal{M}(k) = g_y(k) - \max_{y_i \in \mathcal{Y}, y_i \neq y} g_{y_i}(k)$ , where  $g_y(\cdot)$  and  $\max_{y_i \in \mathcal{Y}, y_i \neq y} g_{y_i}(\cdot)$  refers to the logits of the ground truth class and the runner-up class, respectively. We then calculate the margin given perturbed input  $z + \epsilon$  as  $\mathcal{M}(z + \epsilon)$ . We take its negative as an optimization term for minimization:
$$
\ell_ {c r} = - \max  (\mathcal {M} (z + \epsilon), 0). \tag {5}
$$
$\ell_{cr}$  has a max operation that outputs 0 when  $\mathcal{M}(z + \epsilon) < 0$ . This operation acts as feedback for asking the optimization to concentrate on improving the representation space and the discriminative ability through the other terms when  $\mathcal{M}(z + \epsilon) < 0$ , which will not regularize  $g$  if  $z$  is not well learned.
Modular Regularization Term. After that, these two regularizers are combined as
$$
\ell_ {m r} = \max  \left(\ell_ {f r} + \alpha \cdot \ell_ {c r}, \xi\right), \tag {6}
$$
which we refer to as the modular regularization (MR) term. In Eq. (6), the MR term is calculated from the linear combination of  $\ell_{fr}$  and  $\ell_{cr}$ , where  $\alpha$  is a weight hyperparameter greater than 0. Additionally, we introduce relaxation during optimization with max operation, and  $\xi > 0$  is the relaxation hyperparameter. Putting them together provides a regularization on how small  $||z' - z||$  should be given the margin information. This geometric interpretation is detailed in Sec. A.8.
DR loss. Given the MR term, we can combine it with the CE loss as the DR loss, which is
$$
\ell_ {d r} = \ell_ {c e} + \beta \cdot \ell_ {m r}, \tag {7}
$$
where  $\beta$  is a weight hyperparameter for  $\ell_{mr}$  term. Given that, in the Type I training, we can directly use  $\ell_{dr}$  for training. For Type II training, since the certification is related to  $\hat{R}$ , we also have a regularizer for  $\hat{R}$  with  $\ell_{\hat{R}} = \max (\hat{R} - R + \mu, 0)$  as CISS does, where  $R$  is calculated from Eq. (3) and  $\mu$  is the relaxation hyperparameter. Thus, the training loss in the Type II scenario is  $\ell_{tr} = \ell_{dr} + \gamma \cdot \ell_{\hat{R}}$ , where  $\beta$  and  $\gamma$  are hyperparameters.
Remark. From Eq. (7), we observe that the DR loss is composed of the CE loss and the MR term, which are indispensable to each other. If we only use the CE loss, the individual regularization will be ignored. With the help of the MR term, the introduced regularizers in Eq. (6) directly request the representation to be stable and the classifier to be robust under the perturbation, which helps the CE loss to improve the base model robustness. Meanwhile, since the MR term does not penalize the wrongly classified cases, the CE loss is necessary for improving the discriminative ability.
As for computation cost, we acknowledge that compared to CE loss calculating the proposed DR loss needs to increase the computation complexity due to the calculation of getting the features of perturbed input and original input during training. Such a tradeoff is natural: to obtain better robustness it is generally unavoidable to increase computation cost, and the increase of computation cost also occurs for existing training techniques for improving robustness such as using PGD [13] loss and TRADES [28] loss. We want to point out that the computation cost of calculating DR loss is relatively lower which is approximately twice of the cost of caluating CE loss, which is similar to that of single-step PGD and is acceptable in implementation.
Given the proposed method above, the questions of interest include: (Q1) Does the unified framework UniT outperform existing state-of-the-art approaches in both scenarios? (Q2) How does UniT improve the robustness of the base model in certified robust training? (Q3) Is the design of the UniT reasonable? We now answer them with the following experimental results.
We describe the experimental setup first. Due to the space limit, we kindly ask the readers to refer to the Appendix for the details of hyperparameters, dataset description, and metric calculation.
Datasets. We conduct experiments on four widely used text classification datasets: (1) IMDB [12], (2) SST2 [20], (3) Yelp [19] and (4) AG [29].
Baselines. Because of different certification conditions used in different works, to gain a fair comparison, we compare the proposed method with baselines individually in Type I and Type II scenarios. To illustrate: (1) Since the code of WordDP [23] is unavailable, the Type I baseline we use is SAFER [25], which is the ground-breaking method that trains a model by inputting perturbed text samples constructed by randomly sampled synonyms and certifies the prediction results directly by the prediction margin. (2) The Type II baseline is CISS [31], a recently proposed method that trains the model and certifies the prediction with IBP in a randomized smoothing manner. Their implemented base models are BERT because it yields the best performance. Since [31] has shown that IBP methods [4, 5] cannot certify large networks like BERT, we hereby do not discuss them. We also include adversarial training losses like PGD loss [13] and TRADES loss [28] as baselines.
Metric. Naturally, certified robust accuracy (CRA) is the used comparison metric, and certified robust accuracy = natural accuracy × certification ratio. It denotes the fraction of predictions that are not only correctly classified but also certified robust. As for its calculation, we follow the same process used in the baseline methods to have a fair comparison. Further details are in Appendix.
Comparison of Certified Robustness in Type I Scenario. To answer Q1, we first conduct the comparison experiments in Type I scenario. For SAFER and UniT, both of them use BERT as the base model, and the difference is that SAFER is trained with the CE loss and UniT is with the DR loss. The comparison results are in Table 2, which show that the DR loss increases the certified robust accuracy up to  $3.68\%$ ,  $1.37\%$ ,  $0.68\%$ , and  $0.53\%$  in IMDB, SST2, Yelp, and AG datasets, respectively. These results indicate that the MR term enlarges the prediction margin from the CE loss under perturbation to yield a higher certified robust accuracy. Hence, the DR loss can improve the robustness of the base model and help UniT achieve better performance in Type I scenario.
We also include adversarial training losses as baselines for Type I Scenario. From Table 2, DR loss can consistently perform better than PGD loss and TRADES loss in four datasets. Especially, DR loss is able to obtain a relatively high increment in terms of CRA on the datasets with long text samples such as IMDB and SST2. Thus, DR loss is a more suitable choice for certified robust training for its decomposed regularization, which requires each module to be robust under perturbation.
Base Model	Loss	IMDB	SST2	Yelp	AG
BERT	CE Loss (SAFER)	85.36	91.65	97.19	93.78
PGD Loss	87.52	90.28	97.86	93.98
TRADES Loss	86.80	90.44	97.56	93.96
DR Loss (UniT)	89.04	93.02	97.87	94.31
Comparison of Certified Robustness in Type II Scenario.
We further conduct the comparison experiments in the Type II scenario to answer Q1, and the results are shown in Table 3. During implementation, we find that the training for randomized smoothing in CISS can only work in relatively large datasets (100k+ training samples). Therefore, we can only report comparison results on two large datasets satisfying this condition, i.e., Yelp and AG, with CISS as the baseline. As we have seen before, directly replacing the CE loss (first row) with the DR loss (second row) can help improve CRA.
Table 2: Comparison of certified robust accuracy (%) in the Type I scenario.
Method	Loss	Yelp	AG
CISS	CE	88.60	82.47
CISS	DR	89.22	82.93
UniT	DR	91.24	84.32
Table 3: Comparison of certified robust accuracy  $(\%)$  in the Type II scenario.
In this setting, we are more interested in the analysis of whether using the structure of UniT can make Type II training achieve stronger certification without using extra modules. By comparing the last two rows in Table 3, we find that UniT achieves an increase of  $2.02\%$  and  $1.39\%$  in Yelp and AG in terms of certified robust accuracy than the CISS structure trained with the DR loss, respectively. This improvement shows that by directly conducting randomized smoothing in the BERT embedding space, UniT not only adapts to both Type I and Type II training but also improves the certified robust accuracy by avoiding the loose bound problem of the IBP module.
Comparison of Empirical Robustness. We also conduct experiments on comparing empirical robustness. We employ BERT as the base model and use PGD loss and TRADES loss as baselines. In our experiments, we conduct adversarial attacks by utilizing TextBugger [9], a commonly used method for generating text adversarial examples. Following the existing adversarial attack setup, we randomly draw 1,000 test samples for adversarial example generation. The defense results of different methods are shown in Table 4, which demonstrate the empirical robust accuracy of DR loss is higher than those of PGD loss and TRADES loss. Compared to baselines, the improvement is obtained by the individual robustness supervision for each module provided by DR loss.
Base Model	Loss	IMDB	SST2	Yelp	AG
BERT	PGD Loss	62.0 (87.1)	87.6 (91.5)	92.4 (96.9)	85.5 (93.7)
TRADES Loss	57.8 (84.8)	86.3 (91.7)	91.9 (97.0)	84.4 (94.3)
DR Loss	72.7 (86.9)	89.8 (93.3)	96.9 (98.4)	87.6 (92.9)
Table 4: Comparison of empirical robust accuracy (%) with adversarial training losses. We also show the corresponding natural accuracy indicated by the parentheses.
4.3 Insight Analysis on UniT
Analysis on Certification Condition. For Q2, we first find that the improvement is achieved by improving the certification ratio in the Type II scenario. As shown in Theorem 1, UniT includes a stronger certification condition in the Type II scenario that does not require the calculation from IBP. We now verify it by comparing the certification ratio between CISS and UniT. The certification ratio results shown in Table 5 indicate that UniT has a larger ratio of correctly classified test samples that increases up to  $2.84\%$  and  $4.23\%$  compared to CISS, and  $2.06\%$  and  $3.85\%$  compared to CISS trained with
Method	Loss	Yelp	AG
CISS	CE	94.37	94.48
CISS	DR	95.15	94.86
UniT	DR	97.21	98.71
Table 5: Comparison of certification ratio  $(\%)$  in the Type II scenario.
the DR loss, in Yelp and AG, respectively. This is mainly because the calculation of  $\hat{R}$  in UniT does not need IBP and is able to avoid the loose bound problem caused by IBP. Consequently, directly certifying in the BERT embedding space leads to more correctly classified samples that satisfy the certification condition. Note that the certified robust accuracy equals to the natural accuracy times the certification ratio. Thus, UniT can contribute to higher certified robust accuracy in the Type II setting, as it can have a higher ratio of test samples that satisfy the certification condition.
Analysis on the Effect of the DR Loss. Another aspect of the answer to Q2 is that the DR loss helps improve the robustness of the base model. To obtain a better understanding of the improvement achieved by the DR loss, w.l.o.g., we analyze how the DR loss improves the perturbed sample representation and classification in IMDB in the Type I scenario. We calculate the average  $l_{2}$  norm between the representa
tions of the perturbed and original samples, i.e.,  $||z^{\prime} - z||$ , learned from the CE loss and the DR loss in the test set, which are the inputs of the last FC layer. We also evaluate how the classifier (FC layer) is affected by calculating the margin change  $\mathcal{M}(z') - \mathcal{M}(z)$  given perturbed samples.
Measure	||z' - z||	M(z') - M(z)
Loss	CE	DR	CE	DR
Mean	1.752	0.382	-0.060	-0.032
Variance	1.374	0.494	0.877	0.282
Table 6: Comparison of representational  $l_{2}$  norm and margin change between the CE and DR loss in IMDB.
For both of them, the closer to 0, the better. In Table 6, it is shown that the mean and variance of  $||z^{\prime} - z||$  obtained from the DR loss is significantly smaller than the ones obtained from the CE loss, which means that the DR loss effectively improves the representation compactness between perturbed samples and original ones. Besides, if we replace the CE loss with the DR loss, on average, the margin change will reduce from -0.060 to -0.032, given the perturbed samples
in testing. Therefore, the DR loss can improve the quality of representation and certified results. To further demonstrate, we take one test sample as a visualization example for demonstration. Due to the space limit, we present the visualization result in Sec. A.9 in the Appendix.
Regularization	\( \ell_{\hat{R}} \)	\( \ell_{\hat{R}_{ub}} \)	\( \ell_{\hat{R}_{lb}} \)
Natural Acc. (%)	93.86	90.37	96.72
Cert. Ratio (%)	97.21	99.15	1.08
CRA (%)	91.24	89.60	1.05
Table 7: Trade-off between natural accuracy and certification ratio of UniT in Yelp.
Analysis on Trade-Off Balance. To answer Q3, we conduct the trade-off experiment in Yelp again but on UniT. Comparing Table 7 and Table 1, we can see that UniT does not suffer as much drop in certification ratio as CISS does when it adopts  $\ell_{\hat{R}}$  compared to  $\ell_{\hat{R}_{ub}}$ , and  $\ell_{\hat{R}}$  leads to the best performance. It is because UniT avoids the restriction caused by IBP. Therefore, the design of UniT for conducting certification in the embedding space can achieve a balance for the mentioned trade-off.
Analysis on the Design of DR Loss. For Q3, w.l.o.g., we further justify the design of the DR loss with ablation studies in IMDB in the Type I scenario. First, we illustrate the necessity of designing separate regularization terms for different modules. A straightforward choice that regularizes these two modules simultaneously is  $||g(z^{\prime}) - g(z)||$ , which is referred to as the adversarial logit pairing (ALP) [6], and  $g(\cdot)\in \mathbf{R}^{c}$  is the corresponding logit vector. From Table 8, if we adopt ALP that penalizes two modules together, the CRA only improves from  $85.36\%$  to  $86.64\%$ . Thus, the decoupled regularization brings optimization flexibility.
CE	ALP	lfr	lcr	CRA (%)
✓	85.36 (86.03)
✓	✓	86.64 (86.36)
✓	✓	88.08 (87.62)
✓	✓ (z + ε)	85.92 (87.42)
✓	✓	✓ (z)	87.68 (87.45)
✓	✓	✓ (z')	85.76 (85.71)
✓	✓	✓ (z + ε)	89.04 (87.87)
Meanwhile, from Table 8, we can see that only using  $\ell_{fr}$  or  $\ell_{cr}$  causes a drop of  $0.96\%$  and  $3.12\%$ , respectively, compared to the DR loss. Thus, these two terms are indispensable for modular regularization. Additionally, we also validate our design of modeling perturbation for classifier by Gaussian noise. We incorporate the negative margin calculated from  $z^{\prime}$  with  $\ell_{fr}$ , but this regularization does not collaborate with the representational  $l_{2}$  norm and worsens the regularization effectiveness. Thus, we should use  $z + \epsilon$  to model the perturbation for the classifier module to enhance its robustness.
We further justify the design of each term in the MR term. For clarity, we here denote  $t = \mathcal{M}(z + \epsilon)$ . For the choice of the classifier regularization, it could be the ALP term, the margin loss  $\max(-t, 0)$ , or the function without max operation  $-t$ . Table 9 indicates that these choices all undermine the improvement brought by the representational  $l_{2}$  norm. This is because the calculation of  $||g(z') - g(z)||$  depends on  $||z' - z||$  too, and it attaches too much
Table 8: Ablation study results on loss design in IMDB.  ${\ell }_{fr}$  and  ${\ell }_{fr}$  are indispensable. The natural accuracy is indicated by the parentheses.
Feature Extraction Reg.	Classifier Reg.	CRA (%)
||z' - z||	||g(z') - g(z)||	85.44
||z' - z||	max(-t,0)	86.24
||z' - z||	-t	86.64
||z' - z||2	-max(t,0)	87.36
||z' - z||	-max(t,0)	89.04
Table 9: Ablation study results on the regularization (Reg.) design in IMDB. We denote  $t = \mathcal{M}(z + \epsilon)$ .
weight to  $||z' - z||$ . Additionally, compared to  $\max(-t, 0)$  or  $-t$ , the used  $-\max(t, 0)$  has a selected gradient activation area for  $t > 0$ , so it adjusts the optimization direction when  $z$  has not been trained well yet and sends a clearer optimization guidance than the others. We also note that the improvement decreases if we change the  $l_2$  norm to the squared  $l_2$  norm. Thus, the design in Figure 3 is reasonable.
The development of certified robust training in text is restricted by the lack of a unified framework and modular regularization. To alleviate these problems, we propose a unified framework named UniT to strengthen the certification of prediction results for text data. Under this framework, we introduce the DR loss combining the CE loss with the modular regularization term for different modules specifically to improve the base model robustness. Experiments show that UniT allows more correctly classified samples to satisfy the certification condition by avoiding IBP. Additionally, the DR loss can generate more similar representations for perturbed samples to the original ones and reduce the drop of prediction margin. Higher certified robust accuracy is achieved consequently.
This work is partially supported by the National Science Foundation under Grant No. 1951729, 1953813, 2119331, 2212323, and 2238275.
[31] Haiteng Zhao, Chang Ma, Xinshuai Dong, Anh Tuan Luu, Zhi-Hong Deng, and Hanwang Zhang. Certified robustness against natural language attacks by causal intervention. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, 2022.
The datasets used in experiments are mainly the ones used in the text classification task. We notice that there are other types of text datasets, such as question answering and text generation ones. In our experiments, we follow the existing works and test our methods in text classification datasets only. We expect future works to include a discussion on defending adversarial synonym substitution on text datasets of other types.
A.2 Broader Impacts
The proposed UniT has a positive societal impact because it is able to generate not only accurate but also certified robust predictions on text-related tasks. Nowadays, large language models (LLMs) such as ChatGPT have attracted great attention due to their powerful performance and intriguing interface. However, they are mostly deployed as a black-box service for users to use and lack reliability for their predictions. For UniT, it can provide a robustness guarantee for predictions for tasks such as text classification, and it can be scalable for LLMs. As a result, UniT can positively impact the employment of LLMs in real-world applications and improve the trust between users and LLM service providers.
The used datasets are all in English. Each text sample is tokenized by the "BertTokenizer" provided by the Transformers [24] library. The detailed descriptions of the datasets are as follows:
1. IMDB [12] is a sentiment analysis dataset for movie reviews with either positive or negative sentiment. Its text samples have comparatively longer lengths than the ones of the rest datasets. It has 25,000 train and test samples, respectively.
2. SST2 [20] is another binary text classification dataset for movie reviews. It has 67,349 train samples and 1,821 test samples.‡ Its license is CC0.
3. Yelp [19] is a large-scale sentiment analysis dataset collected from restaurant reviews written by Yelp users with two classes, i.e., positive and negative ones. It has 444,101 samples for training and 126,670 samples for testing. Its license is the Apache-2.0 license.
4. AG [29] is a comparably large-scale news classification dataset with 4 classes, including world, sports, science/technology and business. $^{\S}$  The number of train and test samples are 120,000 and 7,600, respectively.
A.4 Calculation of Certified Robust Accuracy