How Does Diffusion Influence Pretrained Language Models on Out-of-Distribution Data?

Following [34], we use four classification-task datasets as ID data, including SST2 [28], IMDB [22], 20NG [17] and TREC-10 [20]. Any pair of the above datasets can be regarded as OOD to each other, except for IMDB and SST2 which belong to the same task category. Besides, we also select four additional datasets as the OOD datasets, including RTE [4], MNLI [31], WMT16 [3] and Multi30K [6]. We take the test splits in those OOD datasets for testing and report the average sequence lengths of the test splits. The statistics of the datasets are shown in Table 1.

A diffusion model [13] is a latent variable model that builds a transformation from Gaussian distribution to data distribution through a multi-step denoising process. Given a data distribution $\boldsymbol{x}_{0}\sim q(\boldsymbol{x}_{0})$, at timestep $t$, $\boldsymbol{x}_{t}$ is produced by a diffusion process which satisfies a Markov chain according to a variance schedule $\beta_{1},\beta_{2},...,\beta_{T}$ $\in(0,1)$,

where $0\leq t\leq T$, $\alpha_{t}:=1-\beta_{t}$, $\overline{\alpha}_{t}:=\prod_{s=1}^{t}\alpha_{s}$, $T$ is the max iteration step. A closed form of $\boldsymbol{x}_{t}$ can be calculated for any arbitrary $t\geq 1$

where $\epsilon\in N(0,1)$. When the data is fully noised, $\boldsymbol{x}_{T}$ is close to an isotropic Gaussian, $\boldsymbol{x}_{T}\sim N(0,1)$. Note that the predefined forward process $q$ contains no trainable parameters.

In the denoising process, given $\boldsymbol{x}_{t}$ and $t$, the model is trained to reverse the diffusion process iteratively and to reconstruct data. Each denoising transition $\boldsymbol{x}_{t}\rightarrow\boldsymbol{x}_{t-1}$ is parametrized by

where $\boldsymbol{\mu}_{\theta}$ and $\boldsymbol{\Sigma}_{\theta}$ are predicted by a transformer-based PLMs. Following [19], instead of predicting the mean posterior $\mu_{\theta}(\boldsymbol{x}_{t},t)$, we train a neural network to predict $\boldsymbol{x}_{0}$ in every term and simplify the loss to a sum of mean-squared errors between the ground truth data $\boldsymbol{x}_{0}$ and its estimates $\hat{\boldsymbol{x}_{0}}$

To apply diffusion on transformer-based PLMs, we define the diffusion process on the continuous embedding space. Considering a sequence $z=\{w_{1},w_{2},...,w_{n}\}$, where each token $w_{i}\in\mathbb{R}^{v}$ has an associated embedding $e_{i}\in\mathbb{R}^{d}$, the discrete-to-continuous step is then defined as $q_{\phi}(\boldsymbol{x}_{0}|\boldsymbol{w})=N(\boldsymbol{E}\boldsymbol{w},\sigma_{0}^{2}\boldsymbol{I})$, where $d$ is the dimension of input space, $\boldsymbol{E}$ is a matrix of all embeddings and $\sigma_{0}$ is a constant scale factor with a similar order of magnitude as $\beta_{1}$. To map the predicted vector back to tokens, we define a reverse continuous-to-discrete step $p_{\theta}(\boldsymbol{w}|\boldsymbol{x}_{0})=\prod_{i=1}^{n}p_{\theta}(\boldsymbol{w}_{i}|\boldsymbol{x}_{i})$, where $p_{\theta}(\boldsymbol{w}_{i}|\boldsymbol{x}_{i})$ is the softmax probability of token $i$ with logit $\boldsymbol{x}_{i}$. A simple cross-entropy loss is then added to maximise $p_{\theta}(\boldsymbol{w}|\boldsymbol{x}_{0})$

The reconstruction loss function is then defined by

We train the diffusion model on the training splits of the four aforementioned ID datasets respectively and use validation splits for evaluation to avoid overfitting. When it comes to inference, the other test splits of the remaining datasets are treated as OOD data. We implement the model upon pretrained RoBERTa-Large [32]. The learning rate is set to be $5e-5$ with a linear decay rate towards 0. We set $\beta$ to constants as [19] and [13] did. It increases linearly from $1e-4$ to $0.02$, standing for the amount of noise added at diffusion step $t$. And the batch size is 16. Each model is trained for $s=80k$ steps, optimized with Adam [15]. $T$ is set to be 1000. Reconstruction loss is calculated as the average loss of each word in the whole sentence.

OOD data is more sensitive to higher noise level. In the process of diffusion, the added noise level, i.e. the diffusion step $t$, is pivotal in restoring noised inputs. However, under which $t$ the model starts to produce dissimilar outputs of ID and OOD data has not yet been studied. Hence, in this experiment, we explore how reconstruction loss varies with the increase of $t$.

As shown in Fig. 1, the reconstruction loss of both ID and OOD data keep rising as $t$ increases. More specifically, the loss of OOD data rises faster than ID data, indicating that OOD data is more sensitive to higher noise level. The loss gap between ID/OOD is also increasing with $t$ and achieves the maximum when $t=700\sim 800$, where model outputs start to look dissimilar from the inputs. For instance, only the high-frequency word “what” can be correctly restored in ID dataset when $t=800$ (Table 2). As for OOD data, even the common words like “it”, “just” and “as” are wrongly restored. The reason is that highly-noised inputs reserve very little information from the original sentences and the model begins to output unconditioned samples more than doing reconstructions [8], which results in high reconstruction loss. It shows that diffusion models with PLMs have worse OOD robustness when $t$ is large.

Models with more training steps are more resilience to higher noise level. The number of training steps $s$ is another important factor influencing the sensitivity of the model to noise. We explore the model performance on ID/OOD data under different $t$ with the increase of $s$.

As noted in Fig.2, when the added noise level is relative small ($t=500$), reconstruction loss of ID and OOD data reaches the minimum in small training steps ($s=20k\sim 40k$) and then starts to rise and fluctuate. When tested on higher noise level ($t=700$), the reconstruction loss of ID data basically declines and then remains stable. However, reconstruction loss of OOD data goes up with the increase of $s$. Such different behaviors verify that fully-trained models are more robust to higher noise level on ID data but the ability to restore OOD data is greatly degraded.

To verify the generation ability of fully-trained models, we test our models on SST2 dataset with different training steps and compare the score with MUCOLA-DISC [16]. The results are competitive with MUCOLA-DISC when reconstruction loss is relatively small ($s=160k$). This suggests that training models fully with reconstruction loss can guarantee good generation ability.

Larger models are not good at reconstruction and have bad OOD robustness. It is proved that RoBERTa exhibits greater robustness than BERT on OOD data since it is pretrained on more diverse corpora [12]. To test this hypothesis, we also evaluate the reconstruction ability of diffusion models with more PLMs.

As shown in Fig.3, the reconstruction loss of ID and OOD data when using RoBERTa is up to 2.8 times higher than that using BERT. Besides, the loss of OOD data is up to 4.2 times higher than that of ID data when using BERT, while 4.7 times when using RoBERTa. It shows that diffusion model with RoBERTa presents a greater gap of reconstruction loss between ID/OOD data and have worse OOD robustness than BERT.

Besides, diffusion models with BERT/RoBERTa-Base have lower reconstruction loss than BERT/RoBERTa-Large. The way in which the diffusion process adds noise determines this result. Since the latent space dimension of the base model is lower than that of the large model, the influence of the added noise is smaller and the model can make better reconstructions conditioned on the remaining inputs.

Finetuning on more diverse data improves reconstruction quality, but not always. Since more diverse data improves OOD generalization of PLMs [12], we question if this assumption applies to diffusion as well. To test it, we compare the OOD reconstruction performances when training with different ID datasets. We set a small $t=500$ in this experiment where the performance gap between different ID datasets is the most obvious.

As shown in Table 4, compared to TREC-10, the model finetuned on SST2 dataset has a lower OOD reconstruction loss and a greater ability to restore the unfamiliar words that do not appear in ID dataset. Since SST2 is a large dataset with more diverse data and covers 72.5% OOD tokens, it suggests that finetuning on more diverse ID data improves OOD reconstruction performance [14].

However, for the datasets IMDB and 20NG, although the value of “overlap” reaches 99.2% and 97.9% respectively (Table. 4), models finetuned on these two datasets do not have better reconstruction performance. By carefully investigating which aspect of data causes this anomalies, we compare the statistical characteristics of these data and observe that the average sentence length of IMDB and 20NG is 10$\sim$40 times larger than that of SST2 and TREC-10. So we speculate that the OOD performance of diffusion model is also related to sentence length.

Diffusion models with PLMs finetuned on long sentences are sensitive to length. To further test the hypothesis of length dependence, we explore the influence of sentence length on reconstruction quality.

As shown in Fig.4, when the model is finetuned on short sentences (TREC-10), the reconstruction loss of long sentence is overall higher than that of short sentence. The reason is that long sentences tend to contain more words not seen during training.

However, when the model is finetuned on long sentences (IMDB), the reconstruction loss of shorter sentence is at most 155 times higher than that of long sentence. Even when tested on SST2 dataset which is not considered as OOD dataset from IMDB, the model fails to produce low-loss reconstructions on short sentences as well. We call this phenomenon “Length Bias”, meaning that diffusion models with PLMs finetuned on long sentences are more sensitive to sentence length.

We give more examples to illustrate the “length bias” problem using 20NG as ID dataset which has an average sentence length of 865. As shown in Table 5, sentence from ID dataset can be perfectly restored with an average loss of 0.0008. It can also make near-perfect reconstructions on OOD dataset with long sentences (IMDB). However, when tested on short OOD sentences, even the common word “in” and “the” are not reconstructed correctly. Besides, the loss of long OOD sentence is up to 27 times higher than that of short OOD sentences, indicating serious “Length Bias” problem.

We introduce the following methods as baselines of OOD detection.

$\bullet$ Cosine Similarity serves as a scoring function to measure the similarity of input representations. It is defined by the maximum cosine similarity of $\boldsymbol{h}$ to samples of the ID validation set $\boldsymbol{h}^{(dev)}$

$\bullet$ MLM represents continuing training the pretrained transformers on ID data with Masked Language Modeling task [5]. Then we use the corresponding loss as a threshold to detect OOD data in a same way of diffusion-based OOD detection method.

$\bullet$ Hendrycks et al. [11] propose Maximum Softmax Probability (MSP) method by using the maximum class probability of a classifier with a softmax layer. The less confident classifier is, the higher the OOD score will be. It is defined by

where $C$ is the number of classes.

$\bullet$ Liu et al. [21] propose Energy Score method by using the softmax function as the ratio of the joint probability in $X\times Y$ to the probability in $X$, where $X$ and $Y$ denote the training corpus and its label set. A higher $d(x)$ means lower probability density in ID dataset, implying higher OOD likelihood. It calculates the following probability density

where $\boldsymbol{w}_{j}$ is the weight of class $j$ in the softmax layer and $\boldsymbol{h}$ is the input of the softmax layer.

$\bullet$ Lee et al. [18] propose to use Mahalanobis Distance [26] to determine the closeness of a sample to a set of samples belonging to the class $c$ by modeling the ID features with class-conditional multivariate Gaussian distributions

where $\boldsymbol{h(x)}$ is the vector representation of the sample $x$, $\boldsymbol{\mu}_{c}$ is the mean vector of the training set of class $c$, $\boldsymbol{\Sigma}$ is a shared covariance matrix and $\boldsymbol{\Sigma}^{+}$ is the pseudo-inverse of $\boldsymbol{\Sigma}$ defined by

where $N$ is the total number of the samples and $N_{c}$ is the number of samples belonging to class $c$.

$\bullet$ Zhou et al. [34] propose to finetune the Transformers with a contrastive loss. By contrasting samples to those from different ID classes, it improves the compactness of representations. The highest accuracy occurs when combining it with Mahalanobis Distance.

We propose a diffusion-based detector which uses a score function $f(x)=L_{recon}(x)$ to measure the similarity of any sample $x$ with ID dataset by fixing a threshold $\gamma$. If $f(x)\leq\gamma$, we classify $x$ as ID data. Otherwise, $x$ is regarded as OOD data. Formally, we distinguish OOD data using the decision function:

Since models learn to pay more attentions on token level after finetuned with diffusion, we further leverage Mahalanobis Distance method to add sentence level techniques. Following  [18], we define Mahalanobis Distance without using the class labels of the training set, i.e, treating the whole set as one class. Eq. 11 changes to

Hence, we define our score function with a mixed of loss and Mahalanobis Distance score.

where $\lambda$ is a hyperparameter treated as a constant with regards to optimization.

We use “Diffusion” to represent only using reconstruction loss as score function while “Diffusion+Maha” represents using both loss and Mahalanobis Distance.

We apply two metrics that are commonly used in evaluating OOD detection performance [11] and both of them are threshold-independent.

$\bullet$ AUROC: AUROC is the area under the Receiver Operating Characteristic curve. Higher AUROC value indicates better OOD detection performance. A random guessing detector model has an AUROC of 50%.

$\bullet$ FAR95: FAR95 is the probability that a negative example (OOD) to be mistakenly classified as positive (ID) when the true positive rate is 95%. In this case, a lower value indicates better performance.

Overall. As shown in Table 6, the AUROC accuracies of “Diffusion” on most datasets are above 99% and become SOTA. It proves the effectiveness of diffusion models with PLMs on OOD detection task. Specifically, the AUROC accuracy of “Diffusion” is up to 18% higher than that of “MLM”, indicating that diffusion models with PLMs degrade the generalization on OOD data. One exception is that the accuracy when trained on SST2 and tested on Multi30K has a significant decline. We suspect the reason is that SST2 contains more diverse data and has a similar sentence length with Multi30K. So the reconstruction loss between ID and OOD data are too close to distinguish OOD data from ID data.

Impact of $\lambda$. We test the impact of $\lambda$ on TREC-10 dataset while keeping $t=700$ and $s=80K$. As shown in Table 7, the accuracy reaches a peak at $\lambda=0.99$ and it declines by up to 2.27% as the decrease of $\lambda$. It indicates that the reconstruction loss plays a key role in OOD detection.

Impact of $\beta$. We test the impact of $\beta$ on TREC-10 dataset while keeping $t=700$ and $s=80K$. As shown in Table 9, the reconstruction loss of ID and OOD data decreases as $\beta$ gets smaller. But the gap between ID and OOD remains obvious. The model achieves the best results when $\beta$ increases linearly from $1e-4$ to $0.02$.

Visualization. By applying Gaussian Mixture Models (GMMs) to the learned representations, we draw the domain clusters in a 2D visualization as shown in Fig.5. Models are trained on TREC-10 as ID dataset. After finetuning, ID data is scattered throughout the output space while OOD data tends to aggregate into a cluster. Comparing with MLM, diffusion makes OOD data cluster in a more compact way like a Gauss. ID data is more scattered since all sentences can be perfectly restored. It further proves that diffusion models with PLMs are better OOD detectors.

Few-shot scenarios. In reality, many applications have very limited ID data. Hence, we test the robustness of diffusion-based detector finetuned on low-resource ID data. As shown in Fig.6, the AUROC accuracy remains high even when the number of training data is 10. It indicates that a few number of ID data is enough for diffusion-based detector to detect OOD data.

More PLMs. We further explore the OOD detection performance of T5-Large and BART finetuned with diffusion. TREC-10 is used as the ID dataset while keeping $s=80k$ and $t=800$. Results in Table 8 indicate that the OOD detection performances of T5-Large and BART outperform RoBERTa and become the new state-of-the-art, further illustrating that larger models finetuned with diffusion have worse OOD robustness.

We conclude that diffusion models with PLMs are effective on OOD detection, even in few-shot scenarios, though it shows a limitation when two datasets have high lexical coverage and similar sentence length.