[Paper Review] R-Drop: Regularized Dropout for Neural Networks

Xiaobo Liang, Lijun Wu+7 more|arXiv (Cornell University)|Jun 28, 2021

Advanced Neural Network Applications74 references306 citations

TL;DR

R-Drop regularizes dropout by enforcing consistency between two dropout-generated sub-model outputs using bidirectional KL-divergence, improving generalization across NLP and CV tasks and achieving state-of-the-art results on some translation benchmarks.

ABSTRACT

Dropout is a powerful and widely used technique to regularize the training of deep neural networks. In this paper, we introduce a simple regularization strategy upon dropout in model training, namely R-Drop, which forces the output distributions of different sub models generated by dropout to be consistent with each other. Specifically, for each training sample, R-Drop minimizes the bidirectional KL-divergence between the output distributions of two sub models sampled by dropout. Theoretical analysis reveals that R-Drop reduces the freedom of the model parameters and complements dropout. Experiments on $\bf{5}$ widely used deep learning tasks ($\bf{18}$ datasets in total), including neural machine translation, abstractive summarization, language understanding, language modeling, and image classification, show that R-Drop is universally effective. In particular, it yields substantial improvements when applied to fine-tune large-scale pre-trained models, e.g., ViT, RoBERTa-large, and BART, and achieves state-of-the-art (SOTA) performances with the vanilla Transformer model on WMT14 English$ o$German translation ($\bf{30.91}$ BLEU) and WMT14 English$ o$French translation ($\bf{43.95}$ BLEU), even surpassing models trained with extra large-scale data and expert-designed advanced variants of Transformer models. Our code is available at GitHub{\url{https://github.com/dropreg/R-Drop}}.

Motivation & Objective

• Motivate and address training-inference inconsistency caused by dropout in deep networks.
• Propose a simple regularization method that enforces consistency between two dropout-induced sub-model outputs.
• Theoretically analyze how R-Drop reduces training-inference inconsistency.
• Empirically demonstrate universal effectiveness across NLP and CV tasks, including large pre-trained models.

Proposed method

• For each training sample, run two forward passes with different dropout instantiations to obtain P1(y|x) and P2(y|x).
• Minimize the bidirectional KL-divergence between P1 and P2 in addition to the standard negative log-likelihood loss.
• Final objective combines NLL losses from both passes with an alpha-weighted KL term: L = L_NLL1 + L_NLL2 + (alpha/2)[KL(P1||P2) + KL(P2||P1)].
• Provide a single training step by duplicating each input in the batch to compute both passes within the same mini-batch.
• Algorithmically, train with two dropout-submodels per sample and update parameters by minimizing the combined loss.
• Theoretical analysis shows the constraint bounds the inconsistency between sub-model losses and full-model loss in a linear model.

Experimental results

Research questions

• RQ1 Does enforcing output-consistency between two dropout-induced sub-models improve generalization across tasks?
• RQ2 How does R-Drop affect training-inference inconsistency and regularization strength?
• RQ3 Can R-Drop yield state-of-the-art results for vanilla Transformer and large pre-trained models without extra data or architecture changes?
• RQ4 What are the stability and cost implications of applying R-Drop during training across diverse domains?

Key findings

• R-Drop yields substantial improvements across 5 tasks and 18 datasets, including NLP, language modeling, and image classification.
• On WMT14 En→De and En→Fr translation with vanilla Transformer, R-Drop achieves 30.91 BLEU and 43.95 BLEU respectively, surpassing previous SOTA results.
• In GLUE, RoBERTa-large with RD achieves 89.73 average and outperforms several strong baselines (e.g., XLNet-large, ELECTRA-large).
• In CNN/Daily Mail summarization, BART+RD achieves state-of-the-art ROUGE-L, and ROUGE-1/2 improved by about 0.3 points over BART.
• On Wikitext-103 language modeling, RD improves perplexity for both Transformer and Adaptive Input Transformer baselines (e.g., Transformer: valid 25.76 to 23.97; test 26.62 to 24.94).
• On image classification, ViT models gain accuracy improvements with RD (e.g., ViT-B/16: CIFAR-100 92.64→93.29; ImageNet 83.97→84.38).