【深度学习】NLP中的对抗训练

在NLP中，对抗训练往往都是针对嵌入层（包括词嵌入，位置嵌入，segment嵌入等等）开展的，思想很简单，即针对嵌入层添加干扰，从而提高模型的鲁棒性和泛化能力，下面结合具体代码讲解一些NLP中常见对抗训练算法。

1.Fast Gradient Method（FGM）

FGM的思想是针对词嵌入加入梯度方向的干扰，至于干扰的大小是我们可以调节的，增加干扰后的样本可以作为额外的对抗样本进行训练，以此提高模型的效果。由于我们在训练时会针对每个样本都进行一次额外的增加干扰后的训练，所以使用FGM后训练时间理论上也会大概增加一倍。

FGM在原训练代码的基础上，主要增加了以下几个额外的操作：针对嵌入层添加干扰并备份参数，计算添加干扰后的损失，梯度回传从而累积添加干扰后的梯度，恢复原来的嵌入层参数。

1.1 算法流程

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

1.2 具体代码

import torch
class FGM():
    def __init__(self, model):
        self.model = model
        self.backup = {}

    def attack(self, epsilon=1., emb_name='word_embeddings'):
        # emb_name这个参数要换成你模型中embedding的参数名
        for name, param in self.model.named_parameters():
            if param.requires_grad and emb_name in name:
                #print('增加扰动的对象是', name)
                #print(type(param.grad))
                self.backup[name] = param.data.clone()
                norm = torch.norm(param.grad)
                if norm != 0 and not torch.isnan(norm):
                    r_at = epsilon * param.grad / norm
                    param.data.add_(r_at)

    def restore(self, emb_name='word_embeddings'):
        # emb_name这个参数要换成你模型中embedding的参数名
        for name, param in self.model.named_parameters():
            if param.requires_grad and emb_name in name: 
                assert name in self.backup
                param.data = self.backup[name]
        self.backup = {}

1.3 具体用法

fgm = FGM(model) # （#1）初始化
for batch_input, batch_label in data:
    loss = model(batch_input, batch_label) # 正常训练
    loss.backward() # 反向传播，得到正常的grad
    # 对抗训练
    fgm.attack() # （#2）在embedding上添加对抗扰动
    loss_adv = model(batch_input, batch_label) # （#3）计算含有扰动的对抗样本的loss
    loss_adv.backward() # （#4）反向传播，并在正常的grad基础上，累加对抗训练的梯度
    fgm.restore() # （#5）恢复embedding参数
    # 梯度下降，更新参数
    optimizer.step()
    model.zero_grad()

2.Projected Gradient Descent （PGD）

Project Gradient Descent（PGD）是一种迭代攻击算法，相比于普通的FGM 仅做一次迭代，PGD是做多次迭代，每次走一小步，每次迭代都会将扰动投射到规定范围内。其中r为扰动约束空间（一个半径为r的球体），原始的输入样本对应的初识点为球心，避免扰动超过球面。迭代多次后，保证扰动在一定范围内，如下图所示：

2.1 算法流程

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

2.2 具体代码

import torch
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='word_embeddings', is_first_attack=False):
        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='word_embeddings'):
        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.3 具体用法

pgd = PGD(model)
K = 3
for batch_input, batch_label in data:
    # 正常训练
    loss = model(batch_input, batch_label)
    loss.backward() # 反向传播，得到正常的grad
    pgd.backup_grad()
    # 累积多次对抗训练——每次生成对抗样本后，进行一次对抗训练，并不断累积梯度
    for t in range(K):
        pgd.attack(is_first_attack=(t==0)) # 在embedding上添加对抗扰动, first attack时备份param.data
        if t != K-1:
            model.zero_grad()
        else:
            pgd.restore_grad()
        loss_adv = model(batch_input, batch_label)
        loss_adv.backward() # 反向传播，并在正常的grad基础上，累加对抗训练的梯度
    pgd.restore() # 恢复embedding参数
    # 梯度下降，更新参数
    optimizer.step()
    model.zero_grad()