Categorical Normalizing Flows via Continuous Transformations

1 INTRODUCTION . Normalizing Flows have been popular for tasks with continuous data like image modeling ( Dinh et al. , 2017 ; Kingma and Dhariwal , 2018 ; Ho et al. , 2019 ) and speech generation ( Kim et al. , 2019 ; Prenger et al. , 2019 ) by providing efficient parallel sampling and exact density evaluation . The concept that normalizing flows rely on is the rule of change of variables , a continuous transformation designed for continuous data . However , there exist many data types typically encoded as discrete , categorical variables , like language and graphs , where normalizing flows are not straightforward to apply . To address this , it has recently been proposed to discretize the transformations inside normalizing flows to act directly on discrete data . Unfortunately , these discrete transformations have shown to be limited in terms of the vocabulary size and layer depth due to gradient approximations ( Hoogeboom et al. , 2019 ; Tran et al. , 2019 ) . For the specific case of discrete but ordinal data , like images where integers represent quantized values , a popular strategy is to add a small amount of noise to each value ( Dinh et al. , 2017 ; Ho et al. , 2019 ) . It is unnatural , however , to apply such dequantization techniques for the general case of categorical data , where values represent categories with no intrinsic order . Treating these categories as integers for dequantization biases the data to a non-existing order , and makes the modeling task significantly harder . Besides , relations between categories are often multi-dimensional , for example , word meanings , which can not be represented with dequantization . In this paper , we investigate normalizing flows for the general case of categorical data . To account for discontinuity , we propose continuous encodings in which different categories correspond to unique , non-overlapping and thus close-to-deterministic volumes in a continuous latent space . Instead of pre-specifying the non-overlapping volumes per category , we resort to variational inference to jointly learn those and model the likelihood by a normalizing flow at the same time . This work is not the first to propose variational inference with normalizing flows , mostly considered for improving the flexibility of the approximate posterior ( Kingma et al. , 2016 ; Rezende and Mohamed , 2015 ; Van Den Berg et al. , 2018 ) . Different from previous works , we use variational inference to learn a continuous representation z of the discrete categorical data x to a normalizing flow . A similar idea has been investigated in ( Ziegler and Rush , 2019 ) , who use a variational autoencoder structure with the normalizing flow being the prior . As both their decoder and normalizing flow model ( complex ) dependencies between categorical variables , ( Ziegler and Rush , 2019 ) rely on intricate yet sensitive learning schedules for balancing the likelihood terms . Instead , we propose to separate the representation and relation modeling by factorizing the decoder both over the categorical variable x and the conditioning latent z . This forces the encoder and decoder to focus only on the mapping from categorical data to continuous encodings , and not model any interactions . By inserting this inductive bias , we move all complexity into the flow . We call this approach Categorical Normalizing Flows ( CNF ) . Categorical Normalizing Flows can be applied to any task involving categorical variables , but we primarily focus on modeling graphs . Current state-of-the-art approaches often rely on autoregressive models ( Li et al. , 2018 ; Shi et al. , 2020 ; You et al. , 2018 ) that view graphs as sequences , although there exists no intrinsic order of the node . In contrast , normalizing flows can perform generation in parallel making a definition of order unnecessary . By treating both nodes and edges as categorical variables , we employ our variational inference encoding and propose GraphCNF . GraphCNF is a novel permutation-invariant normalizing flow on graph generation which assigns equal likelihood to any ordering of nodes . Meanwhile , GraphCNF efficiently encodes the node attributes , edge attributes , and graph structure in three consecutive steps . As shown in the experiments , the improved encoding and flow architecture allows GraphCNF to outperform significantly both the autoregressive and parallel flow-based state-of-the-art . Further , we show that Categorical Normalizing Flows can be used in problems with regular categorical variables like modeling natural language or sets . Our contributions are summarized as follows . Firstly , we propose Categorical Normalizing Flows using variational inference with a factorized decoder to move all complexity into the prior and scale up to large number of categories . Secondly , starting from the Categorical Normalizing Flows , we propose GraphCNF , a permutation-invariant normalizing flow on graph generation . On molecule generation , GraphCNF sets a new state-of-the-art for flow-based methods outperforming one-shot and autoregressive baselines . Finally , we show that simple mixture models for encoding distributions are accurate , efficient , and generalize across a multitude of setups , including sets language and graphs . 2 CATEGORICAL NORMALIZING FLOWS . 2.1 NORMALIZING FLOWS ON CONTINUOUS DATA . A normalizing flow ( Rezende and Mohamed , 2015 ; Tabak and Vanden Eijnden , 2010 ) is a generative model that models a probability distribution p ( z ( 0 ) ) by applying a sequence of invertible , smooth mappings f1 , ... , fK : Rd → Rd . Using the rule of change of variables , the likelihood of the input z ( 0 ) is determined as follows : p ( z ( 0 ) ) = p ( z ( K ) ) · K∏ k=1 ∣∣∣∣det ∂fk ( z ( k−1 ) ) ∂z ( k−1 ) ∣∣∣∣ ( 1 ) where z ( k ) = fk ( z ( k−1 ) ) , and p ( z ( K ) ) represents a prior distribution . This calculation requires to compute the Jacobian for the mappings f1 , ... , fK , which is expensive for arbitrary functions . Thus , the mappings are often designed to allow efficient computation of its determinant . One of such is the coupling layer proposed by Dinh et al . ( 2017 ) which showed to work well with neural networks . For a detailed introduction to normalizing flows , we refer the reader to Kobyzev et al . ( 2019 ) . 2.2 NORMALIZING FLOWS ON CATEGORICAL DATA . We define x = { x1 , ... , xS } to be a multivariate , categorical random variable , where each element xi is itself a categorical variable of K categories with no intrinsic order . For instance , x could be a sentence with xi being the words . Our goal is to learn the joint probability mass function , Pmodel ( x ) , via a normalizing flow . Specifically , as normalizing flows constitute a class of continuous transformations , we aim to learn a continuous latent space in which each categorical choice of a variable xi maps to a stochastic continuous variable zi ∈ Rd whose distribution we learn . Compared to variational autoencoders ( Kingma and Welling , 2014 ) and latent normalizing flows ( Ziegler and Rush , 2019 ) , we want to ensure that all modeling complexity is solely in the prior , and keep a lossless reconstruction from latent space . To implement this , we simplify the decoder by factorizing the decoder over latent variables : p ( x|z ) = ∏ i p ( xi|zi ) . Factorizing the conditional likelihood means that we enforce independence between the categorical variables xi given their learned continuous encodings zi . Therefore , any interaction between the categorical variables x must be learned inside the normalizing flow . If in this setup , the encoding distributions of multiple categories would overlap , the prior would be limited in the dependencies over x1 , ... , xS it can model as it can not clearly distinguish between all categories . Therefore , the encoder q ( z|x ) is being optimized to provide suitable representations of the categorical variables to the flow while separating the different categories in latent space . Meanwhile , the decoder is incentivized to be deterministic , i.e . precisely reconstructing x from z , in order to minimize the overlap of categories . Overall , our objective becomes : Ex∼Pdata [ logPmodel ( x ) ] ≥ Ex∼PdataEz∼q ( ·|x ) [ log pmodel ( z ) ∏ i p ( xi|zi ) q ( z|x ) ] ( 2 ) We refer to this framework as Categorical Normalizing Flows . In contrast to dequantization , the continuous encoding z is not bounded by the domain of the encoding distribution . Instead , the partitioning is jointly learned with the model likelihood . Furthermore , we can freely choose the dimensionality of the continuous variables , zi , to fit the number of categories and their relations . Modeling the encoder The encoder q ( z|x ) and decoder p ( xi|zi ) can be implemented in several ways . The first and main setup we consider is to encode each category by a logistic distribution with a learned mean and scaling . Therefore , our encoding distribution q ( zi ) is a mixture of K logistics , one per category . With g denoting the logistic , the encoder becomes q ( z|x ) = ∏S i=1 g ( zi|µ ( xi ) , σ ( xi ) ) . In this setup , the decoder likelihood can actually be found correspondingly to the encoder by applying Bayes : p ( xi|zi ) = p̃ ( xi ) q ( zi|xi ) ∑ x̂ p̃ ( x̂ ) q ( zi|x̂ ) with p̃ ( xi ) being a prior over categories . Hence , we do not need to learn a separate decoder but can calculate the likelihood based on the encoder ’ s parameters . The objective in Equation 2 simplifies to the following : Ex∼Pdata [ logPmodel ( x ) ] ≥ Ex∼PdataEz∼q ( ·|x ) [ log ( pmodel ( z ) S∏ i=1 p̃ ( xi ) ∑ x̂ p̃ ( x̂ ) q ( zi|x̂ ) ) ] ( 3 ) Note that the term q ( zi|xi ) in the numerator of p ( xi|zi ) cancels out with the denominator in Equation 2 . Given that the encoder and decoder are sharing the parameters , we remove any possible mismatch between p ( xi|zi ) and q ( xi|zi ) . This allows changes in the encoding distribution to directly being propagated to the decoder , and further moves the focus of the training to the prior . Besides , the mixture encoding introduces a dependency of the true posterior p ( z|x ) on the approximate posterior q ( z|x ) , which potentially tightens the variational gap compared to a separately learned decoder . During testing , we can use importance sampling ( Burda et al. , 2016 ) to further reduce the gap . Details on the posterior dependency in the variational gap , and training and test steps can be found in Appendix A.1 . The mixture model is simple and efficient , but might be limited in the distributions it can express . To test whether greater encoding flexibility is needed , we experiment with adding flows conditioned on the categories which transform each logistic into a more complex distribution . We refer to this approach as linear flows . Taking a step further , we can also represent the encoder q ( z|x ) with a flow across categorical variables , similar to variational dequantization ( Ho et al. , 2019 ) . Experiments presented in Section 5 show however that a simple mixture of logistics usually suffices .

Summary: This work uses the idea of variational inference to map categorical data to continuous space affording the use of normalizing flows. Authors use several ideas to increase their framework's applicability--factorized distribution assumption, use of multi-scale architecture for step-generation, and permutation invariant components—achieving favorable results on several problems. The approach seems to be especially useful when data is non-sequential.