[Paper Review] Argmax Flows and Multinomial Diffusion: Towards Non-Autoregressive Language Models.

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

Topic Modeling25 citations

TL;DR

This paper introduces Argmax Flows and Multinomial Diffusion, two non-autoregressive generative models for categorical data like text and image segmentation. By using continuous normalizing flows with argmax operations or a diffusion process with categorical noise, the models enable parallel generation and achieve competitive performance on language modeling and segmentation tasks.

ABSTRACT

The field of language modelling has been largely dominated by autoregressive models, for which sampling is inherently difficult to parallelize. This paper introduces two new classes of generative models 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 models perform competitively on language modelling and modelling of image segmentation maps.

Motivation & Objective

• Address the inherent sequential bottleneck in autoregressive language models by proposing non-autoregressive alternatives.
• Enable parallel generation of discrete sequences by modeling categorical data through continuous normalizing flows and diffusion processes.
• Develop differentiable, invertible mappings for discrete data via probabilistic inverses of argmax functions.
• Design a diffusion process that gradually adds categorical noise and learns to reverse the corruption for generation.
• Demonstrate competitive performance on language modeling and image segmentation tasks without autoregressive decoding.

Proposed method

• Define Argmax Flows as a composition of a continuous distribution (e.g., normalizing flow) and an argmax operation to map continuous latent variables to discrete tokens.
• Learn a probabilistic inverse of the argmax function to lift discrete data into a continuous space for training and optimization.
• Formulate Multinomial Diffusion as a Markov process that gradually corrupts discrete data with categorical noise using a diffusion schedule.
• Train a denoising network to reverse the corruption process, enabling conditional generation from noise.
• Use normalizing flows and score-based modeling to ensure differentiability and tractable likelihood estimation in both frameworks.
• Optimize both models via maximum likelihood using variational inference and reparameterization gradients.

Experimental results

Research questions

• RQ1 Can argmax flows provide a differentiable, invertible, and scalable way to model discrete sequences using continuous normalizing flows?
• RQ2 Can a diffusion process effectively model categorical data by gradually adding and reversing categorical noise?
• RQ3 Do these non-autoregressive models achieve competitive performance on language modeling and image segmentation compared to autoregressive baselines?
• RQ4 How well can the probabilistic inverse of argmax enable end-to-end training of discrete flows?
• RQ5 Can the denoising process in Multinomial Diffusion generate high-quality sequences without autoregressive generation?

Key findings

• Argmax Flows enable parallel generation of discrete sequences by learning a continuous, invertible mapping from latent space to discrete tokens.
• The probabilistic inverse of argmax allows effective backpropagation through discrete operations, enabling end-to-end training of the flow model.
• Multinomial Diffusion successfully models categorical data through a diffusion process with categorical noise, achieving stable training and generation.
• Both models achieve competitive performance on language modeling benchmarks, demonstrating strong generative quality without autoregressive decoding.
• The models show strong results on image segmentation tasks, indicating generalization beyond natural language.
• The proposed methods provide viable non-autoregressive alternatives to autoregressive models, reducing generation time while maintaining competitive likelihood and quality.