Source-linked AI summary
Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Emiel Hoogeboom, Didrik Nielsen, Priyank Jaini, Patrick Forré, Max Welling
TL;DR
Categorical data has been less extensively modeled by flows and diffusion than ordinal data, despite applications such as language and image segmentation. The paper introduces Argmax Flows with probabilistic argmax inverses and Multinomial Diffusion defined directly on categorical variables. Experiments show that the methods outperform comparable models in negative log-likelihood, while the paper notes scope limitations for direct applications.
Problem
Generative flows and diffusion models have mainly been trained on ordinal data, leaving categorical data such as language and image segmentation less extensively treated.
Method
The paper introduces Argmax Flows, which lift categorical data through a probabilistic inverse of argmax, and Multinomial Diffusion, which directly models categorical diffusion and denoising.
Results
The experiments show that the proposed methods outperform comparable models in negative log-likelihood, including traditional dequantization approaches in language modeling.
Takeaways & Limitations
The methods enable generative normalizing flows and diffusion models to learn categorical data directly within the demonstrated text and image-segmentation settings.
Takeaways & Limitations
The methods are still distant from direct applications such as text generation, automated chat bots, or improved self-driving-car segmentation.
Abstract
from arXiv · showhide
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.
1 Introduction
Categorical data such as language and image segmentation has received less treatment from flows and diffusion models than ordinal data, while autoregressive models are expensive to sample. The paper addresses this gap with categorical extensions of both model families.
- Categorical sources include language and image segmentation, but generative flows and diffusion models have mainly been applied to ordinal data.
- Autoregressive models primarily model categorical data, but they are expensive to sample from.
- Argmax Flows connect categorical variables to continuous normalizing flows through an argmax transformation and probabilistic inverses.
- Multinomial Diffusion is a diffusion model defined directly on categorical variables, using a fixed categorical-noise trajectory and learned denoising process.
- Direct discrete diffusion avoids gradient approximations because its diffusion trajectory is fixed.
2 Background
Continuous normalizing flows require continuous densities, so discrete data is commonly handled by dequantization, while diffusion models gradually destroy information through predefined noise distributions and learn reverse denoising distributions.
- Normalizing Flows: Normalizing flows learn bijective differentiable transformations and exact densities, but their change-of-variables formulation is restricted to continuous densities.
- Dequantization: Dequantization adds noise to ordinal discrete data and uses a probabilistic right-inverse of rounding to lift observations into continuous space.
- Dequantization: The variational dequantization objective is an evidence lower bound involving log p(v) minus log q(v|x), under support restricted to the rounding region.
- Diffusion Models: Diffusion models use predefined distributions q(x_t|x_t−1) that gradually add noise until the final state carries almost no information about the original data.
- Diffusion Models: Conditioning on x_0 makes q(x_t−1|x_t,x_0) tractable, enabling a lower-variance reformulation using KL divergences.
3 Argmax Flows
Argmax Flows model categorical variables by partitioning continuous space with an argmax layer and optimizing a constrained probabilistic inverse. Several inverse constructions are considered, including thresholding, Gumbel-based sampling, and Cartesian products.
- Argmax Flow Construction: Argmax Flows combine a continuous density p(v) with an argmax layer mapping v ∈ R^D×K to categorical x ∈ {1,...,K}^D.
- Argmax Flow Construction: The argmax partitions continuous space into symmetric category regions, making sampling straightforward, but likelihood computation is intractable.
- Probabilistic Inverse: Variational inference uses q(v|x), whose support must satisfy x = arg max v to avoid zero likelihood contributions and negative-infinite ELBO values.
- Probabilistic Inverse: Thresholding samples an unconstrained u and maps non-observed components below v_x using a bijective threshold function, enabling change-of-variables evaluation of log q(v|x).
- Probabilistic Inverse: Gumbel-based inverses sample the maximum and then the remaining values from truncated Gumbel distributions, while Gumbel thresholding can replace uniform noise with a normalizing flow.
- Cartesian Products of Argmax Flows: Cartesian products trade symmetry for dimensionality by representing a K-category variable with multiple lower-cardinality variables, including binary encodings.
4 Multinomial Diffusion
Multinomial Diffusion defines a categorical diffusion process that uniformly resamples categories with increasing noise, then learns a denoising process through predictions of the original signal.
- Multinomial Diffusion represents each categorical variable in one-hot format and independently applies a categorical noise process across dimensions.For category k, the one-hot vector has x_k = 1 and all other entries equal to 0.
- At timestep t, the forward process resamples a category uniformly with probability β_t and retains the previous category with probability 1 − β_t.The process introduces a small amount of uniform noise over the K classes at each step.
- The Markov structure permits computing q(x_t|x_0) and the categorical posterior q(x_{t−1}|x_t, x_0) in closed form.These posterior probabilities support construction of the learned reverse process.
- The generative model predicts an approximation x̂_0 = μ(x_t, t), then uses the posterior parameterized by x̂_0 to define p(x_{t−1}|x_t).The reverse distributions are p(x_0|x_1) = C(x_0|x̂_0) and p(x_{t−1}|x_t) = C(x_{t−1}|θpost(x_t, x̂_0)).
- Because the distributions are discrete, the variational objective’s KL terms can be computed by enumerating probabilities and evaluating discrete KL divergences.The one-hot representation also permits direct computation of log p(x_0|x_1).
5 Related Work
Related work situates the paper among major deep generative-model families, prior continuous-flow methods requiring dequantization, and diffusion models that trade fast training for slower sampling.
- Deep generative models include autoregressive models, variational autoencoders, generative adversarial networks, normalizing flows, energy-based models, and diffusion models.
- Normalizing flows learn continuous distributions and therefore require dequantization for ordinal discrete data such as images.Prior work develops more expressive continuous flows and adds uniform noise between ordinal values.
- Diffusion models have been developed for Gaussian, Bernoulli, and high-dimensional image data, with architectural and prediction improvements supporting image generation.
- Diffusion models are relatively fast to train but slow to sample because generation traditionally iterates over many timesteps.Later work showed that samples can be generated using significantly fewer steps, and importance-weighting can improve log-likelihood performance.
- Table 2 compares coupling and autoregressive flows using uniform dequantization, variational dequantization, and Argmax Flows.
6 Experiments
The experiments evaluate Argmax Flows and Multinomial Diffusion on language modelling and unconditional image-segmentation learning, comparing them with dequantization and other generative models. Argmax Flows consistently outperform traditional dequantization, while Multinomial Diffusion offers competitive non-autoregressive text modelling and supports single-pass spell-checking.
- Experimental setup: Experiments cover text8, enwik8, and unconditional cityscapes segmentation-map modelling.text8 has 27 categories and enwik8 has 256 byte categories; cityscapes maps are rescaled to 32 × 64 images with global categories as targets.
- Language data: Argmax Flows use autoregressive or coupling-based continuous flows with conditional thresholded inverses for categorical variables.For language experiments, the probabilistic inverse is a conditional diagonal Gaussian q(u|x), and categorical spaces are represented as binary Cartesian products.
- Language data: Argmax Flows consistently outperform uniform and variational dequantization on language data.The reported comparison attributes the advantage to lifting variables into the entire Euclidean space rather than leaving gaps between dequantized intervals.
- Language data: Autoregressive Argmax Flows outperform VAE approaches, AF/AF, and CategoricalNF, while non-autoregressive Argmax Flows outperform IAF / SCF.Multinomial text diffusion performs better than the argmax coupling flow but worse than the autoregressive version.
- Segmentation maps: Multinomial Text Diffusion can spell-check character-level corruption in a single forward pass, but it does not handle insertions.The corrupted sequence is denoised to predict the most-likely original sequence.
- Segmentation maps: Multinomial Diffusion achieves 0.37 bpp on the cityscapes test set versus 0.33 bpp on train, with overfitting requiring data augmentation.Argmax Flows outperform traditional uniform and variational dequantization in the cityscapes comparison; softplus has the best ELBO, while Gumbel has the better IWBO.
7 Social Impact and Conclusion
The paper presents Argmax Flows and Multinomial Diffusion as extensions for learning categorical data, with numerically stable log-space implementation details for Multinomial Diffusion. It reports improved negative log-likelihood relative to comparable models and identifies performance gaps for future work.
- 7 Social Impact and Conclusion: The authors note that the methods remain distant from direct applications such as text generation, segmentation, or fake-content detection.
- 7 Social Impact and Conclusion: Argmax Flows and Multinomial Diffusion extend normalizing flows and diffusion models to categorical data.
- 7 Social Impact and Conclusion: The experiments show that the methods outperform comparable models in negative log-likelihood and reveal gaps between autoregressive and non-autoregressive approaches.
- 7 Social Impact and Conclusion: The Multinomial Diffusion implementation computes categorical forward-process probabilities and posterior distributions in log-space.
- 7 Social Impact and Conclusion: The implementation predicts p(x_t−1|x_t) with a neural network parameterized through the categorical posterior.
- 7 Social Impact and Conclusion: Training uses KL divergence between the true categorical posterior and the model-predicted posterior as the loss term L_t.
B Experimental details
The experiments use a cosine diffusion schedule with importance sampling of timesteps, alongside numerically stable implementation choices and documented Gumbel properties.
- B Experimental details: The diffusion experiments use a cosine schedule for α_t adapted from Nichol and Dhariwal’s Gaussian formulation to categorical parameters.
- B Experimental details: Timesteps are importance-sampled rather than sampled uniformly because this reduces variance in the objective.
- B Experimental details: The appendix summarizes Gumbel properties used to sample and compute likelihoods for truncated Gumbel distributions.
B.1 Language Modelling
Language-modeling experiments use standardized text8 and enwik8 settings, train Argmax Flows and Multinomial Text Diffusion for different durations, and vary transformer capacity for diffusion models.
- B.1 Language Modelling: The language experiments use text8 sequences of length 256 and enwik8 sequences of length 320 with standard 90,000,000/5,000,000/5,000,000 train/validation/test splits.
- B.1 Language Modelling: Multinomial Text Diffusion models are trained for 300 epochs, while Argmax Flows use 40 epochs, except the enwik8 Argmax Coupling Flow, trained for 20 epochs.
- B.1 Language Modelling: Text-model architecture details are documented in optimization and architecture tables, including separate summaries for training settings and model designs.
- B.1 Language Modelling: Cityscapes models use a shared architecture, while the segmentation preprocessing and model details specify the corresponding experimental setup.
- B.1 Language Modelling: Transformer experiments vary depth from 1 to 20 and hidden size from 128 to 1024; depth 12 with hidden size 512 balances training time and performance.
C Reproducing Discrete Flows
The paper revisits discrete-flow experiments and reports reproducibility difficulties, including failed language-model reproductions and unavailable original implementation dependencies.
- C Reproducing Discrete Flows: The reproduction study focuses on discrete flows with factorized base distributions and avoids autoregressive bases to isolate the flow’s contribution.
- C Reproducing Discrete Flows: The quantized eight-Gaussians experiment compares probability mass functions with weighted sample counts that converge toward the pmf.
- C Reproducing Discrete Flows: Other researchers failed to reproduce the discrete-flow language experiments, and prior work also reported optimization difficulties.
- C Reproducing Discrete Flows: The authors implemented an expressive autoregressive discrete-flow layer after verifying the original repository’s one-hot operations.
- C Reproducing Discrete Flows: Table 8 reports discrete-flow results on text8 and notes that autoregressive flows are more expressive than coupling flows.
- C Reproducing Discrete Flows: The authors conclude that the discrete-flow language experiments are not reproducible because the original code was unavailable and no successful public implementation was found.
D Additional experiments
Additional experiments examine base choices, sampling speed, run-to-run variability, and a controlled comparison with CategoricalNF. They show small base numbers are preferable, non-autoregressive sampling is much faster, reported variability is generally small, and Argmax Flows outperform CategoricalNF in the matched text8 setting.
- Cartesian products: Performance differences between Cartesian-product bases are relatively small, but larger bases slightly reduce performance and increase overall dimensionality.The experiment used a somewhat smaller architecture than the main text.
- Sampling speed: Autoregressive models require a couple of orders of magnitude more sampling time than non-autoregressive alternatives.Multinomial Diffusion falls between autoregressive and coupling models in sampling time.
- Sampling speed: Reducing diffusion steps reduces sampling time, as expected.The comparison includes a complete 1000-timestep chain, for which improvements are possible by skipping steps.
- Run variability: Standard deviations across three runs are small compared with the reported performance differences between models.Coupling models have larger standard deviations, but also larger performance differences between model types.
- Controlled comparison: 1.43 bpc versus 1.45 bpc: Argmax Flows outperform CategoricalNF on text8 when their density models are exactly matched.This comparison uses an equal experimental setting and the same underlying density model.
E Samples from the text models
The paper presents samples from models trained on text8 and shows intermediate generation steps for Multinomial Text Diffusion. The displayed examples are intended to represent the proposed models without cherry-picking.
- Samples: Samples from the proposed models are presented in Table 12 for models trained on text8.The table is identified as containing samples from models trained on text8.
- Evaluation context: The reported samples and generation-chain visualization were not cherry-picked.This qualification applies to the samples in Table 12 and the Multinomial Text Diffusion train shown in Figure 7.
- Generation chain: Figure 7 shows intermediate steps in the generation chain of the Multinomial Text Diffusion model trained on text8.The figure focuses on the progression of the generation chain rather than only its final output.