i-Mix: A Strategy for Regularizing Contrastive Representation Learning

ICLR (2021)

Kibok Lee (University of Michigan), Yian Zhu (University of Michigan), Kihyuk Sohn (Google Cloud AI), Chun-Liang Li (Google Cloud AI), Jinwoo Shin (KAIST), Honglak Lee

Abstract

Contrastive representation learning has shown to be effective to learn representations from unlabeled data. However, much progress has been made in vision domains relying on data augmentations carefully designed using domain knowledge. In this work, we propose i-Mix, a simple yet effective domain-agnostic regularization strategy for improving contrastive representation learning. We cast contrastive learning as training a non-parametric classifier by assigning a unique virtual class to each data in a batch. Then, data instances are mixed in both the input and virtual label spaces, providing more augmented data during training. In experiments, we demonstrate that i-Mix consistently improves the quality of learned representations across domains, including image, speech, and tabular data. Furthermore, we confirm its regularization effect via extensive ablation studies across model and dataset
sizes.