[论文速览] R-Drop@ Regularized Dropout for Neural Networks-526互联

Pre

title: R-Drop: Regularized Dropout for Neural Networks
accepted: NeurIPS 2021
paper: https://arxiv.org/abs/2106.14448
code: https://github.com/dropreg/R-Drop

关键词： regularization, consistency learning, dropout
阅读理由：简单有效的正则化方法

Idea

一种训练思想，约束同一输入在dropout作用下过两次同一模型得到的输出尽量一致（SimCSE?）

Motivation&Solution

dropout的随机性导致了训练和推理时的不一致性 —— R-Drop

Background

略

Method（Model）

Overview

图1 R-Drop总体框架，以Transformer为例。左边是一个输入过两次得到两个不同分布，右侧是dropout产生的两个不同子模型

每步训练中，同一个输入过两次模型，由于dropout或其他正则手段，输出会有不同，R-Drop通过双向KL散度来约束相同，这样就能减少训练推理的不一致性。不同于之前的大多数方法，R-Drop同时作用域隐层单元和子模型输出，更加有效。

公式2 KL损失计算

公式3 NLL损失

公式4 数据对$x_i, y_i$的总体损失

算法1 R-Drop 训练算法

公式和算法1中的$\mathcal{P}$是模型的输出分布，算法1通过将输入重复并拼接，一次过模型来加速计算，两个输入都算交叉熵，然后两个输出之间再互相算KL散度来拉近。直接看它github中README的代码展示更好理解（但这份的输入分两次过模型）：

import torch.nn.functional as F

# define your task model, which outputs the classifier logits
model = TaskModel()

def compute_kl_loss(self, p, q, pad_mask=None):
    
    p_loss = F.kl_div(F.log_softmax(p, dim=-1), F.softmax(q, dim=-1), reduction='none')
    q_loss = F.kl_div(F.log_softmax(q, dim=-1), F.softmax(p, dim=-1), reduction='none')
    
    # pad_mask is for seq-level tasks
    if pad_mask is not None:
        p_loss.masked_fill_(pad_mask, 0.)
        q_loss.masked_fill_(pad_mask, 0.)

    # You can choose whether to use function "sum" and "mean" depending on your task
    p_loss = p_loss.sum()
    q_loss = q_loss.sum()

    loss = (p_loss + q_loss) / 2
    return loss

# keep dropout and forward twice
logits = model(x)

logits2 = model(x)

# cross entropy loss for classifier
ce_loss = 0.5 * (cross_entropy_loss(logits, label) + cross_entropy_loss(logits2, label))

kl_loss = compute_kl_loss(logits, logits2)

# carefully choose hyper-parameters
loss = ce_loss + α * kl_loss

Theoretical Analysis

作者在这里从理论上（公式推导）分析为何R-Drop能够降低训练推理不一致性。简单看了下是说Dropout在训练的时候，同一个输入相当于会经过不同的子模型处理，损失则是这些子模型损失的平均，推理时用的却是完整的模型，而子模型结构跟完整模型是不一样的，因此有不一致。但是R-Drop通过约束子模型输出一致，从而让它的结构也相似，减少了这种不一致性（那不就是既要模型不一致Dropout，又要一致）

Discussion

跟R-Drop最相关的是 ELD 和 FD，都研究用了dropout的一致性训练。然而R-Drop与它们有关键不同：

差异控制（gap control）来自不同视角。ELD直接减少了dropout导致的子模型和完整模型间的不同（gap），而R-Drop和FD通过惩罚子模型间的差异来实现，对子模型进行正则化的好处已经由FD证实了
正则化效率不同。ELD只反向传播了子模型的梯度，R-Drop同时更新两个子模型更高效。
正则化效果不同。ELD和FD都在隐层状态间算L2距离来作为正则化损失函数，但这跟最小化负对数似然的主要训练目标差距较大。隐层状态的距离跟概率分布不在一个空间中，因为log-softmax会极大地影响优化。

Experiment

Settings

简单起见后面用RD来代指R-Drop

Neural Machine Translation

表1 8种IWSLT机器翻译任务上的BLEU分数

表2 WMT14英译德和英译法的BLEU分数

如果把R-Drop结合到表格里其他方法，换上更强的backbone还能涨点

Language Understanding

表3 GLUE语言理解benchmark上微调模型表现

Summarization

拿预训练BERT作为backbone，用R-Drop来微调

表4 CNN/Daily Mail 摘要数据集上ROUGE结果。RG-x代表ROUGE-x分数

Language Modeling

数据集：Wikitext-103
模型：基础Transformer解码器、Adaptive Input Transformer
公式4 $\alpha$：1.0

表5 Wikitext-103语言建模任务的困惑度结果，Adaptive指的是 Adaptive Input Transformer

这个任务上提升还算明显，特别是如果baseline不是很强，涨点更显眼。

Image Classification

表6 分类任务准确率

Study

通过不同视角的广泛研究来理解R-Drop

Regularization and Cost Analysis

图2 随模型训练的Loss/BLEU曲线

可以看到：

训练时Transformer会迅速过拟合，训练和验证的损失差距很大，而R-Drop有更低的验证损失。
训练早期Transformer迅速提升了BLEU，但马上收敛到局部最优。R-Drop需要更多时间，但能得到更好的性能。同时R-Drop由于要前向同一个样本两次，训练开销会大一些。

k-step R-Drop

图3 不同step的R-Drop

每k步再用一次R-Drop来提升训练效率，图3能看出虽然k大一点收敛快，但是效果不行，迅速过拟合。

m-time R-Drop

默认是一个输入过两次模型来正则化，也就是m=2，如果加大m呢？首先KL损失计算更麻烦，得算$m(m-1)$种两两组合再平均。作者试了m=3的IWSLT14德译英测试集，BLEU分数是37.30，但m=2的时候也有37.25，因此认为m=2的正则化已经足够强。

Two Dropout Rates

图4 不同dropout概率组合的R-Drop，该表格是对称的，25个格子只有15种不同组合

默认两个分布用的dropout概率都一样，如IWSLT翻译用0.3，现在试训练时一个子模型用不同概率去dropout，从$\{0.1, 0.2, 0.3, 0.4, 0.5\}$挑选概率值共15种设置（$C^2_5$两两组合，$C^1_5$两个应用相同概率值）结果如图4。能看到两个都取0.3效果最好，相差不大（都在0.3~0.5）也能有不错结果，而当两个都0.5（一半单元被丢弃），R-Drop仍有不错效果