mlsanityData to 7 Oct 2026

Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions

arXiv 2102.05379 ↗First posted 10 Feb 2021stat.ML cs.CL cs.LGcited by 566

On the map: Diffusion models → Medical imaging and segmentation

Abstract (arXiv)

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.

Key facts. Cited by 566 papers in the map (data through 7 Oct 2026). Placed in Medical imaging and segmentation, one of 14 directions in Diffusion models. Builds on 32 of its 39 references that are in the map. First posted on arXiv 10 Feb 2021.

Builds on

32 of 39 references.

References · 32

  1. 2020Denoising Diffusion Probabilistic Modelscited by 20,851
  2. 2020Denoising Diffusion Implicit Modelscited by 9,655
  3. 2020Score-Based Generative Modeling through Stochastic Differential Equationscited by 8,964
  4. 2015Deep Unsupervised Learning using Nonequilibrium Thermodynamicscited by 7,005
  5. 2016Densely Connected Convolutional Networkscited by 6,691
  6. 2016The Cityscapes Dataset for Semantic Urban Scene Understandingcited by 4,427
  7. 2013Auto-Encoding Variational Bayescited by 3,980
  8. 2018Neural Ordinary Differential Equationscited by 3,733
  9. 2015Variational Inference with Normalizing Flowscited by 2,161
  10. 2016Density estimation using Real NVPcited by 2,029
  11. 2019Transformer-XL: Attentive Language Models Beyond a Fixed-Length Contextcited by 1,830
  12. 2014NICE: Non-linear Independent Components Estimationcited by 1,052
20 more
  1. 2017Masked Autoregressive Flow for Density Estimationcited by 680
  2. 2016Improving Variational Inference with Inverse Autoregressive Flowcited by 636
  3. 2015Importance Weighted Autoencoderscited by 454
  4. 2015MADE: Masked Autoencoder for Distribution Estimationcited by 306
  5. 2015A note on the evaluation of generative modelscited by 291
  6. 2019Flow++: Improving Flow-Based Generative Models with Variational Dequantization and Architecture Designcited by 258
  7. 2019Stochastic Beams and Where to Find Them: The Gumbel-Top-k Trick for Sampling Sequences Without Replacementcited by 205
  8. 2019Learning Likelihoods with Conditional Normalizing Flowscited by 170
  9. 2014A* Samplingcited by 150
  10. 2018Character-Level Language Modeling with Deeper Self-Attentioncited by 138
  11. 2019Integer Discrete Flows and Lossless Compressioncited by 82
  12. 2020SurVAE Flows: Surjections to Bridge the Gap between VAEs and Flowscited by 71
  13. 2019Discrete Flows: Invertible Generative Models of Discrete Datacited by 68
  14. 2019Latent Normalizing Flows for Discrete Sequencescited by 66
  15. 2013RNADE: The real-valued neural autoregressive density-estimatorcited by 65
  16. 2020Categorical Normalizing Flows via Continuous Transformationscited by 36
  17. 2016Recurrent Batch Normalizationcited by 35
  18. 2019MintNet: Building Invertible Neural Networks with Masked Convolutionscited by 35
  19. 2020IDF++: Analyzing and Improving Integer Discrete Flows for Lossless Compressioncited by 35
  20. 2020Closing the Dequantization Gap: PixelCNN as a Single-Layer Flowcited by 7

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2021: 4; 2022: 42; 2023: 106; 2024: 119; 2025: 150; 2026: 145.

566 papers.

Sort:

By area:

  1. 2022Diffusion Models: A Comprehensive Survey of Methods and Applicationscited by 1,107
  2. 2025Large Language Diffusion Modelscited by 872
  3. 2023Discrete Diffusion Modeling by Estimating the Ratios of the Data Distributioncited by 712
  4. 2022Equivariant Diffusion for Molecule Generation in 3Dcited by 666
  5. 2024Simplified and Generalized Masked Diffusion for Discrete Datacited by 606
  6. 2022Fast Sampling of Diffusion Models with Exponential Integratorcited by 526
  7. 2025Dream 7B: Diffusion Large Language Modelscited by 515
  8. 2022DiGress: Discrete Denoising diffusion for graph generationcited by 503
  9. 2022All are Worth Words: A ViT Backbone for Diffusion Modelscited by 498
  10. 2022DiffuSeq: Sequence to Sequence Text Generation with Diffusion Modelscited by 439
  11. 2022A Continuous Time Framework for Discrete Denoising Modelscited by 437
  12. 2025Fast-dLLM: Training-free Acceleration of Diffusion LLM by Enabling KV Cache and Parallel Decodingcited by 404
38 more
  1. 2025Block Diffusion: Interpolating Between Autoregressive and Diffusion Language Modelscited by 385
  2. 2022TabDDPM: Modelling Tabular Data with Diffusion Modelscited by 373
  3. 2024Your Absorbing Discrete Diffusion Secretly Models the Conditional Distributions of Clean Datacited by 371
  4. 2022Analog Bits: Generating Discrete Data using Diffusion Models with Self-Conditioningcited by 362
  5. 2024Discrete Flow Matchingcited by 301
  6. 2024Generative Flows on Discrete State-Spaces: Enabling Multimodal Flows with Applications to Protein Co-Designcited by 294
  7. 2025LLaDA 1.5: Variance-Reduced Preference Optimization for Large Language Diffusion Modelscited by 273
  8. 2024Scaling Diffusion Language Models via Adaptation from Autoregressive Modelscited by 268
  9. 2022Diffsound: Discrete Diffusion Model for Text-to-sound Generationcited by 256
  10. 2021Deep Generative Modelling: A Comparative Review of VAEs, GANs, Normalizing Flows, Energy-Based and Autoregressive Modelscited by 251
  11. 2025MMaDA: Multimodal Large Diffusion Language Modelscited by 241
  12. 2022A Survey on Generative Diffusion Modelcited by 238
  13. 2022Diffusion models as plug-and-play priorscited by 230
  14. 2023DIFUSCO: Graph-based Diffusion Solvers for Combinatorial Optimizationcited by 216
  15. 2021Autoregressive Diffusion Modelscited by 208
  16. 2025Train for the Worst, Plan for the Best: Understanding Token Ordering in Masked Diffusionscited by 208
  17. 2024Masked Diffusion Models are Secretly Time-Agnostic Masked Models and Exploit Inaccurate Categorical Samplingcited by 204
  18. 2022Continuous diffusion for categorical datacited by 202
  19. 2024Scaling up Masked Diffusion Models on Textcited by 195
  20. 2022SSD-LM: Semi-autoregressive Simplex-based Diffusion Language Model for Text Generation and Modular Controlcited by 191
  21. 2025DiffuCoder: Understanding and Improving Masked Diffusion Models for Code Generationcited by 184
  22. 2023Mixed-Type Tabular Data Synthesis with Score-based Diffusion in Latent Spacecited by 183
  23. 20233D Equivariant Diffusion for Target-Aware Molecule Generation and Affinity Predictioncited by 176
  24. 2022Score-based Continuous-time Discrete Diffusion Modelscited by 167
  25. 2022Generative Modelling With Inverse Heat Dissipationcited by 164
  26. 2023Protein Design with Guided Discrete Diffusioncited by 154
  27. 2023A Reparameterized Discrete Diffusion Model for Text Generationcited by 152
  28. 2024Simple Guidance Mechanisms for Discrete Diffusion Modelscited by 148
  29. 2025LLaDA-V: Large Language Diffusion Models with Visual Instruction Tuningcited by 145
  30. 2022Concrete Score Matching: Generalized Score Matching for Discrete Datacited by 144
  31. 2023Crystal Structure Prediction by Joint Equivariant Diffusioncited by 142
  32. 2025dKV-Cache: The Cache for Diffusion Language Modelscited by 140
  33. 2024Beyond Autoregression: Discrete Diffusion for Complex Reasoning and Planningcited by 135
  34. 2025Fast-dLLM v2: Efficient Block-Diffusion LLMcited by 128
  35. 2024Unlocking Guidance for Discrete State-Space Diffusion and Flow Modelscited by 124
  36. 2025The Diffusion Dualitycited by 124
  37. 2022Equivariant 3D-Conditional Diffusion Models for Molecular Linker Designcited by 121
  38. 2023TRACT: Denoising Diffusion Models with Transitive Closure Time-Distillationcited by 115

50 most cited of 566 All 566 (JSON)

Similar, not cited

Close in wording, no citation either way

  1. 2025An Introduction to Flow Matching and Diffusion Models
  2. 2026From Diffusion to Rectified Flow: Rethinking Text-Based Segmentation
  3. 2026Jumping up and down: Denoiser diffusion models for discrete ordinal data
  4. 2021Label-Efficient Semantic Segmentation with Diffusion Models
  5. 2025Diffusion Models: A Mathematical Introduction
  6. 2023Denoising Diffusion Semantic Segmentation with Mask Prior Modeling
  7. 2026Diffusion Model as a Generalist Segmentation Learner
  8. 2026Making Training-Free Diffusion Segmentors Scale with the Generative Power