一文详解对抗训练方法

2022-07-01 本文已影响0人晓柒NLP与药物设计

对抗训练方法

Adversarial learning主要是用于样本生成或者对抗攻击领域，主要方法是通过添加鉴别器或者根据梯度回传生成新样本，其主要是为了提升当前主干模型生成样本的能力或者鲁棒性

一. 对抗训练定义

==对抗训练是一种引入噪声的训练方式，可以对参数进行正则化，提升模型鲁棒性和泛化能力==

1.1 对抗训练特点

相对于原始输入，所添加的扰动是微小的
添加的噪声可以使得模型预测错误

1.2 对抗训练的基本概念

就是在原始输入样本 $x$ 上加上一个扰动 $\Delta x$ 得到对抗样本，再用其进行训练，这个问题可以抽象成这样一个模型：
$\max _{\theta} P(y \mid x+\Delta x ; \theta)\tag{1}$
其中， $y$ 是ground truth, $\theta$ 是模型参数。意思就是即使在扰动的情况下求使得预测出 $y$ 的概率最大的参数，扰动可以被定义为：
$\Delta x=ε \cdot \operatorname{sign}\left(\nabla_{x} L(x, y ; \theta)\right)\tag{2}$
其中， $sign$ 为符号函数， $L$ 为损失函数

最后，GoodFellow还总结了对抗训练的两个作用：

提高模型应对恶意对抗样本时的鲁棒性
作为一种regularization，减少overfitting，提高泛化能力

1.3 Min-Max公式

Madry在2018年的ICLR论文Towards Deep Learning Models Resistant to Adversarial Attacks中总结了之前的工作，对抗训练可以统一写成如下格式：
$\min _{\theta} \mathbb{E}_{(x, y) \sim \mathcal{D}}\left[\max _{\Delta x \in \Omega} L(x+\Delta x, y ; \theta)\right]\tag{3}$
其中 $\mathcal{D}$ 代表输入样本的分布， $x$ 代表输入， $y$ 代表标签， $\theta$ 是模型参数， $L(x+y; \theta)$ 是单个样本的loss， $\Delta x$ 是扰动， $\Omega$ 是扰动空间。这个式子可以分布理解如下：

内部max是指往 $x$ 中添加扰动 $\Delta x$ ， $\Delta x$ 的目的是让 $L(x+\Delta x, y ; \theta)$ 越大越好，也就是说尽可能让现有模型预测出错。但是， $\Delta x$ 也是有约束的，要在 $\Omega$ 范围内. 常规的约束是 $|| \Delta x|| \leq ε$ ，其中 $ε$ 是一个常数
外部min是指找到最鲁棒的参数 $\theta$ 是预测的分布符合原数据集的分布

这就解决了两个问题：如何构建足够强的对抗样本、和如何使得分布仍然尽可能接近原始分布

1.4 NLP领域的对抗训练

对于CV领域，图像被认为是连续的，因此可以直接在原始图像上添加扰动；而对于NLP，它的输入是文本的本质是one-hot，而one-hot之间的欧式距离恒为 $\sqrt{2}$ ，理论上不存在微小的扰动，而且，在Embedding向量上加上微小扰动可能就找不到与之对应的词了，不是真正意义上的对抗样本，因为对抗样本依旧能对应一个合理的原始输入，既然不能对Embedding向量添加扰动，可以对Embedding层添加扰动，使其产生更鲁棒的Embedding向量

二. 对抗训练方法

2.1 FGM(Fast Gradient Method) ICLR2017

FGM是根据具体的梯度进行scale，得到更好的对抗样本：
$r_{adv}=εg/\|g\|_2\tag{4}$
整个对抗训练的过程如下，伪代码如下：

计算x的前向loss、反向传播得到梯度
根据embedding矩阵的梯度计算出 $r$ ，并加到当前embedding上，相当于 $x+r$
计算 $x+r$ 的前向loss，反向传播得到对抗的梯度，累加到(1)的梯度上
将embedding恢复为(1)时的值
根据(3)的梯度对参数进行更新

class FGM:
    def __init__(self, model: nn.Module, eps=1.):
        self.model = (model.module if hasattr(model, "module") else model)
        self.eps = eps
        self.backup = {}
    # only attack word embedding
    def attack(self, emb_name='word_embeddings'):
        for name, param in self.model.named_parameters():
            if param.requires_grad and emb_name in name:
                self.backup[name] = param.data.clone()
                norm = torch.norm(param.grad)
                if norm and not torch.isnan(norm):
                    r_at = self.eps * param.grad / norm
                    param.data.add_(r_at)
    def restore(self, emb_name='word_embeddings'):
        for name, para in self.model.named_parameters():
            if para.requires_grad and emb_name in name:
                assert name in self.backup
                para.data = self.backup[name]
        self.backup = {}

2.2 FGSM (Fast Gradient Sign Method) ICLR2015

FGSM的全称是Fast Gradient Sign Method. FGSM和FGM的核心区别在计算扰动的方式不一样，FGSM扰动的计算方式如下:
$r_{adv}=ε \cdot \operatorname{sign}\left(\nabla_{x} L(x, y ; \theta)\right)\tag{5}$

def FGSM(image, epsilon, data_grad):
    """
    :param image: 需要攻击的图像
    :param epsilon: 扰动值的范围
    :param data_grad: 图像的梯度
    :return: 扰动后的图像
    """
    # 收集数据梯度的元素符号
    sign_data_grad = data_grad.sign()
    # 通过调整输入图像的每个像素来创建扰动图像
    perturbed_image = image + epsilon*sign_data_grad
    # 添加剪切以维持[0,1]范围
    perturbed_image = torch.clamp(perturbed_image, 0, 1)
    # 返回被扰动的图像
    return perturbed_image

2.3 PGD(Projected Gradient Descent)

FGM直接通过epsilon参数算出了对抗扰动，这样得到的可能不是最优的。因此PGD进行了改进，通过迭代慢慢找到最优的扰动

$r_{adv|t+1}=\alpha g_t/\|g_t\|_2\tag{6}$

并且 $\|r\|_2≤ε$

PGD整个对抗训练的过程如下

计算 $x$ 的前向loss、反向传播得到梯度并备份
对于每步 $t$ :
1. 根据embedding矩阵的梯度计算出r，并加到当前embedding上，相当于 $x+r$ (超出范围则投影回epsilon内)
2. if t不是最后一步: 将梯度归0，根据(1)的 $x+r$ 计算前后向并得到梯度
3. if t是最后一步: 恢复(1)的梯度，计算最后的 $x+r$ 并将梯度累加到(1)上
将embedding恢复为(1)时的值
根据(5)的梯度对参数进行更新

在循环中 $r$ 是逐渐累加的，要注意的是最后更新参数只使用最后一个 $x+r$ 算出来的梯度

class PGD():
    def __init__(self, model):
        self.model = model
        self.emb_backup = {}
        self.grad_backup = {}
    def attack(self, epsilon=1., alpha=0.3, emb_name='emb.', is_first_attack=False):
        # emb_name这个参数要换成你模型中embedding的参数名
        for name, param in self.model.named_parameters():
            if param.requires_grad and emb_name in name:
                if is_first_attack:
                    self.emb_backup[name] = param.data.clone()
                norm = torch.norm(param.grad)
                if norm != 0 and not torch.isnan(norm):
                    r_at = alpha * param.grad / norm
                    param.data.add_(r_at)
                    param.data = self.project(name, param.data, epsilon)
    def restore(self, emb_name='emb.'):
        # emb_name这个参数要换成你模型中embedding的参数名
        for name, param in self.model.named_parameters():
            if param.requires_grad and emb_name in name: 
                assert name in self.emb_backup
                param.data = self.emb_backup[name]
        self.emb_backup = {}
    def project(self, param_name, param_data, epsilon):
        r = param_data - self.emb_backup[param_name]
        if torch.norm(r) > epsilon:
            r = epsilon * r / torch.norm(r)
        return self.emb_backup[param_name] + r
    def backup_grad(self):
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                self.grad_backup[name] = param.grad.clone()
    def restore_grad(self):
        for name, param in self.model.named_parameters():
            if param.requires_grad:
                param.grad = self.grad_backup[name]

2.4 FreeAT(Free Adversarial Training)

从FGSM到PGD，主要是优化对抗扰动的计算，虽然取得了更好的效果，但计算量也一步步增加。对于每个样本，FGSM和FGM都只用计算两次，一次是计算 $x$ 的前后向，一次是计算 $x+r$ 的前后向。而PGD则计算了K+1次，消耗了更多的计算资源。因此FreeAT被提了出来，在PGD的基础上进行训练速度的优化

FreeAT的思想是在对每个样本 $x$ 连续重复 $m$ 次训练，计算 $r$ 时复用上一步的梯度，为了保证速度，整体epoch会除以 $m$ 。 $r$ 的更新公式为：
$r_{t+1}=r_t+ε \cdot sign(g)\tag{7}$
FreeAT的训练过程如下：

初始化 $r=0$
对于epoch=:
1. 对于每个:
  1. 对于每步:
    1. 利用上一步的 $r$ ，计算 $x+r$ 的前后向，得到梯度
    2. 根据梯度更新参数
    3. 根据梯度更新 $r$

FreeAT的问题在于每次的 $r$ 对于当前的参数都是次优的(无法最大化loss)，因为当前 $r$ 是由 $r_{t-1}$ 和 $\theta_{t-1}$ 计算出来的，是对于 $\theta_{t-1}$ 的最优

2.5 YOPO(You Only Propagate Once)

YOPO的出发点是利用神经网络的结构来降低梯度计算的计算量。从极大值原理PMP(Pontryagin’s maximum principle)出发，对抗扰动只和网络的第0层有关，即在embedding层上添加扰动。再加之层之间是解耦合的，那就不需要每次都计算完整的前后向传播

基于这个想法，复用后面几层的梯度，减少非必要的完整传播。可以将PGD的 $r$ 次攻击拆成 $m\times n$ 次：
$p=\nabla_{g_\tilde \theta}(l(g_\tilde \theta(f_0(x_i+r_i^{j,0},\theta_0)),y_i))\cdot \nabla_{f_0}(g_\tilde \theta(f_0(x_i+r_i^{j,0},\theta_0)))\tag{8}$
则对r的更新就可以变为:
$r_i^{j,s+1}=r_i^{j,s}+\alpha_1\cdot\nabla_{r_i}f_0(x_i+r_i^{j,s},\theta_0)\tag{9}$

其算法流程为：

对于每个样本 $x$ ，初始化 $r(1,0)$ ，对于 $j=1,2,…,m$ :

根据 $r(j,0)$ ,计算 $p$ 对于 $s=0,1,…,n-1$ :
计算 $r(j,s+1)$
另 $r(j+1,0)=r(j,n)$

2.6 FreeLB (Free Large-Batch)

YOPO的假设对于ReLU-based网络来说是不成立的，因为YOPO要求损失是两次可微的，于是，FreeLB在FreeAT的基础上将每次inner-max中更新模型参数这一操作换掉，利用 $K$ 步之后累积的参数梯度进行更新，于是总体任务的目标函数就记为：
$\underset{\theta}{min}\mathbb E_{(Z,y)\sim \mathcal D}\left[\frac{1}{K}\sum_{t=0}^{K-1}\underset{\delta_t\in\mathcal I_t}{max}\ L(f_\theta(X+\delta_t),y)\right]\\ \mathcal I_t=\mathcal B_{X+\delta_0}(\alpha t)\cap\mathcal B_X(\epsilon)\tag{10}$
$X+\delta_t$ 可以看成两个球形邻域的交上局部最大的近似。同时，通过累积参数梯度的操作，可以看作是输入了 $[X+\delta_0,\cdots,X+\delta_{K-1}]$ 这样一个虚拟的 $K$ 倍大小的batch。其中input subwords的one-hot representations记为 $Z$ ，embedding matrix记为 $V$ ，subwords embedding记为 $X = V Z$

依据下面算法中的数学符号，PGD需要进行 $N_{ep}\cdot(K+1)$ 次梯度计算，FreeAT需要进行 $N_{ep}$ 次，FreeLB需要 $N_{ep}\cdot K$ 次。虽然FreeLB在效率上并没有特别大的优势，但是其效果十分不错

另外，论文中指出对抗训练和dropout不能同时使用，加上dropout相当于改变了网络的结果，影响扰动的计算。如果一定要加入dropout操作，需要在K步中都使用同一个mask

2.7 SMART(SMoothness-inducing Adversarial Regularization)

SMART放弃了Min-Max公式，选择通过正则项Smoothness-inducing Adversarial Regularization完成对抗学习。为了解决这个新的目标函数作者又提出了优化算法Bregman Proximal Point Optimization，这就是SMART的两个主要内容

SMART的主要想法是强制模型在neighboring data points上作出相似的预测，加入正则项后的目标函数如下所示：
$\underset{\theta}{min}\ \mathcal F(\theta)=\mathcal L(\theta)+\lambda_s\mathcal R_s(\theta))\\ \mathcal L(\theta)=\frac{1}{n}\sum_{i=1}^{n}\ell\left(f(x_i;\theta),y_i\right)\\ \mathcal R_s(\theta)=\frac{1}{n}\sum_{i=1}^{n}\underset{||\tilde x_i-x_i||_p\leq\epsilon}{max}\ \ell_s\left[f(\tilde x_i;\theta),f(x_i;\theta)\right]\tag{11}$

$\ell$ 是具体任务的损失函数， $\tilde x_i$ 是generated neighbors of training points， $\ell_s$ 在分类任务中使用对称的KL散度，即 $\ell_s(P,Q)=\mathcal D_{KL}(P||Q)+D_{KL}(Q||L)$ ；在回归任务中使用平方损失， $\ell_s(p,q)=(p-q)^2$ 此时可以看到对抗发生在正则化项上，对抗的目标是最大扰动前后的输出

Bregman Proximal Point Optimization也可以看作是一个正则项，防止更新的时候 $\theta_{t+1}$ 和前面的 $\theta_t$ 变化过大。在第 $t+1$ 次迭代时，采用vanilla Bregman proximal point (VBPP) method：
$\theta_{t+1}=argmin_{\theta}\mathcal F(\theta)+\mu\mathcal D_{Breg}(\theta,\theta_t)\tag{12}$
其中 $\mathcal D_{Breg}$ 表示Bregman divergence定义为：
$\mathcal D_{Breg}(\theta,\theta_t)=\frac{1}{n}\sum_{i=1}^n\ell_s\left(f(x_i;\theta),f(x_i;\theta_t)\right)\tag{13}$
$\ell_s$ 是上面给出的对称KL散度

使用动量来加速VBPP，此时定义 $\beta$ 为动量，记 $\tilde\theta=(1-\beta)\theta_t+\beta\tilde\theta_{t-1}$ 表示指数移动平均，那么momentum Bregman proximal point (MBPP) method就可以表示为：
$\theta_{t+1}=argmin_{\theta}\mathcal F(\theta)+\mu\mathcal D_{Breg}(\theta,\tilde\theta_t)\tag{14}$
下面是SMART的完整算法流程：

对于轮迭代：
1. 备份 $theta$ ，作为Bregman divergence计算的 $\theta_t$
2. 对于每一个：
  1. 使用正态分布随机初始化扰动，结合 $x$ 得到 $x\_tilde$
  2. 循环小步：计
    1. 算扰动下的梯度 $g\_tilde$
    2. 基于 $g\_tilde$ 和学习率更新 $x\_tilde$
  3. 基于 $x\_tilde$ 重新计算梯度，更新参数 $\theta$
3. 更新 $\theta_t$

三. Reference

Madry A, Makelov A, Schmidt L, et al. Towards deep learning models resistant to adversarial attacks[J]. arXiv preprint arXiv:1706.06083, 2017.
Goodfellow I J, Shlens J, Szegedy C. Explaining and harnessing adversarial examples[J]. arXiv preprint arXiv:1412.6572, 2014.
Miyato T, Dai A M, Goodfellow I. Adversarial training methods for semi-supervised text classification[J]. arXiv preprint arXiv:1605.07725, 2016.
Shafahi A, Najibi M, Ghiasi A, et al. Adversarial training for free![J]. arXiv preprint arXiv:1904.12843, 2019.
Zhang D, Zhang T, Lu Y, et al. You only propagate once: Accelerating adversarial training via maximal principle[J]. arXiv preprint arXiv:1905.00877, 2019.
Zhu C, Cheng Y, Gan Z, et al. Freelb: Enhanced adversarial training for natural language understanding[J]. arXiv preprint arXiv:1909.11764, 2019.
Jiang H, He P, Chen W, et al. Smart: Robust and efficient fine-tuning for pre-trained natural language models through principled regularized optimization[J]. arXiv preprint arXiv:1911.03437, 2019.