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