ROIC-DM: Robust Text Inference and Classification via Diffusion Model

Abstract

While language models have made many milestones in text inference and classification tasks, they remain susceptible to adversarial attacks that can lead to unforeseen outcomes. Existing works alleviate this problem by equip** language models with defense patches. However, these defense strategies often rely on impractical assumptions or entail substantial sacrifices in model performance. Consequently, enhancing the resilience of the target model using such defense mechanisms is a formidable challenge. This paper introduces an innovative model for robust text inference and classification, built upon diffusion models (ROIC-DM). Benefiting from its training involving denoising stages, ROIC-DM inherently exhibits greater robustness compared to conventional language models. Moreover, ROIC-DM can attain comparable, and in some cases, superior performance to language models, by effectively incorporating them as advisory components. Extensive experiments conducted with several strong textual adversarial attacks on three datasets demonstrate that (1) ROIC-DM outperforms traditional language models in robustness, even when the latter are fortified with advanced defense mechanisms; (2) ROIC-DM can achieve comparable and even better performance than traditional language models by using them as advisors.

1 Introduction

Text inference and classification are two fundamental and significant tasks in Natural Language Processing (NLP) (Li et al. 2022a). In recent years, large language models have made impressive advancements in these two tasks, however, these models are pointed out to be extremely vulnerable to textual adversarial attacks (Papernot et al. 2016), wherein adversaries can easily compromise the performance of these models by crafting deceptive inputs. To address this issue, many defense methods have been proposed to improve the adversarial robustness of these language models. These defense methods can generally be classified into three categories: adversarial detection (Mozes et al. 2021; Zhou et al. 2019a), adversarial training (Wang et al. 2021a; Zhu et al. 2020; Madry et al. 2018), and adversarial purification (Li, Song, and Qiu 2023a).

Although these defense approaches can slightly alleviate the threats of adversarial attacks to some extent, they either rely on strong assumptions or sacrifice too much model performance, limiting their practical usage (Li, Song, and Qiu 2023a). Specifically, adversarial detection (Mozes et al. 2021; Mosca et al. 2022) and adversarial training methods (Zhu et al. 2020; Madry et al. 2018; Wang et al. 2021b) require prior knowledge of adversarial samples. However, in real-world scenarios, these adversarial samples are unavailable before potential attacks are launched. Even when attacks have been executed, differentiating and collecting the adversarial samples from input data remains challenging since the perturbations in these samples are imperceptible. Adversarial purification methods can avoid this dilemma as they can purify adversarial text without requiring knowledge of the specific attacks. Nevertheless, since purification operations have to be applied to all input data, these methods will deteriorate the model’s performance, as the original information of clean data will be inevitably modified during the purification process. Achieving the removal of adversarial perturbations while preserving the original meanings of the data remains a challenging task for these purification methods.

Considering the limitations of existing defense methods, it is not easy to improve the adversarial robustness of existing language models by equip** them with certain defending patches. In light of this, this paper proposes to explore a new kind of text inference and classification model that is robust enough against most adversarial attacks.

Diffusion models, as a new kind of generative models, have exhibited powerful learning ability in the computer vision community (Kong et al. 2021; Ramesh et al. 2022; Dhariwal and Nichol 2021). Naturally, some researchers attempt to transplant them into natural language processing areas (Li et al. 2022c; Gong et al. 2023; Savinov et al. 2022; Austin et al. 2021). But all these works are focusing on text generation tasks since diffusion models are generally used as generative models.

In this paper, we take the first step to investigate the usage of diffusion models in text inference and classification contexts. The basic motivation for replacing the traditional language models with diffusion models is that the diffusion models contain the diffusion-denoising steps so that they themselves would be more robust as they can more accurately estimate the whole data space (Chen et al. 2023). The left part of Figure 1 explains our motivation. During the reverse process, the model $f_{\theta}$ will “match” the text $\mathbf{x}$ with many polluted $\mathbf{y}_{t}$ until it removes all the noises. As a result, the data space of $(x,y)$ pairs is large and $f_{\theta}$ would be robust to adversarial perturbations. However, applying diffusion models in text inference and classification tasks faces many challenges. Firstly, diffusion models are originally generative models, how to employ them for text classification and inference is non-trivial. Besides, the classical language models have been developed for a long duration in text classification and inference tasks with tons of researchers’ efforts, therefore, naively leveraging the new style of models in these fields may fail to achieve comparable performance. It is significant to incorporate the previous language models’ abilities in the diffusion models to gain better performance.

To address the above challenges, we propose ROIC-DM (Robost text inference and classification diffusion model), which is the first text inference and classification diffusion model. ROIC-DM first modifies the original diffusion models to make them suitable for text classification and inference tasks. Then, to further improve the effectiveness and efficiency, ROIC-DM incorporates the traditional language models as advisors to provide advice during the denoising process. It is worth noting that ROIC-DM is much different from text purification (Li, Song, and Qiu 2023a), as the latter only uses diffusion models for data preprocessing while ROIC-DM directly utilizes diffusion models to solve the tasks. Extensive experiments on three real-world datasets demonstrate that (1) ROIC-DM is more robust to adversarial attacks than existing language models even when they are equipped with advanced defense methods. (2) ROIC-DM can achieve comparable and even better performance than traditional language models by using them as advisors. The main contributions of this paper are as follows.

• To the best of our knowledge, we are the first to introduce diffusion models in text classification and inference fields.

• We propose ROIC-DM which achieves better performance and robustness than language models.

• We conduct extensive experiments on three real-world datasets including AG NEWS, SST, and MRPC. The experimental results showcase the effectiveness and robustness of our proposed methods.

Figure 1: Illustration of our proposed Robust text Inference and Classification Diffusion Model (ROIC-DM). The figure on the left illustrates the diffusion phase and the reverse phase. The figure on the right illustrates the architecture of noise estimator $f_{\theta}$. $\mathbf{x}$ is the input text; $\mathbf{y}$ is the categorical label; $t$ is the time step number; $\epsilon_{\theta}$ is the predicted noise. The $\odot$ is the element-wise product. The $\oplus$ indicates element-wise summation.

2 Related Work

2.1 Textual Adversarial Attacks

Although language models have made great achievements, they are revealed to be vulnerable to adversarial attacks, i.e., these language models can be easily manipulated by adversaries via minor revisions of normal samples. Generally, most existing textual adversarial attacks achieve their malicious goals by replacing specific words in input texts (Alzantot et al. 2018; ** et al. 2020; Ren et al. 2019). These methods usually assume that the model is black-box but the logits of the output prediction are available. Then, they attempt to design some strategies to find appropriate words to replace. For example,  (** et al. 2020; Mrkšić et al. 2016; Ren et al. 2019) replace words with synonyms, while (** et al. 2020; Li et al. 2020) uses greedy-search methods to get substitutes.

2.2 Textual Adversarial Defenses

To improve the robustness of these language models, many defenses are proposed. These defense approaches can be generally classified into adversarial training, adversarial detection, and adversarial purification. The line of adversarial training (Wang et al. 2021a, 2020; Zhu et al. 2020; Zhou et al. 2019b; Madry et al. 2018) is to incorporate perturbations during a model’s training process so that the model can be robust to the potential risks. The works of adversarial detection (Mozes et al. 2021; Zhou et al. 2019a) aim to filter out the adversarial samples. The purification methods (Samangouei, Kabkab, and Chellappa 2018; Li, Song, and Qiu 2023b) employ generative models to purify adversarial inputs before feeding them to a model. However, all these defense methods have certain limitations. To be specific, the adversarial training methods and detection approaches require prior knowledge of attacks which is infeasible in practice, while the adversarial purification methods cannot avoid the modification of normal inputs. As a result, using defense methods to improve language models’ robustness is challenging.

3 Preliminaries: Diffusion Models

In this section, we introduce the theory of diffusion model (Sohl-Dickstein et al. 2015; Ho, Jain, and Abbeel 2020a; Song, Meng, and Ermon 2020). Generally, a diffusion model processes data in two steps. First, it gradually transforms raw data into a Gaussian distribution through a diffusion process. Then, It learns the reverse procedure to reconstruct the original data from Gaussian white noise. (Sohl-Dickstein et al. 2015). The following paragraphs formally describe these two processes.

In diffusion process, the diffusion model incrementally corrupts the original representation $\mathbf{x}_{0}$ into a Gaussian noise $\mathbf{x}_{t}$ via a Markov Chain with fixed parameters (Ho, Jain, and Abbeel 2020a) in T steps:

$$q(\mathbf{x}_{t}|\mathbf{x}_{t-1})=\mathcal{N}\left(\mathbf{x}_{t};\sqrt{1-\beta_{t}}\mathbf{x}_{t-1},\beta_{t}\mathbf{I}\right)$$ (1)

where $\mathcal{N}\left(x;\mu,\sigma^{2}\right)$ represents $x$ sampled from a Gaussian distribution with a mean $\mu$ and variance $\sigma^{2}$. The value of $\beta_{t}$ is determined by a pre-defined noise schedule $\beta$ that regulates the amount of noise injected at each step. The common noise schedules encompass square-root (Li et al. 2022b), cosine (Ho, Jain, and Abbeel 2020a), and linear (Nichol and Dhariwal 2021) functions. According to (Ho, Jain, and Abbeel 2020b), $\mathbf{x}_{t}$ can be directly computed conditioned on $\mathbf{x}_{0}$ with the following transformation:

$$q(\mathbf{x}_{t}|\mathbf{x}_{0}) =\mathcal{N}\left(\mathbf{x}_{t};\sqrt{\overline{\alpha}_{t}}\mathbf{x}_{0},(1-\overline{\alpha}_{t})\mathbf{I}\right)$$ (2)

$$\overline{\alpha}_{t} =\prod_{i=1}^{t}\alpha_{i},~{}~{}\alpha_{i}=1-\beta_{i}.$$ (3)

With the trick of re-parameter, we can get $\mathbf{x}_{t}$ by adding Guassian noise $\epsilon\sim\mathcal{N}(0,\mathbf{I})$ to $\mathbf{x}_{0}$ as follows:

$$\mathbf{x}_{t}=\sqrt{\bar{\alpha}_{t}}\mathbf{x}_{0}+\sqrt{1-\bar{\alpha}_{t}}\epsilon$$ (4)

The reverse process is a Markov Chain with a learnable $\theta$ to denoise $\mathbf{x}_{T}$ to $\mathbf{x}_{0}$. Specifically, when provided with the current representation $\mathbf{x}_{s}$, the subsequent representation $\mathbf{x}_{s-1}$ after denoising is calculated as follows:

$$p(\mathbf{x}_{s-1}|\mathbf{x}_{s},\mathbf{x}_{0}) =\mathcal{N}\left(\mathbf{x}_{s-1};\tilde{\mathbf{\mu}}_{s}(\mathbf{x}_{s},\mathbf{x}_{0}),\tilde{\beta}_{s}\mathbf{I}\right)$$ (5)

$$\tilde{\mathbf{\mu}}_{s}(\mathbf{x}_{s},\mathbf{x}_{0}) =\frac{\sqrt{\overline{\alpha}_{s-1}}\beta_{s}}{1-\overline{\alpha}_{s}}\mathbf{x}_{0}+\frac{\sqrt{\alpha_{s}}(1-\overline{\alpha}_{s-1})}{1-\overline{\alpha}_{s}}\mathbf{x}_{s}$$ (6)

$$\tilde{\beta_{s}} =\frac{1-\overline{\alpha}_{s-1}}{1-\overline{\alpha}_{s}}\beta_{s}$$ (7)

Since $\mathbf{x}_{0}$ is unknown during the reverse phase, the diffusion model utilizes a noise estimator $f_{\theta}(\mathbf{x}_{t},t)$, which is typically modelled by a deep neural network, such as Transformer (Vaswani et al. 2017) or U-Net (Ronneberger, Fischer, and Brox 2015). We optimize the reverse phase, by optimizing the variational lower bound(VLB)

$$\mathcal{L}_{vlb} =\underbrace{\mathbb{E}_{q}\left[D_{KL}\left(q(\mathbf{x}_{t}|\mathbf{x}_{0})||p(\mathbf{x}_{t})\right)\right]}_{L_{t}}$$

$$+\underbrace{\mathbb{E}_{q}\left[\sum_{s=2}^{t}D_{KL}(q(\mathbf{x}_{s-1}|\mathbf{x}_{s},\mathbf{x}_{0})||p_{\theta}(\mathbf{x}_{s-1}|\mathbf{x}_{s}))\right]}_{L_{t-1}}$$

$$-\underbrace{\textrm{log}p_{\theta}(\mathbf{x}_{0}|\mathbf{x}_{1})}_{L_{0}}$$ (8)

Finally, after simplifying (Kingma et al. 2021), the diffusion loss is as follows:

$$\mathcal{L}_{simple}=\mathbb{E}_{t,\mathbf{x}_{0},\mathbf{\epsilon}}\left[||\mathbf{\epsilon}-f_{\theta}(\mathbf{x}_{t},t)||^{2}\right]$$ (9)

More details of the general diffusion model can be found in previous works (Ho, Jain, and Abbeel 2020a; Song, Meng, and Ermon 2020).

4 Methodology

In this section, we first introduce how to apply diffusion models in text classification and inference, i.e., ROIC-DM. Then, we describe how to incorporate the knowledge from pre-trained models to ROIC-DM to improve the model performance. Finally, we present the detailed model architecture of ROIC-DM. The overview of our ROIC-DM is illustrated in Figure 1.

4.1 Diffusion Model for Text Classification and Inference

The skeleton of ROIC-DM is the diffusion model. As diffusion models exhibit strong learning ability in Computer Vision, many works attempt to employ them to solve NLP tasks (Li et al. 2022b; Gong et al. 2022). However, all of these works mainly utilize diffusion models to solve natural language generation problems since diffusion models are a kind of generative model. This paper takes the first step to modifying diffusion models as robust classifiers to solve text classification and inference problems. The technical details are as follows.

In the diffusion process, ROIC-DM gradually adds noise to the label $\mathbf{y}_{0}$. Note that since the label $\mathbf{y}_{0}$ usually can be denoted as a one-hot vector, ROIC-DM can directly add noise to it (Hoogeboom et al. 2021; Han, Zheng, and Zhou 2022). Therefore, the only difference in the diffusion process between ROIC-DM and other standard diffusion models is that our input is a label while others are images or sentences. As a result, the diffused $\mathbf{y}$ in each timestep can be obtained as follows:

$$\mathbf{y}_{t}=\sqrt{\bar{\alpha}_{t}}\mathbf{y}_{0}+\sqrt{1-\bar{\alpha}_{t}}\epsilon\quad\epsilon\sim\mathcal{N}(\mathbf{0},\mathbf{I})$$ (10)

Algorithm 1 Training procedure of ROIC-DM

0:  text $x$, learning epochs $E$, maximum diffusion steps $T$, linear schedule $\beta$, model $f_{\theta}(\cdot)$, advisor prediction $\mathbf{y}^{\prime}$ (optional), $\dots$

0:  well trained $f_{\theta}(\cdot)$;

1:  while $j<E$ do

2:     $t\sim\mathrm{uniform}(\{0,\ldots,T\})$;

3:     $\mathbf{y}_{t},\epsilon\leftarrow$ E.q. 10

4:     if use advisor then

5:        $\epsilon_{\theta}=f_{\theta}(\mathbf{x},\mathbf{y}_{t},t,\mathbf{y}^{\prime})$

6:     else

7:        $\epsilon_{\theta}=f_{\theta}(\mathbf{x},\mathbf{y}_{t},t)$

8:     end if

9:     update $f_{\theta}(\cdot)$ using $\nabla_{\theta}\left\|\epsilon-\epsilon_{\theta}\right\|^{2}$ 

10:     $j=j+1$;

11:  end while

The critical difference between ROIC-DM and the generative diffusion models is the reverse process. Specifically, for traditional diffusion models, during their reverse process, they utilize a noise estimator $f_{\theta}(\mathbf{x}_{t},t)$ to predict the noise (note that the $\mathbf{x}_{t}$ in this case is equal to $\mathbf{y}_{t}$ in ROIC-DM). However, in ROIC-DM, directly denoising $\mathbf{y}_{t}$ without any conditions is meaningless since it is just a certain category label (Hoogeboom et al. 2021; Han, Zheng, and Zhou 2022). Thus, ROIC-DM constructs a trainable model $f_{\theta}(\mathbf{x},\mathbf{y}_{t},t)$ whose goal is to generate a noise $\epsilon_{\theta}$ to recover $\mathbf{y}_{t}$ to $\mathbf{y}_{t-1}$ considering the corresponding text context $\mathbf{x}$.

Algorithm 1 shows how to train the model $f_{\theta}(\mathbf{x},\mathbf{y}_{t},t)$. To be specific, ROIC-DM randomly selects a pair of data $(\mathbf{x},\mathbf{y})$ and calculates $\mathbf{y}_{t}$ using E.q. 10. Then, it predicts the noise using $f_{\theta}(\mathbf{x},\mathbf{y}_{t},t)$ and calculates the loss with E.q. 11. Above steps iteratively continue until model convergence.

$$\mathcal{L}=\mathbb{E}_{t,\mathbf{x}_{0},\mathbf{\epsilon}}\left[||\mathbf{\epsilon}-f_{\theta}(\mathbf{x},\mathbf{y}_{t},t)||^{2}\right]$$ (11)

The inference procedure of ROIC-DM is displayed in Algorithm 2. In the inference phase, ROIC-DM randomly samples a Gaussian noise $\mathbf{y}_{T}$ as the start point. Then, ROIC-DM recovers $\mathbf{y}_{T}$ to $\mathbf{y}_{0}$ by repeating the following equations:

$$\mathbf{y}_{t-1} =\tilde{\mu}(\hat{\mathbf{y}_{0}},\mathbf{y}_{t})+\tilde{\beta}_{t}*\zeta$$ (12)

$$\tilde{\mu}_{t}\left(\mathbf{y}_{t},\mathbf{\hat{y}}_{0}\right) =\frac{\sqrt{\overline{\alpha}_{t-1}}\beta_{t}}{1-\overline{\alpha}_{t}}\mathbf{\hat{y}_{0}}+\frac{\sqrt{\alpha_{t}}(1-\overline{\alpha}_{t-1})}{1-\overline{\alpha}_{t}}\mathbf{y}_{t}$$

$$\tilde{\beta_{t}} =\frac{1-\overline{\alpha}_{t-1}}{1-\overline{\alpha}_{t}}\beta_{t}$$

$$\hat{\mathbf{y}}_{0} =\frac{1}{\sqrt{\bar{\alpha}_{t}}}\left(\mathbf{y}_{t}-\sqrt{1-\bar{\alpha}_{t}}\mathbf{\epsilon}_{\theta}\right)$$

where $\zeta\sim\mathcal{N}(\mathbf{0},\mathbf{I})$, $\epsilon_{\theta}$ is calculated by well-trained $f_{\theta}(\cdot)$. After obtaining $\mathbf{y}_{0}$, ROIC-DM selects the position with the maximum value as the final predicted class for $\mathbf{x}$.

Algorithm 2 Inference procedure of ROIC-DM

0:  text $x$, total reverse steps $T$, linear schedule $\beta$, model $f_{\theta}(\cdot)$, advisor prediction $\mathbf{y}^{\prime}$ $\dots$

0:  $label$

1:  sample $\mathbf{y}_{T}\sim\mathcal{N}(\mathbf{0},\mathbf{I})$

2:  $t=T$

3:  while $t>0$ do

4:     sample $\zeta\sim\mathcal{N}(\mathbf{0},\mathbf{I})$

5:     if use advisor then

6:        $\epsilon_{\theta}=f_{\theta}(\mathbf{x},\mathbf{y}_{t},t,\mathbf{y}^{\prime})$

7:     else

8:        $\epsilon_{\theta}=f_{\theta}(\mathbf{x},\mathbf{y}_{t},t)$

9:     end if

10:     $\mathbf{y}_{t-1}\leftarrow$ calculate E.q. 12 with $\zeta$

11:     $t=t-1$

12:  end while

13:  $label\leftarrow$find the position with max value in $\mathbf{y}_{0}$

4.2 Pre-trained Advisor Improved ROIC-DM

Although language models have been revealed to be vulnerable to adversarial attacks, they have developed for a long time and have achieved remarkable results in text classification and inference tasks, especially the large pre-trained language models (Devlin et al. 2018; Minaee et al. 2021). Therefore, it is necessary to build our ROIC-DM based on the power of pre-trained language models.

To achieve that, before training ROIC-DM, we first fine-tune a large pre-trained language model (e.g., BERT (Devlin et al. 2018)) on the target classification dataset. Then, during the reverse process, ROIC-DM generates $\epsilon_{\theta}$ not only considering $\mathbf{x}$, but also the prediction from the fine-tuned pre-trained model. Specifically, we transfer the knowledge of fine-tuned model to ROIC-DM by utilizing its soft-label $\mathbf{y}^{\prime}$ based on the input $\mathbf{x}$, which is inspired by the observations from knowledge distillation (Gou et al. 2021) said that model’s soft label contains many useful auxiliary information. In ROIC-DM, we incorporate the fine-tuned model’s knowledge by directly adding ${\mathbf{y}^{\prime}}$ to $\mathbf{y}_{t}$. The left part of Figure 1 shows the framework of our ROIC-DM.

4.3 Model Architecture of ${f_{\theta}(\cdot)}$

The right part of Figure 1 presents the model architecture of the noise estimator $f_{\theta}(\cdot)$ in ROIC-DM. Generally, $f_{\theta}$ contains an encoder to extract features from the context $\mathbf{x}$, a time embedding table to encode the timestep, a normalization layer to make the output be more smooth, and a down projector to output the noise.

Encoder

We use the encoder block of the BERT (Devlin et al. 2018) as the feature extractor to convert the words to hidden states. Then, we leverage the hidden state of $[CLS]$ token, which is appended at the start of the text, as the feature vector of text $\mathbf{x}$, i.e., $\mathbf{h}_{1}$ in Figure 1.

$$\mathbf{h}_{1}\leftarrow Encoder(\mathbf{x})$$ (13)

Time Embedding and Linear Layer

The diffusion step $t$ is uniformly sampled randomly from $\{0,\ldots,T\}$. Then, we leverage a time embedding table to learn the features of time steps $\mathbf{t}$ (Xiao, Kreis, and Vahdat 2021). For $\mathbf{y}_{t}$, we utilize a full-connected layer to transform it to the same size as $\mathbf{t}$ and then conduct element-wise product to fuse the information of $\mathbf{t}$ and $\mathbf{y}_{t}$.

$$\mathbf{e}_{t}=\mathbf{t}\odot Linear(\mathbf{y}_{t})$$ (14)

Smoother

Smoother is consisted of a softmax function and a layer normalization. Such a combination can be used in certain layers of neural networks to enhance the network’s representational capacity and convergence performance and lead to improved performance and faster training speed in learning tasks (Huang et al. 2023). Here, we use this module to process the fused vector $\mathbf{e}_{t}$.

$$\mathbf{d}_{t}=LN(Softplus(\mathbf{e}_{t}))$$ (15)

where $LN(\cdot)$ is layer normalization.

Down Projector

In the down projector, we first conduct element-wise product to fuse $\mathbf{d}_{t}$ and text feature vector $\mathbf{h}_{1}$. Then, a stack of linear layers followed by softmax and layer normalization are leveraged to predict the noise $\epsilon_{\theta}$.

$$\epsilon_{\theta}=projector(\mathbf{d}_{t},\mathbf{h}_{1})$$ (16)

Table 1: Detailed statistics of the datasets.

Dataset	Training Set	Test Set	#Avg. words
AG NEWS	120K	7.6K	43
SST-2	67K	1.8K	19
MRPC	3.7K	1.7K	44

5 Experiments

Table 2: Experimental results of adversarial robustness evaluation. The best performance is marked in bold. For the MRPC task, we attack both hypothesis and premise. Methods labelled by ${\dagger}$ are fine-tuning baselines without considering adversarial defence.

Dataset	Method	Clean$\%$	TextFooler		BERT-Attack
			Aua$\%$	Suc$\%$	Aua$\%$	Suc$\%$
AG NEWS	BERT${}^{{\dagger}}$	$94.0$	$20.5$	$78.9$	$14.6$	$84.3$
	PGD (Li and Qiu 2021)	$94.8$	$36.2$	$61.3$	$32.8$	$65.7$
	Free LB (Zhu et al. 2020)	$94.7$	$34.8$	$63.5$	$12.7$	$86.7$
	InfoBERT (Ishida et al. 2020)	$94.9$	$30.4$	$65.9$	$20.4$	$78.0$
	Text Purification (Li, Song, and Qiu 2023b)	$93.0$	$51.0$	42.0	$44.5$	48.5
	ROIC-DM	$\mathbf{95.1}$	$\mathbf{78.7}$	$\mathbf{18.4}$	49.0	47.0
SST-2	BERT${}^{{\dagger}}$	$92.2$	$13.9$	$80.4$	$13.3$	$80.6$
	PGD	$93.2$	$13.4$	$82.1$	$13.4$	$84.5$
	Free LB	$92.2$	$19.4$	$80.3$	$12.1$	$87.1$
	InfoBERT	$92.9$	$20.4$	$76.7$	$16.6$	$82.7$
	Text Purification	$91.8$	$42.6$	$53.4$	$33.5$	$62.4$
	ROIC-DM	$\mathbf{94.1}$	$\mathbf{55.4}$	$\mathbf{39.5}$	$\mathbf{46.4}$	$\mathbf{49.7}$
MRPC	BERT${}^{{\dagger}}$	$83.8$	$6.4$	$92.8$	$9.4$	$89.5$
	PGD	$84.3$	$6.9$	$92.2$	$11.5$	$82.3$
	Free LB	$83.8$	$8.2$	$91.0$	$10.3$	$87.7$
	InfoBERT	$87.9$	$13.9$	$84.1$	$17.4$	$78.9$
	Text Purification	$82.5$	$32.5$	$54.5$	$28.6$	$61.2$
	ROIC-DM	$\mathbf{90.8}$	$\mathbf{53.3}$	$\mathbf{39.8}$	$\mathbf{38.3}$	$\mathbf{58.5}$

Table 3: The comparison of the accuracy of ROIC-DM with different pre-trained advisors on the AG NEWS dataset.

Method	Accuracy %	Method	Accuracy %	Method	Accuracy %
BERT	94.0	DistilBERT	93.4	ALBERT	94.5
ROIC-DM(-advisor)	91.8	ROIC-DM(-advisor)	91.8	ROIC-DM(-advisor)	91.8
ROIC-DM	95.1	ROIC-DM	93.8	ROIC-DM	95.5

In this section, we first introduce the basic experimental settings, and then, present the experimental results with comprehensive analysis to showcase the superiority of our proposed methods.

5.1 Datasets

In this paper, we conduct experiments on three widely used text classification and inference datasets: AG NEWS (Zhang, Zhao, and LeCun 2015), SST-2 (Socher et al. 2013), and MRPC (Dolan and Brockett 2005). The statistics of these involved datasets are illustrated in Table 1, including the size of the training/test set and the average word count of the training samples. The training and test set division is following (Liu et al. 2022). We use the whole test set to evaluate model accuracy and randomly select $500$ samples for robustness evaluation since the attack process is seriously slow (Morris et al. 2020; Moon et al. 2023).

5.2 Baselines

We compare our ROIC-DM with the following defense baselines including two adversarial training algorithms (PGD and FreeLB), one regularization method (InfoBERT), and a text purification method.

PGD

(Madry et al. 2018) formulates adversarial training as a minimax problem, which aims to minimize the empirical loss on adversarial examples that could potentially lead to adversarial risk.

FreeLB

(Zhu et al. 2020) attempts to improve language models’ robustness by enhancing their generalization abilities. Specifically, FreeLB generates some virtual adversarial samples by injecting adversarial perturbations into word embeddings. Then, FreeLB mixes them with normal training data to improve the tolerance of target models to adversarial samples.

InfoBERT

(Wang et al. 2021a) consists of two mutual-information-based regularizers to improve the robustness of the learned representations by suppressing noisy mutual information.

Text Purification

(Li, Song, and Qiu 2023b) is a textual adversarial purification algorithm. It utilizes the mask-infill ability of pre-trained models to recover noisy texts and use these purified texts to make predictions.

5.3 Attack Methods and Evaluation Metrics

Textfooler (** et al. 2020) and BERT-Attack (Li et al. 2020) have exhibited capable of effectively deceiving robust models in text classification and inference tasks with limited perturbations. These two attack methods have been widely used in adversarial robustness research (Liu et al. 2022; Morris et al. 2020). Therefore, in this paper, we leverage these two adversarial attack methods to assess the robustness of our ROIC-DM.

TextFooler identifies crucial words within the input text for the target model and iteratively replaces them with synonyms until the model’s prediction is modified. BERTAttack utilizes BERT in a manner that preserves semantics when generating substitute words for the identified vulnerable words in the input text.

We utilize the following metrics to evaluate models’ resistance to above mentioned adversarial attacks:

Clean%

is the victim model’s accuracy tested on the clean test set, which represents the original performance of the victim model.

Aua%

is short for accuracy under attacks(Aua). It measures the prediction accuracy of victim models on the adversarial data generated by certain attacks. This metric reflects the defensive capability of the model against adversarial attacks. A higher $Aua\%$ indicates a more robust model.

Suc%

is the attack success rate. i.e., it is the ratio of the number of texts that have been successfully perturbed by certain attack methods to the number of involved texts. A lower $Suc\%$ means a more robust model.

5.4 Implementation Details

For ROIC-DM, we set the number of diffusion timesteps to $T=1000$, and employ a linear noise schedule with $\beta_{1}=1e-4$ and $\beta_{T}=0.02$. The optimizer is AdamW (Loshchilov and Hutter 2017) with a linear decay learning rate starting at 1e-4. The batch size is 64. For the pre-trained advisor, we directly use the pre-trained version in the toolkit (Morris et al. 2020) provided by Huggingface.

For the baselines, we directly use the hyper-parameter settings in their released codes (PGD, FreeLB, InforBERT) to generate experimental results. For Text purification (Li, Song, and Qiu 2023b), we re-implement it according to their paper description since there is no public code available.

For the adversarial attack methods, we directly use the corresponding attack’s official codes in Openattack (Zeng et al. 2021) framework.

All the code related to the experiments will be available at https://github.com/ after the double-blind review.

5.5 The Robustness of ROIC-DM

In this part, we showcase that our ROIC-DM outperforms traditional language models in robustness, even when the latter are equipped with advanced defense methods.

Table 2 presents the experimental results of ROIC-DM compared with BERT equipped with baseline defense methods. Note that our ROIC-DM also utilizes BERT as the advisor for the fair comparison. According to the results on clean data ($Clean\%$), we can observe that ROIC-DM outperforms its advisor, BERT, on all three datasets for at most $2.9$ accuracy scores ($Clean\%$). Besides, ROIC-DM is also better than BERT with defenses on the clean data.

Table 4: Comparison of the performance between ROIC-DM(-advisor) and ROIC-DM on the AG News test dataset under BERTAttack and Textfooler adversarial attacks. BERT is used as the advisor.

Method	Clean %	TextFooler		BERT-Attack
		Aua %	Suc %	Aua %	Suc %
ROIC-DM(-advisor)	91.8	60.3	33.4	46.4	47.8
ROIC-DM	95.1	78.7	18.4	49.0	47.0

Table 5: A case study on the AG NEWS dataset

	Text	BERT	Text Purification	Ours
Original	though the violence is far less sadistic than usual, the film is typical miike: fast, furious and full of off-the-cuff imaginative flourishes	$\checkmark$	$\checkmark$	✓
Attacked by BERT-Attack	since the brutal is very smaller sad graphic than usual, the movie is traditional miike way rapid becomes frantic and - of off the cuff creative flourishs.	$\times$	$\times$	$\checkmark$

When attacked by TextFooler, the vanilla BERT’s performance is dramatically dropped on all three datasets. Text Purification achieves the best performance among baselines, but our ROIC-DM still outperforms it by a large margin (e.g., 27.7 $Aua\%$ scores on AG NEWS). Moreover, it is worth mentioning that, Text Purification is harmful to the model’s performance on clean data as shown in $Clean\%$.

BERT-Attack is stronger than TextFooler and all the models suffer performance deterioration but our ROIC-DM still consistently outperforms all these baseline methods.

5.6 ROIC-DM v.s. Traditional Classifiers

Except for the robustness, another advantage of ROIC-DM is that it can incorporate knowledge from traditional classifiers and achieve better performance.

Table 3 presents the results of ROIC-DM using different pre-trained advisors. Due to the space limitation, we only include the results on AG NEWS. A similar conclusion can be obtained from the other two datasets. As shown in the results, when removing the advisor, ROIC-DM’s performance is slightly worse than traditional classifiers. This may be because these classifiers are based on pre-trained models that have been trained on large-scale datasets with long-term developed model architecture, while our ROIC-DM is a new kind of classification model trained from scratch.

When ROIC-DM uses counterpart advisors, it achieves better performance than its advisors. Specifically, ROIC-DM obtains $1.1$, $0.4$, and $1.0$ higher scores compared to BERT, DistillBERT, and ALBERT, respectively.

Figure 2: The training loss trend for ROIC-DM and ROIC-DM(-advisor) on the AG NEWS dataset.

5.7 The Impact of Advisor

In this section, we investigate the impacts of advisors for ROIC-DM from the robustness perspective and the training aspect. We compare ROIC-DM with BERT as advisor and ROIC-DM(-advisor) on the AG News dataset under BertAttack and Textfooler adversarial attacks. As displayed in Table 4, without an advisor, ROIC-DM’s performance declined on both clean and adversarial data. However, if combined the results in Table 4 and Table 2, we can observe that without an advisor, ROIC-DM still exhibits strong robustness, as its performance ($60.3$ and $46.4$ with $Aua\%$ for TextFooler and BERT-Attack respectively) is still much better than BERT with the best baseline defender ($51.0$ and $44.5$ for TextFooler and BERT-Attack respectively). This observation supports our argument that the diffusion model will be more robust than conventional language models.

From the training aspects, Figure 2 shows the training loss curve for ROIC-DM and ROIC-DM(-advisor). The training loss curve for ROIC-DM is smoother than ROIC-DM(-advisor) and the prior’s decreasing speed is also faster than the latter. This phenomenon implies that the advisor can be positive to ROIC-DM’s training process.

5.8 Case Study

Table 5 presents the data sample from the AG NEWS dataset. We utilize BERT-Attack to perturb the original data since it is one of the strongest textual adversarial attacks. We show the comparison with Text Purification as it has the best performance among baselines. For the original text, the vanilla BERT, Text Purification based BERT, and our ROIC-DM can make the correct prediction. After being perturbed, the overall meaning of the sentence is the same as the original text but BERT and Text Purification based BERT cannot make correct predictions. Only our ROIC-DM keeps correct, which indicates its robustness.

6 Conclusion

In this paper, we introduce an innovative model for robust text inference and classification named ROIC-DM, which is built upon the foundational framework of the diffusion model. To enhance the performance of ROIC-DM, we strategically integrate conventional large language models as advisors within the reverse process. The experimental results on three datasets with two competitive textual adversarial attacks indicate that ROIC-DM is more robust to adversarial attacks and can achieve better performance compared with conventional language models.