Adversarially Robust Distillation

In plain words

Small neural networks are cheap to run, but they are easy to fool: a tiny, almost invisible change to an image can flip the answer. Making a small network resist such changes usually means adversarial training, which is slow because every training batch needs extra attack computations.

This paper asks whether a small network can inherit robustness from a big, already robust network, the same way it inherits accuracy in ordinary distillation. The student learns to copy the teacher's answers, but its copy is checked on attacked versions of the images. The paper also finds that distilling a robust teacher onto clean images alone passes on a good share of the robustness.

The method is named Adversarially Robust Distillation (ARD). It makes a small network robust, and in the paper's experiments it beats a small network that was trained adversarially from scratch.

TL;DR

• Distilling a robust teacher onto a small student transfers much of the teacher's robustness, even when the student sees only clean images §4.2.
• ARD trains the student to match the teacher's softmax on adversarial inputs (the student is attacked) §5 eq. 1.
• Author-reported CIFAR-10 result: ARD MobileNetV2 from a TRADES WideResNet teacher gets 50.42% robust accuracy, versus 46.90% for adversarially trained MobileNetV2 tab. 6.
• ARD costs about as much training time as adversarial training, so the cheap part is the student, not the training tab. 10.
• Official PyTorch code is MIT-licensed and matches most of the method, but its test attack and learning-rate schedule differ from the paper (see Paper vs code).

Problem

Knowledge distillation trains a small student to match a large teacher's outputs and gives accurate, compact models. The student is usually not robust to adversarial examples, even when the teacher is robust §1 §2. Adversarial training makes a model robust but needs many gradient computations per batch, so it is slow §5.4. Defensive distillation is reported as broken by stronger attacks §2. The paper focuses on ℓ∞ attacks §3.

The question is whether robustness moves from a robust teacher to a student during distillation §4, and whether a distillation loss built on adversarial inputs can produce better robust students than adversarial training of the same architecture §5.

Method

Robustness transfer under plain distillation §4. The authors distill a non-robust ResNet18 teacher, which gives a student with 0% robust accuracy under 20-step PGD §4.1 tab. 2. Distilling an adversarially trained ResNet18 gives a student that keeps most of the teacher's robust accuracy §4.2 tab. 3. Robust WideResNet teachers transfer less robustness than the weaker ResNet18 teacher, and on CIFAR-100 a robust WideResNet transfers almost none §4.3 tab. 4 tab. 5.

ARD objective §5 eq. 1. The student S_θ with teacher T and temperature t is trained with

\min_{\theta}\ \mathbb{E}_{(X,y)\sim\mathcal{D}}\Big[\alpha t^2\, \mathrm{KL}\big(S_\theta^t(X+\delta_\theta),\,T^t(X)\big) + (1-\alpha)\,\ell\big(S_\theta^t(X),\,y\big)\Big] \qquad \delta_\theta=\arg\max_{\|\delta\|_p<\epsilon}\ell\big(S_\theta^t(X+\delta),y\big) eq. 1 §5

Here ℓ is cross-entropy, S^t and T^t are logits divided by the temperature before the softmax, and α sets the balance. The attack δ_θ is computed against the student, while the teacher sees clean X. The t² factor offsets the gradient shrinkage caused by dividing the logits by t §5. The paper sets α = 1 in most experiments, which removes the clean cross-entropy term §5.

Training procedure alg. 1 §5. For each batch: build adversarial examples for the student with K-step PGD on cross-entropy against the true labels, then take one SGD step on the ARD loss alg. 1. Teacher logits on clean inputs are the targets §5.

# ARD training step  [alg. 1] [§5] [eq. 1]
# Input: student, frozen teacher, images (B, 3, H, W), labels (B,)
# Output: scalar loss for one SGD step

def ard_loss(student, teacher, images, labels, temp, alpha):
    # Attack the student; adversarial inputs depend on the student  [alg. 1]
    adv_images = pgd_attack(student, images, labels)

    # Teacher soft targets computed on the clean images  [§5] [eq. 1]
    teacher_probs = softmax(teacher(images) / temp)

    # Student soft predictions on the adversarial images  [§5] [eq. 1]
    student_log_probs = log_softmax(student(adv_images) / temp)

    # KL(teacher || student), scaled by alpha and temp squared  [§5] [eq. 1]
    distill = kl_divergence(teacher_probs, student_log_probs)
    distill = alpha * temp * temp * distill

    # Clean cross-entropy term, weighted by (1 - alpha)  [§5]
    xent = cross_entropy(student(images), labels)

    return distill + (1 - alpha) * xent

# K-step l-infinity PGD on cross-entropy  [§8] [alg. 1]
# Input: model, images (B, 3, H, W), labels (B,)
# Output: adversarial images inside the eps-ball around images

def pgd_attack(model, images, labels, eps, step, num_steps):
    # Start from a uniform random point in the eps-ball  [§8] [code: main.py#L70]
    adv = images + uniform(-eps, eps)
    for _ in range(num_steps):
        # Ascend the cross-entropy with a signed gradient step  [§8]
        grad = grad_cross_entropy(model(adv), labels, wrt=adv)
        adv = adv + step * sign(grad)
        # Project back into the eps-ball around the clean images  [alg. 1]
        adv = clip(adv, images - eps, images + eps)
        # Keep pixel values in [0, 1]  [code: main.py#L78]
        adv = clip(adv, 0, 1)
    return adv

The method's parameter values are in the table below §8.

Fast-ARD §5.4 alg. 2. The authors adapt "free" adversarial training to distillation. Each mini-batch is replayed m times. On each replay, one backward pass gives both the input gradient, used for a 1-step attack that keeps a running perturbation δ per training image, and the parameter gradient for the update. The epoch count is divided by m so the total training time matches plain distillation alg. 2 §5.4. The paper reports m = 4 and m = 8 tab. 9.

Extensions noted by the paper. Combining ARD with a non-adversarial KL term, using teacher outputs at adversarial points, and maximising the KL instead of cross-entropy in the attack all lowered robust accuracy or natural accuracy app. B. Passing adversarial examples through a large teacher slows training and raises memory use app. B.

Hyperparameters

Results

Author-reported. CIFAR-10 and CIFAR-100. Teacher and student architectures as in the Hyperparameters table. Robust accuracy is under 20-step PGD unless the row says otherwise. Values are natural accuracy / robust accuracy, in percent.

Training time, hours, CIFAR-10, one RTX 2080 Ti, 200 epochs of parameter updates per image tab. 10: AT MobileNetV2 41.09; clean distillation from TRADES WRN 5.13; TRADES WRN ARD 41.06; TRADES WRN Fast-ARD m = 4 5.15; m = 8 5.10.

Attacks on MobileNetV2 (CIFAR-10; all numbers are robust validation accuracy, ε = 8/255) tab. 11: AT MobileNetV2 gets 50.82 (MI-FGSM 20 steps), 57.74 (DeepFool), 46.51 (1000-step PGD), 46.79 (20-step PGD, 100 restarts). ARD from TRADES WRN gets 55.16, 64.61, 49.98 and 50.30. Free-trained MobileNetV2 (m = 4) gets 30.60, 41.09, 22.23 and 22.94; Fast-ARD (m = 4) gets 44.78, 60.03, 36.01 and 36.88.

Ablations, author-reported app. A app. C app. D app. E:

The paper's own summary: ARD students exceed adversarially trained students of the same architecture in robust accuracy and often in natural accuracy; they also exceed the reported state of the art on CIFAR-10 and CIFAR-100 §7. These are the authors' claims, not verified here.

Paper vs code

Official code: https://github.com/goldblum/AdversariallyRobustDistillation (MIT, Copyright 2019 MicahGoldblum), commit 779ade36d612bddb0a5a705fdd74efe7571a1381. The repo owner is an author. The paper links this repo §8.

• Test attack is 10 steps, not 20. Robust accuracy in the paper's tables uses 20-step PGD tab. 1 tab. 6. The code builds a single attack object with num_steps 10 and uses it for both training and evaluation main.py main.py. Numbers from this code are therefore not comparable to the paper's 20-step columns.
• Learning-rate decay does not compound. The paper decreases the LR by a factor of 10 at epochs 100 and 150 §8. The code's adjust_learning_rate multiplies a base value that is passed in unchanged, so both decay epochs set the LR to 0.1 × 0.1 = 0.01 and it never reaches 0.001 main.py main.py. Epochs are also 0-based in the loop, so the first decay happens at index 100 main.py.
• Adversarial examples are generated with the student in training mode. The attack calls the student inside train(), so each PGD forward pass updates the BatchNorm running statistics main.py main.py. The paper does not describe BatchNorm mode during attack generation.
• KL reduction and direction. The paper writes KL(S(X+δ), T(X)) eq. 1. The code uses nn.KLDivLoss() with the student's log-probabilities as input and the teacher's probabilities as target, so it computes KL(teacher ‖ student) main.py main.py. The default reduction averages over every element (batch × classes), which is not a per-example sum. The paper does not state a reduction.
• Teacher runs on every step. The teacher is applied to the clean batch in every iteration main.py. Appendix B notes that stored logits would avoid these forward passes app. B; the code does not use stored logits.
• Clean cross-entropy is computed even when α = 1. The clean forward pass runs regardless and is multiplied by zero main.py. Its BatchNorm side effect still happens.
• CIFAR-100 temperature is not in the code. The default temperature is 30 main.py. The paper uses 5 for CIFAR-100 §8. The code has a CIFAR-100 data branch but no temperature setting for it main.py.
• No teacher training and no adversarial-training baseline. The script loads a teacher checkpoint unconditionally main.py. Adversarial-training teachers and the AT baseline are not in the repo, although the paper says the authors re-implemented them §3.
• No Fast-ARD. The repo has only the standard ARD loop. The "free" replay of Algorithm 2 alg. 2 is absent.
• Pixel clamp. The attack clamps to [0, 1] main.py. The paper does not state an input range.
• Attack initialisation and steps match. Uniform random start, 10 steps, ε = 8/255, step 2/255, and cross-entropy against the true label match §8 alg. 1 main.py main.py.
• Loss weights and temperature squared match. The loss is α t² times the KL plus (1 − α) cross-entropy, with logits divided by t §5 main.py.
• Optimiser, momentum, weight decay, initial LR and epoch count match §8 main.py main.py main.py.
• Batch size is not in the paper. Code uses 128 for CIFAR-10 and 256 for CIFAR-100 main.py main.py.

Pitfalls

• Expect lower numbers if you use the code's test loop. The code evaluates with a 10-step attack, the paper with 20 steps tab. 1 main.py.
• Implement the LR schedule as written in the paper, or document the deviation. The code's schedule keeps LR at 0.01 after epoch 150.
• BatchNorm. Attacks in training mode change BatchNorm statistics. The paper says nothing about this, so it is a choice you must make explicitly.
• Cost is the main limit. ARD trains as slowly as adversarial training: 41.06 hours against 5.13 for clean distillation in the authors' setup tab. 10. Only the Fast-ARD variant brings the cost down, and the paper reports it as less robust than full ARD §5.4.
• α trades robustness for accuracy and matters a lot for clean distillation. For clean distillation robustness falls fast as α drops; ARD is less sensitive to α and temperature app. E tab. 13 tab. 18.
• A robust teacher is required. Naturally trained teachers give students with lower robust accuracy than AT students of the same architecture app. C tab. 15.
• Robust teachers do not always help. WideResNet teachers transfer less robustness than ResNet18 through clean distillation, and on CIFAR-100 almost none §4.3 tab. 5. ARD is what the paper proposes for these cases §4.3.
• Only ℓ∞ attacks are studied. The experiments focus on ℓ∞ attacks, and other norms are not evaluated §3.
• Notation is ambiguous. The order in KL(S(·), T(·)) is written two ways across the paper. Follow the code's convention (teacher as the target distribution) unless you have reason to switch eq. 1 main.py.
• Hidden tuning. The paper's robust results depend on α = 0.95 or 0.93 for some teachers, and on a specific temperature, whose best value differs by dataset tab. 6 tab. 7 tab. 12 tab. 17.
• Single-run results. No variance or seed information is given in the paper. Robust accuracy differences of about one point should be treated with caution.
• Teacher pretraining is not described in terms of the adversarial-training recipe for the teacher beyond "as described in" prior work §3.
• The reported Fast-ARD numbers depend on m, which changes the epoch count §5.4 tab. 9.

Lineage

• Builds on: knowledge distillation §1 §3; adversarial training with PGD §3; TRADES and VAT, whose KL consistency terms ARD adapts §3; "free" adversarial training and a related accelerated method, which Fast-ARD adapts §5.4 (no card for these papers).
• Compared against: defensive distillation §2; defensive quantisation and pruning-based compression §2; adversarial training of same-architecture students §5.
• Discussed as a possible extension: knowledge distillation with feature maps and adaptive data augmentation such as AutoAugment §7.
• Superseded by: not stated in this card.

Attribution

Title, authors and the arXiv link are in the frontmatter and the page footer. This card is not endorsed by the authors.