Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions

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.

arXiv commentAccepted at Neural Information Processing Systems (NeurIPS 2021)

References in corpus (16)

• Auto-Encoding Variational Bayes
• Denoising Diffusion Probabilistic Models
• NICE: Non-linear Independent Components Estimation
• Variational Inference with Normalizing Flows
• Score-Based Generative Modeling through Stochastic Differential Equations
• Improved Denoising Diffusion Probabilistic Models
• MADE: Masked Autoencoder for Distribution Estimation
• Flow++: Improving Flow-Based Generative Models with Variational Dequantization and Architecture Design
• Denoising Diffusion Implicit Models
• Learning Likelihoods with Conditional Normalizing Flows
• Latent Normalizing Flows for Discrete Sequences
• SurVAE Flows: Surjections to Bridge the Gap between VAEs and Flows
• Discrete Flows: Invertible Generative Models of Discrete Data
• IDF++: Analyzing and Improving Integer Discrete Flows for Lossless Compression
• Categorical Normalizing Flows via Continuous Transformations
• Invertible DenseNets

Cited by in corpus (11)

• Structured Denoising Diffusion Models in Discrete State-Spaces
• Diffusion Schrödinger Bridge with Applications to Score-Based Generative Modeling
• On Fast Sampling of Diffusion Probabilistic Models
• Vector Quantized Diffusion Model for Text-to-Image Synthesis
• A Variational Perspective on Diffusion-Based Generative Models and Score Matching
• Self-Supervision Improves Diffusion Models for Tabular Data Imputation
• SegDiff: Image Segmentation with Diffusion Probabilistic Models
• Geometric-Facilitated Denoising Diffusion Model for 3D Molecule Generation
• Autoregressive Diffusion Models
• Zero-Shot Translation using Diffusion Models
• Beyond In-Place Corruption: Insertion and Deletion In Denoising Probabilistic Models