FD-Align: Feature Discrimination Alignment for Fine-tuning Pre-Trained Models in Few-Shot Learning

Code  Poster  知乎

Background

CLIP demonstrates exceptional performance across various visual tasks. When applying it to downstream tasks, it often requires fine-tuning on the downstream data. However, in few-shot learning, directly fine-tuning CLIP can easily lead to overfitting and adversely affect its generalization on out-of-distribution (OOD) data. Consequently, previous methods have mostly attempted to fine-tune only the classification head or introduce additional structures, but these approaches do not fully exploit the potential of CLIP’s visual encoder. Therefore, our research explores how to fine-tune CLIP with a sample number of samples without decreasing its performance on OOD data.

Motivation

We explain the decline in OOD performance of fully fine-tuned CLIP through the robustness to spurious correlations. This robustness refers to the model's ability to distinguish between causally relevant information (related to the class) and irrelevant information (such as background or style), which are spurious. Previous work has shown that CLIP exhibits strong robustness against spurious correlations, thus ensuring impressive OOD performance [1]. However, fully fine-tuning of CLIP results in a decrease in its OOD performance [2]. Visualizing the attention maps of both CLIP and its fully fine-tuned counterpart reveals that the fully fine-tuned CLIP tends to focus more on localized features. This heightened attention to local areas weakens the model's robustness against spurious correlations [3]. In other words, while fully fine-tuning enables the model to better learn the causal features of the fine-tuning samples, it diminishes its capacity to recognize spurious correlations, which leads to the casual features extracted by the model failing to generalize to unseen samples, resulting in overfitting and impaired generalization on OOD data. Therefore, we proposes a fine-tuning methodology that does not compromise the model's ability to recognize spurious correlations, thereby preserving the robustness of the fine-tuned model against spurious correlations.

Method

We aim to fine-tune CLIP’s visual encoder $f_0$ on a few-shot dataset $\mathcal{D}\subset\mathcal{X}\times\mathcal{Y}$ to obtain $f_t$. Our approach is divided into two parts: model fine-tuning and spurious feature constraint.

Model fine-tuning

During the fine-tuning process, CLIP's text encoder $g_0$ is frozen. We input images $x$ along with their corresponding label $y$ and $M$ prompt templates $(P_1, \cdots, P_M)$. For any given category $y$, its associated prompts are denoted as $[P_1, y], \cdots, [P_M,y]$. For each category $y$, we utilize the text encoder to extract features from all prompts, which are then used to compute the prototype of that category.

We calculate the similarity between the image and the category prototypes using cosine similarity, and then compute the probability density accordingly. Subsequently, we employ cross-entropy loss as the classification loss during the fine-tuning process

Where $\mathcal{Y}$ is the set of label.

Spurious Feature Constraint

Ensuring the constancy of spurious features during fine-tuning is most directly achieved by decoupling causal and spurious features, while maintaining the stability of the spurious features. However, decoupling features in images is an extremely challenging task. In contrast, it is much simpler to decouple spurious and causal features in text. For example, in the prompt 'a photo of a dog,' 'dog' is the causal feature, whereas 'a photo of a' represents spurious features. Leveraging CLIP's strong alignment capabilities between vision and text, we can use the spurious features in text as prototypes for the spurious information in images

We can obtain the probability distribution of image features extracted by the fine-tuned model on spurious information:

Meanwhile, We can obtain the probability distribution of image feature extracted by CLIP on spurious information:

Although it is not feasible to directly ensure that the spurious features in the image features extracted by the fine-tuned model and the original CLIP are identical, we can indirectly ensure consistency of the extracted spurious features by constraining their probability distributions on spurious information to be the same.

Finally, we constrain both the classification loss and the spurious consistency loss to ensure the OOD generalization of the model during fine-tuning.

Spurious Prototype Correction

Currently, most prompt templates are either manually designed or generated by language models, which can lead to unreasonable or redundant scenarios, resulting in inaccurate prototypes of spurious information. To address this issue, we initially employ anomaly detection algorithms to remove unreasonable prompt features.

Subsequently, we utilize the $k-Means$ clustering algorithm to deal with the redundant features in the prompts.

Experiment

Results on OOD

As shwon in the table below, we conducted 16-shot fine-tuning of CLIP on ImageNet and evaluated its performance on two variant datasets of ImageNet. Compared to fully fine-tuning, FD-Align achieves a comprehensive improvement in OOD performance. Additionally, by directly applying the fine-tuned visual encoder to Tip and APE, it is evident that our fine-tuned model can be applied to existing methods without the need for re-tuning, thereby enhancing OOD performance.

Similarly, we fine-tuned CLIP on miniImageNet following the $N$-way $K$-shot few-shot learning, and evaluated its performance on various downstream datasets. The following figure illustrates the performance changes of the fine-tuned model across different datasets. FD-Align brings significant improvements in OOD performance on most of the datasets.

Results on ID

We also evaluated the in-distribution (ID) performance of our method across 11 datasets. Our approach demonstrates a significant improvement in performance, which becomes even more pronounced as the number of shots increases.

Similarly, we applied the fine-tuned visual encoder directly to existing methods. The table below shows the performance on ImageNet. It is evident that our fine-tuned model can also directly enhance the ID performance of existing methods.

The Necessity of Correcting Spurious Prototypes

As shown in the table below, we compare the ID performance using all prompt features as spurious prototypes, prompt features manually filtered by Tip as spurious prototypes, and spurious prototypes corrected using Spurious Prototype Correction (SPC). As illustrated, prototypes corrected with SPC achieve higher performance compared to using all prompts. Notably, the performance significantly drops when using prototypes from manually filtered prompts by Tip. We attribute this to the retention of the template ‘itap of a {class}’ in Tip, which is removed as an outlier in SPC. Therefore, SPC module can avoid the irrationality of manual filtering.

Visualization

We also visualized the image features extracted by different models. FD-Align can distinguish different classes and domains better.

Reference

[1] Self-supervision on images and text reduces reliance on visual shortcut features.

[2] Fine-tuning can distort pretrained features and underperform out-of-distribution.

[3] Are vision transformers robust to spurious correlations?