关于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% 左右——下图汇总了训练前后的三组对比。

ResNet-18 / CIFAR-10 对抗训练前后在干净数据、FGSM 与 PGD 攻击下的准确率对比(%)

对抗训练用可接受的干净集精度代价,换来了数量级的鲁棒性提升,这是后续推理层防御的基础。

防御策略:不只是训练,还要在推理层下功夫

对抗训练能显著提升模型鲁棒性,但不是万能药。有时候你没法重新训练模型(比如用别人预训练好的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/) 转载或引用必须申明原指尖魔法屋来源及源地址!