Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions

By Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forr'e, Max Welling. Neural Information Processing Systems, 2021

710 citations recorded by the source corpus.

Generative flows and diffusion models have been predominantly trained on ordinal data, for example natural images. This paper introduces two extensions of flows and diffusion for categorical data such as language or image segmentation: Argmax Flows and Multinomial Diffusion. Argmax Flows are defined by a composition of a continuous distribution (such as a normalizing flow), and an argmax function. To optimize this model, we learn a probabilistic inverse for the argmax that lifts the categorical data to a continuous space. Multinomial Diffusion gradually adds categorical noise in a diffusion process, for which the generative denoising process is learned. We demonstrate that our method outperforms existing dequantization approaches on text modelling and modelling on image segmentation maps in log-likelihood.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Signal Processing

View Paper: https://www.semanticscholar.org/paper/1913d3edcc00d0aba097a9df190dd16f6fdfbf0c

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

In Computer Vision and Pattern Recognition

In Signal Processing

Citations in the corpus

Discrete Flow Matching (2024, background): “D3PM (Austin et al., 2021a) and Argmax flows (Hoogeboom et al., 2021) introduced diffusion in discrete spaces by proposing a corruption process for categorical data.”
DiffuSeq: Sequence to Sequence Text Generation with Diffusion Models (2022, methodology): “For text generation using the diffusion models, Hoogeboom et al. (2021) introduce the multinomial diffusion for character-level text generation, the forward categorical noise is applied through the Markov transition matrix.”
Continuous diffusion for categorical data (2022, methodology): “…and diﬀusion-inspired approaches have been proposed for non-autoregressive iterative reﬁnement of discrete data, and especially for language in particular (Austin et al., 2021; Chang et al., 2022; Ghazvininejad et al., 2019; Hoogeboom et al., 2021a,b; Reid et al., 2022; Savinov et al., 2021).”
Deep Generative Modelling: A Comparative Review of VAEs, GANs, Normalizing Flows, Energy-Based and Autoregressive Models (2021, methodology): “Diffusion models have also been applied to categorical data; multinomial diffusions [93] define a forward process where each discrete variable switches randomly to a different value and the reverse process is trained to approximate the noise.”
Self-conditioned Embedding Diffusion for Text Generation (2022, methodology): “It was extended to multinomial text modeling (Hoogeboom et al., 2021) and further scaled up in the D3PM work (Austin et al.”
Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Data (2024, neutral): “Several discrete-state diffusion models have been proposed (Sohl-Dickstein et al., 2015; Hoogeboom et al., 2021; Austin et al., 2021).”
Diffusion Models: A Comprehensive Survey of Methods and Applications (2022, methodology): “It generalize the multinomial diffusion model [96] through going beyond corruption processes with uniform transition probabilities.”
Equivariant 3D-conditional diffusion model for molecular linker design (2022, neutral): “To illustrate this difference, we adapted DeLinker and 3DLinker to iteratively connect pairs of fragments in molecules where more than two fragments should be connected and tested all the methods on the GEOM dataset.”
Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioning (2022, methodology): “Due to the discreteness, existing work has to re-formulate continuous diffusion models by adopting a discrete data space and state space (Sohl-Dickstein et al., 2015; Hoogeboom et al., 2021; Austin et al., 2021).”
Fast Sampling of Diffusion Models with Exponential Integrator (2022, background): “…Rombach et al., 2021; Dhariwal & Nichol, 2021), text conditioned image generation (Nichol et al., 2021; Ramesh et al., 2022), text generation (Hoogeboom et al., 2021; Austin et al., 2021), 3D point cloud generation (Lyu et al., 2021), inverse problem (Kawar et al., 2021; Song et al., 2021b),…”
3D Equivariant Diffusion for Target-Aware Molecule Generation and Affinity Prediction (2023, methodology): “Thanks to recent progress in probabilistic diffusion models (Ho et al., 2020; Hoogeboom et al., 2021) and equivariant neural networks (Fuchs et al., 2020; Satorras et al., 2021b), our proposed model can generate molecules in continuous 3D space based on the context provided by protein atoms, and…”
Diffusion models as plug-and-play priors (2022, methodology): “Alternatively, we could use diffusion models for categorical data [14] with the appropriate modifications to our inference procedure.”
Diffsound: Discrete Diffusion Model for Text-to-Sound Generation (2022, background): “[38] extended the model to categorical random variables with transition matrices characterized by uniform transition probabilities.”
InstructTTS: Modelling Expressive TTS in Discrete Latent Space With Natural Language Style Prompt (2023, background): “Diffusion models with discrete state spaces are ﬁrst investigated by [43], [44], [45].”
Large Language Diffusion Models (2025, neutral): “Another approach replaces continuous diffusion with discrete processes featuring new forward and reverse dynamics, leading to numerous variants [58–71].”
Simplified and Generalized Masked Diffusion for Discrete Data (2024, neutral): “While diffusion models have been extended to discrete state spaces [1, 14, 15] and have been successfully applied to applications ranging from graph generation [16], text-to-sound generation [17] or protein design [18], they remain not as widely used as their continuous counterparts as they are not…”
BLOCK DIFFUSION: INTERPOLATING BETWEEN AU-TOREGRESSIVE AND DIFFUSION LANGUAGE MODELS (2025, neutral): “Autoregressive diffusion models (Hoogeboom et al., 2021b;a) extend any-order AR models (AO-ARMs; Uria et al. (2014)) to support parallel sampling.”