fzjcdt/co-weight-driftPyTorch reference implementation of a fast adversarial training method that prevents catastrophic overfitting via coupled objectives and an EMA-based weight-drift guard, with AutoAttack evaluation.
Reference PyTorch implementation of the method described in the paper above. Fast adversarial training (FAT) suffers from catastrophic overfitting (CO): the FGSM robust accuracy stays high while PGD robust accuracy collapses within a couple of epochs. This codebase implements the paper's two contributions:
Robust accuracy (%) under the complete standard AutoAttack (AA) and the constraint-relaxation attack (CR), ε = 8/255, ResNet18 (PreActResNet18 on Tiny-ImageNet), mean ± sample std. Baseline rows are the released results of the FastAT_Benchmark under the identical protocol.
| Method | C10 AA | C10 CR | C100 AA | C100 CR | Tiny AA | Tiny CR |
|---|---|---|---|---|---|---|
| PGD-AT-WA (multi-step ref.) | 50.49±0.15 | 50.51±0.16 | 26.90±0.17 | 26.87±0.17 | 19.77±0.16 | 19.75±0.13 |
| NuAT | 50.10±0.18 | 50.11±0.19 | 14.41±3.29 | 17.66±0.42 | 10.20±0.44 | 10.19±0.44 |
| LIET | 49.99±0.15 | 49.93±0.15 | 26.64±0.13 | 26.63±0.14 | 19.88±0.09 | 19.86±0.07 |
| FGSM-PGI | 49.89±0.08 | 49.89±0.12 | 26.29±0.16 | 26.31±0.12 | 19.55±0.11 | 19.52±0.12 |
| FGSM-FUAP | 49.73±0.12 | 49.71±0.12 | 25.95±0.08 | 25.94±0.05 | 19.10±0.18 | 19.09±0.19 |
| FGSM-CUAP | 49.55±0.22 | 49.54±0.22 | 25.85±0.50 | 25.83±0.49 | 19.19±0.09 | 19.20±0.11 |
| Base + guard (ours) | 50.45±0.20 | 50.44±0.16 | 27.15±0.04 | 27.20±0.07 | 20.47±0.11 | 20.54±0.15 |
Strongest AutoAttack and CR accuracy among fast methods on all three datasets, at FAT cost (~0.7–1.1 h/run on one A100, vs ~2 h for the multi-step reference). On CIFAR-10 the unguarded base is a coin flip: 2 of 4 seeds collapse to CO artifacts (AutoAttack 7.5) while the survivors match the guarded row (50.67±0.08) — the guard's effect is exactly the removal of this failure mode. Notably, the FAB-T and square components of AutoAttack find zero additional points on all of our released checkpoints (no gradient masking signature). Per-seed numbers and raw logs: logs/.
pip install -r requirements.txt
Tested with Python 3.12, torch 2.5.1, torchvision 0.20.1, torchattacks 3.5.1, CUDA 12.4, on A100-PCIE-40GB. The autoattack/ directory in this repo is the vendored copy the released numbers were produced with (Francesco Croce's AutoAttack); running from the repo root makes it take precedence over any pip-installed copy.
train.py implements the paper's method; defaults are the CIFAR-10 recipe.
# CIFAR-10 (γ_max=4.0, m_max=2.5, λ_d=0.3, guard θ=0.08) — 4 seeds
bash scripts/train_cifar10.sh 0 # GPU 0, seeds 0 1 2 3
# CIFAR-100 (γ_max=1.25, m_max=4.5, λ_d=0)
bash scripts/train_cifar100.sh 0
# Tiny-ImageNet (PreActResNet18, γ_max=1.0, m_max=2.5, λ_d=0)
bash scripts/train_tiny.sh 0
# or a single run directly (equivalent to the CIFAR-10 script)
python train.py --dataset cifar10 --seed 0 --gpu 0 \
--nuc-max 4.0 --mse-extra 2.5 --lam-d 0.3 --guard val --guard-thr 0.08 \
--out-dir results/cifar10/guard_s0
Each run writes train.log (per-epoch val PGD-10 of raw and EMA weights, guard-fire lines), probe.csv (optional batch-level CO dynamics), and best.pth / last.pth / ema_last.pth. best.pth is the epoch with the highest validation PGD-10 (raw vs EMA weights, tie-broken by clean accuracy).
--guard none disables the guard (the unguarded base; expect the CO coin flip on CIFAR-10).--init-style bu switches the perturbation init from the paper's fresh half-sign at 0.5ε to 70% ±ε Rademacher / 30% U[−ε, ε].--probe-every 0 disables the diagnostic probes. Note the probes consume RNG and therefore perturb the training trajectory (the "instrumentation lottery" reported in the paper) — match this flag to the released run you compare against (see logs/README.md).# complete standard AutoAttack + CR attack, full test set, bs=200
bash scripts/eval.sh results/cifar10/guard_s0 0 # GPU 0, ResNet18
bash scripts/eval.sh results/tiny-imagenet/guard_s0 0 preact18
# or individually
python aa_eval.py --run-dir results/cifar10/guard_s0 --gpu 0 --full # autoattack_full.log
python cr_eval.py --run-dir results/cifar10/guard_s0 --gpu 0 # cr_attack.log
aa_eval.py --full runs the complete standard suite (apgd-ce, apgd-t, FAB-T, square) — the protocol of every AA number in the paper. Without --full it runs the benchmark's 2-attack development variant (apgd-ce + apgd-t), which on our models differs by at most 0.02.
train.py # coupled FAT + EMA self-anchor guard (the method)
aa_eval.py # AutoAttack evaluation (full suite / 2-attack variant)
cr_eval.py # CR attack evaluation
models/ # ResNet18 / PreActResNet18 (from FastAT_Benchmark)
utils/ # datasets, seeding, logger (from FastAT_Benchmark)
attacks/ # CR attack (from FastAT_Benchmark)
autoattack/ # vendored AutoAttack (Francesco Croce, fra31)
scripts/ # per-dataset training launchers + eval
logs/ # raw logs of every run behind the paper's tables
Identical to FastAT_Benchmark: ResNet18 (CIFAR-10/100) and PreActResNet18 (Tiny-ImageNet); ε = 8/255; 100 epochs; SGD (lr 0.1, momentum 0.9, weight decay 5e-4) with OneCycleLR; batch size 128; 49 000 / 1 000 train / val split; label smoothing 0.4; model selection by validation PGD-10; evaluation on the full test set with PGD-{10,20,50}, AutoAttack and the CR attack. All runs use set_seed with cudnn deterministic mode.
models/, utils/, attacks/ CR attack) and all baseline results are from FastAT_Benchmark — please cite it if you use the comparison numbers.autoattack/ is AutoAttack by Francesco Croce (vendored unmodified).The paper is under review. If you use this code, please cite it by its title until a BibTeX entry is available:
MIT (see LICENSE). The vendored third-party files keep their upstream licenses and copyright headers.