Adversarial examples with FGSM

One gradient step on the input flips a classifier

Lab 6 of 10 · about 30 minutes · runs on a laptop. Related chapters: 20. Attacks on language models and their defenses · 3. Backpropagation and optimizers

pip install torch torchvision

import torch, torch.nn as nn, torch.nn.functional as F
from torchvision import datasets, transforms
tf = transforms.ToTensor()
train = datasets.MNIST('.', train=True, download=True, transform=tf)
test = datasets.MNIST('.', train=False, download=True, transform=tf)
model = nn.Sequential(nn.Flatten(), nn.Linear(784,256), nn.ReLU(), nn.Linear(256,10))
opt = torch.optim.Adam(model.parameters(), 1e-3)
loader = torch.utils.data.DataLoader(train, batch_size=128, shuffle=True)
for epoch in range(2):
    for x, y in loader:
        opt.zero_grad(); F.cross_entropy(model(x), y).backward(); opt.step()

def fgsm(x, y, eps):
    x = x.clone().requires_grad_(True)
    F.cross_entropy(model(x), y).backward()
    return (x + eps * x.grad.sign()).clamp(0, 1).detach()

x, y = next(iter(torch.utils.data.DataLoader(test, batch_size=2000)))
for eps in [0, 0.05, 0.1, 0.2, 0.3]:
    adv = fgsm(x, y, eps)
    print(eps, (model(adv).argmax(1) == y).float().mean().item())

In one run the accuracy was about 95% at ε = 0 and fell to about 54% at ε = 0.05, 8% at 0.1 and 0% at 0.2 and 0.3. Your numbers will differ a little between runs, but the collapse should look the same: a perturbation of 0.1 on pixels that range from 0 to 1 is hard to see yet leaves almost nothing correct.

Extensions: turn FGSM into PGD (loop with a small step, project back into the $\epsilon$-ball); adversarially train and re-plot accuracy vs $\epsilon$; test whether attacks crafted on this model fool a second independently trained one (transferability).