Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Authors: Emiel Hoogeboom
Didrik Nielsen
Priyank Jaini
Patrick Forré
Max Welling
Venue: NeurIPS 2021 Poster (2021)
Keywords: categorical
normalizing flows
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. 

Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Emiel Hoogeboom's ' Didrik Nielsen Priyank Jaini', Patrick Forré Max Welling' UvA-Bosch Delta Lab, University of Amsterdam', Technical University of Denmark?, University of Amsterdam didni@dtu dk, e h.oogeboom@uva. ml, p. p.jaini@uva.ni nl,
p d. f.rreeruauua.nl m. m.welling@uva.r
Abstract
Generative flows and diffusion models have been predominantly trained on ordinal data, for example natural images. This paper introduces extensions of flows and diffusion for calegorical 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 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 log-likelihood.
Introduction
Many sources of high-dimensional data are cat- egorical, for example language and image seg- mentation. Although natural images have been studied to a large extent with generative flows and diffusion models, categorical data has not had the same extensive treatment. Currently they are primarily modelled by autoregressive models, which are expensive to sample from Cooijmans al. 2017 Dai et al. 2019).
gltte)
Px(#)
P(x)
Composition of a flow p(v) and argmax transformation which gives the model P(w). The flow maps from base distribution p(z) using bijection g.
(a) Argmax Flow:
Normalizing flows are attractive because they be designed to be fast both in the evaluation and sampling direction. Typically, normalizing flows model continuous distributions. As re- sult, directly optimizing a flow on discrete data lit-
(b) Multinomial Diffusion: Each step p(x ( de- noises the signal starting from a uniform categorical base distribution which gives the model p(zo).
may lead arbitrarily high likelihoods. erature this problem is resolved for ordinal data by adding noise unit interval around the dis- crete value (Uria et al.] 2013) Theis et al., 2016 Ho et al. 2019). However, because these methods have been designed for ordinal data, they do not work well on categorical data.
Figure 1: Overview of generative models.
Other attractive generative models are diffusion models (Sohl-Dickstein et al. 2015). which are fast train due to objective that decomposes over time steps (Ho al. 2020). Diffusion models
Equal contribution.
35th Conference on Neural Information Processing Systems (NcurIPS 2021).

---

Table 1: Surjective flow layers for applying continuous flow models to discrete data. The layers are deterministic in the generative direction, but stochastic in the inference direction. Rounding corresponds to the commonly-used dequantization for ordinal data.
Layer
Generation
Inference
Applications
t ~ (((()) with support
Ordinal Data
Rounding
[c]
S(x)) = {ulæ = [v]]
e.g. images, audio
g (1))) with support
Categorical Data
Argmax
x
arg max
S(x) {0læ arg max v)
e.g. text, segmentation
typically have fixed diffusion process that gradually adds noise. This process is complemented by learnable generative process that denoises the signal. [Snn et al. (2000); NNcholl and Dhariwal (00TT) have shown that diffusion models also be designed for fast sampling. Thus far, diffusion models have been primarily trained to learn ordinal data distributions, such as natural images.
Therefore, in this paper we introduce extensions of flows and diffusion models for categorical variables (depicted in Figure M: i) Argmax Flows bridge the gap between categorical data and continuous normalizing flows using argmax transformation and corresponding family of probabilistic inverses for the argmax. In addition ti) we introduce Multinomial Diffusion, which is a diffusion model directly defined categorical variables. Opposed to normalizing flows, defining diffusion for discrete variables directly does not require gradient approximations, because the diffusion trajectory fixed. As result of our work, generative normalizing flows and diffusion models can directly learn categorical data.
Background
Normalizing Flows malizing flows (Rezende and Mohamed, 2013) learn a bijective and differentiable transformation g: 2 V such that the thange-of-vee -o-vviabll formula gives the density at any point v E V:
Given V = Rd and Z =
Rd with densities py and pz respectively, nor-
dz de
pv (v) = p=) det
g(z),
where Pz can be any density (usually chosen as a standard Gaussian). Thus, normalizing flows provide powerful framework to learn exact density functions. However, Equation T restricted to continuous densities.
To learn densities ordinal discrete data (such natural images), typically dequantization noise added (Uria et al. 2013 Theis et al., 2016; Ho et al. 2019). Nielsen et al. (2020) reinterpreted dequantization surjective flow layer + that is deterministic in one direction (z round(v) and stochastic in the other (v = x + u where u ~ ((ulz)). Using this interpretation, dequantization can be probabilistic right-inverse for the rounding operation in the latent variable model
given by:
P(æ) = 1000=((((xlllll P(zlo)p(U) de,
where round applied elementwise. this case, the density model p(v) is modeled using normalizing flow. Learning proceeds by introducing the variational distribution g(v ar) that models the probabilisticc right-inverse for the rounding surjection and optimizing the evidence lower bound (ELBO):
log P(æ) 2 Eusgtria m) [log P(xlv) + log log ((v[==) = Bus m) (logr(v) log (v)=)) (2)
The last equality holds under the constraint that the support of q(vlx) is enforced to be only over the region& S = Rd (d:ax - ro)) which ensures that P(z//v) 1.
Diffusion Models Given data æg. diffusion model (Sohl- et al. 2015 consists of prede- fined variational distributions (xxx(=;-1) that gradually add noise over time steps t E (T)...,T). The diffusion trajectory is defined such that g(=;/<r,-1) adds a small amount of noise around a, This way, information is gradually destroyed such that at the final time step, Øт carries almost no information about 2p. Their generative counterparts consists of learnable distributions p( (a:)) that learn to denoise the data. When the diffusion process adds sufficiently small amounts of noise,

---

Algorithm Sampling from Argmax Flows
Algorithm 2 Optimizing Argmax Flows
Input: p(v)
Input: æ, p(v). ((v(<)
Output: Sample
Output: ELBO C
Sample v p(0)
Sample v ~ (z)z))
Compute arg max
Compute log logp(e) log ((v/2)
suffices to define the denoising trajectory using distributions that are factorized (without correlation) over the dimension axis. The distribution p(ær) chosen to be similar to the distribution that the diffusion trajectory approaches. Diffusion models can be optimized using variational inference:
Σ log p(x1-1|22)) g (x,
T
log P(xo) 2 E Prog logp(<r) +
An important insight in diffusion is that by conditioning on Zo. the posterior probability ((xx--1 Ft, mo) = 11 / is tractable and straightforward compute, permitting reformulation 2015). Note that KT. (g (()r)())))))))))))))))))))))(( (p()))))))) = 0f the diffusion trajectory q is defined well:
terms of KI. divergences that has lower variance (Sohl-Dickstein et al.
log P(æo)= 2 Eg ( - KL (2o (o)ooppp(())
(L( 2o) p( [z)))
(3)
Argmax Flows
Argmax flows define discrete distributions using 1) density model p(v), such as a normalizing flow, and 2) an argmax layer that maps the continuous v E Dx to a discrete a E (1, ..., K)D using
a arg max t where zd = arg max vikk
This is natural choice to model categorical variables, because it divides the entire continuous space of o into symmetric partitions corresponding to categories in æ. To sample from argmax flow sample e ~ p(v) and compute æ arg max e (Algorithm 1. generate reasonable samples, it is up to the density model p(v) capture any complicated dependencies between the different dimensions. While sampling from an argmax flow is straightforward, the main difficulty lies in optimizing this generative model. To compute the likelihood of a datapoint we have to compute
P(x)) = P(zlv)p(w)de, P(zlu)= 5 = arg max(u)),
(5)
which intractable. Consequently, we resort to variational inference and specify a variational distribution q(v(z)) (z). We note that naïvely choosing any variational distribution may lead to samples ~ q(v(²) where 6(x = arg max v) = O, which yields an ELBO negative infinity. To avoid this, need variational distribution q æ)) that satisfies what term the argmax constraint:
arg max for all q((xx)
That is, the variational distribution q(v))z)) should have support limited to S(x) fo DDX I
arg max Recall that under this condition, the ELBO simplifies to
En-g(()))) [og)(v) log (w)[))) as shown Algorithm 2 For an illustration of the method see Figure Ta
3.1
Probabilistic Inverse
The argmax layer may be viewed as surjective flow layer Niclsen et al. 2020 With this view, the variational distribution g(v(x) ()) specifies a distribution over the possible right- of the argmax function, also known dequantization layer for ordinal data corresponds to the probabilistic inverse rounding operation. As summarized in Table this layer may thus be viewed analogous the argmax layer, where the round is for ordinal data while the argmax for categorical data.
stochastic inverse probabilistic inverse. Recall that the commonly-used
We are free to specify any variational distribution g(vlar a))):) that satisfies the argmax constraint. In the next paragraphs we outline three possible approaches. Since operations are performed independently across dimensions, omit the dimension axis and let E gK and æ E {1....,K K}.

---

Algorithm 3 Thresholldng-bbsee ((v/z)
Algorithm 4 Gumbel-based ((v(z)
Input: 2; ((u(az))
Input: æ, Φ
Output: v, log 21)2))
Output: u, log g ((>/2z)
((u(z))
Omax log 2; exp O;
Ux Ur
Uz ~ Gumbel (max)
v- threshold(u
TruncCumbel(
Un)
log g(v) ()) log ((ul (r) - log det /o/ dul
log ((v)((())) æ)) = log Gumbel Uv x
log TruncGumbel (v-w p- Ux)
A straightforward method to construct distribution g(vlz)) satisfying the argmax constraint is to use thresholding. That is, we first sample an unbounded variable R from ((v(z) (z), which can be for example a conditional Gaussian or normalizing flow. Next, we map u to such that element is the largest:
Thresholding (Alg-
3
Up Uv and V- thresholdr(w-s)
(6)
where the thresholding is applied elementwise with threshold value T' = U2. This ensures that element Us is the largest, and consequently that q(v))) 2) satisfies the argmax constraint. Note that we require the threshold function to be bijective, thresholdy IR - (-o0, T), so that we can use the change-of-vardee implemented using softplus such that all values are mapped below limit T:
formula to compute log q (v)z)) In our implementation, thresholding is
v = thresholdr(u (u) = T' softplus(T - uw)
where softplus(2) log(1 and for which it guaranteed that v E (-00, T).
Gumbel (Alg.
An alternative approach is to let q(v[2) = Gumbel(vlo) restricted to
arg v I. where the location parameters & + NN(x) are predicted using a neural network NN. The Gumbel distribution has favourable properties: The arg max and max are independent and the max is also distributed as Gumbel:
max ~ Gumbel
(8)
where Qmax
log E, exp d;. For more extensive introduction see Maddison et al. 2014 Kool
al. 2019). To sample (012)), we thus first sample the maximum ex according to Eq. g Next,
given the sample Uzz, the remaining values can be sampled using truncated Gumbel distributions:
1/ ~ TruneGumbel(()) T) where +
where the truncation value T given by er which ensures that the argmax constraint e, > v; for i A x is satisfied. Recall that to optimize Eq. 2 log q( ( is also required, which can be computed using the closed-form expressions for the log density functions (see Table 5). Another property of Gumbel distributions is that
P(arg max v = =) = exp d./
Σ exp
(10)
which we use to initialize the location parameters / match the empirical distribution of the first minibatch the data.
Gumbel Thresholding. This method unifies the methods from the previous two sections: Gumbel distributions and thresholding. The key insight that the Gumbel sampling procedures defined above can be seen reparametrization of a uniform noise distribution 2/(0. 1) which is put through the inverse CDF of the Gumbel distributions (see Table 5) From the perspective of change- of-variables, the log likelihood denotes the log volume change of this transformation. To increase expressitivity the uniform distribution can be replaced by normalizing flow (((u)2))) that has support on the interval (0, 1)* which can be enforced using a sigmoid transformation. This section shows that a large collection thresholding functions can be found by studying (truncated) inverse CDFs. In practice find that performance is reasonably similar as long the underlying noise is learned.

---

Behavior of the Variational Posterior Although several methods to learn have been proposed, it is unclear what expressitivity is required. In the following, the interactions between g( æ)) and the density model p(v) are discussed. Recall that the variational bound that is optimized under expectation of data distribution 'D can be seen as minimizing the KT. distance between the aggregated posterior q(v) = Be (1) æ)) and the density model p(v), so (((((())(()))) There are two distinct reasons which can cause this distance to be large: Firstly, the density model p(v) may not have the right probability mass in each argmax region. These desired probabilities solely depend on the data distribution D. Secondly, the variational posterior ((vlz) may not have the correct shape compared to p(v), within an argmax region. At initialization, the thresholding within q can create low density regions at argmax boundaries.
In theory, if p(v) is universal density approximator, then the model can be fitted for any well- behaved g(v) ax). Then p(v) can even fit the low density regions in the boundaries. This argument is trivial, as one can simply set p(v) to g(v)) B==Dq(((2:) In practice, over training steps we find that &g does smooth out these boundary artifacts, and counteracts the thresholding so that the aggregated posterior becomes smoother.
3.2 Cartesian Products of Argmax Flows
In the current description, Argmax Flows require the same number of dimensions in v as there are classes in æ. To alleviate this constraint we introduce Cartesian products of Argmax Flows. To illustrate our method, consider a 256 class problem. One class can be represented using a single number in f1...., ., 256], but also using two hexadecimal numbers {1.... 16)² 2 or alternatively using eight binary numbers. Specifically, any base K variable (K) E (1, (1...., K)D can be converted to a base . variable x (MM) E (1...., M) x where dm = [og K]. Then the variable (M) with dimensionality M dm represents the variable (K) with dimensionality K D, trading off symmetry for dimensionality. Even though this may lead to some unused additional classes, the ELBO objective in Equation P can still be optimized using an M-categorical Argmax Flow. Finally, note that Cartesian products of binary spaces are special case where the variable can be encoded symmetrically into single dimension to the positive and negative part using binary dequantization (Winkler et al. 2019). In this case, by trading-off symmetry the dimensionality increases only proportional to log, K
4 Multinomial Diffusion
In this section we introduce an alternative likelihood-based model for categorical data: Multinomial Diffusion. In contrast with previous sections, 2, will be represented in one-hot encoded format 2; E {0, 1}* Specifically, for category k, I'k = 1 and r = 0 for j * k. Note that again the dimension axis is omitted for clarity as all distributions are independent over the dimension axis. We define the multinomial diffusion process using a categorical distribution that has 3, chance of resampling a category uniformly:
q(=:/<=--1) = C(x,1(1 - 3.)a.-1 + 3,/K),
(11)
where C denotes a categorical distribution with probability parameters after |. Further addition (and subtraction) between scalars and vectors is done elementwise. This convention kept throughout this section. Since these distributions form a Markov chain, we can express the probability of any x given Tn
q( ( C(x. (o + (1 (1-aa a=)/
(12)
where ag = 1 3, and ā
II'= a, Intuitively, for each next timestep, little amount of uniform
noise 3e over the K classes is introduced, and with a large probability (1 - B:) the previous value
T;-1 1 is sampled. Using Equation II and 12 the categorical posterior q [z,, xp) can be computed
in closed-form:
K
((#1-112,, vo) = C(x; 10 (a (xz, œo)); where (post
(2,, x00) = 0/
0x
(13)
k=1
and 0 = [a,, + at) ④ fā-- 20 + -ār-))/K
5

---

nelal
1nnnn
o..
do floe de?
Figure 2: Overview of multinomial diffusion. A generative model p(z; (z;,) learns to gradually denoise signal from left to right. An inference diffusion process g(zr//x;-) 1) gradually adds noise form right to left.
One of the innovations in Ho et al. (2020) was the insight to not predict the parameters for the generative trajectory directly, but rather to predict the noise using the posterior equation for Although predicting the noise is difficult for discrete data, we predict a probability vector for wo from 2; and subsequently parametrize p( (zx) using the probability vector from g(xt- [2,, &o), where 25 is approximated using neural network 2o p( (). Equation 13 will produce valid probability vectors that are non non-neeative and sums to one under the condition that the prediction do non-negative and sums to one, which is ensured with a softmax function in μ.. To summarize:
p( = ((rolão) and p(æ,-1læx,) = C(#:-110
p post (Iz, (b)) where ao = p(æ,, t)
(14)
The KL terms in Equation 3 cca be simply computed by enumerating the probabilities in Equation B
and 14 and computing the KL divergence for discrete distributions in L with 2 2:
(L mo)) [o))p((
(T. popt (Ze, zo)) (()pp ka))),
(15)
which can be computed using Σx Post (we, 2o ②o)),, log 0 ( wa). log p( (1) use that wo is onehot:
Furtermore, to compute
post (2) (2),
p(pol())))) = Σ Lo,k log Lo.
(16)
5 Related Work
Deep generative models broadly fall into the categories autoregressive models ARMs (Germain
et al. 2015), Variational Autoencoders (VAEs) (Kingma and Welling 2014 Rezende et al. 2014).
Adversarial Network (GANs) (Goodfellow et al. 2014), Normalizing Flows (Rezende and Mohamed
2015 Energy-Based Models (EBMs) and Diffusion Models (Sohl-Dickstein al. 2015).
Normalizing Flows typically learn continuous distribution and dequantization is required to train these methods on ordinal data such as images. large body of work dedicated to building more
expressive continuous normalizing flows (Dnh et al. 2017: Germain et al., 2015; Kingma et al.
2016 Papamakarios et al. 2017 Chen et al. 2018 (ong al. 2019 Perugachi Diaz al. 2020). To
learn ordinal discrete distributions with normalizing flows, adding uniform noise in in-between ordinal
classes was proposed in (Uria et al. 2013) and later theoretically justified in Theis et al. 2016).
An extension for more powerful dequantization based on variational inference was proposed in (Ho
al. 2019), and connected to autoregressive models in (Nielsen and Winther, 2020). Dequantization for binary variables was proposed in (Winkler et al. 2019). Tran et al. (2019) propose invertible transformations for categorical variables directly. However, these methods can be difficult to train because of gradient bias and results images have thus far not been demonstrated. In addition flows for ordinal discrete data (integers) have been explored in (Hoogeboom et al. 2019: van den Berg al. 2020). In other works, VAEs have been adapted to learn normalizing flow for the latent space ZZigler and Rush 2019: Lippe and Gavves 2020). However, these approaches typically still utilize though this not the distribution specified during training.
argmax heuristic to sample,
Diffusion models were first introduced in Sohl-Dickstein al. (205)), who developed diffusion for Gaussian and Bernoulli distributions. Recently. Denoising Diffusion models Ho er al. 2020) have been shown capable of generating high-dimensional images by architectural improvements and reparametrization of the predictions. Diffusion models are relatively fast to train, but slow to sample

---

Table 2: Comparison of coupling and autoregressive generative flows with uniform (Uria al. and variational (Ho et al. 2019) dequantization and our proposed Argmax flows.
2013
Dequantization
Flow type
text8 (bpe)
enwik8 (bits per raw byte)
Uniform dequantization
1.90
2.14
Variational dequantization
Autoregressive
1.43
1.44
Argmax Flow (ours)
1.38
1.42
Uniform dequantization
2.01
2.33
Variational dequantization
Coupling
2.08
2.28
Argmax Flow (ours)
1.82
1.93
from as they require iterations over the many timesteps in the chain. [Song al. (2020); Nichol and Dhariwal (202T) showed that in practice samples can be generated using sigmificantly fewer steps. Nichol and Dhariwal 202T) demonstrated that importance-weightng the objective components greatly improves log-lik log-likelihood performance. In SSng et al. (2020) a continuous cittimmuuu--meee extension of denoising diffusion models was proposed. After initial release of this paper we discovered that [Song et al. (2020) concurrently also describe framework for discrete diffusion, but without empirical evaluation.
Experiments
In our experiments we compare the performance of our methods on language modelling tasks and learning image segmentation maps unconditionally.
6.1 Language data
In this section we compare our methods on two language datasets, text8 and enwik8. text8 contains 27 categories ('a' through 'z' and "" and for enwik8 the bytes are directly modelled which results in 256 categories.
Model description Two versions of generative argmax flows tested: using autoregressive (AR) flow and a coppling-based flow for p(v). In these experiments the probabilistic inverse is based on the thresholding approach. Specifically, conditional diagonal Gaussian q ar) is trained and thresholded which gives the distribution ((v))))) The argmax flow is defined on binary Cartesian products. This means that for K 27, 5-dimensional binary space is used and for K 256 an 8-dimensional binary space. The argmax flow is compared to the current standard of training generative flows directly on discrete data: dequantization. We compare to both uniform and variational dequantization, where noise on a (0, 1) interval is added to the onehot representation of the categorical data. The autoregressive density model is based the model proposed in (Lippe and Gavves 2020). The coupling density model consists of 8 flow layers where each layer consists of a x I convolution and mixture of logistics transformations Ho al. (2019). the multinomial text diffusion model, the network is modeled by 12-layer Transformer. For more extensive details about the experiment setup see Appendix
Table 3: Comparison of different methods on text8 and enwik8. Results are reported in negative log-likes likelihood with units bits per character (bpe) for text8 and bits per raw byte (bpb) for enwik8.
Model type
Model
text8 (bpc)
enwik8 (bpb)
64 Layer Transformer (Al-Rfou et al. 2019 TransformerXI.
1.13
1.06
ARM
Dai ct al. 2019
1.08
0.99
AF/AF* (AR) Ziegler and Rush 2019
1.62
1.72
VAE
IAF SCF* Ziegler Rush, 2019
1.88
2.03
CategoricalNN (A (Lappe and Gavves 2020
1.45
Argmax Flow, AR (ours)
1.39
1.42
Generative Flow
Argmax Coupling Flow (ours)
1.82
1.93
Diffusion
Multinomial Text Diffusion (ours)
1.72
1.75
Results obtained by ramning code from the official repository for textō and emviks datasets.
7

---

di on ionatteacie ahead conllitia tarat ze closest to the expe and di veness cussunctng Lotal party exte
t. thines the cweet wee seven zive ain ditel: zoom or conteir you are Lse eilon words thereis greatly usuged up not
(a) Samples from Multinomial Text Diffusion.
(a) Samples from the Argmax Flow.
ceurcegnase Srepe chenaticelly intered by the teecor excatence of to del asy specaticts edtisioned sise their cunctiomaiity sussot ewed thore
otat cast of leaening axjeccs frep Jasgrage to platonse itspuaatics ets situessa azirals bisd aeisnces mesievel bioiogy bip Iooy and nens vil andintice Starovery ve
(b) Samples from Argmax AR Flow.
(b) Samples from the Multinomial Diffusion model.
sosatie ope morean aeans as syagathesic desch the sece biliedhaan sccuiced ted vaviaus harey
insencertia
danibah nancetta hot nalgere of the wolntere ncods and drintt brows es cotpdage unseoce shle ous bochere to borel opin poppppiii
de acceapt causs
(c) Samples from Argmax Coupling Flow.
(c) Cityscapes data.
Figure 3: Samples from models, text8.
Figure 4: Samples from models, cityscapes.
Comparison with Generative Flows Firstly we compare the performance of generative flows directly trained on language data (Table 2). These experiments are using the same underlying normalizing flow: either coupling based flow or an autoregressive flow. Note that Argmax Flows consistently outperform both uniform and variational dequantization. This indicates that it is easier for generative flow to learn the lifted continuous distribution using an argmax flow. An advantage of Argmax flows that may explain this difference that they lift the variables into the entire Fuclidean space, whereas traditional dequantization only introduce probability density on (0, 1) intervals, leaving gaps with no probability density. The performance improvements of Argmax flows are even more pronounced when comparing coupling-baid based approaches. Also note that coupling flows have worse performance than autoregressive flows, with a difference that generally smaller for images. This indicates that designing more expressive coupling layers for text is an interesting future research direction.
Comparison with other generative models The performance compared to models in literature is presented BableeBalongsidd the performance of our Argmax Flows and Multinomial Diffusion. The latent variable approaches containing autoregressive components are marked using (AR). Although autoregressive flows still have the same disadvantages as ARMs, they provide perspective on where performance deficiencies are coming from. We find that our autoregressive Argmax Flows achieve better performance than the VAE approaches, they outperform AF/AF (Ziegler and Rush 2019) and
CategoricalNF (Lippe and Gavves 2020
When comparing non-autoregreesivv models, Argmax Flows also outperforms the method that lifts the categorical space to continuous space: IAF SCF (Ziegler and Rush 2019). Interestingly, the multinomial text diffusion is a non-autoreeeessive model that performs even better than the argmax coupling flow, but performs worse than the autoregressive version. For this model it is possible that different diffusion trajectories for q would result in even better performance, because the current form the denoising model has to be very robust to input noise. These experiments also highlight that there is still a distinct performance gap between standard ARMs and (autoregressive) continuous density model on text, possibly related to the dequantization gap (Nielsen and Winther 2020). Samples from different models trained on text8 are depicted in Figure 3 Because of difficulties in reproducing results from Discrete Flows, comparison and analysis of discrete flows are left out of this section. Instead they are extensively discussed in Appendix C For additional experiments
regarding Cartesian products and sampling time see Appendix D
8

---

Unsupervised spell-checking An interesting by-product of the text diffusion model is that can be used to spell-check text using single forward pass. To demonstrate this, a sentence taken from the test data is corrupted by chang- ing a few characters. This corrupted sequence is given as 1 to the generative denoising model, which is close to the data at step O. Then the de- noising model predicts p( o/zo/ar and the most- likely 2y can be suggested. Note that this model only works for character-lewe Corruption, not insertions. An example is depicted in Figure [ Since the model chooses the most-likely match- ing word, larger corruptions will at some point lead to word changes.
Ineeion city the ester staelicy cotadioo nakeos nome tib swae se one zeze seco saso raze tane mexcce bearted the tootbell world cup in
(a) Ground truth sequence from text8.
nexiin wist the nstedl atalliie rosadioo sakens hore Clip swari se one zate zeco zesd pare Exnx pexics bosted the cootgall wolld cup In
appenee
(b) Corrupted sentence.
roucice city the uctoc atadiup axtadio axtace club se one sase seco sasd sase Eaos sexics boated the football would cop An
(c) Suggested, prediction by the model.
Figure 5: Spell checking with Multinomial Text Diffusion.
6.2 Segmentation maps
For image-type data, we introduce categorical image dataset: the cityscapes dataset is repurposed for unconditional image segmentation learning. In contrast with the standard setting. the distribution over the segmentation targets needs to be learned without conditioning on the photograph. To reduce computational cost, we rescale the segmentation maps from cityscapes to 32 x 64 images using nearest neighbour interpolation. We utilize the global categories as prediction targets which results in an 8-class problem.
Model description The Argmax Flows are de- fined directly the K 8 categorical space. The density model p(w) is defined using affine coupling layers parametrized by DenseNets (Huan et al.] 2017). For the probabilistic in- verse we learn conditional flow g(ularr) which is also based on the affine coupling structure. Depending on the method, either softplus or Gumbel thresholding is applied to obtain v. Re- call that for our first Gumbel approach is equiv- alent to set g( g(u) to the unit uniform distri- bution, whereas g(ul is learned for Gumbel thresholding. We compare to existing dequantization strategies in literature: uniform UUria et al. 2013) and variational dequantization (Ho et al.] 2019) which are applied on the onehot representation. All models utilize the same underlying flow architectures and thus the number of parameters is roughly the same. The exception are uniform dequantization and the Gumbel distribution, since no additional variational flow distribution is needed. For more extensive details see Appendix B
Table 4: Performance of different dequantization methods on squares and cityscapes dataset, bits per pixel, lower is better. Cityscapes
ELBO
IWBO
Round/ Unif. UUria al. 2013
1.010
0.930
Round/ / Var. Ho et al., 2019
0.334
0.315
Argmax / Softplus thres. (ours)
0.303
0.290
Argmax / Gumbel dist. (ours)
0.365
0.341
Argmax Gumbel thres. (ours)
0.307
0.287
Multinomial Diffusion (ours)
0.305
The results of this experiment are shown in Table 4 in terms of ELBO and if available the IWBO (importance weighted bound) (Burda et al.) 2016) with 1000 samples measured in bits per pixel. Consistent with the language experiments, the traditional dequantization approaches (uniform/ / variational) are outperformed by Argmax Flows. Interestingly, although argmax flows with softplus thresholding achieves the best ELBO, the argmax flow with Gumbel thresholding approach achieves better IWBO. The Multinomial Diffusion model performs somewhat worse with 0.37 bpp on test whereas it scored 0.33 bpp on train. Interestingly, this the only model where overfitting was an issue and data augmentation was required, which may explain this portion of the performance difference. For all other models training performance was comparable to test and validation performance. Samples from the different models trained on cityscapes are depicted in Figure 4 Another interesting point is that coupling flows had difficulty producing coherent text samples (Figure Fp but do not suffer from this problem on the cityscapes data which is more image- like. As coupling layers where initially designed for images (Dinh et al. 2015), they may require adjustments to increase their expressiveness on text.
Comparison
9

---

7 Social Impact and Conclusion
Social Impact The methods described in this paper can be used learn categorical distributions. For that reason, they can potentially be used to generate high-dimensional categorical data, such as text or image segmentation maps, faster than iterative approaches. Possibly negative influences are the generation of fake media in the form of text, or very unhelpful automated chat bots for customer service. Our work could positively influence new methods for text generation, or improved segmentation for self-driving cars. In addition, our work may also be used for outlier detection to flag fake content. Also, we believe the method in its current form still distant from direct applications the ones mentioned above.
Conclusion In this paper we propose two extensions for Normalizing Flows and Diffusion models to learn categorical data: Argmax Flows and Multinomial Diffusion. Our experiments show that our methods outperform comparable models in terms of negative log-likelihood. In addition, our experiments highlight distinct performance gaps in the field: Between standard ARMs, continuous autoregressive models and non-autorre -utorergggree continuous models. This indicates that future work could focus two sources of decreased performance: 1) when discrete variables are lifted to continuous space and further 2) when removing autoregressive components.
Funding Disclosure
There are no additional sources of funding to disclose, beyond the affiliations of the authors.
10

---

References
Cooijmans, T.; Ballas, N.; Laurent, C.: Gülçehre, Ç.: Courville, C. Recurrent Batch Normalization. 5th International Conference on Learning Representations, ICLR. 2017.
Dai, Z.; Yang, Z.; Yang. Y.; Carbonell, J. G.; Le, Q. V.; Salakhutdinov, R. Transformer-XL: Attentive Language Models beyond Fixed-Length Context. Proceedings of the 57th Conference of the Association for Computational I Linguistics, ACI 2019. 2019.
Uria, B.; Murray, L,; Larochelle, H. RNADE: The Real-valued Neural Autoregressive Density-
estimator. Advances in Neural Information Processing Systems. 2013; pp 2175-2188
Theis, L.; van den Oord, A.; Bethge, M. A note on the evaluation of generative models. International Conference on Learning Representations. 2016.
Ho, J.; Chen, X.; Srinivas, A.; Duan, Y.; Abbeel, P. Flow++: Improving Flow-Based Generative
Models with Variational Dequantization and Architecture Design. 36th International Conference on Machine Learning 2019,
Sohl-Dickstein, J.: Weiss, E. A.; Maheswaranathan, N.; Ganguli, S. Deep Unsupervised Learning using Nonequilibrium Thermodynamies, Proceedings of the 32nd International Conference on Machine Learning, ICML. 2015.
Ho, J.; Jain, A.; Abbeel, P. Denoising Diffusion Probabilistic Models. CoRR 2020, abs/2006. 1239.
Song, J.; Meng, C.; Ermon, S. Denoising Diffusion Implicit Models. CoRR 2020, abs/2010.02502.
Nichol, A. Q.: Dhariwal, P. Improved Denoising Diffusion Probabilistic Models. 2021; https:
Rezende, D.; Mohamed, S. Variational Inference with Normalizing Flows. Proceedings of the 32nd
International Conference on Machine Learning. 2015; pp 1530-1538.
Nielsen, D.; Jaini, P:; Hoogeboom, E.; Winther, O.: Welling, M. SurVAE Flows: Surjections to Bridge the Gap between VAEs and Flows. CoRR 2020, abx/2007.02731.
Maddison, C. J.; Tarlow, D.; Minka, T. A* Sampling. Advances in Neural Information Processing Systems 27: Annual Conference on Neural Information Processing Systems. 2014.
Kool, W.; Van Hoof. H.: Welling, M. Stochastic Beams and Where To Find Them: The Gumbel- Top-k Trick for Sampling Sequences Without Replacement. Proceedings of the 36th International Conference on Machine Learning. 2019.
Winkler, C.; Worrall, D. E.; Hoogeboom, E.; Welling, M. Learning Likelihoods with Conditional Normalizing Flows. CoRR 2019, abs/1912.00044..
Germain, M.; Gregor, K.; Murray, L.; Larochelle, H. Made: Masked autoencoder for distribution estimation. International Conference on Machine Learning. 2015; pp 881 881-889.
Kingma, D. P.; Welling, M. Auto-Encoding Variational Bayes. Proceedings of the 2nd International
Conference on Learning Representations. 2014.
Rezende, D. J.; Mohamed, S.; Wierstra, D. Stochastic Backpropagation and Approximate Inference in Deep Generative Models. Proceedings of the 31th International Conference on Machine Learning, ICML. 2014.
Goodfellow, L; Pouget-Abadie, J.; Mirza, M.; Xu, B.; Warde-Farley, D.; Ozair, S.; Courville, A.; Bengio, Y. Generative adversarial nets. Advances in neural information processing systems. 2014; pp 2672-2680.
Dinh, L.; Sohl-Dickstein, J.; Bengio, S. Density estimation using Real NVP. Sth International Conference Learning Representations, ICLR 2017,
Kingma, D. P.; Salimans, T.; Jozefowicz, R.; Chen, X.; Sutskever, I.; Welling, M. Improved variational inference with inverse autoregressive flow. Advances Neural Information Processing Systems. 2016; pp 4743-4751.
11

---

Papamakarios, G.; Murray, L.; Pavlakou, T. Masked autoregressive flow for density estimation. Advances in Neural Information Processing Systems. 2017: pp
Chen, T. Q.; Rubanova, Y.; Bettencourt, J.; Duvenaud, D. K. Neural ordinary differential equations. Advances in Neural Information Processing Systems. 2018; pp 6572-6583.
Song, Y.; Meng. C.; Ermon, S. MintNet: Building Invertible Neural Networks with Masked Convo- lutions. Advances in Neural Information Processing Systems 32: Annual Conference on Neural Information Processing Systems 2019, NeurIPS 2019. 2019.
Perugachi-Diate Y.; Tomczak, M.; Bhulai, S. Invertible DenseNets. CoRR 2020, abs/2010.02125..
Nielsen, D.; Winther, O. Closing the Dequantization Gap: PixeICNN as Single-Layer Flow. Advances in Neural Information Processing Systems 33: Annual Conference on Neural Information Processing Systems 2020, NeurIPS. 2020.
Tran, D.; Vafa, K.; Agrawal, K.; Dinh, L.: Poole, B. Discrete Flows: Invertible Generative Models of Discrete Data. ICLR 2019 Workshop DeepGenStruct 2019,
Hoogeboom, E.; Peters, J. W. T.; den Berg, R.; Welling, M. Integer Discrete Flows and Lossless Compression. Neural Information Processing Systems 2019, NeurIPS 2019. 2019; pp 12134 12144.
van den Berg, R.; Gritsenko, A. A.; Dehghani, M.; Sønderby, C. K.; Salimans, T. IDF++: Analyzing and Improving Integer Discrete Flows for Lossless Compression. CoRR 2020, abs/2006.12455 12459,
Ziegler, 7. M.; Rush, A. M. Latent Normalizing Flows for Discrete Sequences. Proceedings of the 36th International Conference on Machine Learning, ICML. 2019.
Lippe, P.; Gavves, E. Categorical Normalizing Flows via Continuous Transformations. CoRR 2020, abs/2006.09790.
Song. Y.; SShl-Dickstedd J.; Kingma, D. P.; Kumar, A.; Ermon, S.; Poole, B. Score-Based Generative Modeling through Stochastic Differential Equations. CoRR 2020, 0tbb/2011.134555
Al-Rfou, R.; Choe, D.; Constant, N.; Guo, M.; Jones, L. Character-Level Language Modeling with Deeper Self-Attention. The Thirty-Third AAAI Conference on Artificial Intelligence, AAAI 2019. 2019.
Huang, G.; Liu, Z.; Van Der Maaten, L., Weinberger, K. Q. Densely connected convolutional networks. Proceedings of the IEEE conference on computer vision and pattern recognition. 2017; pp 4700-4708.
Burda, Y.; Grosse, R. B.; Salakhutdinov. R. Importance Weighted Autoencoders. 4th International Conference on Learning Representations. 2016.
Dinh, L.; Krueger, D.; Bengio, Y. NICE: Non-linear independent components estimation. 3rd International Conference on Learning Representations, ICLR, Workshop Truck Proceedings 2015,
Cordts, M.; Omran, M.; Ramos, S.; Rehfeld, Enzweiler, M.; Benenson, R.; Franke, U.; Roth, S.; Schiele, B. The Cityscapes Dataset for Semantic Urban Scene Understanding. 2016 IEEE Confer- ence on Computer Vision and Pattern Recognition, CVPR. 2016; pp 3213-3223.
12


Reviews:
[{"review_id": "xRIGBgQzwWn", "reviewer": "Reviewer_YpDK", "summary": "This paper proposes two methods for generative modeling of categorical data, one based on normalizing flows and the other on diffusions. The first method models observed data as the result of a continuous normalizing flow, followed by an argmax operation rather than a rounding one sometimes used for ordinal data. The model is then trained by introducing a carefully constructed approximate posterior (of the continuous variable given the categorical observation), and maximizing the ELBO. The second method is a fairly straightforward extension of commonly used diffusions to the categorical setting.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Overall I quite liked the paper, it's very clearly written, well motivated, and showed good results. I found the posterior constructions from section 3.1 ensuring that the approximate posterior has the correct support to be elegant solutions; and the proposed argmax flows do beat competing dequantization schemes empirically, as well as VAE-based approaches. While argmax flows are still outperformed by purely autoregressive approaches, I think this is OK, and is actually a similar behaviour as what is observed with images.\n\nWhile I found showing that categorical diffusion models can get reasonable performance interesting (Figure 5 is nice), the results for diffusions are weaker than for argmax flows, and the paper reads like two papers pasted together. Aside from the application, there really isn't that much in common between the proposed approaches and although the paper remains clear, I think this weakens the narrative of the paper.\n\nMinor notational issue: It looks from the unnumbered equation between equations 1 and 2 that \"P\" is reserved for probabilities, and \"p\" for continuous densities, but for diffusion models \"P(x_0)\" is used, but so are \"p(x_{t-1}|x_t)\" and \"p(x_T)\".\n\n======================================================================================================\n\nEDIT 1 AFTER REBUTTAL\n\nI have read the other reviews as well as the author's rebuttal. Although I still think the two proposed methods are mostly disconnected, and find the contributions for argmax flows more exciting and novel than the ones for multinomial diffusions, I also still think the proposed argmax flow model is an elegant approach with good empirical results deserving publication, and thus have kept my original score.\n\n======================================================================================================", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Limitations and potential negative societal impacts are adequately discussed.", "review_text": "", "extra_scores": {"ethical_concerns": "None.", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "4"}, "extra_text": {}}, {"review_id": "gEjje2ed0LU", "reviewer": "Reviewer_rJxa", "summary": "The paper proposed two methods to learn categorical distributions. \n\nMethod 1: ArgMax Flows. The sampling procedure is simply taking argmax of a continuous random variable. The learning procedure involves a variational distribution that satisfies the argmax constraint and optimizing ELBO. Several variation distributions are introduced. In addition, the Cartesian product of argmax flows is introduced to trade-off between symmetry and dimension.\n\nMethod 2: Multinomial DDPM. The diffusion process is defined by categorical distributions. The paper derived ELBO for these models, and introduced a special parameterization of the reverse process.\n\nThe paper evaluated the methods in several experiments including text generation and image segmentation. They show better performances than baseline methods.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "The paper is original to my knowledge. It is well written and easy to follow. The research problem is significant because many tasks involve discrete data generation. \n\n\nPros:\n- The ArgMax Flow is generally an elegant framework. It is natural to optimize ELBO when likelihood cannot be computed. Then, the argmax condition for the variational distribution is a simple but smart solution that validates the ELBO objective. The several proposed variational distributions are elegant. \n- The Cartesian product is very useful especially when data is supported on a big, discrete domain. \n- The multinomial DDPM defines the diffusion process with categorical distributions, and therefore the model can be applied to categorical data. The closed-form ELBO is derived so we can directly optimize it. \n- The experiments are generally persuasive, and the character-level spell check is interesting. \n\n\nCons / main concerns:\n- Except that both ArgMax flows and multinomial DDPM are designed for categorical data, what are some inherent connections between these two types of models? Why do we want two models? Is there anything that one model can do while the other cannot? If one model is significantly more advantageous than the other in certain cases, is there any reason / recipe? \n- In eq(6), what does $v_{-x}$ mean? A clear definition is needed.\n- Line 109, it seems the threshold function should be $\\mathbb{R}^2\\rightarrow(-\\infty,T)$. This still looks strange: $T$ is an input; how can an input determine the range of the output? A re-formulation is needed here: e.g., $\\mathrm{threshold}(\\cdot;T)$ or $\\mathrm{threshold}_T(\\cdot)$.\n- Regarding the Cartesian products, the authors had an additional experiment in the appendix, where the data domain size is 27. Theoretically, this method should be more useful when the data domain is very large. For instance, in 64*64 image generation, the data domain is [0,...,255] * [1,..,3] * [1,...,64] * [1,...,64]. In 16kHz 1-sec audio generation, the data domain is [-32768,...,32768] * [1,...,16000]. Therefore, it is more persuasive to show the effect of the Cartesian products if this method is evaluated on one of these tasks (although Table 1 indicates the rounding operation is more suitable). \n- The diffusion process defined by categorical distribution is a natural extension of the Bernoulli distribution in the original diffusion model. Although the paper mentioned it in the related work, I would strongly suggest the authors to write down the original Bernoulli version in the background section so that readers can have a clear sense of the background and the contribution of this paper. In addition, a clear definition for the $\\mathcal{C}(\\cdot|\\cdot)$ in equation (11) is needed. There should be the unit vector $\\textbf{1}$ after $\\beta_t/K$. \n- Line 164-165: it is a good idea to parameterize $\\mu(x_t,t)$, but why do you call it $\\hat{x}_0$? Why can it approximate $x_0$? Either a theoretical or experimental validation is necessary. \n- Note: most of these concerns are about validation and explanation, and they do not harm the innovation or contribution of the paper. I would claim again that this paper is generally very elegant. \n\n\nMinor concerns:\n- There should be commas or dots after each equation. There should be a number for the eq at line 56. \n- Some math symbols are bold somewhere but not bold elsewhere. Some symbols are confusing, e.g. in eq (4), I would prefer $v_{d,k}$ rather than $v_{dk}$.\n\n\nOther questions just for curiosity:\n- Is there a way to estimate the likelihood of ArgMax flows, even with a specially designed variational distribution?\n- The spell check is interesting but seems a little too heuristic. If there are more mistakes then you probably need to start from an earlier step. Do you have a solution for that? E.g., one way is to learn a network that predicts the noise level of the input, and then determine which step to begin accordingly. \n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "Social impact is discussed in paper. Concerns on the limitations are in the main review. ", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "5 hours. "}, "extra_text": {}}, {"review_id": "_0kAPYm_k2-", "reviewer": "Reviewer_Cd2a", "summary": "This paper proposes two distinct methods for learning categorical distributions. The first is based on an extension to variational dequantization with argmax operators which allows normalizing flow models to learn categorical distributions. The second is based on an extension to Bernoulli diffusion probabilistic models so that it can handle larger than 2 categories. Proof-of-concept experiments demonstrate that the proposed methods can successfully handle high dimensional categorical distributions.", "questions": "", "limitations": "", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Significance**: Normalizing flow models and diffusion models are mostly designed for continuous distributions. Extending them to the realm of  discrete distributions has important practical values and is highly relevant to the NeurIPS community. This paper proposes new methods to achieve this.\n\n**Clarity**: The submission is clearly written and very easy to follow.\n\n**Quality and originality**: Here is a list of concerns:\n1. Argmax flows and multinomial diffusions are two completely different model families. The ways of extending them to categorical distributions are also very different. It is therefore unclear why authors are combining these two hardly related ideas into a single paper. Authors should either separate these ideas into two different piece of publications, or have a more unified story.\n\n2. I am not impressed by the novelty of argmax flows. It is a simple extension to variational dequantization. The argmax flow layers cannot be composed to form deeper architectures. In fact, it is a simple \"preprocessing\" layer attached to standard normalizing flows, same as variational dequantization. \n\n3. It is unclear why argmax flows can perform better than VAEs when priors are learnable and parameterized by normalizing flows. Authors argue that argmax flows are better than variational dequantization because the continuous variables are not restricted to the unit interval [0, 1). However, if this is the right intuition for why argmax flows are better, we should instead use VAEs and train a normalizing flow model on the latent space as the prior. The latent space of VAEs is even more flexible than argmax flows, since it is not restricted by the argmax constraint. Although authors have compared VAE + IAF with their argmax flows in one experiment, I find it unsatisfying as the difference in model architectures is unclear, and comparison in other experimental settings is also needed.\n\n4. I am not impressed by the novelty of multinomial diffusions. It is a straightforward generalization to the Bernoulli diffusion models in [Sohl-Dickstein et al. 2015], and as the author correctly noted, [Song et al. 2020] has already included a same formulation in the appendix. Moreover, in the language data experiment, authors use binary Cartesian products to represent multiple categories. Doesn't this reduce the problem to modeling binary data? If so, isn't multinomial diffusion the same as Bernoulli diffusion in the original 2015 paper?\n\n**Post rebuttal update**:\nThanks for the response. Authors have addressed some of my questions, so I'm happy to increase my score. However, my major concerns still remain valid:\n\nFirst, as noted by other reviewers as well, the connection between argmax flows and multinomial diffusion models are underwhelming. Authors responded that the biggest connection of them is that they both model discrete distributions. This is not convincing. Modeling discrete distributions is not an unsolved problem—autoregressive models, VAEs, energy-based models, and sometimes even GANs can model discrete distributions. We can't just publish paper by putting together separate ideas on VAEs, EBMs, etc, if they both model discrete distributions.\n\nSecond, I'm still not convinced that the proposed argmax flows are useful. The proposed method is a special case of a VAE. The paper should focus more on why a special variant of a VAE, with a hand-designed decoder, is necessary for good performance on discrete data. In the paper, authors motivate argmax flows as a more flexible variant to variational dequantization. However, VAEs are even more flexible and should be better from that reasoning. Authors also mentioned in the response that argmax flows are easier to optimize due to the fixed decoder, which seems to say that limiting the flexibility can contribute to the performance. This contradicts with the motivation from variational dequantization to argmax flows.\n", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "No obvious concerns.", "review_text": "", "extra_scores": {"needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": "3"}, "extra_text": {}}, {"review_id": "DH8GXEcXRTP", "reviewer": "Reviewer_4fJX", "summary": "The paper proposes two new generative models for non-ordinal discrete random variables. \n\nFirst, the paper proposes Argmax Flows, which apply argmax operation on continuous random variables, modeled by flow-based models. Note that estimating the likelihood of a model with a surjection requires a valid stochastic inverse, which should satisfy right-inverse (argmax constraint in the paper) and absolute continuity. For stochastic inverses of the argmax operation, the authors propose three novel methods: (1) thresholding, (2) Gumbel, and (3) Gumbel thresholding. The proposed stochastic inverses and the Argmax Flows are jointly trained by maximizing the ELBO.\n\nSecond, the authors propose a diffusion-based generative model for non-ordinal discrete data. Consider a pre-defined Markov chain, by which a data distribution is transformed to a prior distribution. The model learns a reverse Markov chain from the prior to the data. In this context, the paper proposes to diffuse non-ordinal discrete data to a uniform categorical distribution. For each transition, with a small probability, every value is independently decided to be either resampled or remaining the same; when we resample, a new value will be drawn from a uniform categorical distribution. For the reverse Markov chain, the authors acknowledge that a closed-form (conditional) reverse transition exists for the proposed diffusion, and they follow categorical distributions. Based on that, the paper proposes an efficient parameterization of the reverse transitions, which reduces the variance of the ELBO estimation.\n\nIn the experiments, the paper demonstrates that the proposed methods improve modeling language and segmentation maps (of cityscapes images) compared to previous generative models designed for ordinal data.", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "(Originality)  \nThe originality of the paper is clear:  \nFirst, the paper proposes novel dequantization methods for non-ordinate discrete data, which are stochastic inverses of the argmax operation—noting that it is nontrivial to satisfy the necessary conditions to design valid stochastic inverses of a given surjection. The authors discuss how to achieve the constraints for a given surjection (aka argmax) to design its probabilistic inverse.  \n\nSecond, the paper also introduces a diffusion-based model for non-ordinate discrete data.  \n\nFinally, the paper demonstrates the practicalities of the proposed methods via language modeling experiments and density estimation tasks on segmentation maps.  \n\n(Significance & Quality)  \nIn general, the results of the paper are important. However, the followings aspects of the paper can be improved.\n\nFirst, while one of the main contributions is the novel stochastic inverses of argmax operation, the discussions about their properties are limited. Aside from the proposed inverses satisfy the necessary constraints, I consider it important to discuss the expressivity of the proposed methods or potential failure cases. For instance, due to the truncations in the variational distributions, the proposed methods may not be optimal when prior $p(u_x) > 0$ for $u_x > T, x \\\\in \\\\{1, \\dots, K\\\\}$ (considering $T > 0$). What can we do instead? In other instances, it would be helpful to provide a simple example of how the proposed models represent dequantizations for uniform categorial distribution or other toy cases.\n\nSecond, the connection between the Argmax Flows and diffusion-based models is limited. I believe that they are both important contributions. Based on the current version, however, I found that they are somewhat different topics. I believe that clarifying their connection improves the significance of the paper.\n\n(Clarity)  \nThe paper is self-contained and has a well-organized structure to present relevant topics. In particular, the paper concisely raises the difficulty of previous methods regarding modeling non-ordinal data and explains how the proposed models address the issue.", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "(Limitations)  \nPlease find comments on (Significance & Quality) section in the main review. \n\n(Social Impact)  \nN/A", "review_text": "", "extra_scores": {"ethical_concerns": "N/A", "needs_ethics_review": "No", "ethics_review_area": "", "time_spent_reviewing": ">12hrs"}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper presents two approaches for learning expressive distributions over discrete variables. Although there are valid criticisms regarding the novelty (especially for the multinomial diffusion), the majority of the reviewers found the introduction of argmax flows an interesting and non-trivial extension of prior works. Given the community's interest in discrete variables, I recommend this paper for acceptance. For the final camera-ready version, I'd appreciate it if the authors could expand the discussion of limitations as suggested by 4fJX (in addition to other promised changes). 
Decision: Accept (Poster)