Generalization Beyond Feature Alignment: Concept Activation-Guided Contrastive Learning

TL;DR

This paper identifies that traditional feature alignment in contrastive learning can hinder domain generalization by reducing feature diversity, and proposes Concept Contrast (CoCo) to enhance invariant feature learning and improve generalization.

Contribution

It introduces CoCo, a neuron activation-based contrastive approach that relaxes feature alignment, promoting diversity and better generalization in contrastive learning methods.

Findings

CoCo improves generalization across four contrastive methods.

It increases the diversity of learned feature representations.

CoCo activates more meaningful neurons during training.

Abstract

Learning invariant representations via contrastive learning has seen state-of-the-art performance in domain generalization (DG). Despite such success, in this paper, we find that its core learning strategy -- feature alignment -- could heavily hinder model generalization. Drawing insights in neuron interpretability, we characterize this problem from a neuron activation view. Specifically, by treating feature elements as neuron activation states, we show that conventional alignment methods tend to deteriorate the diversity of learned invariant features, as they indiscriminately minimize all neuron activation differences. This instead ignores rich relations among neurons -- many of them often identify the same visual concepts despite differing activation patterns. With this finding, we present a simple yet effective approach, Concept Contrast (CoCo), which relaxes element-wise feature alignments by contrasting high-level concepts encoded in neurons. Our CoCo performs in a plug-and-play fashion, thus it can be integrated into any contrastive method in DG. We evaluate CoCo over four canonical contrastive methods, showing that CoCo promotes the diversity of feature representations and consistently improves model generalization capability. By decoupling this success through neuron coverage analysis, we further find that CoCo potentially invokes more meaningful neurons during training, thereby improving model learning.