Identifiable Contrastive Learning with Automatic Feature Importance Discovery

TL;DR

This paper introduces tri-factor contrastive learning (triCL), a novel method that enhances feature interpretability and identifiability by incorporating a learnable importance matrix, improving data representations in contrastive learning.

Contribution

The paper proposes triCL, extending contrastive learning with a learnable importance matrix to achieve identifiable and interpretable features, applicable to existing methods like SimCLR and CLIP.

Findings

TriCL produces identifiable features eliminating randomness.

Features with high importance are more interpretable and class-specific.

TriCL improves image retrieval performance with fewer features.

Abstract

Existing contrastive learning methods rely on pairwise sample contrast $z_x^\top z_{x'}$ to learn data representations, but the learned features often lack clear interpretability from a human perspective. Theoretically, it lacks feature identifiability and different initialization may lead to totally different features. In this paper, we study a new method named tri-factor contrastive learning (triCL) that involves a 3-factor contrast in the form of $z_x^\top S z_{x'}$, where $S=\text{diag}(s_1,\dots,s_k)$ is a learnable diagonal matrix that automatically captures the importance of each feature. We show that by this simple extension, triCL can not only obtain identifiable features that eliminate randomness but also obtain more interpretable features that are ordered according to the importance matrix $S$. We show that features with high importance have nice interpretability by capturing common classwise features, and obtain superior performance when evaluated for image retrieval using a few features. The proposed triCL objective is general and can be applied to different contrastive learning methods like SimCLR and CLIP. We believe that it is a better alternative to existing 2-factor contrastive learning by improving its identifiability and interpretability with minimal overhead. Code is available at https://github.com/PKU-ML/Tri-factor-Contrastive-Learning.