关于AI模型鲁棒性的几点记录
「模型在测试集上 92%,上线应该稳了吧?」
去年 10 月,朋友把图像分类模型推上生产。两天后群里炸锅:有人用精心构造的对抗样本喂图,准确率从 92% 掉到 3%,连猫狗都开始瞎猜。那之后我才认真把鲁棒性从"论文里的边角料"当成上线必查项。
先搞明白:模型为什么会被"骗"
先说个简单的例子。假设你有一个训练好的图像分类模型,它在一张熊猫的图片上预测"熊猫"的概率是99.9%。但是,只要在原图上加一层肉眼几乎看不见的噪声,模型就会开始坚定地认为这是"长臂猿",而且置信度高达99.8%。
这种噪声就是对抗样本。有意思的是,同样的噪声模式换到另一张熊猫图片上,依然能让它误分类;但是同样的噪声加到其他模型上,可能就没用了。这说明对抗样本是针对特定模型结构和参数的"定制攻击",而不是图像本身有什么特殊之处。
下面这个简单的代码示例能让你快速理解对抗攻击是怎么回事:
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
# 快速构建一个简单的3层全连接网络(仅用于演示)
class SimpleNet(nn.Module):
def __init__(self, input_dim=784, num_classes=10):
super().__init__()
self.fc1 = nn.Linear(input_dim, 256)
self.fc2 = nn.Linear(256, 128)
self.fc3 = nn.Linear(128, num_classes)
def forward(self, x):
x = x.view(x.size(0), -1)
x = F.relu(self.fc1(x))
x = F.relu(self.fc2(x))
return self.fc3(x)
# FGSM对抗攻击(Fast Gradient Sign Method)
def fgsm_attack(model, x, y, epsilon=0.1):
x_adv = x.clone().detach()
x_adv.requires_grad = True
logits = model(x_adv)
loss = F.cross_entropy(logits, y)
loss.backward()
# 计算梯度符号
data_grad = x_adv.grad.data
sign_data_grad = data_grad.sign()
# 在输入上添加扰动
x_adv = x_adv + epsilon * sign_data_grad
x_adv = torch.clamp(x_adv, 0, 1)
return x_adv
这段代码跑起来非常快,但效果相当吓人。我在MNIST上测试的时候,用一个epsilon=0.1的扰动,就能让原本99.2%准确率的模型降到17.3%。关键是这种扰动肉眼几乎看不见,它不是随便乱加噪声,而是沿着让模型损失函数最大的方向"微调"像素值。
# 实际测试效果
python attack_demo.py
原始测试集准确率: 99.2%
FGSM攻击后准确率: 17.3% # epsilon=0.1
FGSM攻击后准确率: 3.7% # epsilon=0.3
这时候我才发现,模型在训练集和测试集上表现再好,也根本不能说明它在真实环境里到底有多"可靠"。
鲁棒训练:让模型学会"抵抗"干扰
理解了对抗攻击的原理,第一反应肯定是"那我们就训练一个不怕攻击的模型"。这个思路叫对抗训练(Adversarial Training),本质上就是在训练过程中不断用对抗样本攻击模型,逼它学会抵抗。
听起来很合理,实践起来坑不少。我最开始的尝试就踩了个大坑:直接用FGSM生成的对抗样本混入训练数据,结果发现模型虽然在FGSM攻击下表现好了,但遇到其他攻击方式(比如PGD)依然一塌糊涂。更糟糕的是,干净数据上的准确率也下降了5-6个百分点。
后来才搞明白,对抗训练不能只用一种攻击方法,也不能只在一个epsilon值上训练。生产环境里遇到的攻击手段多种多样,扰动程度也完全不可预测。
下面这个是我后来调整后的训练流程,效果明显好很多:
# 对抗训练流程(使用PGD攻击)
def adversarial_train(model, train_loader, optimizer, device,
epochs=10, epsilon=0.3, num_steps=7, step_size=0.01):
model.train()
for epoch in range(epochs):
total_loss = 0
correct = 0
total = 0
for batch_idx, (data, target) in enumerate(train_loader):
data, target = data.to(device), target.to(device)
# 生成PGD对抗样本
data_adv = pgd_attack(model, data, target, epsilon,
num_steps, step_size, device)
# 混合原始样本和对抗样本
mixed_data = torch.cat([data, data_adv])
mixed_target = torch.cat([target, target])
optimizer.zero_grad()
output = model(mixed_data)
loss = F.cross_entropy(output, mixed_target)
loss.backward()
optimizer.step()
total_loss += loss.item()
pred = output.argmax(dim=1)
correct += pred.eq(mixed_target).sum().item()
total += mixed_target.size(0)
if batch_idx % 100 == 0:
print(f'Epoch: {epoch}, Batch: {batch_idx}, Loss: {loss.item():.4f}')
acc = 100. * correct / total
print(f'Epoch {epoch}: Loss: {total_loss/len(train_loader):.4f}, Acc: {acc:.2f}%')
# PGD攻击(比FGSM更强的迭代攻击)
def pgd_attack(model, x, y, epsilon, num_steps, step_size, device):
x_adv = x.clone().detach() + torch.zeros_like(x).uniform_(-epsilon, epsilon)
x_adv = torch.clamp(x_adv, 0, 1)
for _ in range(num_steps):
x_adv.requires_grad = True
logits = model(x_adv)
loss = F.cross_entropy(logits, y)
loss.backward()
data_grad = x_adv.grad.data
x_adv = x_adv + step_size * data_grad.sign()
x_adv = torch.clamp(x_adv, x - epsilon, x + epsilon)
x_adv = torch.clamp(x_adv, 0, 1)
x_adv = x_adv.detach()
return x_adv
实际训练的时候,我用了一个混合策略:每个batch里70%用原始样本,30%用PGD生成的对抗样本,并且epsilon值在[0.1, 0.3]之间随机采样。这样训练出来的模型在面对不同攻击强度时都表现得比较稳定。
训练成本确实高了不少——原本训练一个ResNet-18在CIFAR-10上要3小时,加上对抗训练后变成了8小时多。但我后来测试发现,这种投入是值得的:干净数据上准确率从92%降到89%,但面对FGSM攻击时准确率从17%升到了82%,面对PGD攻击时从5%升到了78%。
# 训练效果对比
原始模型:
干净数据准确率: 92.1%
FGSM攻击后: 17.3%
PGD攻击后: 5.1%
对抗训练后:
干净数据准确率: 89.4%
FGSM攻击后: 82.7%
PGD攻击后: 78.3%
干净数据只牺牲约 3 个百分点,但 FGSM/PGD 下的准确率从个位数拉到 80% 左右——下图汇总了训练前后的三组对比。

对抗训练用可接受的干净集精度代价,换来了数量级的鲁棒性提升,这是后续推理层防御的基础。
防御策略:不只是训练,还要在推理层下功夫
对抗训练能显著提升模型鲁棒性,但不是万能药。有时候你没法重新训练模型(比如用别人预训练好的BERT),或者训练成本实在太高。这时候就要在推理层动点心思。
我实际用过的几个防御策略,按实用程度排序:
1. 输入预处理:防御蒸馏最有效
防御蒸馏(Defensive Distillation)这个方法一开始我还怀疑过——不就是再训练一个模型吗,能有多大用?但实际测试下来,它确实是最稳定的防御手段之一。基本思路是用一个训练好的"教师"模型(可以是经过对抗训练的)来训练一个"学生"模型,学生模型学习的是教师模型的概率分布,而不是硬标签。
# 防御蒸馏训练
def distillation_train(student_model, teacher_model, train_loader,
optimizer, device, temperature=3.0, alpha=0.7):
teacher_model.eval()
student_model.train()
for epoch in range(epochs):
for data, target in train_loader:
data, target = data.to(device), target.to(device)
optimizer.zero_grad()
with torch.no_grad():
teacher_logits = teacher_model(data)
teacher_probs = F.softmax(teacher_logits / temperature, dim=1)
student_logits = student_model(data)
student_probs = F.log_softmax(student_logits / temperature, dim=1)
# 蒸馏损失 + 传统交叉熵损失
distill_loss = F.kl_div(student_probs, teacher_probs,
reduction='batchmean') * (temperature ** 2)
ce_loss = F.cross_entropy(student_logits, target)
loss = alpha * distill_loss + (1 - alpha) * ce_loss
loss.backward()
optimizer.step()
我在一个图像分类项目里用这个方法,把一个原始ResNet-50蒸馏到一个更小的ResNet-18上,结果学生模型的对抗鲁棒性反而比教师模型还好。这说明蒸馏过程本身就有平滑决策边界的作用,让模型不会因为输入的微小变化就剧烈改变预测。
2. 随机平滑:给推理过程加点噪声
随机平滑(Randomized Smoothing)的核心思想很简单:在推理时对输入添加随机噪声,然后多次预测取平均。这样即使某个特定的对抗样本能骗过一次预测,也很难在多次随机噪声下持续骗过模型。
# 随机平滑预测
def randomized_smoothing_predict(model, x, num_samples=100, sigma=0.1, device='cuda'):
model.eval()
x = x.to(device)
# 生成多个噪声版本
noisy_samples = []
for _ in range(num_samples):
noise = torch.randn_like(x) * sigma
noisy_samples.append(x + noise)
noisy_samples = torch.cat(noisy_samples, dim=0)
with torch.no_grad():
logits = model(noisy_samples)
probs = F.softmax(logits, dim=1)
# 平均概率分布
avg_probs = probs.mean(dim=0)
# 返回平均概率最高的类别
return avg_probs.argmax().item()
这个方法的优点是不需要重新训练模型,而且理论上能提供可证明的鲁棒性保证(在一定噪声范围内)。缺点是推理速度会变慢——每次预测要做100次前向传播,延迟会增加两个数量级。我在一个实时性要求不高的API里用过,效果还不错:面对FGSM攻击,准确率从12%提升到了65%,但响应时间从80ms增加到了7.2s。
# 随机平滑效果对比
原始推理:
干净数据: 92.1%
FGSM攻击: 12.3%
平均推理时间: 80ms
随机平滑 (100 samples):
干净数据: 88.7%
FGSM攻击: 65.4%
平均推理时间: 7200ms
所以这个方法更适合对实时性要求不高的场景,比如离线批量处理或者后台分析任务。
3. 特征压缩:让对抗攻击"迷路"
特征压缩(Feature Squeezing)这个方法想法很巧妙:压缩输入数据的特征空间,让对抗扰动在压缩过程中丢失。具体做法可以是降低图像分辨率、减少颜色深度,或者用中值滤波去掉高频噪声。
# 特征压缩实现
def feature_squeeze(x, squeeze_type='bit_depth', bits=4):
if squeeze_type == 'bit_depth':
# 减少颜色深度(比如从8-bit降到4-bit)
factor = 2 ** (8 - bits)
squeezed = torch.round(x * 255 / factor) * factor / 255
elif squeeze_type == 'median_filter':
# 中值滤波(需要PIL或者OpenCV)
squeezed = x.clone()
for i in range(x.size(0)):
for j in range(x.size(1)):
# 简化的中值滤波(实际用cv2.medianBlur更高效)
channel = x[i, j].cpu().numpy()
squeezed[i, j] = torch.from_numpy(scipy.ndimage.median_filter(channel, size=3))
return torch.clamp(squeezed, 0, 1)
# 防御预测:用压缩特征做二次验证
def defended_predict(model, x, device='cuda'):
# 原始预测
original_pred = model(x.to(device)).argmax(dim=1)
# 压缩特征预测
squeezed_x = feature_squeeze(x, squeeze_type='bit_depth', bits=4)
squeezed_pred = model(squeezed_x.to(device)).argmax(dim=1)
# 如果两者不一致,可能被攻击,返回不确定标记
if not torch.equal(original_pred, squeezed_pred):
return -1 # 表示不确定
return original_pred
我在一个项目里把8-bit图像降到4-bit,虽然图像质量肉眼看着还行,但大部分对抗扰动都被"挤"掉了。测试结果显示,面对FGSM攻击时准确率从12%提升到了57%,而且对正常样本的影响很小——准确率只下降了1.2%。
这个方法的成本很低,不需要重新训练模型,推理延迟也几乎没有增加(只需要多做一次前向传播)。缺点是面对强攻击(比如PGD)时效果会减弱,这时候可能需要和其他防御策略组合使用。
踩过的几个坑
折腾这半年,踩过的坑比做对的决策多。这里挑几个最有代表性的,免得你重复走弯路。
坑1:过度依赖单一评估指标
最开始我只看对抗样本下的准确率,觉得准确率上去就万事大吉了。但后来发现,模型可能会通过一种"偷懒"的方式提升准确率:它不再认真对抗攻击,而是干脆对所有可疑输入都输出最常见的那一类。这样准确率确实提高了(因为对抗样本也可能被分类到这个最常见类别),但模型实际上变得更不可用了。
后来我引入了更多评估指标:除了准确率,还看置信度分布、类别多样性、以及对抗样本和干净样本的预测一致性。一个健康的鲁棒模型,不应该通过"摆烂"来提升对抗准确率。
坑2:忽视了分布外数据
对抗样本只是模型鲁棒性问题的一个侧面。实际生产环境里,模型更常遇到的是分布外(Out-of-Distribution)数据——比如一个在ImageNet上训练的模型突然看到一张X光片,它大概率会胡乱猜一个类别,还可能给出很高的置信度。
我在一个医疗图像项目里就遇到过这个问题:模型训练集中全是正常组织的图像,推理时却经常碰到病变组织。模型不但没法识别病变,还会用高置信度把它们强行分类到训练集中的某个类别,完全误导了临床决策。
# 简单的OOD检测示例
def detect_ood(model, x, temperature=1.0, threshold=0.5):
with torch.no_grad():
logits = model(x)
probs = F.softmax(logits / temperature, dim=1)
# 如果最大概率低于阈值,可能是OOD
max_prob = probs.max(dim=1)[0]
# 计算预测熵
entropy = -torch.sum(probs * torch.log(probs + 1e-10), dim=1)
# 综合判断
is_ood = (max_prob < threshold) | (entropy > 2.0)
return is_ood
后来我加了一个OOD检测模块:在模型预测前先判断输入是否在训练分布内,如果是就正常预测,如果不是就返回"无法识别"或触发人工审核。这个简单的改动让实际生产中的误判率下降了40%。
坑3:防御策略的组合不是简单堆叠
我一直以为既然单个防御策略都有用,把它们组合起来肯定效果更好。但实践发现,某些防御策略之间会互相抵消效果。比如随机平滑和特征压缩一起用的时候,压缩后的特征空间变小,随机平滑需要的噪声范围反而变得难以控制。
更糟糕的是,有些组合还会产生新的漏洞。我在一个项目里把对抗训练和输入预处理组合起来,结果发现预处理后的对抗样本反而更容易骗过模型——因为预处理改变了原始对抗样本的梯度方向,让对抗训练学到的防御模式失效了。
所以防御策略组合需要仔细设计,最好先在小规模数据上测试,确认没有互相抵消或产生新漏洞后再上生产。
实际落地的一些经验
这些理论和实践都摆在这了,但真正落地到生产环境还有不少细节要处理。
首先是训练数据的问题。很多项目可能只有几千条标注数据,根本不够做对抗训练(对抗训练通常需要更多数据来学习更平滑的决策边界)。这种情况下,可以考虑用半监督学习或者数据增强来扩充数据集。我在一个只有3000条标注数据的项目里,通过MixUp和CutMix数据增强,把有效数据量扩大了10倍,对抗训练的效果也跟着上来了。
其次是模型选择的权衡。鲁棒性更好的模型通常更大、更慢。如果你的应用对延迟很敏感(比如实时视频流处理),可能需要在鲁棒性和延迟之间做取舍。一个实用策略是:对关键业务(比如支付验证)用鲁棒性强的模型,对非关键业务(比如内容推荐)用效率更高的普通模型。
最后是监控和快速响应。即使做了各种防御,模型在生产环境里还是可能遇到新的攻击方式或分布外数据。建议建立一个监控系统,实时跟踪模型的预测置信度分布、异常输入比例、错误率变化等指标。一旦发现异常,可以快速回滚到旧版本或者触发人工审核。
# 简单的监控指标收集
class ModelMonitor:
def __init__(self, alert_threshold=0.1):
self.confs = []
self.error_rate = []
self.alert_threshold = alert_threshold
def log_prediction(self, conf, is_correct):
self.confs.append(conf)
if not is_correct:
self.error_rate.append(1)
else:
self.error_rate.append(0)
# 保持最近1000条记录
self.confs = self.confs[-1000:]
self.error_rate = self.error_rate[-1000:]
def check_health(self):
avg_conf = np.mean(self.confs)
recent_error = np.mean(self.error_rate[-100:])
if recent_error > self.alert_threshold or avg_conf < 0.6:
return "ALERT"
elif recent_error > self.alert_threshold * 0.7:
return "WARNING"
else:
return "HEALTHY"
写在最后
模型鲁棒性不是一次性的工程任务,而是需要持续关注和迭代的过程。新攻击方式在不断出现,新防御策略也在持续被提出,这本质上是一场持续的攻防博弈。
但这不意味着我们要陷入永无止境的军备竞赛。大部分实际应用里,合理的防御策略已经能覆盖95%以上的真实风险。关键是要找到平衡点:在保证鲁棒性的同时,不过度牺牲准确率、效率或开发成本。
我现在的实践策略通常是:先评估业务风险等级,对高风险场景用对抗训练+防御蒸馏的组合,对中风险场景用特征压缩或随机平滑,对低风险场景则主要依赖监控和快速响应。这样能在成本和效果之间找到一个相对平衡点。
模型鲁棒性可能不是最性感的技术方向,但它决定了你训练的模型到底能不能真正用起来。毕竟,一个在训练集上表现完美但在生产环境里一碰就崩的模型,除了给benchmark刷数据之外,实在没什么实际意义。
版权声明: 本文首发于 指尖魔法屋-关于AI模型鲁棒性的几点记录(https://blog.thinkmoon.cn/post/259-ai-model-robustness-training-defense/) 转载或引用必须申明原指尖魔法屋来源及源地址!
评论
使用 GitHub 账号登录后即可留言,支持 Markdown。