Masked Images Are Counterfactual Samples for Robust Fine-tuning

TL;DR

This paper introduces a novel fine-tuning approach using masked images as counterfactual samples to enhance model robustness against distribution shifts, effectively balancing in-distribution accuracy and out-of-distribution robustness.

Contribution

The proposed method leverages masked images as counterfactuals for fine-tuning, explicitly addressing OOD robustness and improving the ID-OOD performance trade-off.

Findings

Outperforms previous methods on OOD robustness benchmarks

Improves the trade-off between in-distribution and out-of-distribution performance

Uses class activation maps to generate effective counterfactual samples

Abstract

Deep learning models are challenged by the distribution shift between the training data and test data. Recently, the large models pre-trained on diverse data have demonstrated unprecedented robustness to various distribution shifts. However, fine-tuning these models can lead to a trade-off between in-distribution (ID) performance and out-of-distribution (OOD) robustness. Existing methods for tackling this trade-off do not explicitly address the OOD robustness problem. In this paper, based on causal analysis of the aforementioned problems, we propose a novel fine-tuning method, which uses masked images as counterfactual samples that help improve the robustness of the fine-tuning model. Specifically, we mask either the semantics-related or semantics-unrelated patches of the images based on class activation map to break the spurious correlation, and refill the masked patches with patches from other images. The resulting counterfactual samples are used in feature-based distillation with the pre-trained model. Extensive experiments verify that regularizing the fine-tuning with the proposed masked images can achieve a better trade-off between ID and OOD performance, surpassing previous methods on the OOD performance. Our code is available at https://github.com/Coxy7/robust-finetuning.