Argmax Flows and Multinomial Diffusion: Learning Categorical Distributions

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 . 2 Background . 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 ) 3 Argmax Flows . 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 . 3.1 Probabilistic Inverse . 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 .

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.