{'fraction_non_alphanumeric': 0.03696612665684831, 'fraction_numerical': 0.0095139911634757, 'mean_word_length': 4.69934530804096, 'pattern_counts': {'":': 0, '<': 2, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
Mixture of Discrete Normalizing Flows for Variational Inference

Advances in gradient-based inference have made distributional approximations for posterior distribution of latent-variable models easy, but only for continuous latent spaces. Models with discrete latent variables still require analytic marginalization, continuous relaxations, or specialized algorithms that are difficult to generalize already for minor variations of the model. Discrete normalizing flows could, in principle, be used as approximations while allowing efficient gradient-based learning, but as explained in this work they are not sufficiently expressive for representing realistic posterior distributions even for simple cases. We overcome this limitation by considering mixtures of discrete normalizing flows instead, and present a novel algorithm for modeling the posterior distribution of models with discrete latent variables, based on boosting variational inference.


Introduction
Variational inference is increasingly based on gradient-based algorithms [15,16], using reparameterized representation of the posterior approximation q(x) that enables passing gradients through Monte Carlo estimates of expectations, such as E q(x) [log p(y, x)] ≈ 1 S s log p(y, x s ) for some set of S samples x ∼ q(x). This has made it easier to step outside the rigid model and approximation families designed based on computational tractability, and to switch to flexible distributions parameterized by neural networks, including normalizing flows relying on invertible transformation of simple base distributions [21]. Models with discrete latent variables, however, remain problematic due to non-differentiability of the sampling operation that prevents efficient optimization of expectations over the approximation. Hence in practice, we still largely resort to model-specific algorithms that are tedious to extend already for minor variants of the model, analytic marginalization of the discrete latent variables (e.g., mixture models and LDA in Stan [7]), continuous relaxations like the concrete distribution [19,13], or (semi-)implicit approximations that do not support probability evaluation [34,30].
Normalizing flows [25,15,21], suitable for learning flexible posterior approximations for continuous variables, have recently been generalized also for discrete categorical [31] and ordinal [12] variables. However, the discrete variants have only been applied in generative modeling of discrete observations. Even though discrete normalizing flows for categorical distributions (DNF) retain the property of differentiable Monte Carlo estimates, they are not very suitable for variational approximation due to their limited expressive power. As we will show later, a DNF can only move probability mass around and hence relies extremely strongly on use of base distributions that are already expressive (and of the same dimension as the final distribution). For generative modeling this can be satisfied, for example by using recurrent neural networks as base distributions in the case of language modeling [31], but for modeling latent variables there are no easy ways of learning strong base distributions.
We improve the expressive power of DNFs by constructing mixtures of categorical discrete normalizing flows (MDNF). Instead of a fixed mapping from the base distribution to the observations (forward sampling), we consider a mixture of multiple flows where each sample is generated by one of the underlying component flows. Since this choice is in itself a latent variable, the probability evaluation (backward transformation) of the normalizing flow then passes the sample through all of the flows, weighting them accordingly. We show that MDNFs do not suffer from the limited modeling capacity of DNFs, and in particular can model arbitrarily complex distributions even when using naive base distributions, such as delta distributions that assign all probability mass for a single observation. This makes MDNFs suitable also for variational approximations. We demonstrate them in context of relatively simple probabilistic models (discrete Bayes networks, mixture models, and variational autoencoders) to provide quantifiable comparisons, but note that the approach is directly applicable also for models with more complex dependencies between the variables.
Finally, we also analyse the question of how MDNF can be learned. The mixture is not inherently identifiable and stochastic optimization becomes increasingly more difficult for mixtures with large number of components. To overcome these issues we propose a practical learning algorithm based on the idea of variational boosting [20,18]. We start by first training a single DNF and then iteratively refine the approximation by training more flows while keeping the old ones intact, except for the mixture weights. This typically results in better posterior approximation compared to naively training the mixture directly, and always improves from the baseline of stacked DNFs.

Problem formulation
Normalizing flows can be used for modeling arbitrary discrete probability distributions, either for modeling density of observed data p y (Y = y), or for modeling the probability of latent variables p x (X = x|Y = y) conditional on the observed data.
Generative model for discrete data Direct representation of an arbitrary joint distribution of D discrete variables y d of K d categories has exponential complexity (K D if all K d are equal). DNF models the distribution using y := f λ (u), where f λ (·) is an invertible transformation (parametrized by λ) of samples from a suitable base distribution p u (u). By using DNF to model dependencies between the variables, we can use, e.g., factorized or autoregressive base distributions, dramatically reducing the number of parameters. However, as explained in Section 3, the base distribution still needs to be rich and learned for the task. We contribute to this use-case by introducing a new flow type in Section 4.
Discrete posterior approximations We concentrate on using DNFs for modeling posterior distributions of latent categorical variables X for some model p(y|x), where the observations y may be continuous or discrete. We are not aware of any earlier works that would have used discrete normalizing flows for this purpose. In particular, as explained in more detail in Section 5, we extend reparameterized Monte Carlo estimates for variational lower bound of the marginal likelihood for discrete data, but note that other divergences and gradient estimators could also be used.

Discrete normalizing flows
Normalizing flows (NFs) [28,27,6,22,32,21] transform (typically simple) base distributions p u (u) into almost arbitrary complex distributions in an invertible manner. Samples x from the target distribution are generated by transforming samples u from the base using x := f λ (u), and the corresponding density is denotes the inverse transformation. The core design challenge typically is in crafting a flow with tractable Jacobian du dx . Discrete flows [31] build on the same principle, but with some technical differences. The Jacobian is no longer needed and the density is simply p x (x) = p u (invf λ (x)), but the choice of f λ (u) is considerably more limited -we need a discrete transformation that is invertible and can be trained w.r.t. the parameters λ. Tran et al. [31] proposed one that treats each dimension d separately as where K denotes a number of possible categories (K d different for each d can be used, but we drop d for brevity) and σ λd and K need to be coprime. The transformation parameters µ λd and σ λd are modeled with neural networks with K outputs, scaled with a temperature hyperparameter τ and passed through softmax and the straight-through (ST) operation [2] to obtain discrete one-hot encoded values differentiable w.r.t. parameters λ. That is, we have µ λd = ST(softmax(net λ (. . . )/τ )) for a suitably chosen neural network, where the temperature τ controls 'magnitude' of the gradients so that larger τ typically makes optimization easier while introducing more gradient bias. Those networks also model dependencies between the dimensions d, typically either using autoregressive or bipartite design. In autoregressive networks, the transformation of the dth dimension depends on outputs of the preceding dimensions 1 . . . d − 1, i.e., net λ (x 1 , . . . , x d−1 ), implemented using a recurrent network such as LSTM [11] or a masked autoencoder [8]. We use the latter in this work. Bipartite networks, e.g. [5], divide the dimensions d into two disjoint sets so that the second set depends only on the outputs for the first set. This simplifies computation (no iterating over d is required) but does not allow for modeling arbitrary dependencies, although stacking multiple flows with different subsets helps. Finally, we note that assuming independence between dimensions, it is sufficient to treat the parameters λ directly as outputs (logits) of the network, e.g., net λ (∅) = λ; this is not useful for generative modeling, but will be helpful in our use-cases.
The main limitations of discrete flows are illustrated on Figure 1. Since f λ is invertible and the samples x and u have the same shape, we have one-to-one mapping between x and u. This implies that the possible values for target probabilities p x (x) are directly determined by the values p u (u) of the base distribution, such that for each unique x such that p x (x) = a for any a, we must have p u (u) = a for some unique u. Consequently, the modeling capability of a discrete flow is determined directly by the flexibility of a selected base distribution, and the flow can only shuffle values in p u (u) around but cannot modify the probabilities themselves. For specific constructions there can also be other limitations -for example as mentioned above (1) cannot achieve all permutations -but these can usually be alleviated by stacking multiple flows. Stacking, however, does not help with the more fundamental limitation of inability to change the probability values.
The need to have sufficiently rich base distribution p u (u) limits the applications of DNFs. For generative modeling Tran et al. [31] overcame the problem by first training strong autoregressive or factorized models and only applying DNFs for re-organizing the values (e.g., to break up the factorization assumption), reaching consistent but somewhat small improvement over the baseline. Unfortunately, training such flexible base models is not always possible, which has thus far excluded, e.g., use of DNFs for variational approximation, even though for continuous variables normalizing flows are highly suitable for that [15].

Mixtures of discrete normalizing flows
To overcome the limitations of ordinary DNFs, we construct a mixture of discrete normalizing flows (MDNF) that uses multiple transformations. Probability estimation with MDNF is performed by where the mixing weights π b ≥ 0 s.t. b π b = 1 delegate partial responsibility of the total probability mass to individual component flows f b λ . For B = 1 the mixture reduces to standard DNF, whereas for B > 1 both sampling and probability evaluation are algorithmically somewhat more involved but remain tractable; the details are described in the next subsection.
The mixture resembles the construct of Papamakarios et al. [21] for continuous normalizing flows, but both the motivation and technical details differ notably. Even though they mix a collection of flows with expression similar to (2), they use flows with restricted and non-overlapping support in the u-space, and hence the mixture actually corresponds to piece-wise application of B separate flows. We use a proper mixture instead, where every component operates on the same domain.

Operations
Forward sampling for MDNF consists of three stages: (1) choosing a flow b ∼ Categ(π 1 , . . . , π B ), (2) sampling u ∼ p b u (u) from the base distribution of the bth flow (which may be the same for all flows), and (3) passing the sample through the flow to obtain the final sample: We assume that only the last step is differentiable, and hence the particular sample x is a function of λ associated only with the particular f b λ . Direct implementation of this scheme, however, requires a separate dynamic computation graph associated with every sample x. For a more efficient implementation using only a single graph, we can mask the outputs, conceptually similar to how masking is used in transformers or masked autoencoders. First, a whole batch consisting of n samples {u} is passed through all flows f b λ . Then, flow outputs are multiplied with B masks {M b } responsible for choosing which flow each individual sample is assigned to (1 appears exactly for one of the B positions). Finally, the masked outputs are summed up using

Choice of flows
The mixture formulation is in principle compatible with any choice of flows f b λ . In practice, the formulation of Tran et al. [31] in (1) provides a good starting point, despite not being able to model all possible transformations without stacking multiple flows (and even then, we still lack guarantees).
We also propose an alternative of partial flows that apply the same transformation as in (1) (usually with σ ≡ 1) but only on a subset of positions of a one-hot encoded vector u d . Instead of acting on vector of length K the transformation acts on vector of length K < K. The remaining category positions are passed untouched. Note that the 'subset' vectors do need to be valid one-hot encodings anymore and may consist of only 0 in all positions. Such partial flows make, e.g., swapping elements considerably easier; in the extreme case of K = 2 two elements can be swapped by shifting them by one position, without influencing any of the other probabilities (that would happen when obtaining the same swap with (1)). By simply stacking O(K log 2 (K)) partial flows with appropriately chosen subsets of size K = 2, we can construct a flow that could perform any reordering (=relabeling) of categories by swapping individual pairs, analogous to bitonic sorter [1]. Partial flows can be implemented with a slight modification of code for (1), e.g., by passing only the correct subset for an ordinary discrete flow transformation. Note however that transformation parameters µ need to receive full vectors x 1 , . . . , x d−1 when learning dependencies between dimensions.  x. Likelihood of the samples x is evaluated by passing them in a reverse direction (invf λ ) through a stack consisting of a number of either partial of location-scale flows to get samples u for which base probabilities p u are known. For partial flows we always used K = 2 and organized them in a way resembling a bubble-sort sorting network ( Figure 2) so we used 10 flows for K = 5 and 21 flows for K = 7. For location-scale flows we experimented with stacks of 3-25 identical layers. In each of the repeated runs, we created a new base distribution p u by randomly shuffling p x and optimized the stacked flows to maximize the likelihood by using Adam optimizer with learning rate 0.1 and τ = 1.0. The optimal solution is reached only when categories of p u are correctly reordered back, so that samples f λ (u), u ∼ p u follow p x . This happens only for one of K! permutations. Table 1 presents results estimated with 40 repeated runs, where as a success we count only the perfect recovery of the original probability distribution within at most 5000 iterations. Partial flows are superior with both K = 5 and K = 7 categories, but require more iterations for convergence. Additionally, for stacked location-scale flows, we observe that increasing the number of layers makes learning harder, i.e., the success is achieved in fewer cases (but also in fewer iterations).

Base distributions
MDNF is compatible with any joint base distribution over the D variables. Due to the exponential storage cost we would typically use factorized distributions with O(KD) storage, evaluated in log-domain using log p b u (u) = log p b 1 (u 1 ) + · · · + log p b D (u D ), and leaving the task of modeling dependencies for the flow. For ordinary DNFs it is crucial to use expressive base distributions, as explained in Section 3, but for MDNF this is not critical. In fact, it turns out that already delta distributions that allocate all probability for a single category c (that is, p u (u) = 1 ⇐⇒ u = c) are sufficient, offering both computational advantage and easy theoretical analysis. To illustrate their potential, let us first note that any delta distribution can be transformed into another one by renaming c to some k in samples drawn from the distribution, and this can be achieved already with a single shift transformation -by a simplified (1) with σ = 1. In particular, for any choice of k and c, we can always find some µ such that p x (X = k) = 1 for any delta base distribution p u (U = c) using k = (c + µ) mod K. In practice, we use trainable µ = ST(sof tmax(net λ (. . . )/τ )), and multiple flows in a mixture (MDNF) to be able to allocate separately (but not independently) fractions of the total probability mass.
Using sufficiently many component flows with delta distributions allows for modeling arbitrary distributions. For example, Figure 3 illustrates how the example considered in Figure 1 can be represented using only delta distributions as the base. In more detail, already with uniform weights π b = 1/B there exist flows f b λ such that the absolute error for any target distribution p t (x) and any outcome x can be bounded by and hence the approximation converges when B → ∞. Note that the bound is trivial for B = 1 corresponding to a single DNF, providing no information on their accuracy. The bound follows directly from approximating each p t (x) with a subset of the B flows, chosen independently so that the proportion of the flows best matches it. As explained above, a single flow f b λ using delta base distribution p b u (invf b λ (x)) can allocate all probability mass of that flow for any given category. To verify the bound, it hence remains to show that by combining B component flows with uniform weights we can represent any categorical distribution sufficiently accurately. We do this in a constructive manner, and note that the reasoning does not directly say anything about how MDNFs trained with practical learning algorithms would behave.
Let us first assume K ≤ B. We approximate p t (x) (for each x) using P (x) ∈ N 0 flows, each carrying a probability mass of 1/B. Let us first allocate P (x) = p t (x)B component flows for modeling each category, so that p t (x) is approximated byp t (x) = P (x) . For this it already holds that |p t (x) − p t (x)| ≤ 1/B. However, we have B −K ≤ x P (x) ≤ B, and hencep t (x) is not a distribution and there are some unallocated component flows, unless the right inequality is equality (in which casep t (x) = p t (x) for all x). However, since the number of unallocated flows is at most K, we can allocate the remaining ones on arbitrary categories, one for each, so that P (x) = pt(x)B B for these categories. This makesp t (x) a valid distribution, while retaining the maximum approximation error of 1/B.
If K > B we simply set P (x) = 0 for the K − B smallest probabilities p t (x), all of which are guaranteed to be at most 1/K < 1/B, and apply the above procedure to approximate p t (x) for the remaining categories.

Variational inference with MDNFs
We use MDNFs to form variational approximation q x (x) of the posterior p(x|D) of a probabilistic model with parameters x and observed data D = {y}. Monte Carlo estimate for the ELBO objective [29,16] can be then optimized using reparameterized gradients estimated with S samples drawn from the approximation as in For discrete variables, f λ denotes a MDNF but the same objective supports using other transformations for continuous variables as long as the corresponding Jacobian is included. For a given B, learning MDNF requires specifying both the weights π b as well as the flows f b λ , controlled by the parameters λ b of a suitable neural network outputting the parameters µ λd and σ λd (see Section 3 for examples). This can naturally be carried out using multiple algorithms; in the following we describe few alternatives.
Variational inference on flows (VIF) As explained in Section 4.3, the mixture is highly expressive already for fixed π b . By fixing them to 1/B, we can jointly train λ b for all flows with ordinary gradient descent. One optimization step consists of sampling x as explained in Section 4.1 and estimating gradients of the objective following (3), where the Monte Carlo estimate of entropy relies on (2). This algorithm is easy to implement, but simultaneous training of multiple flows is poorly identified.
Boosting variational inference on flows (BVIF) To avoid simultaneous training of competing flows, we turn the attention to existing literature on variational boosting [9,20,17,18,4], originally developed for iteratively increasing flexibility of approximations (in context of continuous variables). Standard variational boosting algorithms construct posterior approximation of the form q where we intentionally match notation with (2), so that in our case q b λ (x) = p b u (invf λ (x)). We adapt the algorithm of Miller et al. [20], which proceeds by iteratively training the π b+1 and λ b+1 by keeping flows and weights up to b fixed, using the objective where q (1...b) denotes a mixture composed from components 1 . . . b. The optimization consists of B steps, each using a number of gradient ascent iterations.
Boosting baseline (BVI) Boosting as explained above is, in principle, applicable to directly learning a flexible approximation without training the individual flows at all. For completeness, we explain also a baseline algorithms doing this, to illustrate the importance of training the individual flows. For fixed component distributions q b , each of which is a delta distribution (other choices could be used as well), we merely train the weights π b iteratively as above. For perfect recovery we need B ≥ K D base distributions, which naturally makes this baseline reasonable only for small problems.

Experiments
We demonstrate the flexibility of MDNFs as general-purpose tool for approximation discrete posteriors by applying them in context of three distinct model families. The first (Bayes network, BN) involves only dicrete variables, the second (mixture model, GMM) corresponds to common case of discrete latent variables for continuous data, and the final example (variational autoencoder (VAE) with discrete latent space) considers a more flexible model. For the main comparison of different algorithms against ordinary DNF, we use base delta distributions with one flow (no stacking; only shift transformation) per component. The shifts we found using standard gradient-based optimizers (RMSprop and Adam). In all cases, following the common practice (see for example, [13]), we also performed annealing of the temperature hyperparameter τ controlling bias of gradients of the straight-through estimator, by slowly decreasing its value in each iteration t with τ t = τ exp(−γt), where γ > 0 controls the rate of annealing. The code used in the experiments is publicly available 1 .

Discrete Bayes networks
A Bayes network represents a joint distribution of random variables factorized according to a directed acyclic graph (DAG) that determines their conditional independence. Often the task is learning the structure of the network [26], but for BNs with latent nodes determining their joint posterior is difficult even if the structure is known. For sufficiently small networks the true posterior can be evaluated by direct enumeration of all configurations, with exponential cost in the number of latent nodes. We use this as ground truth, and hence experiment on small-to-medium sized publicly available 2 networks for which it can still be evaluated, observing values for 1-2 variables and leaving all others latent. We use Asia (8 binary nodes; we fix asia:=yes and xray:=yes), Sachs (11 variables with 3 categories; we fix Akt:=LOW) and Hepar II (70 nodes with up to 6 categories; we fix carcinoma:=present). To approximate the posteriors, we use masked autoencoders [8] (MADE) representing flows' shifts as , where we set the initial temperature τ = 0.1 and annealing rate γ = 0.001. Monte-carlo estimate (with S = 100 samples) of ELBO we optimize w.r.t. the parameters λ b of a MDNF using RMSprop with learning rate 0.001.
The Asia and Sachs have the same cardinality for all variables, but for Hepar II the cardinality depends on the variable and ranges from 2 to 6. To handle the varying dimensionality in an environment designed for processing fixed-size tensors (in our case TensorFlow), to represent S D-dimensional samples, we use tensors of size S × D × K with K = max (K 1 , . . . , K D ), and map excess positions in one-hot encoded vectors (category positions with numbers larger than K d ) down to positions representing valid categories (by summing up zeros and ones from respective positions) only at the end -when evaluating joint probability of observed and latent variables for a model. Entropy term can not be handled this way, but entropy for variables with K categories bounds the entropy for the original set of variables, H K (x) ≥ H K1,...,K d ...K D (x), and with Hepar II we used the approximation. Figure 4 (top row) compares MDNF trained with three learning algorithms explained in Section 5, evaluated with ELBO and whenever possible with Kullback-Leibler divergence between the approximation and the true posterior, so that zero indicates perfect match. MDNF when trained with BVIF outperforms the baseline of standard DNF (B=1) clearly for all three instances, and in general increasing the number of flows improves the result roughly monotonically. For the two smaller networks VIF performs roughly as well, but breaks down for the larger Hepar II network. BVI is clearly inferior, showing that training the component flows is necessary.

Mixture models
As an example of a model with both discrete and continuous latent variables we used Gaussian mixture model (GMM) [3]. GMM is a model allocating D observed data points y d ∈ R N to one of K clusters with D latent K-dimensional categorical variables x d (one per data point; the assignments are assumed to be conditionally independent). Note that we here denote -somewhat unconventionally -by D the number of samples to emphasize that in our context the modeling task concerns learning the D-dimensional distribution of the latent allocations.
The model with multivariate normal component distributions is where we used α 0 = 1 K · I K×K , µ 0n = 1 D D d=1 y n , β 0 = 1, W 0 = I N ×N , and ν 0 = N . The model was trained using closed form Variational EM [3] by alternating updates of the allocations (E-step) and the cluster parameters (M-step). Our implementation and choice of the hyperparameters follow a publicly available implementation 3 , where we replaced the closed-form E-step with stochastic gradient-based optimization of ELBO w.r.t. parameters of a MDNF modeling the distribution of latent allocations x. The optimization we performed using RMSprop optimizer (learning rate 0.1) with S = 100 samples used for MC estimate of the objective. The posterior for x factorizes and therefore we also used factorized flows with µ d = ST (softmax(λ d /τ t )). That is, we passed the trainable parameters λ directly through softmax and the straight-through (ST) operation. In variational EM with gradient-based E-step, stochastic optimization is performed multiple times, each time for slightly different clusters' found in M-step. As a consequence the temperature hyperparmeter we slightly anneal in each step and then in each iteration of the step, e.g., in our schedule t = step + iteration with initial temperature τ = 10 and rate γ = 0.01. Figure 4 (bottom row) compares BVIF and VIF for three data sets: 2-dimensional simulated data with 3 partially overlapping clusters (100 points each) with centers in (0,2),(1.7,-1),(-1.7,-1) and diagonal unit covariances; and for Google 4 (5456 data points with 24 features) and TripAdivsor 5 (980 points with 10 features) travel reviews [24]. In all our experiments we used K = 3 -the same number of clusters that was used to generate the simulated data, and that was suggested for the Google set by the authors [24] who provided the data. They also indicated K = 3 is good, though not necessarily optimal, for TripAdivsor. We drop BVI that performed very poorly already in the previous experiment, and again demonstrate that MDNF works as intended and sometimes improves on top of DNF (B=1), but for these simpler posteriors the difference is marginal because in non-overlapping case all probability mass of a data point posterior is allocated to a single cluster. Nevertheless, we see that MDNF can be used as plug-and-play approximation and works even when modeling a relatively large number of variables.

Variational autoencoder
Finally, we apply MDNF for learning VAEs [14] with discrete latent variables. Various specialized methods have been proposed for this, e.g., based on continuous relaxations [13] or vector quantization [33,23], sometimes coupled with additional modifications for the learning objective, for example, [10]. In contrast to them, we demonstrate that by using MDNF, we can use the ordinary learning objective without additional modifications for learning discrete codes. The model is presented in Figure 5. The encoder's output λ is passed to each of the component flows and parameterizes the flows' transformation µ b (λ; τ ) (the transformations factorize over latent dimensions d). For the base distributions p b u we use delta distributions, making it sufficient to use a single shift transformation µ b with each of the B = 40 used flows.
The encoder and decoder architectures match the ones used in 6 [13]. The encoder has two dense layers with 512 and 256 nodes and ReLU activations, outputting 20 (10 · 2; we used D = 10 latent variables with K = 2 categories) logits λ. The decoder has the same layers in reverse order, taking inputs of size 20 and outputting 768 logits for Bernoulli distributions for the pixels. The flow transformations µ b are obtained by passing the encoder's output through a network with a hidden layer consisting of D · K · B nodes with ReLU activations and outputting tensors of the same shape. The model we trained following VIF scheme and by maximizing standard ELBO with uniform priors p x (x) = 1/K. We also used stochastic gradient decent (Adam with learning rate 0.001) with minibatches of size 128. Finally, similarly as Jang et al. [13], we anneal the temperature hyperparameter in consecutive iterations with initial temperature τ = 100 and rate γ = 0.0005. Figure 6 compares MNIST digits generated by the same VAE architecture when using MDNF, and two variants of the concrete distribution: (1) direct Gumbel relaxation and (2) Straight-Through estimator for Gumbel relaxation, proposed by Jang et al. [13]. Note that ELBO can not be directly compared between the models due to use of the relaxed approximation in baselines, and hence we resort to illustrating that the generated samples are of comparable quality, verifying that MDNF are readily applicable without modifications.

Base distribution and temperature
Besides the flow design, we need to determine the base distribution (Section 4.3) and the temperature parameter τ (Section 3). Figure 7 illustrates that for both choices we can make a safe default choice, of using delta distributions while avoiding very small τ .
For comparing different base distributions we sample them from a symmetric Dirichlet, i.e., Dir(α). For α → 0 all realizations are delta distributions whereas for α → ∞ they are uniform distributions, and by varying α we can systematically create base distributions between these extremes. Note that uniform distributions provide a worst-case baseline, since re-arranging identical values does not help. Figure 7 shows that for both BNs and GMMs very small α is optimal, supporting use of delta distributions. Similarly, large τ is clearly better for both cases, and e.g., τ = 10 is a good practical suggestion. Using a smaller temperature (e.g., τ = 0.1 for BN) may help slightly in reducing gradient bias (which is why we used annealing for the main experiments), but the effort in finding optimal value may not be worthwhile.

Conclucions
For continuous variables normalizing flows have become a routine element in the modeling toolbox, providing both strong generative models as well as flexible representations for latent variables. Categorical discrete flows [31] build on the same conceptual ideas, but as explained in this work are considerably less expressive. This has, until now, restricted their use to generative modeling tasks where strong base distributions can be trained. We are the first ones to demonstrate in practice that DNFs can be used for modeling latent variable posteriors. Besides the core contribution achieving this with MDNF, we also introduced partial flows that lift some of the limitations of earlier discrete flows and are expected to be beneficial in generative modeling applications, even though we did not consider those explicitly in this work.
Even though the mixture formulation is highly expressive, it is typically overparameterized and involves high degree of stochasticity due to discrete sampling. As illustrated in our empirical experiments, this makes learning somewhat challenging, but we also showed how borrowing elements from the variational boosting literature helps in consistently outperforming the baselines, resulting in an algorithm that can be used in plug-and-play manner for learning posteriors of different models.