AFD: Mitigating Feature Gap for Adversarial Robustness by Feature Disentanglement

TL;DR

This work tackles the persistent feature-gap problem observed during adversarial fine-tuning by introducing Adversarial Fine-tuning via Disentanglement (AFD). It explicitly disentangles adversarial features into intrinsic and confused components with a learnable feature disentangler and further aligns the intrinsic features with pre-trained natural features to reduce misclassification under attack. The approach combines a targeted disentanglement loss and a alignment loss into a unified objective, demonstrating superior robustness across CIFAR-10/100 and Tiny-ImageNet on multiple architectures, while maintaining competitive clean accuracy and modest computational overhead. Overall, AFD offers a principled mechanism to mitigate feature-gap-induced vulnerabilities, advancing practical adversarial robustness with efficient fine-tuning.

Abstract

Adversarial fine-tuning methods enhance adversarial robustness via fine-tuning the pre-trained model in an adversarial training manner. However, we identify that some specific latent features of adversarial samples are confused by adversarial perturbation and lead to an unexpectedly increasing gap between features in the last hidden layer of natural and adversarial samples. To address this issue, we propose a disentanglement-based approach to explicitly model and further remove the specific latent features. We introduce a feature disentangler to separate out the specific latent features from the features of the adversarial samples, thereby boosting robustness by eliminating the specific latent features. Besides, we align clean features in the pre-trained model with features of adversarial samples in the fine-tuned model, to benefit from the intrinsic features of natural samples. Empirical evaluations on three benchmark datasets demonstrate that our approach surpasses existing adversarial fine-tuning methods and adversarial training baselines.

Figures (5)

• Figure 1: Illustration of our motivation. (a) $L_{\infty}$ distances of features between natural and adversarial samples during fine-tuning. Previous AFT methods exhibit an increasing trend of the feature gap. (b) Toy illustration of Disentanglement. We model features of adversarial samples and then remove the specific latent features confused by adversarial perturbation to correct the prediction.
• Figure 2: Overview of our AFD. Lock denotes the frozen parameters, $BZ$ denotes the batch size, $Cls$ denotes the class number, and numbers in brackets denote the sizes of features or parameters. During fine-tuning, Adversarial features are disentangled into $f_{1}$ and $f_{2}$. $\mathcal{L}_{CE}$ ensures $f_{2}$ to approximate confused features. We keep $f_{1}$ away from $f_{2}$ to eliminate confused features in $f_{1}$. Besides, we align $f_{1}$ with pre-trained intrinsic features $f_{nat}^{\prime}$ to further correct the prediction. During inference, we multiply $D_{1}$ and ${\omega}^F$ to obtain the robust model without additional modules.
• Figure 3: t-SNE visualization on CIFAR-10 on ResNet18.
• Figure 4: The clean accuracy (NAT) and robust accuracy against PGD-20 (PGD) of AFD with different hyperparameters.
• Figure 5: t-SNE visualization of different methods. Our AFD shows a clearer separation in the area of the red box.