An Improved Semi-Supervised VAE for Learning Disentangled Representations

TL;DR

This paper proposes an improved semi-supervised VAE framework that enhances disentangled representation learning by incorporating label replacement, leading to significant improvements with limited supervision on synthetic and real datasets.

Contribution

It introduces a novel label replacement supervision method within semi-supervised VAEs for better disentanglement learning, extending prior work by Locatello et al. (2019).

Findings

Significant improvement in disentanglement quality.

Effective with very limited supervision.

Validated on synthetic and real datasets.

Abstract

Learning interpretable and disentangled representations is a crucial yet challenging task in representation learning. In this work, we focus on semi-supervised disentanglement learning and extend work by Locatello et al. (2019) by introducing another source of supervision that we denote as label replacement. Specifically, during training, we replace the inferred representation associated with a data point with its ground-truth representation whenever it is available. Our extension is theoretically inspired by our proposed general framework of semi-supervised disentanglement learning in the context of VAEs which naturally motivates the supervised terms commonly used in existing semi-supervised VAEs (but not for disentanglement learning). Extensive experiments on synthetic and real datasets demonstrate both quantitatively and qualitatively the ability of our extension to significantly and consistently improve disentanglement with very limited supervision.