Flow Factorized Representation Learning

TL;DR

This paper introduces Flow Factorized Representation Learning, a new generative model that learns structured, flexible, and robust representations by modeling input transformations through latent probability paths driven by optimal transport.

Contribution

It proposes a novel framework that improves disentanglement and equivariance in representations using flow-based latent paths and optimal transport, outperforming existing methods.

Findings

Achieves higher likelihoods on standard benchmarks.

Produces representations closer to equivariant models.

Transforms are flexibly composable and extrapolate well.

Abstract

A prominent goal of representation learning research is to achieve representations which are factorized in a useful manner with respect to the ground truth factors of variation. The fields of disentangled and equivariant representation learning have approached this ideal from a range of complimentary perspectives; however, to date, most approaches have proven to either be ill-specified or insufficiently flexible to effectively separate all realistic factors of interest in a learned latent space. In this work, we propose an alternative viewpoint on such structured representation learning which we call Flow Factorized Representation Learning, and demonstrate it to learn both more efficient and more usefully structured representations than existing frameworks. Specifically, we introduce a generative model which specifies a distinct set of latent probability paths that define different input transformations. Each latent flow is generated by the gradient field of a learned potential following dynamic optimal transport. Our novel setup brings new understandings to both \textit{disentanglement} and \textit{equivariance}. We show that our model achieves higher likelihoods on standard representation learning benchmarks while simultaneously being closer to approximately equivariant models. Furthermore, we demonstrate that the transformations learned by our model are flexibly composable and can also extrapolate to new data, implying a degree of robustness and generalizability approaching the ultimate goal of usefully factorized representation learning.