Alignment-aware fine-tuning of vision-language models for out-of-distribution generalization
摘要
Out-of-distribution (OOD) generalization enhances the capacity of the trained model to handle distribution shifts. Through large-scale training, pre-trained vision-language models establish well-aligned vision-language spaces, thus possessing good zero-shot generalization capabilities. Adapting these models to downstream tasks often requires fine-tuning for task-specific knowledge. However, independent optimization of vision and language branches employed in traditional fine-tuning disrupts the crucial vision-language alignment. Moreover, the fixed prompts often fail to leverage the rich semantic knowledge embedded in text modality. In this paper, we present AAFT, i.e., Alignment-Aware Fine-Tuning, to improve the OOD generalization of fine-tuned vision-language models on downstream tasks. First, we propose an Alignment Optimization Module (AOM) to maintain alignment by promoting better consistency with the previous step model as a teacher. AOM dynamically guides the optimization by a self-distillation objective from vision and vision-language alignment perspectives. It further implements a model ensemble strategy weighted by alignment scores to ensure reliability. Second, we propose a Collaborative Prompt Module (CPM) to enhance text prompts by refining learnable text embeddings. These embeddings are initialized with class-aware descriptions generated from a Large Language Model (LLM). Extensive experiments across three OOD scenarios demonstrate the outstanding performance of AAFT.