Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions
Year: 2021

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.

text: 1 Introduction
Many sources of high-dimensional data are categorical, for example language and image segmentation. 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 et al., 2017; Dai et al., 2019).
Normalizing flows are attractive because they can be designed to be fast both in the evaluation and sampling direction. Typically, normalizing flows model continuous distributions. As a result, directly optimizing a flow on discrete data may lead to arbitrarily high likelihoods. In literature this problem is resolved for ordinal data by adding noise in a 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.
Other attractive generative models are diffusion models (Sohl-Dickstein et al., 2015), which are fast to train due to an objective that decomposes over time steps (Ho et al., 2020). Diffusion models
⇤Equal contribution.
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
typically have a fixed diffusion process that gradually adds noise. This process is complemented by a learnable generative process that denoises the signal. Song et al. (2020); Nichol and Dhariwal (2021) have shown that diffusion models can 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 1): i) Argmax Flows bridge the gap between categorical data and continuous normalizing flows using an argmax transformation and a corresponding family of probabilistic inverses for the argmax. In addition ii) we introduce Multinomial Diffusion, which is a diffusion model directly defined on categorical variables. Opposed to normalizing flows, defining diffusion for discrete variables directly does not require gradient approximations, because the diffusion trajectory is fixed. As a result of our work, generative normalizing flows and diffusion models can directly learn categorical data.

heading: 2 Background
text: Normalizing Flows Given V = Rd and Z = Rd with densities pV and pZ respectively, normalizing flows (Rezende and Mohamed, 2015) learn a bijective and differentiable transformation g : Z ! V such that the change-of-variables formula gives the density at any point v 2 V:
pV (v) = pZ(z) · det dz dv , v = g(z), (1)
where pZ can be any density (usually chosen as a standard Gaussian). Thus, normalizing flows provide a powerful framework to learn exact density functions. However, Equation (1) is restricted to continuous densities.
To learn densities on ordinal discrete data (such as natural images), typically dequantization noise is added (Uria et al., 2013; Theis et al., 2016; Ho et al., 2019). Nielsen et al. (2020) reinterpreted dequantization as a surjective flow layer v 7! x that is deterministic in one direction (x = round(v)) and stochastic in the other (v = x+ u where u ⇠ q(u|x)). Using this interpretation, dequantization can be seen as a probabilistic right-inverse for the rounding operation in the latent variable model given by:
P (x) = Z P (x|v)p(v) dv, P (x|v)= x= round(v) ,
where round is applied elementwise. In this case, the density model p(v) is modeled using a normalizing flow. Learning proceeds by introducing the variational distribution q(v|x) that models the probabilistic right-inverse for the rounding surjection and optimizing the evidence lower bound (ELBO):
logP (x) Ev⇠q(v|x) [logP (x|v) + log p(v) log q(v|x)] = Ev⇠q(v|x) [log p(v) log q(v|x)] . (2)
The last equality holds under the constraint that the support of q(v|x) is enforced to be only over the region S = {v 2 Rd : x = round(v)} which ensures that P (x|v) = 1.
Diffusion Models Given data x0, a diffusion model (Sohl-Dickstein et al., 2015) consists of predefined variational distributions q(xt|xt 1) that gradually add noise over time steps t 2 {1, . . . , T}. The diffusion trajectory is defined such that q(xt|xt 1) adds a small amount of noise around xt 1. This way, information is gradually destroyed such that at the final time step, xT carries almost no information about x0. Their generative counterparts consists of learnable distributions p(xt 1|xt) that learn to denoise the data. When the diffusion process adds sufficiently small amounts of noise, it
Algorithm 1 Sampling from Argmax Flows Input: p(v) Output: Sample x Sample v ⇠ p(v) Compute x = argmaxv
Algorithm 2 Optimizing Argmax Flows Input: x, p(v), q(v|x) Output: ELBO L Sample v ⇠ q(v|x) Compute L = log p(v) log q(v|x)
suffices to define the denoising trajectory using distributions that are factorized (without correlation) over the dimension axis. The distribution p(xT ) is chosen to be similar to the distribution that the diffusion trajectory approaches. Diffusion models can be optimized using variational inference:
logP (x0) Ex1,...xT⇠q h log p(xT ) + TX
t=1
log p(xt 1|xt) q(xt|xt 1)
i .
An important insight in diffusion is that by conditioning on x0, the posterior probability q(xt 1|xt,x0) = q(xt|xt 1)q(xt 1|x0)/q(xt|x0) is tractable and straightforward to compute, permitting a reformulation in terms of KL divergences that has lower variance (Sohl-Dickstein et al., 2015). Note that KL q(xT |x0)|p(xT ) ⇡ 0 if the diffusion trajectory q is defined well:
logP (x0) Eq h log p(x0|x1) KL q(xT |x0)|p(xT )
TX
t=2
KL q(xt 1|xt,x0)|p(xt 1|xt) i (3)

heading: 3 Argmax Flows
text: Argmax flows define discrete distributions using 1) a density model p(v), such as a normalizing flow, and 2) an argmax layer that maps the continuous v 2 RD⇥K to a discrete x 2 {1, 2, ...,K}D using
x = argmaxv where xd = argmax k vdk. (4)
This is a natural choice to model categorical variables, because it divides the entire continuous space of v into symmetric partitions corresponding to categories in x. To sample from an argmax flow sample v ⇠ p(v) and compute x = argmaxv (Algorithm 1). To generate reasonable samples, it is up to the density model p(v) to 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 x, we have to compute
P (x) = Z P (x|v)p(v)dv, P (x|v)= x=argmax(v) , (5)
which is intractable. Consequently, we resort to variational inference and specify a variational distribution q(v|x). We note that naïvely choosing any variational distribution may lead to samples v ⇠ q(v|x) where (x = argmaxv) = 0, which yields an ELBO of negative infinity. To avoid this, we need a variational distribution q(v|x) that satisfies what we term the argmax constraint:
x = argmaxv for all v ⇠ q(v|x). That is, the variational distribution q(v|x) should have support limited to S(x) = {v 2 RD⇥K : x = argmaxv}. Recall that under this condition, the ELBO simplifies to Ev⇠q(v|x) [log p(v) log q(v|x)], as shown in Algorithm 2. For an illustration of the method see Figure 1a.

heading: 3.1 Probabilistic Inverse
text: The argmax layer may be viewed as a surjective flow layer (Nielsen et al., 2020). With this view, the variational distribution q(v|x) specifies a distribution over the possible right-inverses of the argmax function, also known as a stochastic inverse or probabilistic inverse. Recall that the commonly-used dequantization layer for ordinal data corresponds to the probabilistic inverse of a rounding operation. As summarized in Table 1, this layer may thus be viewed as analogous to the argmax layer, where the round is for ordinal data while the argmax is for categorical data.
We are free to specify any variational distribution q(v|x) that satisfies the argmax constraint. In the next paragraphs we outline three possible approaches. Since operations are performed independently across dimensions, we omit the dimension axis and let v 2 RK and x 2 {1, . . . ,K}.
Algorithm 3 Thresholding-based q(v|x) Input: x, q(u|x) Output: v, log q(v|x) u ⇠ q(u|x) vx = ux v x = threshold(u x,x) log q(v|x) = log q(u|x) log | det dv/ du| Algorithm 4 Gumbel-based q(v|x) Input: x, Output: v, log q(v|x) max = log P i exp i vx ⇠ Gumbel( max) v x ⇠ TruncGumbel( x,vx) log q(v|x) = logGumbel(vx| max)
+ logTruncGumbel(v x| x,vx)
Thresholding (Alg. 3). A straightforward method to construct a distribution q(v|x) satisfying the argmax constraint is to use thresholding. That is, we first sample an unbounded variable u 2 RK from q(u|x), which can be for example a conditional Gaussian or normalizing flow. Next, we map u to v such that element x is the largest:
vx = ux and v x = thresholdT (u x) (6)
where the thresholding is applied elementwise with threshold value T = vx. This ensures that element vx is the largest, and consequently that q(v|x) satisfies the argmax constraint. Note that we require the threshold function to be bijective, thresholdT : R ! ( 1, T ), so that we can use the change-of-variables formula to compute log q(v|x). In our implementation, thresholding is implemented using a softplus such that all values are mapped below a limit T :
v = thresholdT (u) = T softplus(T u), (7)
where softplus(z) = log(1 + ez) and for which it is guaranteed that v 2 ( 1, T ).
Gumbel (Alg. 4). An alternative approach is to let q(v|x) = Gumbel(v| ) restricted to argmaxv = x, where the location parameters NN(x) are predicted using a neural network NN. The Gumbel distribution has favourable properties: The argmax and max are independent and the max is also distributed as a Gumbel:
max i vi ⇠ Gumbel( max), (8)
where max = log P
i exp i. For a more extensive introduction see (Maddison et al., 2014; Kool et al., 2019). To sample v ⇠ q(v|x), we thus first sample the maximum vx according to Eq. 8. Next, given the sample vx, the remaining values can be sampled using truncated Gumbel distributions:
vi ⇠ TruncGumbel( i;T ) where i 6= x (9)
where the truncation value T is given by vx which ensures that the argmax constraint vx > vi for i 6= x is satisfied. Recall that to optimize Eq. 2, log q(v|x) 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 (argmaxv = i) = exp i/ X
i
exp i, (10)
which we use to initialize the location parameters to match the empirical distribution of the first minibatch of the data.
Gumbel Thresholding. This method unifies the methods from the previous two sections: Gumbel distributions and thresholding. The key insight is that the Gumbel sampling procedures as defined above can be seen as a reparametrization of a uniform noise distribution U(0, 1)K which is put through the inverse CDF of the Gumbel distributions (see Table 5). From the perspective of changeof-variables, the log likelihood denotes the log volume change of this transformation. To increase expressitivity the uniform distribution can be replaced by a normalizing flow q(u|x) that has support on the interval (0, 1)K , which can be enforced using a sigmoid transformation. This section shows that a large collection of thresholding functions can be found by studying (truncated) inverse CDFs. In practice we find that performance is reasonably similar as long as the underlying noise u is learned.
Behavior of the Variational Posterior Although several methods to learn q have been proposed, it is unclear what expressitivity is required. In the following, the interactions between q(v|x) and the density model p(v) are discussed. Recall that the variational bound that is optimized under expectation of a data distribution D can be seen as minimizing the KL distance between the aggregated posterior q(v) = Ex⇠Dq(v|x) and the density model p(v), so KL(q(v)|p(v)). 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 q(v|x) 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 a universal density approximator, then the model can be fitted for any wellbehaved q(v|x). 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 q(v) = Ex⇠Dq(v|x). In practice, over training steps we find that q does smooth out these boundary artifacts, and counteracts the thresholding so that the aggregated posterior becomes smoother.

heading: 3.2 Cartesian Products of Argmax Flows
text: In the current description, Argmax Flows require the same number of dimensions in v as there are classes in x. 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 {1, . . . , 256}, but also using two hexadecimal numbers {1, . . . , 16}2 or alternatively using eight binary numbers. Specifically, any base K variable x(K) 2 {1, . . . ,K}D can be converted to a base M variable x(M) 2 {1, . . . ,M}dm⇥D where dm = dlogM Ke. Then the variable x(M) with dimensionality M · dm · D represents the variable x(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 2 can still be optimized using an M -categorical Argmax Flow. Finally, note that Cartesian products of binary spaces are a special case where the variable can be encoded symmetrically into a 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 log2 K .

heading: 4 Multinomial Diffusion
text: In this section we introduce an alternative likelihood-based model for categorical data: Multinomial Diffusion. In contrast with previous sections, xt will be represented in one-hot encoded format xt 2 {0, 1}K . Specifically, for category k, xk = 1 and xj = 0 for j 6= 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 a t chance of resampling a category uniformly:
q(xt|xt 1) = C(xt|(1 t)xt 1 + t/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 xt given x0 as:
q(xt|x0) = C(xt|↵̄tx0 + (1 ↵̄t)/K) (12)
where ↵t = 1 t and ↵̄t = Qt
⌧=1 ↵⌧ . Intuïtively, for each next timestep, a little amount of uniform noise t over the K classes is introduced, and with a large probability (1 t) the previous value xt 1 is sampled. Using Equation 11 and 12 the categorical posterior q(xt 1|xt,x0) can be computed in closed-form:
q(xt 1|xt,x0) = C(xt 1|✓post(xt,x0)), where ✓post(xt,x0) = ✓̃/ KX
k=1
✓̃k
and ✓̃ = [↵txt + (1 ↵t)/K] [↵̄t 1x0 + (1 ↵̄t 1)/K].
(13)
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 q. Although predicting the noise is difficult for discrete data, we predict a probability vector for x̂0 from xt and subsequently parametrize p(xt 1|xt) using the probability vector from q(xt 1|xt, x̂0), where x0 is approximated using a neural network x̂0 = µ(xt, t). Equation 13 will produce valid probability vectors that are non-negative and sums to one under the condition that the prediction x̂0 is non-negative and sums to one, which is ensured with a softmax function in µ. To summarize:
p(x0|x1) = C(x0|x̂0) and p(xt 1|xt) = C(xt 1|✓post(xt, x̂0)) where x̂0 = µ(xt, t) (14)
The KL terms in Equation 3 can be simply computed by enumerating the probabilities in Equation 13 and 14 and computing the KL divergence for discrete distributions in Lt 1 with t 2:
KL q(xt 1|xt,x0)|p(xt 1|xt) = KL C(✓post(xt,x0))|C(✓post(xt, x̂0)) , (15)
which can be computed using P
k ✓post(xt,x0))k · log ✓post(xt,x0))k ✓post(xt,x̂0))k . Furtermore, to compute log p(x0|x1) use that x0 is onehot:
log p(x0|x1) = X
k
x0,k log x̂0,k (16)

heading: 5 Related Work
text: 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 et al., 2015).
Normalizing Flows typically learn a continuous distribution and dequantization is required to train these methods on ordinal data such as images. A large body of work is dedicated to building more expressive continuous normalizing flows (Dinh et al., 2017; Germain et al., 2015; Kingma et al., 2016; Papamakarios et al., 2017; Chen et al., 2018; Song et al., 2019; Perugachi-Diaz et al., 2020). To learn ordinal discrete distributions with normalizing flows, adding uniform noise 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 et 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 on 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 et al., 2020). In other works, VAEs have been adapted to learn a normalizing flow for the latent space (Ziegler and Rush, 2019; Lippe and Gavves, 2020). However, these approaches typically still utilize an argmax heuristic to sample, even though this is not the distribution specified during training.
Diffusion models were first introduced in Sohl-Dickstein et al. (2015), who developed diffusion for Gaussian and Bernoulli distributions. Recently, Denoising Diffusion models Ho et 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
from as they require iterations over the many timesteps in the chain. Song et al. (2020); Nichol and Dhariwal (2021) showed that in practice samples can be generated using significantly fewer steps. Nichol and Dhariwal (2021) demonstrated that importance-weighting the objective components greatly improves log-likelihood performance. In Song et al. (2020) a continuous-time extension of denoising diffusion models was proposed. After initial release of this paper we discovered that Song et al. (2020) concurrently also describe a framework for discrete diffusion, but without empirical evaluation.

heading: 6 Experiments
text: In our experiments we compare the performance of our methods on language modelling tasks and learning image segmentation maps unconditionally.

heading: 6.1 Language data
text: 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 are tested: using an autoregressive (AR) flow and a coupling-based flow for p(v). In these experiments the probabilistic inverse is based on the thresholding approach. Specifically, a conditional diagonal Gaussian q(u|x) is trained and thresholded which gives the distribution q(v|x). The argmax flow is defined on binary Cartesian products. This means that for K = 27, a 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 on the model proposed in (Lippe and Gavves, 2020). The coupling density model consists of 8 flow layers where each layer consists of a 1⇥ 1 convolution and mixture of logistics transformations Ho et al. (2019). In the multinomial text diffusion model, the µ network is modeled by a 12-layer Transformer. For more extensive details about the experiment setup see Appendix B.
? Results obtained by running code from the official repository for the text8 and enwik8 datasets.
that the role of tellings not be required also action characters passe d on constitution ahmad a nobilitis first be closest to the cope and dh ur and nophosons she criticized itm specifically on august one three mo vement and a renouncing local party of exte
nt is in this meant the replicat today through the understanding elemen t thinks the sometimes seven five his final form of contair you are lot ur and me es to ultimately this work on the future all all machine the silon words thereis greatly usaged up not t
(a) Samples from Multinomial Text Diffusion.
heartedness frege thematically infered by the famous existence of a fu nction f from the laplace definition we can analyze a definition of bin ary operations with additional size so their functionality cannot be re viewed here there is no change because its
otal cost of learning objects from language to platonic linguistics exa mines why animate to indicate wild amphibious substances animal and mar ine life constituents of animals and bird sciences medieval biology bio logy and central medicine full discovery re
(b) Samples from Argmax AR Flow.
ns fergenur d alpha and le heigu man notabhe leglon lm n two six a gg opa movement as sympathetic dutch the term bilirubhah acquired the bava rian cheeh segt thmamouinaire vhvinus lihnos ineoneartis or medical iod ine the rave wesp published harsy varb hhgh
danibah or manuccha but calpere that of the moisture soods and dristi ng attempt to cause any moderator called lk brown or totpdngs is usuall y able to nus and hockecrits borel qbisupnias section rybancase untecce mentation anymore the motion of plays on qr
(c) Samples from Argmax Coupling Flow.
Figure 3: Samples from models, text8.
(a) Samples from the Argmax Flow.
(b) Samples from the Multinomial Diffusion model.
(c) Cityscapes data.
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 a 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 a generative flow to learn the lifted continuous distribution using an argmax flow. An advantage of Argmax flows that may explain this difference is that they lift the variables into the entire Euclidean 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-based approaches. Also note that coupling flows have worse performance than autoregressive flows, with a difference that is 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 in Table 3 alongside 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-autoregressive models, Argmax Flows also outperforms the method that lifts the categorical space to a continuous space: IAF / SCF (Ziegler and Rush, 2019). Interestingly, the multinomial text diffusion is a non-autoregressive 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 in 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, a 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.
Unsupervised spell-checking An interesting by-product of the text diffusion model is that it can be used to spell-check text using a single forward pass. To demonstrate this, a sentence taken from the test data is corrupted by changing a few characters. This corrupted sequence is given as x1 to the generative denoising model, which is close to the data at step 0. Then the denoising model predicts p(x0|x1) and the mostlikely x0 can be suggested. Note that this model only works for character-level corruption, not insertions. An example is depicted in Figure 5. Since the model chooses the most-likely matching word, larger corruptions will at some point lead to word changes.

heading: 6.2 Segmentation maps
text: For image-type data, we introduce a 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 ⇥ 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 defined directly on the K = 8 categorical space. The density model p(v) is defined using affine coupling layers parametrized by DenseNets (Huang et al., 2017). For the probabilistic inverse we learn a conditional flow q(u|x) which is also based on the affine coupling structure. Depending on the method, either softplus or Gumbel thresholding is applied to obtain v. Recall that for our first Gumbel approach it is equivalent to set q(u|x) to the unit uniform distribution, whereas q(u|x) is learned for Gumbel
thresholding. We compare to existing dequantization strategies in literature: uniform (Uria 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.
Comparison 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 a 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 3) but do not suffer from this problem on the cityscapes data which is more imagelike. As coupling layers where initially designed for images (Dinh et al., 2015), they may require adjustments to increase their expressiveness on text.

heading: 7 Social Impact and Conclusion
text: Social Impact The methods described in this paper can be used to 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 is still distant from direct applications as 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-autoregressive continuous models. This indicates that future work could focus on two sources of decreased performance: 1) when discrete variables are lifted to a 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.

Review summaries:
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.

The paper proposed two methods to learn categorical distributions.   Method 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.  Method 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.  The paper evaluated the methods in several experiments including text generation and image segmentation. They show better performances than baseline methods.

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.

The paper proposes two new generative models for non-ordinal discrete random variables.   First, 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.  Second, 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.  In 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.