[论文解读] Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions

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

Music and Audio Processing被引用 79

一句话总结

论文提出 Argmax Flows 和 Multinomial Diffusion，用于学习分类分布，使正规化流和扩散模型能够处理离散数据如文本和分割图，并在对数似然上优于标准去量化方法。

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.

研究动机与目标

• 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.

提出的方法

• 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.

实验结果

研究问题

• 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?

主要发现

• 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.