[Paper Review] Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions

Emiel Hoogeboom, Didrik Nielsen+3 more|arXiv (Cornell University)|Feb 10, 2021

Music and Audio Processing79 citations

TL;DR

The paper introduces Argmax Flows and Multinomial Diffusion to learn categorical distributions, enabling normalizing flows and diffusion models to handle discrete data like text and segmentation maps, and shows improved log-likelihood over standard dequantization methods.

ABSTRACT

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.

Motivation & Objective

• Motivate learning high-dimensional categorical data with generative flows and diffusion models.
• Extend normalizing flows to direct categorical distributions via argmax-based lifting and probabilistic inverses.
• Propose a diffusion model directly on categorical variables (Multinomial Diffusion).
• Demonstrate improved log-likelihood on language modeling and segmentation datasets compared to dequantization baselines.

Proposed method

• Argmax Flows combine a continuous density p(v) with an argmax v -> x mapping and a probabilistic inverse q(v|x) that satisfies the argmax constraint.
• Three probabilistic inverses for q(v|x) are explored: Thresholding, Gumbel-based, and Gumbel Thresholding, each with log-determinant terms for training.
• Optimization uses a variational lower bound (ELBO) where log P(x) >= E_{v~q(v|x)}[log p(v) - log q(v|x)].
• Cartesian products of Argmax Flows expand to base-K spaces via base-M representations to trade symmetry for dimensionality.
• Multinomial Diffusion defines a diffusion process directly on one-hot categorical vectors with a fixed denoising trajectory and closed-form posterior q(x_{t-1}|x_t,x_0).
• The model predicts either p(v) via flows (Argmax Flows) or learns a denoising trajectory p(x_{t-1}|x_t) in the categorical space.

Experimental results

Research questions

• RQ1 Can Argmax Flows effectively lift discrete categorical data into a continuous space for density modeling with normalizing flows?
• RQ2 How can a probabilistic inverse q(v|x) satisfying the argmax constraint be constructed and trained?
• RQ3 Does a diffusion process defined on categorical variables (Multinomial Diffusion) yield better log-likelihoods than dequantization baselines for discrete data?
• RQ4 What are the trade-offs between autoregressive and non-autoregressive (coupling) flow architectures when modeling text and segmentation maps?
• RQ5 How do Cartesian products of Argmax Flows affect performance and model capacity for large category counts?

Key findings

• Argmax Flows consistently outperform uniform and variational dequantization on text8 and enwik8 in log-likelihood (bits per character/byte).
• Autoregressive Argmax Flows achieve better performance than several VAE and discrete-flow baselines on language data.
• Non-autoregressive Argmax Flows (coupling) improve over some discrete-flow baselines but lag behind autoregressive variants on text; Multinomial Diffusion can outperform some non-autoregressive Argmax variants in certain settings.
• In image segmentation maps (Cityscapes), Argmax Flows and Multinomial Diffusion achieve competitive ELBO and IWBO scores, outperforming traditional dequantization methods.
• The results reveal a performance gap between standard autoregressive models, continuous autoregressive models, and non-autoregressive continuous models, suggesting directions for future work in discrete-to-continuous lifting and autoregressive components.