[Paper Review] Self-Supervised Learning Disentangled Group Representation as Feature

Tan Wang, Zhongqi Yue+3 more|arXiv (Cornell University)|Oct 28, 2021

Domain Adaptation and Few-Shot Learning4 citations

TL;DR

This paper proposes Iterative Partition-based Invariant Risk Minimization (IP-IRM), a self-supervised learning method that disentangles visual representations by iteratively partitioning data into subsets corresponding to group elements and enforcing invariance via contrastive learning. It achieves state-of-the-art performance on multiple benchmarks by learning fully disentangled features that separate semantic factors like digit and color, outperforming existing SSL methods in downstream linear evaluation and intervention-based few-shot learning.

ABSTRACT

A good visual representation is an inference map from observations (images) to features (vectors) that faithfully reflects the hidden modularized generative factors (semantics). In this paper, we formulate the notion of "good" representation from a group-theoretic view using Higgins' definition of disentangled representation, and show that existing Self-Supervised Learning (SSL) only disentangles simple augmentation features such as rotation and colorization, thus unable to modularize the remaining semantics. To break the limitation, we propose an iterative SSL algorithm: Iterative Partition-based Invariant Risk Minimization (IP-IRM), which successfully grounds the abstract semantics and the group acting on them into concrete contrastive learning. At each iteration, IP-IRM first partitions the training samples into two subsets that correspond to an entangled group element. Then, it minimizes a subset-invariant contrastive loss, where the invariance guarantees to disentangle the group element. We prove that IP-IRM converges to a fully disentangled representation and show its effectiveness on various benchmarks. Codes are available at https://github.com/Wangt-CN/IP-IRM.

Motivation & Objective

• To formalize 'good' visual representation using group theory and Higgins' definition of disentangled representation.
• To identify the limitation of existing SSL methods in disentangling complex semantics beyond simple augmentations like rotation and colorization.
• To develop a method that systematically disentangles multiple independent semantic factors in visual representations.
• To achieve fully disentangled representations through iterative partitioning and invariant contrastive learning.
• To validate the method's superiority on standard benchmarks and in robust few-shot learning scenarios.

Proposed method

• IP-IRM iteratively partitions training samples into two subsets corresponding to a group element acting on the semantic space.
• At each iteration, it minimizes a subset-invariant contrastive loss to enforce invariance under the group action, thereby disentangling the corresponding semantic factor.
• The method leverages group-theoretic principles to ensure equivariance and decomposability of the learned representation across semantic factors.
• It alternates between partitioning and optimizing the contrastive objective until convergence to a fully disentangled representation.
• The algorithm is grounded in invariant risk minimization (IRM) principles, ensuring robustness to spurious correlations.
• The framework is implemented end-to-end with deep neural networks and trained using standard contrastive learning objectives.

Experimental results

Research questions

• RQ1 Can self-supervised learning methods fully disentangle complex, modularized semantic factors such as digit and color in images?
• RQ2 Why do existing SSL methods fail to disentangle semantics beyond simple augmentations like rotation and colorization?
• RQ3 Can iterative partitioning and group-invariant contrastive learning lead to a fully disentangled representation?
• RQ4 How does the disentangled representation from IP-IRM compare to supervised and other SSL methods in downstream linear evaluation?
• RQ5 Does the disentangled representation improve robustness in intervention-based few-shot learning?

Key findings

• IP-IRM achieves 89.38% accuracy on ImageNet-100 linear evaluation, outperforming MoCo-v2 (88.10%) and other SOTA SSL methods.
• On the IFSL benchmark, IP-IRM achieves 75.57% accuracy, significantly outperforming the best baseline (MoCo-v2 at 74.38%).
• The method shows superior robustness in few-shot learning, with a 2.22% performance gain over MoCo-v2 when using the IFSL classifier.
• Visualization shows that IP-IRM disentangles features into distinct dimensions corresponding to augmentation-related and unrelated semantics, with high classification accuracy on each.
• The ablation study confirms that iterative partitioning and group-invariant contrastive learning are essential for achieving full disentanglement.
• IP-IRM converges to a fully disentangled representation, as proven theoretically, with each semantic factor localized in a separate subspace of the feature space.