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

TL;DR

The paper presents attention head purification as a new paradigm to harness CLIP for domain generalization. It introduces two complementary mechanisms: task-level head-aware LoRA (HA-LoRA) for per-head task adaptation and domain-level domain-invariant gating (DIG) to select heads that generalize across domains, coupled with an MMD-based objective to encourage domain-invariant representations. The approach is evaluated across five DG benchmarks, showing consistent improvements over zero-shot CLIP and prior DG methods, and it is shown to be compatible with various prompt-learning strategies. The findings highlight that selectively purifying attention heads can mitigate domain-specific cues embedded in CLIP and substantially enhance cross-domain performance with efficient training and no extra inference cost.

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.

Figures (7)

• Figure 1: Left: We use different strategies to evaluate the importance of each attention head on domain generalization and drop the least important ones to see how accuracy changes. The strategies we adopt include "randomly drop(--$\lozenge$--)", "manually drop(--$\CIRCLE$--)", "drop by cross-validation(--$\blacktriangle$--)", and "adapt & drop(--$\blacksquare$--)". Details of strategies can be found in Appendix. Via appropriately dropping heads, CLIP's domain generalization performance can be improved. Right: Attention map attentionmap generated by specific heads. The middle columns are from least important heads determined by the cross-validation strategy. They all capture lots of background information. The last column represents the most important one whose attention map mainly focuses on the object itself. The experiments are conducted on OfficeHome officehome. Best viewed in color.
• Figure 2: The architecture of Attention Head Purification. We design head-aware LoRA to perform task-level purification and domain-invariant gating strategy to perform domain-level purification. Further, we minimize $\mathcal{L}_{cls}$ (Section \ref{['ha-lora']}) to update head-aware LoRA and minimize $\mathcal{L}_{cls},\mathcal{L}_{\mathrm{MMD}}$ (Section \ref{['domainpuri']}) to update the gates of heads.
• Figure 3: Attention maps for target samples in DomainNet.
• Figure 4: t-SNE visualization of features from target domain on OfficeHome. We compare the visualization results of the original feature (black dot) and our method (red dot).
• Figure 5: Comparison of the overall attention. Images are sampled from the DomainNet dataset. The "Real" is selected as the target domain while others are selected as source domains to train the CLIP.