Generative Assignment Flows for Representing and Learning Joint Distributions of Discrete Data

TL;DR

This paper presents a new generative model using measure transport via randomized assignment flows to efficiently represent and sample from complex joint distributions of many discrete variables, with applications to large-scale structured data.

Contribution

The paper introduces a novel measure transport approach with assignment flows for modeling joint distributions of discrete data, scalable to large class sets.

Findings

Efficient sampling and likelihood assessment for complex discrete distributions.

Scales better with increasing classes compared to recent methods.

Demonstrated applicability to large-scale structured image labeling problems.

Abstract

We introduce a novel generative model for the representation of joint probability distributions of a possibly large number of discrete random variables. The approach uses measure transport by randomized assignment flows on the statistical submanifold of factorizing distributions, which enables to represent and sample efficiently from any target distribution and to assess the likelihood of unseen data points. The complexity of the target distribution only depends on the parametrization of the affinity function of the dynamical assignment flow system. Our model can be trained in a simulation-free manner by conditional Riemannian flow matching, using the training data encoded as geodesics on the assignment manifold in closed-form, with respect to the e-connection of information geometry. Numerical experiments devoted to distributions of structured image labelings demonstrate the applicability to large-scale problems, which may include discrete distributions in other application areas. Performance measures show that our approach scales better with the increasing number of classes than recent related work.