SMART: 不可错过的通用对抗式训练算法

背景

纯粹的微调下，由于目标函数设置过于简单，神经网络很容易急不可耐地收敛到最近的局部最优点。这些最优点不仅并非全局最优点，同时不抗攻击。只要在模型输入上稍加扰动，输出的概率分布便会产生大幅度的偏移。为此，一些研究潜心于此，希望能让神经网络学习到的映射能力更为 平滑，从而进一步地提高 泛化能力。如果你对深度学习有一定深入的了解，对于这一点一定不陌生。

这些研究里最为经典的当属 VAE (变分自编码器)，通过 reparameterization，将噪声结合到模型前馈的中间流程，基于不变的输入和输出，提高中间语义向量的平滑度。VAE 至今已演化出数百个版本，是 GAN 之外最为成功的生成模型。来到自然语言处理领域，添加噪声的方式则主要集中在 embedding 层，对 embedding 层的输出添加符合正态/均匀分布的扰动。如同 VAE 简单地设置噪声的方差，在每次输入时添加噪声，能获得一定的成效。但这样统一的标准，容易造成训练后期噪声范围过大而导致的参数收敛不平稳，达不到最优点。因此后来的研究则专注于根据梯度的分布，动态调整噪声的方向和范围，统称为对抗式训练。

$\min_\theta \mathbb{E}_{(\mathcal{Z},y)\sim\mathcal{D}}\Big[\max_{||\delta||_p<\epsilon}L(f_\theta(X+\delta),y)\Big]$θmin​E(Z,y)∼D​[∣∣δ∣∣p​<ϵmax​L(fθ​(X+δ),y)]

FGSM 和 FGM 是非常优秀的代表，每批样本进来，沿着梯度的正方向添加扰动，将添加扰动前的梯度和添加扰动后的梯度取平均值更新参数。随后 PGD 被提出，在 FGM 上多次迭代，找到最优的扰动。在 PGD 之后，还出现了 FreeAT、YOPO、FreeLB 等优秀的对抗式训练算法。PGD 被引用了近 2,000 次，而 FreeLB 至今在 GLUE 上依然榜上有名。SMART 正是在这样的背景下诞生。

尽管 SMART 自推出至今，并未受到太多的关注。但实战说话，在我个人的试验下，它在不同任务上都超越了其他一众对抗式训练算法，达到了最优异的成绩。至于为何关于 SMART 的讨论热度不高，我认为是因为作者始终没有开放源码，同时论文存在信息跳跃，对没有接触过对抗式训练的新人不太友好。是金子总会发光，相信 SMART 在未来几个月会逐渐产生热度，成为如同分层学习率一样的通用训练技巧。

关于其他对抗式训练算法的讨论，网上公开的资料很多，我这里就不复述了。

SMART

SMART: Robust and Efficient Fine-Tuning for Pre-trained Natural Language Models through Principled Regularized Optimization
https://arxiv.org/abs/1911.03437

进入正题，SMART 的算法思路与此前基于 FGM、PGD 的一派对抗式算法有些许差异。此前的对抗式训练算法基本没有对目标函数进行修改，仅仅只是对添加扰动后传回的梯度进行处理，而 SMART 则是大刀阔斧地添加了两个正则项。两个正则项涉及到的思想，分别叫做 Smoothness-inducing Adversarial Regularization 和 Bregman Proximal Point Optimization。前者要求模型在一定的扰动范围内，输出完全一致的概率分布；后者要求参数与 epoch 初始时的分布相近。
$\mathcal{R}_s(\theta)=\frac{1}{n}\sum_{i=1}^n\max_{||\tilde{x}_i-x_i||_p\le\epsilon}\mathcal{l}_s(f(\tilde{x}_i;\theta), f(x_i;\theta))$Rs​(θ)=n1​i=1∑n​∣∣x~i​−xi​∣∣p​≤ϵmax​ls​(f(x~i​;θ),f(xi​;θ))

$\mathcal{D}_{Breg}(\theta,\theta_t)=\frac{1}{n}\sum_{i=1}^n\mathcal{l}_s(f(x_i;\theta),f(x_i;\theta_t))$DBreg​(θ,θt​)=n1​i=1∑n​ls​(f(xi​;θ),f(xi​;θt​))

度量是否相似的损失 $\mathcal{l}_s$ls​ 为 对称 KL 散度，相对于 KL 散度，是双向的考量：
$\mathcal{l}_s(P,Q)=\mathcal{D}_{KL}(P||Q)+\mathcal{D}_{KL}(Q||P)$ls​(P,Q)=DKL​(P∣∣Q)+DKL​(Q∣∣P)

在这两项正则项之下，最终的损失函数成为：
$\mathcal{F}(\theta)=\mathcal{L}(\theta)+\lambda_s\mathcal{R}_s (\theta)+\mu\mathcal{D}_{Breg}(\theta,\theta_t)$F(θ)=L(θ)+λs​Rs​(θ)+μDBreg​(θ,θt​)

其中 $\lambda_s$λs​ 和 $\mu$μ 是两个正则项分别对应的权重。最终的优化流程如下：

翻译一下：

对于第 $t$t 轮迭代
 备份 $\theta$θ，作为参与 $\mathcal{D}_{Breg}$DBreg​ 计算的 $\theta_t$θt​
 对于每一个 batch 数据
  使用正态分布随机初始化扰动，结合 $x$x 得到 $\tilde{x}$x~
  对于 $m$m 个小步
   计算扰动下的梯度 $\tilde{g}$g~​
   基于 $\tilde{g}$g~​ 和学习率更新 $\tilde{x}$x~
  基于 $\tilde{x}$x~ 重新计算梯度，更新参数 $\theta$θ
 更新 $\theta_t$θt​，$\theta_t\leftarrow(1-\beta)\theta+\beta\theta_t$θt​←(1−β)θ+βθt​

SMART 在我的各项任务之下都达到了 SOTA 的效果，私以为潜力不可估量。我在使用以下参数配置时达到最好的成绩，供参考：$\lambda_s=0.5$λs​=0.5，$\mu=0.2$μ=0.2，$T_{\tilde{x}}=2$Tx~​=2，$\beta=0.3$β=0.3。

由于代码比较复杂，且必须与我的模型框架结合才能使用，这里就不公开源码了。