FeDi: Feature disentanglement for self-supervised learning

Abstract

Self-supervised learning (SSL) has revolutionized the field of deep learning by enabling the extraction of meaningful representations from unlabeled data. In this work, we introduce FeDi, a novel SSL method that leverages feature disentanglement to enhance the quality and robustness of learned representations. FeDi maximizes the lower bound on mutual information between representation vectors across batch dimensions, effectively disentangling features and preventing representation collapse. Our proposed method serves as a hardness-aware loss function that automatically balances alignment and disentanglement terms, effectively managing the challenges of disentangling high-dimensional representations. Our extensive experiments demonstrate that FeDi consistently outperforms state-of-the-art SSL methods across a variety of tasks, including image classification, object detection, and segmentation. Code is available at: https://github.com/mongeoroo/fedi.