PaperlayerView exact revision on arXiv

Source-linked AI summary

R-Drop: Regularized Dropout for Neural Networks

Xiaobo Liang, Lijun Wu, Juntao Li, Yue Wang, Qi Meng, Tao Qin, Wei Chen, Min Zhang, Tie-Yan Liu

arXiv:2106.14448v2cs.LG

TL;DR

Dropout can produce inconsistent predictions across randomly sampled training submodels and the full inference model. R-Drop addresses this by minimizing bidirectional KL divergence between two dropout submodels’ outputs, and experiments across five tasks and 18 datasets report universal effectiveness, including SOTA translation results.

  • Problem

    Dropout uses random submodels during training but the full model during inference, creating a training–inference inconsistency.

  • Method

    R-Drop performs two dropout-based forward passes per sample and minimizes the bidirectional KL divergence between their output distributions.

  • Results

    Across 5 tasks and 18 datasets, R-Drop is broadly effective, including 30.91 BLEU for WMT14 English→German and 43.95 BLEU for English→French.

  • Takeaways & Limitations

    R-Drop enhances models including ViT, BART, and RoBERTa-large and can achieve SOTA results with vanilla Transformer models.

  • Takeaways & Limitations

    The paper tests R-Drop only during downstream fine-tuning, leaving pre-training experiments for future work because of computational cost.

Abstract

from arXiv · show

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$\to$German translation ($\bf{30.91}$ BLEU) and WMT14 English$\to$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}}.

1 Introduction

R-Drop regularizes dropout by making predictions from independently sampled dropout submodels consistent, addressing training–inference inconsistency. Across five tasks and 18 datasets, it produces broadly strong results, including multiple SOTA outcomes.

  • Dropout regularizes networks by randomly dropping hidden units, but its sampled training submodels differ from the full inference model.
  • R-Drop runs each sample through two dropout submodels and minimizes their bidirectional KL divergence.
  • R-Drop operates on both dropout-hidden units and sampled-submodel outputs, complementing conventional dropout regularization.
  • 5 tasks and 18 datasets show broad effectiveness, including 30.91 BLEU on WMT14 English→German and 43.95 on English→French with vanilla Transformer.
  • The method is presented as a simple regularizer that can be applied across different deep model types.
  • R-Drop theoretically reduces inconsistency between dropout training and full-model inference.

2 Approach

R-Drop adds consistency regularization to dropout by comparing two predictions generated from independently masked versions of the same input. Its objective combines the usual negative log-likelihood with a bidirectional KL term, while theory relates the constraint to reduced training–inference inconsistency.

  • R-Drop Regularization: R-Drop feeds each input through two independently dropped submodels and compares their output distributions.
  • Training Objective: The training objective adds a KL-divergence loss to the negative log-likelihood loss without structural model changes.
  • Training Algorithm: Implementation repeats each input within one mini-batch, computes both likelihood and KL losses, and updates parameters using their combined objective.
  • Theoretical Analysis: R-Drop represents dropout-layer masking with Bernoulli random vectors and optimizes two sampled dropout instantiations stochastically.
  • Motivation: Dropout creates a mismatch because training averages losses over random submodels while inference uses the full model.
  • Theoretical Analysis: For a constrained linear model, the empirical-loss gap between the full model and a random submodel is bounded by c√ϵ.

3 Experiments

Experiments evaluate R-Drop across 5 tasks and 18 datasets spanning machine translation, summarization, language understanding, language modeling, and image classification. Across these settings, R-Drop improves diverse baselines, including strong pre-trained and vision models.

  • Experimental scope: Experiments cover 5 tasks and 18 datasets across four NLP tasks and one computer-vision task.The tasks are neural machine translation, abstractive summarization, language understanding, language modeling, and image classification.
  • Language understanding: 1.21 and 0.80 average-point improvements over BERT-base and RoBERTa-large, respectively, are reported on the GLUE benchmark.RoBERTa-large + RD also surpasses XLNet-large and ELECTRA-large.
  • Abstractive summarization: 0.3-point gains on both RG-1 and RG-2 over fine-tuned BART yield SOTA results on CNN/Daily Mail summarization.The R-Drop model surpasses PEGASUS and BART+R3F without extra parameters or structural changes.
  • Language modeling: R-Drop improves perplexity for both Transformer and Adaptive Input Transformer on Wikitext-103.The gains are 1.79 on validation and 1.68 on test over Transformer, and 0.80 on test over Adaptive Input Transformer.
  • Image classification: 0.65 accuracy improvement over ViT-B/16 and 0.41 points over ViT-L/16 are obtained on CIFAR-100, with consistent ImageNet improvements.These results indicate benefits even when the baseline model is powerful.

4 Study

The studies examine R-Drop’s regularization, training-cost trade-off, and sensitivity to application frequency, distribution count, dropout rates, and KL weight. They find that applying R-Drop at every step and using two distributions provide strong performance, while larger application intervals weaken results.

  • Regularization and cost: R-Drop maintains lower validation loss than Transformer during training, while achieving a better final BLEU optimum despite slower convergence.The Transformer over-fits quickly; R-Drop improves more gradually and requires more training to converge.
  • Regularization and cost: R-Drop increases computation at each step because each input is processed through a second dropout sub model.The authors report negligible overall added cost in an appendix study alongside stronger performance.
  • Application frequency: Applying R-Drop every k steps makes convergence faster but progressively worsens BLEU and leads to over-fitting as k increases.The study evaluates k in {1, 2, 5, 10}; k = 1 is the current strategy.
  • Number of distributions: 37.30 BLEU with m = 3 is similar to 37.25 BLEU with m = 2 on IWSLT14 De→En, indicating that two distributions already provide strong regularization.The authors therefore find no necessity for stronger regularization with more distributions.
  • Dropout-rate combinations: The same dropout rate of (0.3, 0.3) is best, while rates in the 0.3–0.5 range remain strong without large performance differences.The study evaluates 15 combinations of two dropout rates selected from {0.1, 0.2, 0.3, 0.4, 0.5}.
  • KL-divergence weight: KL-divergence weight α = 5 is the best-balanced choice: α = 1 underperforms, whereas α = 10 applies excessive regularization.The tested values are α ∈ {1, 3, 5, 7, 10}.

5 Related Work

The related work positions R-Drop at the intersection of dropout regularization, consistency training, and self-distillation. Its distinguishing choice is to regularize prediction consistency between dropout-sampled sub models.

  • Regularization methods: Dropout regularizes hidden units or model parameters, whereas R-Drop additionally regularizes the outputs of dropout-sampled sub models.The paper describes R-Drop as a dropout variation that operates at both hidden-unit and output levels.
  • Consistency training: R-Drop differs from ELD and FD by targeting consistency between dropout sub-model predictions through bidirectional KL divergence.ELD compares a dropout sub model with the expected full model, while FD operates between sub models.
  • Self-distillation: R-Drop resembles self-knowledge distillation because its teacher and student are dropout instantiations of the same model.The paper relates minimizing KL divergence between model outputs to knowledge distillation.

6 Conclusions and Future Work

The conclusion presents R-Drop as a simple consistency-training method that improves diverse deep-learning models and datasets. It also identifies untested pre-training and non-Transformer architectures as future directions.

  • Conclusion: R-Drop minimizes bidirectional KL divergence between output distributions from pairs of dropout-sampled sub models during training.The method is presented as a consistency-training approach built upon dropout.
  • Conclusion: Across 18 datasets, R-Drop enhances models including ViT, BART, and RoBERTa-large, and reaches SOTA translation results with vanilla Transformer models.The conclusion reports effectiveness on large-scale datasets as well as strong pre-trained models.
  • Future work: Because computational resources limited the study, R-Drop was tested only for downstream fine-tuning rather than pre-training.The authors plan to evaluate R-Drop on pre-training tasks in future work.
  • Future work: The work focuses on Transformer-based models and proposes applying R-Drop to other architectures, including convolutional neural networks.This is stated as a future extension rather than a demonstrated result.

A Detailed Experimental Settings

The experimental settings describe datasets, tokenization, model configurations, evaluation procedures, and selected fine-tuning setups across translation, summarization, language modeling, and GLUE experiments.

  • Neural machine translation: NMT datasets come from IWSLT and WMT competitions and are tokenized with byte-pair encoding, yielding vocabularies near 10k for IWSLT and about 32k for WMT.The source and target sentence pairs jointly build the dictionary except for IWSLT17 En↔Zh.
  • Neural machine translation: IWSLT translation uses a six-layer encoder and decoder with 512-dimensional embeddings, while WMT uses the larger Vaswani Transformer configuration.The settings specify feed-forward sizes of 1,024 for IWSLT and 4,096 for WMT.
  • Evaluation: Translation evaluation uses multi-bleu.perl for IWSLT14 and WMT tasks, sacre-bleu for other NMT tasks, and task-specific beam-search settings.WMT14 En→De uses beam size 4 and length penalty 0.6; other tasks use beam size 5 and penalty 1.0.
  • Summarization: Summarization fine-tunes pre-trained BART on CNN/DailyMail using the hyperparameters from the original BART work.The BART backbone has 12 Transformer encoder layers and 12 decoder layers.
  • Language modeling: Language modeling experiments use Transformer decoder and Adaptive Input Transformer models with configurations named transformer_lm_gpt and transformer_lm_wiki103.The configurations differ in layer counts, embedding sizes, attention heads, and dropout settings.
  • Language understanding: The appendix includes a GLUE fine-tuning hyper-parameter table and results across 8 GLUE tasks with different random seeds.The supplied materials identify these resources but do not specify their individual hyper-parameter values or outcomes.

A.4 Language Understanding

The language-understanding experiments use GLUE tasks with pre-trained Transformer models and extend R-Drop to regression through MSE regularization. Across random seeds, R-Drop consistently improves fine-tuned BERT performance, while α remains task-sensitive.

  • Evaluation setup: GLUE evaluation covers eight tasks spanning single-sentence classification, sentence-pair classification, and sentence-pair regression.The regression task is STS-B.
  • Results: R-Drop fine-tuning produces consistent performance improvements across different random seeds on the GLUE tasks.The comparison uses a pre-trained BERT model and vanilla fine-tuning.
  • Hyperparameter analysis: α = 1.0 is a good choice for most GLUE tasks, although α is sensitive for each task.The paper reports these comparisons in Table 9.
  • Regression extension: For regression, R-Drop replaces the KL-divergence regularization with an MSE-based regularization between two predicted values.The input is forwarded twice to obtain the two predictions.
  • Evaluation scope: The image-classification experiments evaluate R-Drop with pre-trained ViT-B/16 and ViT-L/16 models.These models are pre-trained on ImageNet-21k.

B Theoretical Discussion of R-Drop

The theoretical discussion analyzes how R-Drop constrains dropout-induced variation and relates this constraint to the full-model and random-submodel losses. The analysis uses Lipschitz properties, normalization, and the relationship between KL-divergence and total variation distance.

  • Theoretical discussion: The appendix provides a proof of Proposition 2.1 and related discussions of the regularization analysis.The proof material is presented as supplementary theoretical discussion.
  • Theoretical bound: Proposition B.1 bounds the difference between full-model and random-submodel empirical losses by c√ϵ under a stated constraint.Here, c depends on the Lipschitz constant of the softmax operator.
  • Proof ingredients: The proof uses Lipschitz continuity of the loss and normalizes rows of the weight matrix to have unit norm.The normalization assumption is applied before relating model outputs to parameter changes.
  • Proof ingredients: The argument connects KL-divergence to total variation distance and uses the softmax map together with its inverse under a norm constraint.The inverse map is characterized through a Lipschitz constant c2.

C.1 Batch Size Doubled Training

The doubled-batch experiment separates the benefit of seeing repeated stochastic views from the benefit of R-Drop's KL-divergence loss. R-Drop substantially outperforms both direct batch doubling and removal of the KL term.

  • Batch size comparison: 34.93 BLEU with a doubled batch size remains below R-Drop's 37.25 BLEU on IWSLT14 De→En translation.The doubled-batch result improves only slightly over the 34.64 BLEU baseline.
  • Training procedure: R-Drop training repeats each input and concatenates the duplicate within the same mini-batch for a single forward pass.The duplicated examples receive different dropout submodels.
  • Comparison design: Table 10 compares baseline and R-Drop using BLEU scores and training time for reduced half-batch training on IWSLT14 De→En.The table is specifically devoted to the efficiency comparison.
  • KL-divergence ablation: 34.93 BLEU after removing the KL-divergence loss is far below R-Drop's 37.25 BLEU on the same translation task.The ablation retains only the NLL loss and is slightly above the 34.64 BLEU Transformer baseline.

C.3 Training Time and Efficiency

R-Drop requires more training to converge but reaches a substantially better final model, with comparable time cost under reduced half-batch training.

  • Convergence: R-Drop needs more training to converge, yet its final model is much better than the baseline.The comparison is made on IWSLT14 De→En translation.
  • Efficiency: R-Drop's time cost is comparable to the baseline when trained with a reduced half-batch and matched at the same BLEU score.This comparison concerns reaching an equal BLEU score rather than equal training duration.

C.4 Experiments for ELD and FD

The experiments compare R-Drop with hidden-space regularization methods and with deep ensembling or weight averaging. R-Drop is evaluated across image classification and translation settings, including comparisons of BLEU scores.

  • Comparison with ELD and FD: R-Drop is compared primarily with FD, following FD’s reported advantage over ELD, on CIFAR-10 image classification and IWSLT14 De→En translation.The CIFAR-10 experiments use released code, with R-Drop added as an implementation.
  • Ensembling and weight averaging: Table 12 compares BLEU scores from deep ensembling and weight averaging across models trained with different random seeds or epoch checkpoints.The caption distinguishes different-model ensembles from same-seed checkpoint averaging.
  • Ensembling and weight averaging: R-Drop achieves 37.25 BLEU with a single full model, outperforming weight averaging and deep ensembling in the reported comparison.The authors state that the best result remains R-Drop even with at least six models ensembled.
Loading 2106.14448v2…