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

arXiv: https://arxiv.org/abs/2406.04527 · HTML: https://mlsanity.com/paper/2406.04527/

## Key facts

- Cited by 6 papers in the map (data through 14 Oct 2026).
- Builds on 21 of its 66 references that are in the map.
- First posted on arXiv 6 Jun 2024.

## Builds on (most cited first, 21 of 21)

- [Diffusion Models Beat GANs on Image Synthesis](https://mlsanity.com/paper/2105.05233/) (2021)
- [Flow Matching for Generative Modeling](https://mlsanity.com/paper/2210.02747/) (2022)
- [The Cityscapes Dataset for Semantic Urban Scene Understanding](https://mlsanity.com/paper/1604.01685/) (2016)
- [Neural Ordinary Differential Equations](https://mlsanity.com/paper/1806.07366/) (2018)
- [Density estimation using Real NVP](https://mlsanity.com/paper/1605.08803/) (2016)
- [Normalizing Flows for Probabilistic Modeling and Inference](https://mlsanity.com/paper/1912.02762/) (2019)
- [FFJORD: Free-form Continuous Dynamics for Scalable Reversible Generative Models](https://mlsanity.com/paper/1810.01367/) (2018)
- [Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions](https://mlsanity.com/paper/2102.05379/) (2021)
- [PixelCNN++: Improving the PixelCNN with Discretized Logistic Mixture Likelihood and Other Modifications](https://mlsanity.com/paper/1701.05517/) (2017)
- [A note on the evaluation of generative models](https://mlsanity.com/paper/1511.01844/) (2015)
- [Flow++: Improving Flow-Based Generative Models with Variational Dequantization and Architecture Design](https://mlsanity.com/paper/1902.00275/) (2019)
- [Dirichlet Flow Matching with Applications to DNA Sequence Design](https://mlsanity.com/paper/2402.05841/) (2024)
- [An Introduction to Deep Generative Modeling](https://mlsanity.com/paper/2103.05180/) (2021)
- [Fisher Flow Matching for Generative Modeling over Discrete Data](https://mlsanity.com/paper/2405.14664/) (2024)
- [RNADE: The real-valued neural autoregressive density-estimator](https://mlsanity.com/paper/1306.0186/) (2013)
- [Natural gradient via optimal transport](https://mlsanity.com/paper/1803.07033/) (2018)
- [Semi-Discrete Normalizing Flows through Differentiable Tessellation](https://mlsanity.com/paper/2203.06832/) (2022)
- [Image Labeling by Assignment](https://mlsanity.com/paper/1603.05285/) (2016)
- [Generative Modeling of Discrete Joint Distributions by E-Geodesic Flow Matching on Assignment Manifolds](https://mlsanity.com/paper/2402.07846/) (2024)
- [Assignment Flows for Data Labeling on Graphs: Convergence and Stability](https://mlsanity.com/paper/2002.11571/) (2020)
- [A Geometric Embedding Approach to Multiple Games and Multiple Populations](https://mlsanity.com/paper/2401.05918/) (2024)

## Built on by (most cited first, 6 of 6)

- [PnP-Flow: Plug-and-Play Image Restoration with Flow Matching](https://mlsanity.com/paper/2410.02423/) (2024)
- [α-Flow: A Unified Framework for Continuous-State Discrete Flow Matching Models](https://mlsanity.com/paper/2504.10283/) (2025)
- [Sigma Flows for Image and Data Labeling and Learning Structured Prediction](https://mlsanity.com/paper/2408.15946/) (2024)
- [MolPIF: A Parameter Interpolation Flow Model for Molecule Generation](https://mlsanity.com/paper/2507.13762/) (2025)
- [Riemannian Patch Assignment Gradient Flows](https://mlsanity.com/paper/2504.13024/) (2025)
- [Generative Modeling of Discrete Data Using Geometric Latent Subspaces](https://mlsanity.com/paper/2601.21831/) (2026)

## Abstract (arXiv)

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.