自然言語処理分野で用いられる敵対的学習手法について

# reference: https://www.kaggle.com/c/tweet-sentiment-extraction/discussion/143764
class FGM():
    def __init__(self, model):
        self.model = model
        self.backup = {}

    def attack(self, epsilon=1., emb_name='word_embeddings'):
　　　　　　　"""
        敵対的な摂動を求め、現在のembedding layerに摂動を加える
        """
        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 != 0:
                    r_at = epsilon * param.grad / norm
                    param.data.add_(r_at)

    def restore(self, emb_name='word_embeddings'):
        """
        敵対的な摂動を求める際に変更してしまったembedding layerのパラメータについて
        元のパラメータを代入する
        """
        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 = {}

fgm = FGM(model)
 for batch_input, batch_label in data:
       # オリジナルのサンプルについての損失を計算
       loss = model(batch_input, batch_label)
       loss.backward()  

       # adversarial training
       # embedding layerに敵対的な摂動を加える
       fgm.attack() 
       # 敵対的な摂動を加えられた状態での損失を計算
       loss_adv = model(batch_input, batch_label)
       loss_adv.backward() 
       fgm.restore()  

       optimizer.step()
       model.zero_grad()

# reference: https://www.kaggle.com/code/wht1996/feedback-nn-train/notebook
class AWP:
    def __init__(
        self,
        model,
        optimizer,
        adv_param="weight",
        adv_lr=1,
        adv_eps=0.2,
        start_epoch=0,
        adv_step=1,
        scaler=None
    ):
        self.model = model
        self.optimizer = optimizer
        self.adv_param = adv_param
        self.adv_lr = adv_lr
        self.adv_eps = adv_eps
        self.start_epoch = start_epoch
        self.adv_step = adv_step
        self.backup = {}
        self.backup_eps = {}
        self.scaler = scaler

    def attack_backward(self, x, y, attention_mask,epoch):
    　　　　"""
        敵対的な摂動を加えた損失を計算し、パラメータを更新する
        """
        if (self.adv_lr == 0) or (epoch < self.start_epoch):
            return None

        self._save() 
        for i in range(self.adv_step):
            self._attack_step() 
            with torch.cuda.amp.autocast():
                adv_loss, tr_logits = self.model(input_ids=x, attention_mask=attention_mask, labels=y)
                adv_loss = adv_loss.mean()
            self.optimizer.zero_grad()
            self.scaler.scale(adv_loss).backward()
            
        self._restore()

    def _attack_step(self):
        """
        敵対的な摂動を求め、重みに加える
        重みの範囲をbackup_epsで制限している
        """
        e = 1e-6
        for name, param in self.model.named_parameters():
            if param.requires_grad and param.grad is not None and self.adv_param in name:
                norm1 = torch.norm(param.grad)
                norm2 = torch.norm(param.data.detach())
                if norm1 != 0 and not torch.isnan(norm1):
                    r_at = self.adv_lr * param.grad / (norm1 + e) * (norm2 + e)
                    param.data.add_(r_at)
                    param.data = torch.min(
                        torch.max(param.data, self.backup_eps[name][0]), self.backup_eps[name][1]
                    )

    def _save(self):
        """
        重みのバックアップと、重みの範囲を取得する
        重みの範囲はパラメータの絶対値とadv_epsによって決定する
        """
        for name, param in self.model.named_parameters():
            if param.requires_grad and param.grad is not None and self.adv_param in name:
                if name not in self.backup:
                    self.backup[name] = param.data.clone()
                    grad_eps = self.adv_eps * param.abs().detach()
                    self.backup_eps[name] = (
                        self.backup[name] - grad_eps,
                        self.backup[name] + grad_eps,
                    )

    def _restore(self):
        """
        バックアップを取っていたパラメータを代入するとともに初期化する
        """
        for name, param in self.model.named_parameters():
            if name in self.backup:
                param.data = self.backup[name]
        self.backup = {}
        self.backup_eps = {}