整合了以下五大核心技術(shù)：R-Drop、PGM 對(duì)抗訓(xùn)練、EMA、標(biāo)簽平滑、CosineAnnealing 學(xué)習(xí)率調(diào)度。

1. R-Drop（Regularized Dropout）

原理：同一個(gè)樣本做兩次前向傳播（同 dropout mask），計(jì)算兩次輸出的 KL 散度，作為正則項(xiàng)加入損失中。
目標(biāo)：增強(qiáng)魯棒性，提升泛化能力。
損失組合：

loss = CrossEntropy(logits1, labels) + CrossEntropy(logits2, labels) + α * KL(logits1 || logits2)

2. PGM（Projected Gradient Method）對(duì)抗訓(xùn)練

機(jī)制：

在詞嵌入空間中添加擾動(dòng)，制造“敵意樣本”。
多步迭代（PGM_STEPS=3），每步計(jì)算擾動(dòng)梯度并累積。
作用：
增強(qiáng)模型對(duì)小擾動(dòng)的魯棒性，提高對(duì)抗泛化能力。
干預(yù)時(shí)機(jī)：在每次主 loss 反向傳播后注入對(duì)抗 loss 的梯度。

3. EMA（Exponential Moving Average）

思路：
模型參數(shù)滑動(dòng)平均（shadow weights），推理時(shí)使用這些平滑參數(shù)。
核心優(yōu)勢(shì)：
抑制訓(xùn)練波動(dòng)、緩解過(guò)擬合、穩(wěn)定收斂。

4. 標(biāo)簽平滑（Label Smoothing）

方式：將 one-hot 標(biāo)簽略微“平滑”，防止模型過(guò)度自信。
具體值：label_smoothing=0.1
結(jié)果：能緩解過(guò)擬合、提升模型穩(wěn)定性。

5. Cosine Annealing 學(xué)習(xí)率衰減

調(diào)度策略：余弦退火（cosine），帶 warmup。
優(yōu)勢(shì)：
前期快速學(xué)習(xí)，后期逐步收斂；適合 fine-tuning 場(chǎng)景。

模型訓(xùn)練流程

Trainer 子類化：自定義 AdvancedTrainer，重載 compute_loss 以支持雙前向（R-Drop）訓(xùn)練。
Callbacks 集成：
PGMCallback：注入多步對(duì)抗擾動(dòng)。
EmaCallback：更新并應(yīng)用 shadow 參數(shù)。
EarlyStoppingCallback：監(jiān)控 f1，連續(xù) 3 輪無(wú)改進(jìn)則提前停止。

總體優(yōu)勢(shì)

多重正則和魯棒性增強(qiáng)機(jī)制疊加，極大提升模型泛化能力和抗干擾能力。
適合工業(yè)級(jí) NLP 分類任務(wù)的強(qiáng)化訓(xùn)練。

代碼

# Advanced RoBERTa Sentiment Classifier with R-Drop + PGM + EMA + LabelSmoothing + CosineAnnealingimport os
import numpy as np
import torch
import torch.nn as nn
from transformers import (AutoTokenizer,AutoModelForSequenceClassification,Trainer,TrainingArguments,DataCollatorWithPadding,set_seed,EarlyStoppingCallback,TrainerCallback
)
from datasets import load_dataset
from sklearn.metrics import accuracy_score, f1_score, roc_auc_score# 固定隨機(jī)種子
set_seed(42)# 配置參數(shù)
MODEL_NAME = "roberta-base"
NUM_LABELS = 2
R_DROP_ALPHA = 5.0
LABEL_SMOOTHING = 0.1
PGM_EPSILON = 1.0
PGM_ALPHA = 0.3
# --- PGM 多步迭代次數(shù) ---
PGM_STEPS = 3 # 例如，迭代 3 次來(lái)生成對(duì)抗擾動(dòng)
EMA_DECAY = 0.999
# 加載數(shù)據(jù)
dataset = load_dataset("imdb")
train_dataset = dataset["train"]
test_dataset = dataset["test"]# Tokenizer
tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)def preprocess_function(examples):return tokenizer(examples["text"], truncation=True)train_dataset = train_dataset.map(preprocess_function, batched=True)
test_dataset = test_dataset.map(preprocess_function, batched=True)# 數(shù)據(jù)整理器
data_collator = DataCollatorWithPadding(tokenizer=tokenizer)# 加載模型
model = AutoModelForSequenceClassification.from_pretrained(MODEL_NAME, num_labels=NUM_LABELS
)# --- R-Drop Loss ---
class RDropLoss(nn.Module):def __init__(self, alpha=1.0, label_smoothing=0.0):super().__init__()self.alpha = alphaself.ce = nn.CrossEntropyLoss(label_smoothing=label_smoothing)self.kl = nn.KLDivLoss(reduction="batchmean")def forward(self, logits1, logits2, labels):ce_loss1 = self.ce(logits1, labels)ce_loss2 = self.ce(logits2, labels)ce_loss = 0.5 * (ce_loss1 + ce_loss2)p = torch.log_softmax(logits1, dim=-1)q = torch.log_softmax(logits2, dim=-1)p_softmax = torch.softmax(logits1, dim=-1)q_softmax = torch.softmax(logits2, dim=-1)kl_loss = 0.5 * (self.kl(p, q_softmax) + self.kl(q, p_softmax))return ce_loss + self.alpha * kl_loss# --- PGM ---
class PGM:def __init__(self, model, epsilon=1.0, alpha=0.3, emb_name='embeddings.word_embeddings'):self.model = modelself.epsilon = epsilonself.alpha = alphaself.emb_name = emb_nameself.backup = {}
# 擾動(dòng)詞嵌入可以被理解為在原始單詞的語(yǔ)義空間中進(jìn)行微小的“移動(dòng)”，
# 使其略微偏離原來(lái)的意義，但又不至于完全改變其含義，從而模擬“對(duì)抗性樣本”。def attack(self, is_first_attack=False):for name, param in self.model.named_parameters():if param.requires_grad and self.emb_name in name and param.grad is not None:if is_first_attack:self.backup[name] = param.data.clone()norm = torch.norm(param.grad)if norm != 0:r_at = self.alpha * param.grad / normparam.data.add_(r_at)param.data = self.project(name, param.data, self.backup[name])def restore(self):for name, param in self.model.named_parameters():if name in self.backup:param.data = self.backup[name]self.backup = {}def project(self, param_name, param_data, param_backup):r = param_data - param_backupif torch.norm(r) > self.epsilon:r = self.epsilon * r / torch.norm(r)return param_backup + r# --- EMA ---
class EMA:def __init__(self, model, decay):self.model = modelself.decay = decayself.shadow = {}self.backup = {}def register(self):for name, param in self.model.named_parameters():if param.requires_grad:if name not in self.shadow:self.shadow[name] = param.data.clone()def update(self):for name, param in self.model.named_parameters():if param.requires_grad:if name not in self.shadow:continue  # 保護(hù)：skip 未注冊(cè) param，避免 KeyErrornew_avg = (1.0 - self.decay) * param.data + self.decay * self.shadow[name]self.shadow[name] = new_avg.clone()def apply_shadow(self):for name, param in self.model.named_parameters():if param.requires_grad:if name not in self.shadow:continueself.backup[name] = param.data.clone()param.data = self.shadow[name]def restore(self):for name, param in self.model.named_parameters():if param.requires_grad:if name in self.backup:param.data = self.backup[name]self.backup = {}# --- Callbacks ---
class PGMCallback(TrainerCallback):def __init__(self, pgm, rdrop_loss_fn, pgm_steps=1):self.pgm = pgmself.rdrop_loss_fn = rdrop_loss_fnself.pgm_steps = pgm_steps # 對(duì)抗迭代步數(shù)def on_after_backward(self, args, state, control, model=None, inputs=None, optimizer=None, **kwargs):# 備份原始的梯度狀態(tài)，以便在對(duì)抗訓(xùn)練結(jié)束后恢復(fù)# PyTorch 優(yōu)化器在 step() 時(shí)會(huì)清零梯度，但我們需要在 PGM 內(nèi)部操作時(shí)保留它們# 更安全的做法是使用一個(gè)更細(xì)粒度的梯度累積或在每次 PGM 迭代后清零# 為了簡(jiǎn)化，我們假設(shè)在這個(gè)回調(diào)中梯度是可用的# 開始多步 PGMfor step in range(self.pgm_steps):is_first_attack = (step == 0)# 在第一次攻擊時(shí)備份參數(shù)并施加擾動(dòng)# 在后續(xù)攻擊時(shí)，只施加擾動(dòng)，不備份self.pgm.attack(is_first_attack=is_first_attack)# 在擾動(dòng)后的模型上進(jìn)行前向傳播# 這里需要確保模型處于訓(xùn)練模式，并且梯度是開啟的model.train() # 確保模型處于訓(xùn)練模式model.zero_grad() # 在每次PGM步清零模型梯度adv_outputs1 = model(**{k: v for k, v in inputs.items() if k != "labels"})adv_outputs2 = model(**{k: v for k, v in inputs.items() if k != "labels"})adv_logits1 = adv_outputs1.logitsadv_logits2 = adv_outputs2.logitslabels = inputs["labels"]# 計(jì)算對(duì)抗損失adv_loss = self.rdrop_loss_fn(adv_logits1, adv_logits2, labels)# 對(duì)抗損失的反向傳播# 注意：這里不能簡(jiǎn)單地直接調(diào)用 .backward()# 因?yàn)?Trainer 已經(jīng)處理了主損失的梯度累積和優(yōu)化器步驟# PGM 的梯度應(yīng)該累積到主梯度中，而不是覆蓋它們# 最簡(jiǎn)單的集成方式是讓對(duì)抗損失也產(chǎn)生梯度，并累積到參數(shù)上# 在 Hugging Face Trainer 的 on_after_backward 中，# 已經(jīng)進(jìn)行了一次主損失的 backward，因此這里的 adv_loss.backward() 會(huì)累積梯度。# 但是，為了避免在多步中梯度累積不當(dāng)，需要更細(xì)致的控制。# 通常，PGM 是在優(yōu)化器步驟之前，對(duì)參數(shù)進(jìn)行修改并重新計(jì)算損失。# --- 關(guān)鍵點(diǎn)：如何處理多步 PGM 的梯度 ---# 這里的 `adv_loss.backward()` 會(huì)計(jì)算并累積梯度。# 由于每次 `pgm.attack()` 都會(huì)修改參數(shù)，所以 `adv_loss` 都會(huì)基于當(dāng)前擾動(dòng)后的參數(shù)計(jì)算。# 在每次 `step` 中，我們計(jì)算 `adv_loss` 的梯度并累加到參數(shù)上。# 注意：`model.zero_grad()` 放在循環(huán)內(nèi)部可以確保每次 PGM 步只計(jì)算當(dāng)前擾動(dòng)下的梯度，# 如果放在循環(huán)外部，則所有 PGM 步的梯度會(huì)累積到同一個(gè)梯度值上。# 這里設(shè)置為每次 PGM 步清零梯度，然后計(jì)算當(dāng)前步的對(duì)抗梯度。# 這樣做可以確保 `adv_loss.backward()` 每次計(jì)算的是相對(duì)于當(dāng)前擾動(dòng)參數(shù)的梯度。accelerator = kwargs.get("accelerator", None)if accelerator is not None:accelerator.backward(adv_loss)else:adv_loss.backward()# 多步 PGM 結(jié)束后，恢復(fù)模型參數(shù)到原始狀態(tài)（即未被 PGM 擾動(dòng)前的狀態(tài)）self.pgm.restore()# 此時(shí)，model 的所有 param.grad 中已經(jīng)包含了# (主損失的梯度) + (最后一次 PGM 迭代的對(duì)抗損失的梯度)# HuggingFace Trainer 會(huì)緊接著調(diào)用優(yōu)化器的 step() 方法來(lái)更新模型的參數(shù)。#最終留下并用于優(yōu)化器更新的梯度，是最后一次 PGM 迭代所產(chǎn)生的對(duì)抗損失的梯度。class EmaCallback(TrainerCallback):def __init__(self, ema):self.ema = emadef on_step_end(self, args, state, control, **kwargs):self.ema.update()def on_evaluate(self, args, state, control, **kwargs):self.ema.apply_shadow()def on_evaluate_end(self, args, state, control, **kwargs):self.ema.restore()# --- AdvancedTrainer  ---
class AdvancedTrainer(Trainer):def __init__(self, *args, alpha=1.0, label_smoothing=0.0, ema=None, **kwargs):super().__init__(*args, **kwargs)self.rdrop_loss_fn = RDropLoss(alpha=alpha, label_smoothing=label_smoothing)self.ema = emaif self.ema is not None:self.ema.register()def compute_loss(self, model, inputs, return_outputs=False, **kwargs):labels = inputs["labels"]# 兩次前向傳播用于 R-Dropoutputs1 = model(**{k: v for k, v in inputs.items() if k != "labels"})outputs2 = model(**{k: v for k, v in inputs.items() if k != "labels"})logits1 = outputs1.logitslogits2 = outputs2.logitsloss = self.rdrop_loss_fn(logits1, logits2, labels)return (loss, outputs1) if return_outputs else loss# --- Metrics ---
def compute_metrics(eval_pred):logits, labels = eval_predprobs = torch.softmax(torch.tensor(logits), dim=-1).numpy()predictions = np.argmax(logits, axis=-1)acc = accuracy_score(labels, predictions)f1 = f1_score(labels, predictions)try:auc = roc_auc_score(labels, probs[:, 1])except:auc = 0.0return {"accuracy": acc, "f1": f1, "auc": auc}# --- TrainingArguments  ---
training_args = TrainingArguments(output_dir="./results_adv_rdrop_pgm_ema_multistep", # 更改輸出目錄eval_strategy="epoch",save_strategy="epoch",learning_rate=2e-5,per_device_train_batch_size=16,per_device_eval_batch_size=16,num_train_epochs=5,weight_decay=0.01,warmup_ratio=0.1,lr_scheduler_type="cosine",logging_dir="./logs_adv_rdrop_pgm_ema_multistep", # 更改日志目錄logging_steps=50,load_best_model_at_end=True,metric_for_best_model="f1",fp16=True,save_total_limit=2,
)# --- 初始化模塊 ---
pgm = PGM(model, epsilon=PGM_EPSILON, alpha=PGM_ALPHA)
ema = EMA(model, decay=EMA_DECAY)# --- Trainer ---
trainer = AdvancedTrainer(model=model,args=training_args,train_dataset=train_dataset,eval_dataset=test_dataset,tokenizer=tokenizer, # 使用 tokenizer 而不是 processing_classdata_collator=data_collator,compute_metrics=compute_metrics,alpha=R_DROP_ALPHA,label_smoothing=LABEL_SMOOTHING,callbacks=[PGMCallback(pgm=pgm, rdrop_loss_fn=RDropLoss(alpha=R_DROP_ALPHA, label_smoothing=LABEL_SMOOTHING), pgm_steps=PGM_STEPS),EmaCallback(ema=ema),EarlyStoppingCallback(early_stopping_patience=3, early_stopping_threshold=0.01),],
)# --- 訓(xùn)練 ---
trainer.train()# --- 評(píng)估 ---
trainer.evaluate()

結(jié)果

{'eval_loss': 0.2900645434856415, 'eval_accuracy': 0.95836, 'eval_f1': 0.9586822782298074, 'eval_auc': 0.9911689504000001, 'eval_runtime': 275.0978, 'eval_samples_per_second': 90.877, 'eval_steps_per_second': 5.682, 'epoch': 5.0}                                                                                                             
{'train_runtime': 171019.0699, 'train_samples_per_second': 0.731, 'train_steps_per_second': 0.046, 'train_loss': 0.30841634256749756, 'epoch': 5.0}                      
100%|█████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 7815/7815 [47:30:19<00:00, 21.88s/it]
100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 1563/1563 [04:32<00:00,  5.75it/s]

TensorBoard

注意

實(shí)際操作，這里要保存模型。還要轉(zhuǎn)成ONNX模型，用C++ OnnxRuntime推理等等推理。