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: The paper considers the problem of modeling discrete distributions with normalizing flows. Authors propose a novel framework “Categorical Normalizing Flows”, i.e CNF. By jointly modeling a mapping to continuous latent space, and the likelihood of flows CNF solves some of the bottlenecks in current algorithms. With experiments on some synthetic domains, and benchmarking tasks like Zinc250, the authors empirically demonstrate that CNF-based algorithms perform competitively and often improve significantly on related approaches like Latent NF, discrete flows.