fzjcdt/co-weight-driftImplementação de referência em PyTorch de um método de treinamento adversarial rápido que previne o sobreajuste catastrófico por meio de objetivos acoplados e uma proteção contra deriva de pesos baseada em EMA, com avaliação AutoAttack.
Implementação de referência em PyTorch do método descrito no artigo acima. O treino adversarial rápido (FAT) sofre de overfitting catastrófico (CO): a acurácia robusta ao FGSM mantém-se elevada enquanto a acurácia robusta ao PGD colapsa em poucas épocas. Este código implementa as duas contribuições do artigo:
Acurácia robusta (%) sob o AutoAttack (AA) padrão completo e o ataque de relaxamento de restrições (CR), ε = 8/255, ResNet18 (PreActResNet18 em Tiny-ImageNet), média ± desvio padrão amostral. As linhas de base são os resultados publicados do FastAT_Benchmark sob o protocolo idêntico.
| Método | C10 AA | C10 CR | C100 AA | C100 CR | Tiny AA | Tiny CR |
|---|---|---|---|---|---|---|
| PGD-AT-WA (ref. multi-step) | 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 + guarda (nosso) | 50.45±0.20 | 50.44±0.16 | 27.15±0.04 | 27.20±0.07 | 20.47±0.11 | 20.54±0.15 |
A acurácia mais forte ao AutoAttack e ao CR entre os métodos rápidos nos três conjuntos de dados, a custo de FAT (~0.7–1.1 h/execução numa A100, vs ~2 h para a referência multi-step). Em CIFAR-10 a base sem guarda é uma moeda ao ar: 2 de 4 seeds colapsam em artefactos de CO (AutoAttack 7.5) enquanto os sobreviventes igualam a linha com guarda (50.67±0.08) — o efeito da guarda é exatamente a remoção deste modo de falha. Notavelmente, as componentes FAB-T e square do AutoAttack encontram zero pontos adicionais em todos os nossos checkpoints publicados (sem assinatura de mascaramento de gradiente). Números por seed e registos brutos: logs/.
pip install -r requirements.txt
Testado com Python 3.12, torch 2.5.1, torchvision 0.20.1, torchattacks 3.5.1, CUDA 12.4, em uma A100-PCIE-40GB. O diretório autoattack/ neste repositório é a cópia vendorizada com a qual os números divulgados foram produzidos (o AutoAttack de Francesco Croce); executar a partir da raiz do repositório faz com que ele tenha precedência sobre qualquer cópia instalada via pip.
train.py implementa o método do artigo; os padrões correspondem à receita para CIFAR-10.
# 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
Cada execução grava train.log (PGD-10 de validação por época dos pesos brutos e EMA, linhas de disparo do guard), probe.csv (dinâmica de CO opcional ao nível do lote) e best.pth / last.pth / ema_last.pth. best.pth é a época com o maior PGD-10 de validação (pesos brutos vs EMA, desempate pela acurácia limpa).
--guard none desativa o guard (a base sem guard; espere o cara-ou-coroa de CO no CIFAR-10).--init-style bu alterna a inicialização da perturbação do meio-sinal fresco a 0.5ε do artigo para 70% ±ε Rademacher / 30% U[−ε, ε].--probe-every 0 desativa as sondas de diagnóstico. Note que as sondas consomem RNG e, portanto, perturbam a trajetória de treinamento (a "loteria de instrumentação" relatada no artigo) — combine esta flag com a execução publicada com a qual você compara (veja 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 executa a suíte padrão completa (apgd-ce, apgd-t, FAB-T, square) — o protocolo de todos os números de AA no artigo. Sem --full, ele executa a variante de desenvolvimento de 2 ataques do benchmark (apgd-ce + apgd-t), que nos nossos modelos difere em no máximo 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
Idêntico ao FastAT_Benchmark: ResNet18 (CIFAR-10/100) e PreActResNet18 (Tiny-ImageNet); ε = 8/255; 100 épocas; SGD (lr 0.1, momentum 0.9, weight decay 5e-4) com OneCycleLR; tamanho do lote 128; divisão de 49 000 / 1 000 treino / validação; label smoothing 0.4; seleção do modelo por validação PGD-10; avaliação no conjunto de teste completo com PGD-{10,20,50}, AutoAttack e o ataque CR. Todas as execuções usam set_seed com modo determinístico do cudnn.
models/, utils/, ataque CR em attacks/) e todos os resultados de linha de base são do FastAT_Benchmark — cite-o se utilizar os números de comparação.autoattack/ é o AutoAttack de Francesco Croce (incorporado sem modificações).O artigo está em revisão. Se você usar este código, cite-o pelo título até que uma entrada BibTeX esteja disponível: