Adversarial Attacks — Robustness & Input Manipulation

import torch
import torch.nn as nn
import torch.nn.functional as F
from torchvision import models, transforms
from PIL import Image
import numpy as np

# ── Fast Gradient Sign Method (FGSM) — the simplest adversarial attack ──
# Goodfellow et al. 2014: "Explaining and Harnessing Adversarial Examples"
def fgsm_attack(model, loss_fn, image, label, epsilon=0.03):
    """
    Generate an adversarial example by taking one step in the
    direction that MAXIMISES the loss.
    image:   (1, C, H, W) tensor, requires_grad=True
    epsilon: perturbation budget (max pixel change, in [0,1] scale)
    """
    image.requires_grad = True
    output = model(image)
    loss   = loss_fn(output, label)
    model.zero_grad()
    loss.backward()

    # Perturb in the direction that increases loss
    # sign(gradient) = +1 where gradient is positive, -1 where negative
    perturbation = epsilon * image.grad.sign()
    adversarial  = image + perturbation
    adversarial  = torch.clamp(adversarial, 0, 1)   # Keep pixels in valid range
    return adversarial.detach()

# ── Projected Gradient Descent (PGD) — stronger multi-step attack ──
def pgd_attack(model, loss_fn, image, label,
               epsilon=0.03, alpha=0.01, n_steps=40):
    """
    Multi-step FGSM with random start — much stronger than FGSM.
    Used for adversarial training (Madry et al. 2018)
    """
    # Random start within epsilon ball
    delta = torch.empty_like(image).uniform_(-epsilon, epsilon)
    adversarial = (image + delta).clamp(0, 1).detach().requires_grad_(True)

    for step in range(n_steps):
        output = model(adversarial)
        loss   = loss_fn(output, label)
        model.zero_grad()
        loss.backward()

        # Gradient ascent step
        adversarial = adversarial + alpha * adversarial.grad.sign()
        # Project back to epsilon-ball around original image
        delta = torch.clamp(adversarial - image, -epsilon, epsilon)
        adversarial = torch.clamp(image + delta, 0, 1).detach().requires_grad_(True)

    return adversarial.detach()

# ── Adversarial training: train on adversarial examples to improve robustness ──
def adversarial_training_step(model, optimizer, loss_fn, X, y, epsilon=0.03):
    """
    Madry et al. adversarial training: for each batch, generate adversarial examples
    and train on them INSTEAD OF clean examples.
    """
    model.eval()   # Generate adversarial examples without dropout
    X_adv = pgd_attack(model, loss_fn, X.clone(), y, epsilon=epsilon, n_steps=10)

    model.train()
    optimizer.zero_grad()
    loss = loss_fn(model(X_adv), y)   # Train on adversarial examples
    loss.backward()
    optimizer.step()
    return loss.item()

# ── NLP adversarial attacks ──
# TextFooler: substitute words with synonyms to flip sentiment
def textfooler_attack(text: str, model, tokenizer) -> str:
    """Simplified word substitution attack."""
    from transformers import pipeline
    clf = pipeline("sentiment-analysis", model=model)

    original_pred = clf(text)[0]['label']
    words = text.split()

    # Try substituting each word with synonyms
    synonyms = {
        'good': ['decent', 'adequate', 'satisfactory'],
        'great': ['acceptable', 'reasonable', 'moderate'],
        'amazing': ['okay', 'fine', 'average'],
        'bad': ['challenging', 'complex', 'difficult'],
        'terrible': ['suboptimal', 'imperfect', 'flawed'],
    }

    for i, word in enumerate(words):
        if word.lower() in synonyms:
            for synonym in synonyms[word.lower()]:
                candidate = words.copy()
                candidate[i] = synonym
                candidate_text = ' '.join(candidate)
                new_pred = clf(candidate_text)[0]['label']
                if new_pred != original_pred:
                    print(f"Attack succeeded: '{text}' → '{candidate_text}'")
                    return candidate_text
    return text   # Attack failed