Attention Head Purification: A New Perspective to Harness CLIP for Domain Generalization

TL;DR

This paper introduces attention head purification for CLIP to improve domain generalization by selecting and refining attention heads to focus on domain-invariant features, leading to better cross-domain performance.

Contribution

It proposes a novel attention head purification method involving task-level and domain-level purification to enhance CLIP's domain generalization capabilities.

Findings

Outperforms previous state-of-the-art methods on DG benchmarks.

Head-aware LoRA improves task adaptation of attention heads.

MMD loss encourages domain-invariant feature learning.

Abstract

Domain Generalization (DG) aims to learn a model from multiple source domains to achieve satisfactory performance on unseen target domains. Recent works introduce CLIP to DG tasks due to its superior image-text alignment and zeros-shot performance. Previous methods either utilize full fine-tuning or prompt-learning paradigms to harness CLIP for DG tasks. Those works focus on avoiding catastrophic forgetting of the original knowledge encoded in CLIP but ignore that the knowledge encoded in CLIP in nature may contain domain-specific cues that constrain its domain generalization performance. In this paper, we propose a new perspective to harness CLIP for DG, i.e., attention head purification. We observe that different attention heads may encode different properties of an image and selecting heads appropriately may yield remarkable performance improvement across domains. Based on such observations, we purify the attention heads of CLIP from two levels, including task-level purification and domain-level purification. For task-level purification, we design head-aware LoRA to make each head more adapted to the task we considered. For domain-level purification, we perform head selection via a simple gating strategy. We utilize MMD loss to encourage masked head features to be more domain-invariant to emphasize more generalizable properties/heads. During training, we jointly perform task-level purification and domain-level purification. We conduct experiments on various representative DG benchmarks. Though simple, extensive experiments demonstrate that our method performs favorably against previous state-of-the-arts.