Argmax Flows: Learning Categorical Distributions with Normalizing Flows
Abstract
This paper introduces a new method to deﬁne and train continuous distributions such as normalizing ﬂows directly on categorical data, for example text and image segmentation.
The generative model is deﬁned by a composition of a normalizing ﬂow and an argmax function.
To optimize this model, we dequantize the argmax using a distribution that is a probabilistic right-inverse to the argmax.
This distribution lifts the categorical data to a continuous space on which the ﬂow can be trained.
We demonstrate that applying existing dequantization techniques na¨ıvely to categorical data leads to suboptimal solutions.
In addition, the model is fast both in generative (for sampling) and inference direction (for training), as opposed to autoregressive models.
1. Introduction
Typically, normalizing ﬂows model continuous distributions.
As a result, directly optimizing a ﬂow on discrete data may lead to arbitrarily high likelihoods.
An example of this phenomenon is when ﬂows are trained on ordinal 8-bit image data.
Pixels are discretely valued from 0 to 255 and need to be dequantized (Uria et al., 2013; Theis et al., 2016), i e. noise is added to lift the discrete pixels into a continuous space.
In their framework, dequantization is the inference distribution that is the natural counterpart of generative rounding.
Even though rounding is a natural transformation to obtain ordinal discrete variables, it places an unwanted inductive bias on categorical variables as diﬀerent pairs of categories can be closer or further apart.
In this paper we resolve these issues by proposing a generative model using an argmax surjection and a corresponding family of probabilistic right-inverses for these argmax surjections.
Argmax surjections are a deterministic map to obtain categorical variables from a continuous representation and they do not add uncorrelated noise to the sampling proce- dure.
To learn the underlying density model, we parametrize a probabilistic right-inverse to the argmax surjection, referred to as dequantization .
2. Preliminaries and Problem Setup
Let X = { 1 , 2 , · · · , K } d be a d -dimensional categorical space with probability mass function
P data ( x ) for x ∈ X where x = ( x 1 , x 2 , · · · , x d ) and each x i is a univariate categorical random variable of K categories.
Given V = R d and Z = R d with densities p V and p Z respectively, normalizing ﬂows (Rezende and Mohamed, 2015) learn a bijective and diﬀerentiable transformation g :
Z → V such that the change of variables formula gives the density at any point v ∈ V :
[Equation]
where p Z can be any density (usually chosen as standard Gaussian).
Thus, normalizing ﬂows provide a powerful framework to learn exact density functions in an unsupervised manner.
However, Equation (1) is restricted to continuous densities and cannot be applied in a straight-forward manner to discrete random variables.
Theis et al.
(2016) have shown that modeling this continuous density p ( v ) lower bounds the discrete distribution P ( x ) for uniform distribution.
Ho et al.
(2019) extended this framework for any variational distribution q ( u | x ).
In (Hoogeboom et al., 2020) it is shown that from a variational inference perspective, not only hypercubes but any partitioning of the space V can be optimized using this objective.
Furthermore, Nielsen et al.
(2020) reinterpreted the process of dequantization as a surjective transformation f : X → V that is deterministic in one direction (since 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 family of probabilistic right-inverses for a rounding surjection in the latent variable model given by:
[Equation]
In this case, the density model p ( v ) can be any distribution and is modeled using a normalizing ﬂow.
Learning proceeds by introducing the variational distribution q ( v | x ) that models the family of probabilistic right-inverses for rounding surjection and optimizing the
[Table]
x = (cid:98) v (cid:99)
[Table]
∼
[Table]
following bound:
[Equation]
Under the constraint that the support of q ( v | x ) is enforced to be only over the region
S = { v ∈ R d : x = round ( v ) } which ensures that P ( x | v ) = 1.
In the next section, we will propose a novel surjective transformation called argmax ﬂows that directly extends the ideas of rounding surjection to categorical random variables by designing a probabilistic right-inverse q ( v | x ) for a surjective transformation that maps a categorical random variable to a continuous random variable.
3. Argmax Flows
We propose a novel method to learn categorical data with continuous distributions.
This method consists of two parts: (1) the generative model that comprises of an argmax function i.e. x = arg max v , and (2) an inference model that requires a probabilistic right inverse q ( v | x ) with support over the region S = { v ∈ R d : x = arg max v } where v ∼ p ( v ) can be any underlying continuous distribution.
Such a latent variable model induces the following discrete distribution P model :
[Equation]
where p ( v ) may be modelled by any continuous distribution, such as a normalizing ﬂow.
Importantly, P ( x | v ) denotes a Kronecker delta peak such that P ( x | v )
=
δ ( x = arg max v ).
Intuitively, one can see P ( x | v ) as partitioning the space V for diﬀerent values of x .
To be precise, we deﬁne the (elementwise) argmax operation as:
[Equation]
assigning for each dimension d separately the index
k d such that v d,k d ≥ v d,k for all k = 1 , . . .
, K .
Consequently, in concise notation, the corresponding categorical variable is deﬁned to be x = arg max v , where x ∈ { 1 , . .
. , K } D .
See Fig.
1 for an illustration.
The main diﬃculty lies in optimizing this generative model.
Suppose one would na¨ıvely choose any variational distribution, then a sample v ∼ q ( v | x ) may lead to samples where there is no probability at all because δ ( x = arg max v ) = 0.
Instead, we need to learn the probabilistic right-inverses to the generative argmax function.
In other words, for any sample v ∼ q ( v | x ) it is desired that δ ( x = arg max v ) = 1.
Recall that under this condition, the expected lowerbound (ELBO) can be simpliﬁed as in Equation 2.
3.1. Dequantization by asymptotic thresholding
A relatively straightforward method to construct a distribution satisfying the argmax constraint, is by thresholding values using injective functions.
More concretely, assume a distribution with inﬁnite support q ( u | x ) where u ∈ R D
×
K , for example a Gaussian distribution or a normalizing ﬂow.
We can map u injectively to an argmax partition using a threshold function and x .
The injectivity of the map is important, because in that case the likelihood q ( v | x ) is easily computed via the change of variables formula q ( u | x ) (cid:12)(cid:12) d u d v (cid:12)(cid:12) .
In our implementation thresholding is implemented using a softplus function such that all values are mapped below a limit T .
[Equation]
where softplus( x ) = log(1 + e x ) , (5) for which it is guarenteed that v ∈ ( −∞ , T ).
In particular, the variable u is injectively mapped to v such that v d,k = u d,k if x d equals k and otherwise v d,k = threshold( u d,k , u d,x d ) if k (cid:54) =
x d where x d is used as an index.
In other words, all values except for the argmax indices are thresholded to be below the argmax values.
3.2. Alternative methods for dequantization
For a detailed description of alternative dequantization methods based on Gumbel reparametrization see Appendix A. Further, we outline a method to trade-oﬀ between symmetry and the number of dimensions in Argmax Flows, which we term Cartesian products of Argmax
Flows (further details in Appendix A).
4. Experiments
In this section we compare the performance of our method to alternative dequantization methods, a standard VAE baseline with ﬂexible posterior and prior, and to a VAE-based latent normalizing ﬂow approach.
In the ﬁrst experiment, we ﬁt a toy 50 class problem using maximum log-likelihood using diﬀerent dequantization methods.
As can be seen in
Table 6, our proposed Argmax based methods perform better than existing approaches
[Table]
4.1. Image data
To show that our method also generalized to images
[Table]
we introduce an unconditional image segmentation learning experiment.
In contrast with the standard setting where the [Figure] [Figure]
[Figure]
prediction is conditioned on photographs, this problem is designed to be more complicated by learning a distribution over the target without this photograph.
We introduce a 5-class toy segmentation problem named ”Squares” (see Figure 2).
Further, we rescale segmentation maps from cityscapes to 32 × 64 images and use the global category for an 8-class segmentation problem.
The results of these experiments are depicted in Table 3.
Results depicted are the ELBO and IWBO (with 1000 samples) measured in bits per dimension.
The thresholded-based argmax dequantization peforms best in terms of ELBO values
[Table]
and Gumbel-thresholded dequantization performs best in terms of IWBO evaluation for cityscapes.
[Table]
4.2. Text data
In this section we learn a normalizing ﬂow on the text8 dataset.
Instead of using a K =argmax space, we ﬁnd empirically that using a 5 × 2-binary space with threshold dequantization (see Appendix A) work better.
Further, the same density model as proposed in
(Lippe and Gavves, 2020) is utilized.
This experiment shows that Argmax Flow achieve lower bits per character (bpc) than CategoricalNF on text8 (see Table 4).
[Table]
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), Adversarial Network (GANs) (Goodfellow et al., 2014),
Normalizing Flows Rezende and Mohamed (2015) and Diﬀusion Models (Sohl-Dickstein et al., 2015; Ho et al., 2020).
Normalizing
Flows and Diﬀusion models typically learn a continuous distribution and dequantization is required to train these methods.
A large body of work is dedicated to building more expressive continuous normalizing ﬂows (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 distribution, adding uniform noise in-between ordinal classes was proposed in (Uria et al., 2013) and later theoretically justiﬁed in (Theis et al., 2016).
An extension for more powerful dequantization based on variational inference was proposed in Ho et al.
Dequantization for binary variables was proposed in (Winkler et al., 2019).
In other works, VAEs have been adapted to learn a normalizing ﬂow for the latent space (Ziegler and Rush, 2019; Lippe and Gavves, 2020).
However, these approach typically still utilize an argmax heuristic to sample, even though this is not the distribution speciﬁed during training.
Further, Tran et al.
(2019) propose invertible transformations for categorical variables directly, but results on images have thus far not been demonstrated.
In addition ﬂows for ordinal discrete data (integers) have been explored in (Hoogeboom et al., 2019; van den Berg et al., 2020)
6. Conclusion
In this paper we introduce a principled method to train continuous distributions such as ﬂows on categorical data.
The generative model is deﬁned by an argmax function which can be evaluated using a variational distribution over right-inverses.
Diﬀerent from other approaches is that our method does not require a stochastic decoder in the generative process, and as a result our model does not suﬀer from undesired uncorrelated noise in the model distribution.
We demonstrate that our method performs competitively on similar approaches on unconditional image segmentation and text.
Appendix A. Further methods for dequantization A.1. Dequantization with Gumbel distributions
An alternative method to construct q ( v | x ) is to ﬁrst deﬁne a marginal q ( v ) from which we can sample conditioned on the argmax q ( v | arg max v = x ).
A natural distribution to do such manipulations is the Gumbel distribution because of its favourable properties: the arg max and max are independent and the max is also distributed as a Gumbel distribution.
As a result it is straightforward to sample from a Gumbel distribution conditioned on its argmax.
For simplicity in notation, we will write that v ∈ R K and x ∈ { 1 , . . .
, K } without dimension d , as the distribution is deﬁned independently over the dimension axis d .
We ﬁrst deﬁne a distribution for v as the Gumbel distribution with location parameters φ ∈ R K :
[Equation]
Gumbel have the nice property that the argmax and max are independent distributions, and are accessible in closed-form.
In particular, the variable max i v i is distributed as a
Gumbel distribution itself and importantly does not depend on i :
[Equation]
where φ max = log (cid:80) i exp φ i .
These properties make it very easy to sample conditionally from the Gumbel distribution when the argmax index is given by x .
Conditioned on x , the variable v x is distributed as:
[Equation]
Furthermore, given this sampled maximum, the remaining indices can be directly sampled using truncated Gumbel distributions:
[Equation]
where the truncation value T is given by v x .
By computing ﬁrst Equation 8 and subsequently
Equation 9 we have drawn our sample v ∼ q ( v | x ) for which arg max v = x .
Recall that to optimize Equation 2, the log-likelihood log q ( v | x ) is also required.
This can be computed in closed-form expressions using the log density functions.
Another useful property of the Gumbel distribution is that its argmax is distributed as the categorical distribution
P (arg max v = i ) =
exp φ i / (cid:80)
i exp φ
i .
As such, the location parameters φ can be initialized to match the empirical distribution of the ﬁrst minibatch of the data.
For completeness here follows a quick summary of Gumbel properties:
To sample g ∼
Gumbel( φ ), sample u ∼ U (0 , 1) and compute g = − log( − log( u ))
+ φ .
Further the loglikelihood log Gumbel( g | φ, 1) =
φ − g − exp( φ − g ).
To sample g ∼ TruncGumbel( φ, 1; T ), sample u ∼ U (0 , 1) and compute g = φ − log(exp( φ − T ) − log( u )).
Further the log-likelihood log TruncGumbel( g | φ, 1 , T ) =
exp( φ − T ) − exp( φ − g ) + φ − g under the condition that g < T and otherwise −∞ .
A.2. Unifying Thresholding and Gumbel dequantization into Gumbel Thresholding
This section combines the results from the previous two sections.
The key insight is that the Gumbel sampling procedures as deﬁned above can be seen as a reparametrization of a uniform noise distribution on (0 , 1] d which is put through the inverse CDF of the Gumbel distributions.
Additionally, the log-likelihoods can be seen as log-derivatives of the (forward) CDF.
Consequently, instead of using uniform noise, we may learn any distribution on the interval (0 , 1] d and reparametrize using the same inverse CDF functions, where the log Jacobian determinant is equal to the log likelihood at that point.
The idea is that the smoothness of the Gumbel distribution is retained more while correlations across dimensions can still be learned by the interval distribution.
Compared to the plain Gumbel dequantization, the interval noise can now be conditioned on x and can be further correlated across dimensions d , which leads to more expressive dequantization distributions.
Suggestions to learn an interval distribution are (1) to learn a ﬂow with inﬁnite support composed with a sigmoid that injectively maps to (0 , 1) d or (2) learn a ﬂow starting from a uniform distribution that is transformed using interval preserving transformations such as splines.
A.3. Cartesian products of Argmax Flows
In the current description, Argmax ﬂows require the same number of dimensions in v as there are classes in x .
To alleviate this constraint we introduce Cartesian products of argmax ﬂows.
To illustrate our method, consider a 256 class problem.
One class can be represented using a single 256-onehot vector, but also using two hexadecimal numbers or alternatively using eight binary numbers.
Formally, any categorical variable x ( K ) ∈ { 1 , . . .
, K } d in base K can be converted to x ( M ) ∈ { 1 , . . .
, M } d m × d in base M ≥ 2, where at least d m = (cid:100)
M log K (cid:101)
M -categorical variables are required to model a single K -categorical variable.
Then the variable x ( M ) with dimensionality M · d m · d is dequantized instead of the variable x ( K ) with dimensionality K · d .
Even though this may lead to some unused additional classes, the ELBO objective in Equation 2 remains valid and thus the model can be optimized using an M -categorical argmax ﬂow.
To illustrate the changes in dimensionality, in Table 5.
Finally note that binary class problems are a special case where the variable can be straightforwardly encoded symmetrically into a single dimension.
As a consequence the last row in Table 5 is a special case that can alternatively be encoded using only d m =dimensions. [Table]
A.4. Ablation study: Gumbel distribution versus Gumbel thresholding
[Table]