对抗训练fgm、fgsm和pgd原理和源码分析

当前，在各大NLP竞赛中，对抗训练已然成为上分神器，尤其是fgm和pgd使用较多，下面来说说吧。对抗训练是一种引入噪声的训练方式，可以对参数进行正则化，提升模型鲁棒性和泛化能力。

fgm

FGM的全称是Fast Gradient Method, 出现于Adversarial Training Methods for Semi-supervised Text Classification这篇论文，FGM是根据具体的梯度进行scale，得到更好的对抗样本：

整个对抗训练的过程如下，伪代码如下：

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

fgm代码实现如下：

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 = {}

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26

fgm应用代码如下：

##对应第一步
loss = model(**batch_data)[0]	
loss.backward()
##对应第二步
fgm.attack()
#对应第三步
loss_adv = model(**batch_data)[0]
loss_adv.backward()
#对应第四步
fgm.restore()
#对应第五步
optimizer.step()
1
2
3
4
5
6
7
8
9
10
11
12

fgsm

FGSM的全称是Fast Gradient Sign Method. 如果要说FGSM和FGM的区别，核心区别在计算扰动的方式不一样，FGSM扰动的计算方式如下:

FGSM的其他算法流程跟FGM一样，这里不再赘述。

pgd

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

FGM简单粗暴的“一步到位”，可能走不到约束内的最优点。PGD则是“小步走，多走几步”，如果走出了扰动半径为epsilon的空间，就映射回“球面”上，以保证扰动不要过大

并且

pgd整个对抗训练的过程如下，伪代码如下：

• 1.计算x的前向loss、反向传播得到梯度并备份；
• 2.对于每步t:

•  a.根据embedding矩阵的梯度计算出r，并加到当前embedding上，相当于x+r(超出范围则投影回epsilon内)；
  1

•  if t 不是最后一步，则进行b步骤：将模型梯度归0，根据a的x+r计算前后向并得到梯度,继续a步骤；if t 是最后一步,则进行c步骤：恢复(1)的梯度，根据a的x+r计算前后向得到梯度并将梯度累加到(1)的梯度上，跳出循环；
  1

• 3.将embedding恢复为(1)时的值；
• 4.根据2c的梯度对参数进行更新。

可以看到，在循环中r是逐渐累加的，要注意的是最后更新参数只使用最后一个x+r算出来的梯度。
pgd代码实现如下：

class PGD:
    def __init__(self, model, eps=1., alpha=0.3):
        self.model = (
            model.module if hasattr(model, "module") else model
        )
        self.eps = eps
        self.alpha = alpha
        self.emb_backup = {}
        self.grad_backup = {}

    def attack(self, 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 = self.alpha * param.grad / norm
                    param.data.add_(r_at)
                    param.data = self.project(name, param.data)

    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):
        r = param_data - self.emb_backup[param_name]
        if torch.norm(r) > self.eps:
            r = self.eps * 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 and param.grad is not None:
                self.grad_backup[name] = param.grad.clone()

    def restore_grad(self):
        for name, param in self.model.named_parameters():
            if param.requires_grad and param.grad is not None:
                param.grad = self.grad_backup[name]

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44

pgd应用代码如下：

loss = model(**batch_data)[0]	
loss.backward()
pgd.backup_grad()
for _t in range(pgd_k):
	pgd.attack(is_first_attack=(_t == 0))
	if _t != pgd_k - 1:
		model.zero_grad()
	else:
		pgd.restore_grad()
	loss_adv = model(**batch_data)[0]
	loss_adv.backward()
pgd.restore()
optimizer.step()
1
2
3
4
5
6
7
8
9
10
11
12
13

注：在torch中，每次迭代时，如果不把模型的梯度清零，会默认将模型每次迭代的梯度累加的。

FreeAT

FreeAT (Free Adversarial Training):来源于NIPS2019的一篇论文，从FGSM到PGD，主要是优化对抗扰动的计算，但计算量也一步步增加。对于每个样本，FGSM和FGM都只用计算两次，一次是计算x的前后向，一次是计算x+r的前后向。而PGD则计算了K+1次，消耗了更多的计算资源。因此FreeAT被提了出来，在PGD的基础上进行训练速度的优化。

FreeAT的思想是在对每个样本x连续重复m次训练，计算r时复用上一步的梯度，为了保证速度，整体epoch会除以m。r的更新公式为：

整个对抗训练的过程如下，伪代码如下：

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

以上所述的对抗训练方法在不同的训练数据上表现大同小异，需要根据具体场景和具体数据去选择对应的算法，没有最好的，只有当下场景最适合的算法。
详细代码实现参见github代码实现: 对抗训练.