TWINS: A Fine-Tuning Framework for Improved Transferability of Adversarial Robustness and Generalization

TL;DR

This paper introduces TWINS, a novel fine-tuning framework that enhances both adversarial robustness and generalization of pre-trained models across various image classification tasks by maintaining robust feature transfer and improving training dynamics.

Contribution

The paper proposes TWINS, a statistics-based fine-tuning method that preserves robust features and accelerates training, outperforming existing approaches in robustness and generalization.

Findings

TWINS improves robustness and generalization across multiple datasets.

It accelerates training by increasing effective learning rate.

It alleviates robust overfitting during fine-tuning.

Abstract

Recent years have seen the ever-increasing importance of pre-trained models and their downstream training in deep learning research and applications. At the same time, the defense for adversarial examples has been mainly investigated in the context of training from random initialization on simple classification tasks. To better exploit the potential of pre-trained models in adversarial robustness, this paper focuses on the fine-tuning of an adversarially pre-trained model in various classification tasks. Existing research has shown that since the robust pre-trained model has already learned a robust feature extractor, the crucial question is how to maintain the robustness in the pre-trained model when learning the downstream task. We study the model-based and data-based approaches for this goal and find that the two common approaches cannot achieve the objective of improving both generalization and adversarial robustness. Thus, we propose a novel statistics-based approach, Two-WIng NormliSation (TWINS) fine-tuning framework, which consists of two neural networks where one of them keeps the population means and variances of pre-training data in the batch normalization layers. Besides the robust information transfer, TWINS increases the effective learning rate without hurting the training stability since the relationship between a weight norm and its gradient norm in standard batch normalization layer is broken, resulting in a faster escape from the sub-optimal initialization and alleviating the robust overfitting. Finally, TWINS is shown to be effective on a wide range of image classification datasets in terms of both generalization and robustness. Our code is available at https://github.com/ziquanliu/CVPR2023-TWINS.