Dropout Twice Again! This Time It Achieves SOTA in Supervised Tasks

Classification Problems

In my view, R-Drop is highly related to SimCSE and was likely inspired by it, although the R-Drop paper does not cite SimCSE, which is somewhat puzzling.

Specifically, taking classification as an example, given training data \{x_i, y_i\}_{i=1}^n and a model P_{\theta}(y|x), the loss for each sample is typically cross-entropy: \mathcal{L}_i = -\log P_{\theta}(y_i|x_i) In the "dropout twice" scenario, we can consider the sample as having passed through two slightly different models, denoted as P_{\theta}^{(1)}(y|x) and P_{\theta}^{(2)}(y|x). The R-Drop loss consists of two parts. One part is the conventional cross-entropy: \mathcal{L}_i^{(CE)} = -\log P_{\theta}^{(1)}(y_i|x_i) -\log P_{\theta}^{(2)}(y_i|x_i) \label{eq:ce} The other part is the symmetric KL divergence between the two model outputs, which encourages the outputs of different dropout instances to be as consistent as possible: \mathcal{L}_i^{(KL)} = \frac{1}{2}\big[KL\left(P_{\theta}^{(2)}(y|x_i)\big\Vert P_{\theta}^{(1)}(y|x_i)\right) + KL\left(P_{\theta}^{(1)}(y|x_i)\big\Vert P_{\theta}^{(2)}(y|x_i)\right)\big] \label{eq:kl} The final loss is the weighted sum of the two: \mathcal{L}_i = \mathcal{L}_i^{(CE)} + \alpha\mathcal{L}_i^{(KL)} In other words, it adds a regularization term to the conventional cross-entropy to strengthen model robustness.

General Form

Some readers might ask what the KL term should be replaced with for non-classification problems. In fact, the original paper did not conduct experiments on non-classification tasks, but we can supplement that here. Note that: -\log P_{\theta}(y_i|x_i) = KL\left(\text{one\_hot}(y_i)\big\Vert P_{\theta}(y|x_i)\right) Thus, the aforementioned \mathcal{L}_i is just the repeated use of KL divergence. Its general form is: \mathcal{L}_i = \mathcal{D}\left(y_i, f_{\theta}^{(1)}(x_i)\right)+\mathcal{D}\left(y_i, f_{\theta}^{(2)}(x_i)\right) + \frac{\alpha}{2} \left[\mathcal{D}\left(f_{\theta}^{(2)}(x_i), f_{\theta}^{(1)}(x_i)\right)+\mathcal{D}\left(f_{\theta}^{(1)}(x_i), f_{\theta}^{(2)}(x_i)\right)\right] Therefore, for non-classification problems, we can replace \mathcal{D} with an appropriate metric (instead of KL divergence).

Paper Reports

To be honest, the results reported in the R-Drop paper are quite stunning, which is the main reason I felt compelled to introduce it. The original paper conducted comparative experiments on various tasks including NLU, NLG, and CV classification. Most results can be described as "significant improvements."

Official Implementation: https://github.com/dropreg/R-Drop

Below are some screenshots of the experimental results:

Notably, in machine translation tasks, the simple "Transformer + R-Drop" outperformed other more complex methods:

The paper also includes experiments on automatic summarization, language models, image classification, and ablation studies on hyperparameters. You can refer to the original paper for details. Overall, R-Drop’s "report card" is indeed worthy of praise.

Personal Experiments

Of course, I maintain the view that "a model not tested on Chinese tasks has no soul." Usually, I only write and share after personally trying it on Chinese tasks.

Personal Implementation: https://github.com/bojone/r-drop

On Chinese supervised tasks, I experimented with two text classification tasks (IFLYTEK and TNEWS from the CLUE benchmark):

And one text generation task (CSL title generation, referring to "Analysis and Countermeasures for Exposure Bias in Seq2Seq"):

As we can see, the results of R-Drop are competitive with famous regularization methods like "Adversarial Training" and "Gradient Penalty" introduced in "A Brief Talk on Adversarial Training: Meaning, Methods, and Thoughts".

Implementation Details

Compared to complex regularization methods like adversarial learning, the implementation difficulty of R-Drop is quite low. Here, using bert4keras as an example, I will briefly introduce how to modify a standard training script to include R-Drop.

First, the data generation part is modified as follows:

class data_generator(DataGenerator):
    """Data Generator
    """
    def __iter__(self, random=False):
        batch_token_ids, batch_segment_ids, batch_labels = [], [], []
        for is_end, (text, label) in self.sample(random):
            token_ids, segment_ids = tokenizer.encode(text, maxlen=maxlen)
            # batch_token_ids.append(token_ids)
            # batch_segment_ids.append(segment_ids)
            # batch_labels.append([label])
            for i in range(2):
                batch_token_ids.append(token_ids)
                batch_segment_ids.append(segment_ids)
                batch_labels.append([label])
            # if len(batch_token_ids) == self.batch_size or is_end:
            if len(batch_token_ids) == self.batch_size * 2 or is_end:
                batch_token_ids = sequence_padding(batch_token_ids)
                batch_segment_ids = sequence_padding(batch_segment_ids)
                batch_labels = sequence_padding(batch_labels)
                yield [batch_token_ids, batch_segment_ids], batch_labels
                batch_token_ids, batch_segment_ids, batch_labels = [], [], []

Then, define a new loss function:

from keras.losses import kullback_leibler_divergence as kld

def categorical_crossentropy_with_rdrop(y_true, y_pred):
    """R-Drop Loss for the above generator
    Dividing loss_kl by 4 aligns the quantity with the formula description.
    """
    loss_ce = K.categorical_crossentropy(y_true, y_pred)  # Original loss
    loss_kl = kld(y_pred[::2], y_pred[1::2]) + kld(y_pred[1::2], y_pred[::2])
    return K.mean(loss_ce) + K.mean(loss_kl) / 4 * alpha

Finally, enable Dropout in the model and use this data_generator and categorical_crossentropy_with_rdrop to train the model.

Consistency

R-Dropout can be seen as an improvement over Dropout. So what is the problem with Dropout? Dropout is a classic example of training-inference inconsistency. Specifically, during training, Dropout adds multiplicative noise to (certain layers’) inputs, changing the model from f_{\theta}(x) to f_{\theta}(x, \varepsilon), where each element of \varepsilon has a probability p of being 0 and 1-p of being 1/(1-p). The training objective is: \mathbb{E}_{(x,y)\sim\mathcal{D}}\mathbb{E}_{\varepsilon}[l(y, f_{\theta}(x,\varepsilon))] After training this way, which model should we use for prediction? It’s uncertain, but if the loss function is l_2 distance, we can derive that the optimal prediction model should be: \mathbb{E}_{\varepsilon}[f_{\theta}(x,\varepsilon)]

Derivation: If using l_2 loss, the loss for a single sample is: \mathbb{E}_{\varepsilon}\left[\Vert y - f_{\theta}(x,\varepsilon)\Vert^2\right] = \Vert y\Vert^2 - 2\langle y,\mathbb{E}_{\varepsilon}\left[f_{\theta}(x,\varepsilon)\right]\rangle + \mathbb{E}_{\varepsilon}\left[\Vert f_{\theta}(x,\varepsilon)\Vert^2\right] Note that the question is "what function should be used for prediction after training," so f_{\theta}(x,\varepsilon) is a constant, and y is the variable to optimize. This is just a quadratic minimization problem, easily solved as y=\mathbb{E}_{\varepsilon}[f_{\theta}(x,\varepsilon)].

Assuming this result generalizes, the correct step for a model with Dropout is "model ensemble":

Pass the same input through the model multiple times (without disabling Dropout), and then take the average of the multiple predictions as the final result.

However, in practice, we usually don’t do this. Instead, we disable Dropout for deterministic prediction, which is equivalent to changing the prediction model from "model averaging" to "weight averaging": f_{\theta}(x,\mathbb{E}_{\varepsilon}[\varepsilon])=f_{\theta}(x,1)=f_{\theta}(x) Here, 1 refers to the all-ones vector. Thus, we train an ensemble of different Dropout models but use a single model with disabled Dropout for prediction. The two are not necessarily equivalent, which is the training-inference inconsistency of Dropout.

Now, R-Drop is easy to understand. It adds a regularization term to strengthen the model’s robustness to Dropout, making the outputs under different Dropout instances consistent. This reduces the inconsistency and promotes similarity between "model averaging" and "weight averaging," making the effect of simply disabling Dropout equivalent to the result of multi-Dropout model fusion, thereby improving final performance.

Continuity

As mentioned at the beginning, R-Drop is similar to SimCSE. In fact, it is also quite similar to "Virtual Adversarial Training (VAT)." (R-Drop doesn’t cite VAT either; is it just me who thinks they are similar??)

For an introduction to VAT, you can refer to my previous article "Random Thoughts on Generalization: From Random Noise, Gradient Penalty to Virtual Adversarial Training". Briefly, VAT also uses a regularization term to make the model more robust to perturbations, enhancing the model’s continuity (small changes in input shouldn’t lead to large changes in output). The difference lies in how perturbations are added: VAT only adds perturbations to the input and uses adversarial ideas to make them more targeted; R-Drop’s perturbations can be applied to every layer and are random.

Some readers might have realized that VAT is primarily used for semi-supervised training. Does this mean R-Drop can also do semi-supervised training? The original paper didn’t experiment on this, but I did. The answer is yes. Like VAT, the KL divergence term in R-Drop doesn’t require labels, so it can be used for unsupervised training, and combined with labeled data, it becomes semi-supervised. The results are quite good. Below are my experimental results:

As we can see, R-Drop’s semi-supervised performance is not inferior to VAT, and its implementation is simpler and faster! It seems VAT might be ready for retirement. Intuitively, although R-Drop’s perturbations are random, there are more of them (across all layers), so the resulting perturbation might be amplified enough to match the targeted perturbations of VAT.

Non-target Classes

A direct question is: if my model is complex enough, can’t cross-entropy alone make the model robust to Dropout? What specific difference does the KL divergence term make?

Actually, it can’t. It’s important to note that the training objective of cross-entropy is primarily to make the score of the target class higher than the scores of non-target classes (refer to "Generalizing Softmax+Cross-Entropy to Multi-label Classification"). That is to say, if there is only cross-entropy, the training result is at most:

Under different Dropout instances, the target class score is greater than non-target class scores.

But it cannot achieve:

Under different Dropout instances, the score for every class is consistent.

Therefore, it doesn’t solve the training-inference inconsistency. From the formula, cross-entropy \eqref{eq:ce} only relates to the target class and doesn’t care about the distribution of non-target classes. If the target class is the first one, predictions like [0.5, 0.2, 0.3] or [0.5, 0.3, 0.2] make no difference to it. But for the KL divergence term \eqref{eq:kl}, it’s different; the scores of every class participate in the calculation, and there is a non-zero loss between [0.5, 0.2, 0.3] and [0.5, 0.3, 0.2].