Categorical Normalizing Flows via Continuous Transformations
Authors: Phillip Lippe
Efstratios Gavves
Venue: ICLR 2021 Poster (2021)
Keywords: Normalizing Flows
Density Estimation
Graph Generation

Abstract: Despite their popularity, to date, the application of normalizing flows on categorical data stays limited. The current practice of using dequantization to map discrete data to a continuous space is inapplicable as categorical data has no intrinsic order. Instead, categorical data have complex and latent relations that must be inferred, like the synonymy between words. In this paper, we investigate Categorical Normalizing Flows, that is normalizing flows for categorical data. By casting the encoding of categorical data in continuous space as a variational inference problem, we jointly optimize the continuous representation and the model likelihood. Using a factorized decoder, we introduce an inductive bias to model any interactions in the normalizing flow. As a consequence, we do not only simplify the optimization compared to having a joint decoder, but also make it possible to scale up to a large number of categories that is currently impossible with discrete normalizing flows. Based on Categorical Normalizing Flows, we propose GraphCNF a permutation-invariant generative model on graphs. GraphCNF implements a three step approach modeling the nodes, edges, and adjacency matrix stepwise to increase efficiency. On molecule generation, GraphCNF outperforms both one-shot and autoregressive flow-based state-of-the-art.


Published as a conference paper at ICLR 2021
CATEGORICAL NORMALIZING FLOWS VIA CONTINU- OUS TRANSFORMATIONS
Phillip Lippe
Efstratios Gavves University of Amsterdam egavves@uva.. m1
University of Amsterdam, QUVA lab p. 1 1. w.lippeeuva... nl
ABSTRACT
Despite their popularity, date, the application of normalizing flows on categorical data stays limited. The current practice of using dequantization to map discrete data to a continuous space inapplicable as categorical data has no intrinsic order. Instead, categorical data have complex and latent relations that must be inferred, like the synonymy between words. In this paper, we investigate Categorical Normalizing Flows, that is normalizing flows for categorical data. By casting the encoding of categorical data in continuous space as a variational inference problem, we jointly optimize the continuous representation and the model likelihood. Using a factorized decoder, we introduce an inductive bias to model any interactions in the normalizing flow. As consequence, compared to having a joint decoder, but also make it possible to scale up to a large number of categories that is currently impossible with discrete normalizing flows. Based on Categorical Normalizing Flows, we propose GraphCNF a permutation- invariant generative model on graphs. GraphCNF implements a three-step approach modeling the nodes, edges, and adjacency matrix stepwise to increase efficiency. On molecule generation, GraphCNF outperforms both one-shot and autoregressive flow-based state-of-the-art.
do not only simplify the optimization
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 gradient approximations (Hoogeboom) jt al. 2019 Tran et al. (201)). For the specific case of discrete but ordinal data, like images where integers represent quantized values, popular strategy to add 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 non-existing order, and makes the modeling task significantly harder. Besides, relations between categories are often multi-dimensional, for example, word meanings, which cannot 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 propose variational inference with normalizing flows, mostly considered for improving the flcxibility of the approximate posterior (Kingma et al. 2016)] Rezende and Mohamed, 2015 Van Den Berg et al. 2018). Different from previous works, variational inference learn

---

Published as a conference paper at ICLR 2021
continuous representation of the discrete categorical data normalizing flow. A similar idea has been investigated in (Ziegler and Rush] (019), who use 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 2. 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 applied to any task involving categorical variables, but we primarily focus modeling graphs. Current state-of- of-the-art approaches often rely on autoregressive models (Li et al. 2018 Shi et al. 2020; You al. 2018) that view graphs as sequences, although there exists intrinsic order of the node. In contrast, normalizing flows can perform generation in parallel making definition of order unnecessary. By treating both nodes and edges as categorical variables, employ variational inference encoding and propose GraphCNF. GraphCNF is novel permutation-inet innvaiiannnn 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-arr. 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 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-invarian normalizing flow on graph generation. molecule generation, GraphCNF sets new state-of-thh-arr for flow-hased methods outperforming one-shot and autoregressive baselines. Finally, we show that simple mixture models for encoding distributions are accurate, efficient, and generalize across 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 generative
model that models probability distribution p(z ()) by applying sequence of invertible, smooth
mappings /I,..., fK Rd + Rdd Using the rule of change of variables, the likelihood of the input
(0) is determined as follows:
(1
(() = ((2(K))
II det
(z(k-1)
where (4) - f& (x(-1))))).. and (()( represents a prior distribution. This calculation requires to compute the Jacobian for the mappings 1,...// 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 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 {I1,,.., } to be a multivariate, categorical random variable, where each element I; is itself a categorical variable of K categories with no intrinsic order. For instance, æ could be a sentence with z; being the words. Our goal is to learn the joint probability mass function, (2) via a normalizing flow. Specifically, normalizing flows constitute class of continuous transformations, we aim to learn a continuous latent space in which each categorical choice of a variable x; maps to stochastic continuous variable 2; Rdd whose distribution learn
2

---

Published as a conference paper at ICLR 2021
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 lossless reconstruction from latent space. To implement this, we simplify the decoder by factorizing the decoder over latent variables: p(æ]z) = II p(z; Factorizing the conditional likelihood means that we enforce independence between the categorical variables Iy given their learned continuous encodings 2; Therefore, any interaction between the categorical variables æ must be learned inside the normalizing flow. in this setup, the encoding distributions of multiple categories would overlap, the prior would be limited in the dependencies over .1, .... T5 it can model cannot clearly distinguish between all categories. Therefore, the encoder g(2 optimized to provide suitable representations of the categorical variables to the flow while separating the different categories latent space. Meanwhile, the decoder is incentivized to be deterministic, i.e. precisely reconstructing from z, in order to minimize the overlap of categories. Overall, our objective becomes:
being
log Pmodel (z) II, p(x,(2,))
Pgn log (x)) 2 Eee Paza
(2)
((2(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, Z;, to fit the number of categories and their relations.
Modeling the encoder The encoder (==(az) and decoder p(a;; zzi can be implemented in several ways. The first and main setup we consider to encode each category by a logistic distribution with a encoding distribution q (z,) is mixture of K logistics, one
learned mean and scaling. Therefore, per category. With g denoting the logistic, the encoder becomes q x>) = I g(z (u(2;), o (x;)). In this setup, the decoder likelihood can actually be found correspondingly to the encoder by applying
((11))( with ((z)) being a prior categories. Hence, do not need Bayes: p(z,/zz) m ((:)g( 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:
plaa)
Pum-Puu [log P (x) 2 Pim Eg ((-1a) log Pmmodet (z) i=1 Ee p( ())
II
Note that the term ((z, (z;) in the numerator of (z)) cancels out with the denominator in Equa- tion 2 Given that the encoder and decoder are sharing the parameters, we remove any possible mismatch between p(æ,) xz) and g( (z). 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 dependency of the true posterior p(z] æ)) on the approximate posterior g=))ar), which potentially tightens the variational gap compared to a separately learned decoder. During testing,
use importance sampling (Burda et al. (016) 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.I
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 step further, we can also represent the encoder (z(22) with flow across categorical variables, similar to variational dequantization (Ho et al., (019). Experiments presented in Section 5oohow however that simple mixture of logistics usually suffices.
GRAPH GENERATION WITH CATEGORICAL NORMALIZING FLOWS
A graph G = EV) is defined by a set of nodes V. and set of edges E representing connections between nodes. When modeling a graph, we must take into account the node and edge attributes, often represented by categorical data, as well as the overall graph structure. Moreover, nodes and edges are better viewed as sets and not sequences since any permutation represents the same graph and assigned the same likelihood.
We propose GraphCNF, a normalizing flow for graph generation, which is invariant to the order of nodes by generating all nodes and edges at once. Given graph G, we model each node and edge as
3

---

Published as a conference paper at ICLR 2021
Node type regresentationn
(discrete comlinuuu)
Prior distribution
Filee-GNN
Edoe-GNN
Coupling layers
Edge representation
[discrete - contimuous)
Coupling layers
Addiug viriual adge represiatation
Couping layers
Figure 1: Visualization of GraphCNF for an example graph of five nodes. We add the node and edge attributes, as well the virtual edges stepwise to the latent space while leveraging the graph structure in the coupling layers. The last step considers fully connected graph with features per edge.
separate categorical variable where the categories correspond to their discrete attributes. To represent the graph structure, i.e. between which pairs of nodes does or does not exist an edge, we add an extra category the edges representing the missing virtual edges. Hence, to model arbitrary graph, we consider an edge variable for every possible tuple of nodes.
To apply normalizing flows the node and edge categorical variables, map them into continuous latent space using Categorical Normalizing Flows. Subsequent coupling layers map those represen- tations a continuous prior distribution. Thereby, GraphCNF uses two crucial design choices for graph modeling: (1) we perform the generation stepwise by first encoding the nodes, edges and then the adjacency matrix for improved efficiency, and (2) we introduce an inductive bias that the model assigns equal likelihood to any ordering of the nodes.
3.1 THREE-STER GENERATION
Modeling all edges including the virtual ones requires a significant amount latent variables and is computationally expensive. However, normalizing flows have been shown to benefit from splitting of latent variables at earlier layers while increasing efficiency (Dinh et al. 2017 Kingma and Dhariwal 2018). Thus, we propose to add the node types, edge attributes and graph structure stepwise the latent space as visualized in Figure [ In the first step, we encode the nodes into continuous latent space, (V), using Categorical Normalizing Flows. On those, we apply group of coupling layers, f1, which additionally use the adjacency matrix and the edge attributes, denoted by Eattr as input. Thus, summarize the first step
(4) The second step incorporates the edge attributes, Eattr. into latent space. Hence, all edges of the graph except the virtual edges are encoded into latent variables, z (Eater) representing their attribute. The following coupling layers, denoted by fz, transform both the node and edge attribute variables:
M
f((=6") E,
2g
(V) (Eaur)
f2(21
(V) :00
(Eattr),
(5)
Finally, we add the virtual edges the latent variable model as 2n (E*) Thereby, we need slightly adjust our encoding from Categorical Normalizing Flows as we consider the virtual edges as an additional category of the edges. While the other categories are already encoded by 2 ((uuto)) add a separate encoding distribution for the virtual edges, for which we use a simple logistic. Meanwhile, the decoder needs be applied on all edges, as need to distinguish the continuous representation between virtual and non-virtual edges. Overall, the mapping can be summarized as:
(V)
(E) = [3(=
=0 (E) where =0 =
(E') (Eaut) ,
(6)
23
where the latent variables 2) and 2/
(E) are trained to follow prior distribution. During sampling, we first inverse fa and determine the general graph structure. Next, we inverse f2 and reconstruct the edge attributes. Finally, apply the inverse of fi and determine the node types.
4

---

Published as a conference paper at ICLR 2021
3.2
GRAPH MODELING
To achieve permutation invariance for the likelihood estimate, the transformations of the coupling layers need to be independent of the node order. This includes both the split of variables that will be transformed and the network model that predicts the transformation parameters. We ensure the first aspect by applying a channel masking strategy (Dinh et al., 2017). where the split is performed over the latent dimensions for each node and edge separately making it independent of the node order. For the second aspect, we leverage the graph structure in the coupling networks and apply graph neural networks. In the first step of GraphCNF, f1, we use Relation GCN (Schlichtkrul et al. 2018) which incorporates the categorical edge attributes into the layer. For the second and third steps, we need graph network that supports the modeling of both node and edge features. We implement this by alternating between updates of the edge and the node features. Specifically, given node features v and edge features e' at layer t, update those follows:
++1 fnode (v': e²),
e'+l= fedge o't))
We call this network Edge-GNN, and compare different implementations of fnode and fedge in Appendix B Using both design choices, GraphCNF models a permutation invariant distribution.
4 RELATED WORK
Dequantization Applying continuous normalizing flows on discrete data leads to undesired density
models where arbitrarily high likelihoods are placed on particular values (Theis et al. 2016 Uria et al.) 2013). A common solution to this problem is to dequantize the data by adding noise u E [0, I) Theis et al. (0016) have shown that modeling ( + u) lower-bounds the discrete distribution P The noise distribution q(u x>) is usually uniform learned by second normalizing flow. The latter is referred to as variational dequantization and has been proven to be crucial for state-of-the-art image modeling (Ho et al. 2019 Hoogeboom et al.. 2020). Categories, however, not quantized values, so that ordering them as integers introduces bias to the representation.
Discrete NF Recent works have investigated normalizing flows with discretized transformations. Hoogeboom et al.] (2019) proposed to use additive coupling layers with rounding operators for ensuring discrete output. Tran et al. (2019) discretizes the output by a Gumbel-Softmax approximating an argmax operator. Thereby, the coupling layers resemble a reversible shift operator. While both approaches achieved competitive results, due discrete operations the gradient approximations have been shown to introduce new challenges, such as limiting the number of layers or distribution size.
Latent NF Several works have investigated the application of normalizing flows in variational auto-encoders (VAEs) (Kingma and Welling 2014) for increasing the flexibility of the approximate posterior (Kingma | [2016] Van Den Berg et al. 20188) Tomezak and Welling, 2017). However, VAES model lower bound of the true likelihood. To minimize this gap, Ziegler and Rush (2019) proposed Latent Normalizing Flows that move the main model complexity into the prior by using normalizing flows. In contrast to CNFs, Latent NF have joint decoder over the latent space, p(x|2z) 2²) which allows the modeling of interactions between variables in the decoder. Thus, instead of an inductive bias to push all complexity to the normalizing flow, Latent NF rely on loss-scheduling weighting the decoder loss much higher. This pushes the decoder to be deterministic but can lead to unstable optimization due to neglecting the flow's likelihood. Further, experiments on sequence tasks show Latent NF to be competitive but are still outperformed by an LSTM baseline as we observed in our experiments.
Graph modeling The first generation models on graphs have been autoregressive (Liao et al. 2019 You al. 2018)) generating nodes and edges in sequential order. While being efficient in memory,
they are slow in sampling and assume an order in the set of nodes. The first application of normalizing
flows for graph generation was introduced by [Luu et al. (2019), where a flow modeled the node representations of pretrained autoencoder. Recent works of GraphNVP (Madhawa et al. 2019) and GraphAF (Shi et al. 2020 proposed normalizing flows for molecule generation. GraphNVP consists of two separate flows, one for modeling the adjacency matrix and second for modeling the node types. Although allowing parallel generation, the model is sensitive to the node order due to its masking strategy and feature networks in the coupling layer. GraphAF is an autoregressive normalizing flow sampling nodes and edges sequentially but allowing parallel training. However, both flows use standard uniform dequantization to represent the node and edge categories. VAE have
5

---

Published as a conference paper at ICLR 2021
also been proposed for latent- graph generation (Simonovsky and Komodakis 2018 Ma et al. 2018; Liu et al. 2018;; Jin et al.) 2018). Although those models can be permutation iinittatiionnnnnnnnn,, they model a lower bound and do not provide a lossless reconstruction from latent space.
EXPERIMENTAL RESULTS
We start experiments by evaluating GraphCNF two benchmarks for graph generation, namely molecule generation and graph coloring. Further, to test generality we evaluate CNRs on other categorical problems, specifically language and set modeling. For the normalizing flows, we use sequence of logistic mixture coupling layers (Ho et al. 2019) mapping mixture of logistic distribu- tions back into single mode. Before each coupling layer, we include an activation normalization layer and invertible Ixl convolution (Kingma and Dhariwal [2188) 2201)) For reproducibility. we provide all hyperparameter details in Appendix D and make our code publicly available
5.1 MOLECULE GENERATION
Modeling and generating graphs is crucial in biology and chemistry for applications such as drug discovery, where molecule generation has emerged as a common benchmark (Jin et al., (2018) Shi et al. )020). In molecule graph, the nodes are atoms and the edges represent bonds between atoms, both represented by categorical features. Using a dataset of existing molecules, the goal is to learn a distribution of valid molecules not all possible combinations of atoms and bonds are valid. We perform experiments on the Zinc250k (Irwin et al., (012)) dataset which consists of 250,000 drug-like molecules. The molecules contain up 38 atoms of 9 different types, with three different bond types possible between the atoms. For comparability, we follow the preprocessing of Shi et al. (2020).
We compare GraphCNF to baselines that consider molecules as a graph and not as text representation. As per VAE-based approaches, we consider R VAE (Ma et al. 2018) and Junction-Tres VAE (JT- VAF) (Jin et al. 2018). R-VAF is a one-shot generation model using regularization to ensure semantic validity. JT-VAE represents a molecule as junction tree of sub-graphs that are obtained from the training dataset. We also compare our model to GraphNVP (Madhawa et al., 2209)) and GraphAF (Shi ct al. 2020 The models are evaluated by sampling 10,000 examples and measuring the proportion of valid molecules. We also report the proportion of unique molecules and novel samples that are not in the training dataset. These metrics prevent models from memorizing a small subset of graphs. Finally, the reconstruction rate describes whether graphs can be accurately decoded from latent space. Normalizing Flows naturally score 100% due their invertible mapping, and we achieve the same with our encoding despite no guarantees.
Table shows that GraphCNF generates almost twice many valid molecules than other one- shot approaches. Yet, the validity and uniqueness stay at almost 100%. Even the autoregressive normalizing flow, GraphAF, is outperformed by GraphCNFby by 15% However, the rules for generating valid molecules can be enforced in autoregressive models by masking out the invalid outputs. This has been the case for JT- VAE as it has been trained with those manual rules, and thus achieves a
HC
SHy
HF
(a) Molecule of Zinc250k
(b) Generated molecule
(c) Sub-molecule generations
Figure 2: Visualization of molecules from (a) the Zinc250k dataset and (b,c) generated by GraphCNF. Nodes with black connections and no description represent carbon atoms. Sub-figure (c) shows the failure case of two valid sub-graphs. More example generations can be found in Appendix C
Coode available here: http:::/ggthhbb.cmm///////////ee

---

Published as a conference paper at ICLR 2021
Table 1: Performance on molecule generation trained on Zinc250k (Irwin et al. 2012), calculated on 10k samples and averaged over 4 runs. The standard deviation of those runs can be found in Appendix C
Scores of the baselines are taken from their respective papers.
Method
Validity
Uniqueness
Novelty
Reconstruction
Parallel
Manual Rules
JT-VAE 10
100%
100%
100%
71%
GraphAF 336
68%
99.10%
100%
100%
X
R-VAE 23
34.9%
100%
54.7%
GraphNVP
24
42.60%
94.80%
100%
100%
/
GraphCNF
83.41% 96.35%
99.99%
100%
100%
+ Sub-graphs
99.98%
99.98%
100%
X
validity of 100%. Nevertheless, we are mainly interested in the model's capability of learning the rules by itself and being not specific to any application. While GraphNVP and GraphAF sample with a lower standard deviation from the prior to increase validity, we explicitly sample from the original prior underline that our model covers the whole latent space well. Surprisingly, we found out that most invalid graphs consist of two or more that in isolation are valid, as shown in Figure 2c This can happen as one-shot models have no guidance for generating single connected graph. By taking the largest sub-graph of these predictions, we obtain validity ratio of 96.35% making our model generate almost solely valid molecules without any manually encoded rules. We also evaluated model on the Moses (Polykovskiy et al. 2018) dataset and achieved similar scores shown in Appendix
5.2 GRAPH COLORING
Graph coloring is well-known combinatorial problem (Bondy et al., 1976) where for a graph G. each node is assigned to one of K col- Yet, two adjacent nodes cannot have the same color (see Fig- ure 3. Modeling the distribution of valid color assignments to arbi- trary graphs is No-complete. To train models on such distribution, we generate dataset of valid graph colorings for randomly sampled graphs. To further investigate the effect of complexity, we create two dataset versions, one with graphs of size 10 < [V| < 20 and another with 25 S [V| < 50, as larger graphs are commonly harder to solve.
For graph coloring, we rely GraphCNF and compare to variational autoencoder and an autoregressive model generating one node at a time. As edges are being modeled here, only the first step of GraphCNF's s three-step generation. For all models, we apply the same Graph Attention network (Veličković et al. 2018). As autoregressive models require a manually prescribed node order, we compare the following: a random ordering per graph, largest_firs: which is inspired by heuristics of automated theorem provers that start from the nodes with the most connections, and smallest_fine where we reverse the order of the previous heuristic. We evaluate the models by measuring the likelihood of color assignments to unseen test graphs in bits per node. Secondly, we sample one color assignment per model for each test graph and report the proportion of valid colorings.
Figure 3: Example graph with |V|= = 26 and generated valid coloring by GraphCNF.
The results in Table show that the node ordering has indeed a significant effect on the autoregressive model's performance. While the smallest first ordering leads to only 32% valid solutions on the large dataset, reversing the order simplifies the task for the model such that it generates more than twice as many valid color assignments. generating all nodes in parallel, it outperforms all node orderings the small dataset, while being close to the best ordering on the larger dataset. This invariance property is especially beneficial in tasks where optimal order of nodes is not known, like molccule generation. Although having more parameters, the sampling with GraphCNF is also considerably faster than the autoregressive models. The sampling time can further be improved by replacing the logistic mixture coupling layers with
contrast, GraphCNF is invariant of the order of nodes. Despite
7

---

Published as a conference paper at ICLR 2021
Table 2: Results on the graph coloring problem. Runtimes are measured on NVIDIA TitanRTX
GPU with batch size of 128. The standard deviations of the results can be found in Appendix D.2
10 S [V] <
20
25 S
V S
50
Method
Validity
Bits per node
Time
Validity
Bits per node
Time
VAE
44.95%
0.84bpd
0.05s
7.75%
0.64bpd
0.10s
-NNl-esallett.on first
76.86%
0.73bpd
0.69s
32.27%
0.50bpd
2.88s
RNN+ Random
88.62%
0.70bpd
0.69s
49.28%
0.46bpd
2.88s
RNN
ENN++Largee -Largest_ftii first
93.41%
0.68bpd
0.69s
71.32%
0.43bpd
2.88s
GraphCNF
94.56%
0.67bpd
0.28s
66.80%
0.45bpd
0.54s
Affine coupling
93.90%
0.69bpd
0.12s
65.78%
0.47bpd
0.35s
LSTM
Categorical NF
Latent NF
1.30
655
6.39
133
1,28
1.27
1.44
1.45
5.43
125
50
4.81
120
Penn Treebank
Text8
Wikitext103
Metrie: bits per churacter/word
Figure Results on language modeling. The reconstruction error shown a lighter color corresponding the model. The exact results including standard deviations can be found in Table II
affine ones. Due to the lower complexity, we see a slight decrease in validity and bits per dimension but can verify that logistic mixture couplings are not crucial for CNFs.
5.3 LANGUAGE MODELING
To compare CNF with Latent NF, we test both models on language modeling. We experiment with two popular character-level datasets, Penn Treebank (Marcus al. 1994) and text8 (Mahoney 2011 with a vocabulary size of K = 51 and K = 27 respectively. We also test a word-level dataset, Wikitext103 (Merity al., 2017). with due to its gradient approximations (Tran et al. 2019). We follow the setup of Ziegler and Rush (2019) for the Penn Treebank and train on sequences of 256 tokens for the other two datasets. Both Latent NF and CNF apply single mixture coupling layer being autoregressive across time and latent dimensions and differ only in the encoding/decoding strategy. The applied LSTM in the coupling additional RNN baseline in Figure 4 Categorical Normalizing Flow performs
10, 000 categories, which Discrete NF cannot handle
layers is shown on par with their autoregressive baseline, only slightly underperforming on Wikitext103 due to using single flow layer. Latent NF however performs considerably worse text8 and Wikitext100 due a non -on-metermiiittcc decoding and higher fluctuation in the loss that we experienced during training (see Appendix A.4.2 for visualizations). This shows the importance of factorized likelihood and underlines the two benefits of CNFs. Firstly, CNFs are more stable and simpler to train as no loss scheduling is required, and the likelihood is mainly modeled in the flow prior. Secondly, the encoding of CNFs is much more efficient than Latent NFs. We conclude that CNFs are more widely applicable, and Latent NFs might be the better choice if the decoder can explicitly help, i.e., prior knowledge.
5.4 SET MODELING
Finally, present experiments on sets of categorical variables, for which we create two toy datasets with known data likelihood: set shuffling and set summation. The goal is to assign high likelihood only those possible sets that oceur in the dataset, and shows how accurately an arbitrary, discrete dataset. set shuffling, we model a set of N categorical variables each having out of N categories. Each category has appear exactly once, which leads N! possible
flows model
8

---

Published as a conference paper at ICLR 2021
Table Results on set modeling, Metric used is bits per categorical variable (dimension).
Model
Set shuffling
Set summation
Discrete NF (Tran et al.
2019)
3.87 +0.04
2.51 +0.00
Variational Dequant. (Ho et al. 2019)
3.01 +0002
2.29 +0.01
Latent NF (Ziegler and Rush 2019)
2.78 +.000
2.26 +0.01
CNF Mixture model
2.78 10.00
2.24 =0.00
CNF + Linear flows
2.78 +0.00
2.25 #0.00
CNF + Variational Encoding
2.79 +0.01
2.25 +0.01
Optimal
2.77
2.24
assignments that need to be modeled. In summation, we again consider set of size N with N categories, but those categories represent the actual integers 1, 2,.,,.. N and have to sum to an arbitrary number, In contrast to set shuffling, the data is ordinal, which we initially expected to help dequantization methods. For both experiments we set N 16 and L = 42.
In Table we compare CNFs applying variational dequantization (Ho Normalizing Flows (Ziegler and Rush 2019) and Discrete Normalizing Flows (Tran et al. 2019).
2019). Latent
The results show that CNFs achieve nearly optimal performance. Although we model a lower bound in continuous space, our flows can indeed model discrete distributions precisely. Interestingly, representing categories by a simple mixture model sufficient for achieving these results. We observe the same trend in domains with more complex relations between categories, such as graphs and language modeling, presumably because of both the coupling layers and the prior distribution resting upon logistic distributions well. Variational dequantization performs worse the shuffling dataset, while on set summation with ordinal data, the gap to the optimum is smaller. The same holds for Discrete NFs, although it is worth noting that unlike CNFs, optimizing Discrete NFs had issues due to their gradient approximations. Latent Normalizing Flows with the joint decoder achieve similar performance CNF, which can be attributed close-to deterministic decoding. When looking the encoding space (see Appendix A.4.I for visualization), we see that Latent NF has indeed learned a mixture model well. Hence, the added complexity is not needed on this simple dataset.
CONCLUSION
We present Categorical Normalizing Flows which learn a categorical, discrete distribution by jointly optimizing the representation of categorical data in continuous latent space and the model likelihood of a normalizing flow. Thereby, we apply variational inference with a factorized posterior to maintain almost unique decoding while allowing flexible encoding distributions. We find that plain mixture model is sufficient for modeling discrete distributions accurately while providing an efficient way for encoding and decoding categorical data. Compared to joint posterior, CNF' are more stable, efficient, and have an inductive bias to move all modeling complexity into the prior. Furthermore, GraphCNF, a normalizing flow on graph modeling based on CNFs, outperforms autoregressive and one-shot approaches on molecule generation and graph coloring while being invariant to the node order. This emphasizes the potential of normalizing flows on categorical tasks, especially for such with non-sequential data.
ACKNOWLEDGEMENTS
We thank SURFsara for the support in using the Lisa Compute Cluster.
REFERENCES
John Adrian Bondy, Uppaluri Siva Ramachandra Murty, and others. 1976. Graph theory with applications volume 290. Macmillan London.
Yuri Burda, Roger Grosse, and Ruslan Salakhutdinov, 2016. Importance weighted autoencoders. 4th International Conference on Learning Representations, ICLR 2016 Conference Track Proceedings, pages 1-14.

---

Published as a conference paper at ICLR 2021
Laurent Dinh, Jascha Sohl-Dickstein, and Samy Bengio. 2017. Density estimation using Real NVP In 5th
International Conference on Learning Representations, ICLR 2017 Toulon, France.
Dheeru Dua and Casey Graff. 2019. UCI Machine Learning Repository http://archive.ics.uiieem
Irvine
CA: University of Califoria, School of Information and Computer Science.
Dan Hendrycks and Kevin Gimpel. 2016. arXiv: 00..08./5v3.
Gaussian Error Linear Units (GELUs) arXiv preprint
Jonathan Ho, Xi Chen, Aravind Srinivas, Yan Duan, and Pieter Abbecl. 2019. Flow++: Improving Flow-Based Generative Models with Variational Dequantization and Architecture Design In Proceedings of the 36th International Conference Machine Learning, volume 97, pages 2730,, Long Beach, California, USA.
PMLR.
Emiel Hoogeboom, Taco S. Cohen, and Jakub M. Tomezak. 2020. Learning Discrete Distributions by Dequanti- Zation arXiv preprint arXiv:2001.//235v1.
Emiel Hoogeboom, Jom W. T. Peters, Rianne van den Berg, and Max Welling. 2019. Integer Discrete Flows and Lossless Compression In Advances in Neural Information Processing Systems 32, pages 12134-122 12144.
Vancouver, BC, Canada.
John J Irwin, Teague Sterling, Michael M Mysinger, Erin S Bolstad, and Ryan G Coleman. 2012. ZINC: A Free Tool to Discover Chemistry for Biology Journal of Chemical Information and Modeling, 527771757-1766
Wengong Jin, Regina Barzilay, and Tbmmi Jaakkola. 2018. Junction tree variational autoencoder for molecular ggaah generation 35th International Conference on Machine Learning, ICML 2018, 5:3632 13632-36488.
Sungwon Kim, Sang-Gil Gil Lee, Jongyoon Song, Jaehyeon Kim, and Sungroh Yoon. 2019. FloWaveNet A
Generative Flow for Raw Audio In Proceedings of the 36th International Conference on Machine Learning. volume 97 of Proceedings of Machine Learning Research, pages 3370-3378 Long Beach, Californaa, USA. PMLR.
Diederik P. Kingma and Jimmy Ba. 2015. Adam: A Method for Stochastic Optimization Conference on Learning Representations (ICLR) 2015, San Diego, CA, USA.
In International
Diederik P. Kingma and Prafulla Dhariwal. 2018. Glow: Generative Flow with Invertible 1xl Convolutions In Advances Neural Information Processing Systems, volume 31, pages 10215-10224, Curran Associates, Inc.
Diederik P. Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, llya Sutskever, and Max Welling, 2016, [mproved variational inference with inverse autoregressive flow Advances in Neural Information Processing Systems, 29:4743-4751.
Diederik P Kingma and Max Welling. 2014. Auto-Encoding Variational Bayes In 2nd International Conference on Learning Representations, ICLR, Banff, AB, Canada.
Ivan Kohyzev, Simon Prince, and Marcus A. Brubaker. 2019. Normalizing Flows: Introduction and Ideas arXiv preprint 1rXiv:1908. 09257v1.
Henrique Lemos, Marcelo Prates, Pedro Avelar, and Luis Lamb. 2019. Graph colouring meets deep learning Effective graph neural network models for combinatorial problems Proceedings International Conference on Tools with Artificial Intelligence, ICTAI, volume 2019-Novem, pages 879-885. IEEE Computer Society.
Yujia Li, Oriol Vinyals, Chris Dyer, Razvan Pascanu, and Peter Battaglia. 2018. Learning deep generative models of graphs arXiv preprint arXiv:1803.0332 083..3334..
Renjie Liao, Yujia Li, Yang Song, Shenlong Wang, Will Hamilton, David K Duvenaud, Raquel Urtasun, and Richard Zemel. 2019. Efffiientt Graph Generation with Graph Recurrent Attention Networks In H Wallach, H Larochelle, A Beygelzimer, D' Alche- Buc, E Fox, and R Gameet, editors, Advances in Newral Information Processing Systems 32, pages 4255-4265. Curran Associates, Inc.
Jenny Liu, Aviral Kumar, Jimmy Ba, Jamie Kiros, and Kevin Swersky. 2019. Graph Normalizing Flows In Advances in Neural Information Processing Systems, pages 13556 13566.
Liyuan Liu, Haoming Jiang, Pengcheng He, Weizhu Chen, Xiaodong Liu, Jianfeng Gao, and Jiawei Han. 2020. On the Variance of the Adaptive Learning Rate and Beyond. In International Conference on Learning Representations.
10

---

Published as a conference paper at ICLR 2021
Qi Liu, Miltiadis Allamanis, Mare Brockschmidt, and Alexander Gaunt. 2018. Conntrriieed Graph Variational Autoencoders Molecule Design In S Bengio, H Wallach, H Larochelle, K Grauman, N Desa-Bianchi, and R Garnett, editors, Advances in Neural Information Processing Systems 31, pages 7795-7804, Curran Associates, Inc.
Tengfei Ma, Jie Chen, and Cao Xiao. 2018. Constrained Generation of Semantically Valid Graphs via Reg ularizing Variational Autoencoders) In S Bengio, H Wallach, H Larochelle, K Grauman, N Disa-Bianchi, and R Garnett, editors, Advances in Neural Information Processing Systems 3/, pages 71 711--1124 Curran Associates, Inc.
Kaushalya Madhawa, Katushiko Ishiguro, Kosuke Nakago, and Motoki Abe. 2019. GraphNVP: An Invertible Flow Model for Generating Molecular Graphs arXiv preprint arXiv: 1600vl.
Matt Mahoney, 2011.
Large text
compression
benchmark
Benchmark published at
http:///nattnchhoneeeeeetttttttt
Mitch Marcus, Grace Kim, Mary Ann Marcinkiewicz. Robert MacIntyre. Ann Bies, Mark Ferguson, Karen Katz, and Britta Schasherger. 1994. [The Penn Treehank: Annotating Predicate Argument Structure In Proceedings of the Workshop on Human Language Technology, pages 114-119, Plainsboro, NJ. Association for Computational Linguistics.
Stephen Merity, Caiming Xiong, James Bradbury, and Richard Socher. 2017. Pointer Sentinel Mixture Models In 5th International Conference on Learning Representations, ICLR 2017, Toudon, France, April 24-26, 2077, Conference Track Proceedings. OpenReview.net.
Tomáš Mikolov, Ilya Sutskever, Anoop Deoras, Hai-Son Le, Stefan Kombrink, and Jan Cemooky,, 2012. Subword language modeling with neural networks preprint (http://www.fir wtt.://wwwwt....m/ooooo////nnnna pdf), 8.
Adam Paszke, Sam Gross, Francisco Massa, Adam Lerer, James Bradbury, Gregory Chanan, Trevor Killeen, Zeming Lin, Natalia Gimelshein, Luca Antiga, Alban Desmaison, Andreas Kopf, Edward Yang. Zachary DeVito, Martin Raison, Alykhan Tejani, Sasank Chilamkurthy, Benoit Steiner, Lu Fang. Junjic Bai, and Soumith Chintala. 2019. PoTch:: An Imperative Style, High-Perlornance Deep Learning Library In H Wallach, H Larochelle, A Beygelzimer, F Alché-Buc, E Fox, and R Garnett, editors, Advances in Neural Information Processing Systems 32, pages 1024-8035. Curran Associates, Inc.
Jeffrey Pennington, Richard Socher, and Christopher Manning. 2014. Glove: Global Vectors for Word Repre- sentation In Proceedings of the 2014 Conference Empirical Methods in Natural Language Processing (EMNLP), pages 1532-15433 1543, Doha, Qatar. Association for Computational Linguistics.
Daniil Polykovskiy, Alexander Zhebrak, Benjamin Sanchez-Lengelnn Sergey Golovanov, Oktai Tatanov, Stanislav Belyaev, Rauf Kurbanov, Aleksey Artamonov, Vladimir Aladinskiy, Mark Veselov, Artur Kadurin, Simon Johansson, Hongming Chen, Sergey Nikolenko, Alan Aspuru- Guzik, and Alex Zhavoronkov. 2018. Molecular Sets (MOSES): A Benchmarking Platform for Molecular Generation Models arXiv preprint
1.X.:://11,12223:1 pages 1-17.
Ryan Prenger, Rafacl Valle, and Bryan Catanzaro. 2019. Waveglow: A flow-based generative network for speech synthesis In IEEE International Conference on Acoustics, Speech and Signal Processing (ICASSP), pages IEEE
Afshin Rahimi, Trevor Cohn, and Timothy Baldwin. 2018 Semi-supervised User Geolocation via Graph Convolutional Networks In Proceedings of the 56th Annual Meeting of the Association for Computational Linguistics Volume 1: Long Papers), pages 2009-2019, Melbourne, Australia. Association for Computational Linguistics.
Danilo Jimenez Rezende and Shakir Mohamed. 2015. Variational Inference with Normalizing Flows In Proceedings of the 32nd International Conference on Machine Learning, volume 37, Lille, France.
Michael Schlichtkrull, Thomas N Kipf, Peter Bloem, Rianne van den Berg, Ivan Titov, and Max Welling 2018. Modeling Relational Data with Graph Convolutional Networks In The Semantic Web, pages 593 007, Cham. Springer International Publishing.
Chence Shi, Minkai Xu, Zhaocheng Zhu, Weinan Zhang, Ming Zhang, and Jian Tang. 2020. GraphAF: a Flow-based Autoregressive Model for Molecular Graph Generation In International Conference on Learning Representations.
Martin Simonovsky and Nikos Komodakis. 2018 GraphVAE: Towards Generation of Small Graphs Using Variational Autoencoders. In International Conference Artificial Neural Networks, volume Abs/1802.0,, pages 412-422.
11

---

Published as a conference paper at ICLR 2021
Nitish Srivastava, Geoffre Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. 2014. Dropout Simple Way to Prevent Neural Networks from Overtitting Journal of Machine Learning Research, 15:1929-1958
Esteban Tabak and Eric Vanden Eijnden. 2010. Density estimation by dual ascent log- likeliio Communications in Mathematical Sciences, 8(1):217-233
Lucas Theis, Aäron Van Den Oord, and Matthias Bethge. 2016. A note on the evaluation generative models 4th International Conference on Learning Representations, ICLR 2016 Conference Track Proceedings.
Jakub M. Tomezak and Max Welling. 2017. Improving Variational Auto-Encoders using Householder Flow arXiv preprint Xrviv: 1ee::6111.000 1.066304..
Dustin Tran, Keyon Vafa, Kumar Krishna Agrawal, Laurent Dinh, and Poole. 2019, Discrete Flows
Invertible Generative Models of Discrete Data In Advances in Neural Information Processing Systems, pages
14692-14701.
Benigno Uria, lain Murray, and Hugo Larochellchugo. 2013. RNADE: The real-valued neural autoregressive eattity- In Advances in Neural Information Processing Systems, volume 26, pages 2175-2183.
Rianne Van Den Berg, Leonard Hasenclever, Jakub M. Tomezak, and Max Welling, 2018. Sylvester normalizing flows for variational inference) 34th Conference on Uncertainty in Artificial Intelligence 2018, UAI 2018, 1:393 402.
Ashish Vaswani, Noam Shazeer. Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. 2017. Attenttin is All you Need In I Guyon. U V Luxburg. Bengio. H Wallach, Fergus, Vishwanathan, and R Gareett, editors, Advances in Neural Information Processing Systems 30, pages 5998-6008. Curran Associates, Inc.
Petar Veličković, Guillem Cucurull, Arantxa Casanova, Adriana Romero, Pictro Liò, and Yoshua Bengio. 2018. Graph Attention Networks International Conference on Learning Representations.
Jiaxuan You, Rex Ying, Xiang Ren, William Hamilton, and Jure Leskovec. 2018. GraphRNN: Generating Realistic Graphs with Deep Auto-regressive Models In Proceedings of the 35th International Confer- ence on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pages 5708-5717, Stockholmsmääsan, Stockholm Sweden, PMLR.
Jie Zhou, Ganqu Cui, Zhengyan Zhang, Cheng Yang, Zhiyuan Liu, Lifeng Wang, Changcheng Li, and Maosong Sun. 2018. Graph Neural Networks: A Review of Methods and Applications, arXiv preprint arXiv: 1812.08434.
Zachary M. Ziegler and Alexander M. Rush. 2019. Latent Normalizing Flows for Discrete Sequences In Proceedings the 36th International Conference Machine Learning, volume 97, pages 7673-7682, Long Beach, California, USA. PMLR.
12

---

Published as a conference paper at ICLR 2021
A VISUALIZATIONS AND DETAILS ON ENCODING DISTRIBUTIONS
In the following, we visualize the different encoding distributions we tested in Categorical Normaliz- ing Flows, and outline implementation details for full reproducibility.
A.1 MIXTURE OF LOGISTICS
The mixture model represents each category by an independent logistic distribution in continuous latent space, as visualized in Figure Specifically, the encoder distribution q( ((21 22)=), with æ being the categorical input and z the continuous latent representation, can be written as:
N (<(()) = Il
(())) ==1
(8)
g(vlp,o)
II j=1 (1+ expepp-pep -))) where Ej
exp( (())
M4j Jj
(9)
9 represent the logistic distribution, and d the dimensionality of the continuous latent space per category. Both parameters je and a are learnable parameter, which can be implemented via a simple table lookup. For decoding the discrete categorical data from continuous space, the true posterior is calculated by applying the Bayes rule:
pT,, ((x;)g(Z(aa)) (()
(10)
where the prior categories, p(x;), ), is calculated based the category frequencies in the training dataset. Although the posterior models a distribution over categories, the distribution is strongly peaked for most continuous points in the latent space as the probability steeply decreases the further point is away from a specific mode. Furthermore, the distribution is trained to minimize the posterior entropy which pushes the posterior to be deterministic for commonly sampled continuous points. Hence, the posterior partitions the latent space into fragments in which all continuous points are assigned to one discrete category. The borders between the fragments, where the posterior is not close to deterministic, are small and very rarely sampled by the encoder distribution. We visualized the partitioning for an example of three categories in Figure
(a) Encoding distribution q(z;
(b) Decoder partitioning p(x;)
Figure 5: Visualization of the mixture model encoding and decoding for categories. Best viewed in color. (a) Each category is represented by a logistic distribution with independent mean and scale which are learned during training. (b) The posterior partitions the latent space which we visualize by the background color. The borders show from when on we have an almost unique decoding of the corresponding mixture (> 0.95 decoding probability). Note that these borders do directly correspond to the euclidean distance we use logistic distributions instead of Gaussians.
Notably, the posterior can also be learned by a second, small linear network. While this possibly introduces a difference between encoder and decoder, experienced it to vanish quickly over training iterations and did not observe any significant difference compared to using the Bayes posterior besides slower training in the very early stages of training. Additionally, we were able
13

---

Published as a conference paper at ICLR 2021
to achieve very low reconstruction errors in two dimensions for most discrete distributions of < 50 categories. Nevertheless, higher dimensionality of the latent space is not only crucial for a large number of categories for word-level vocabularies but also be beneficial for more complex problems. Still, using even higher dimensionality rarely caused any problems or showed significantly decreasing performance. Presumably, the flow learns to ignore latent dimensions if those not needed for modeling the discrete distribution. To summarize, the dimensionality of the latent space should be considered as important, but robust hyperparameter which can be tuned in an early stage of hyperparameter search.
In the very first training iterations, it can happen that the mixtures of multiple categories are the exact same spot and have troubles to clearly separate. This can be easily resolved by either weighing the reconstruction loss slightly higher for the first 5500 iterations or initializing the mean of the mixtures with higher variance. Once the mixtures are separated, the model has no incentive to group them together again as has started learn the underlying discrete distribution which results in considerably higher likelihood than plain uniform distribution.
A.1.1 KL DIVERGENCE
The objective for learning categorical normalizing flows with a mixture encoding, shown in Equation 3 constitutes a lower bound. The variational gap between the objective and the evidence log Pmoded is given by the KL divergence between the upproximate posterior (z(2x), and the true posterior p(z|æ); : D [q(z)=)][p(z]aa)
The advantage of using the mixture encoding that by replacing the decoder by Bayes formulation of the encoder, as in Equation T0 we introduce dependency between q( (z) and p(z(2:). Specifically, we can rewrite the true posterior as:
p(((x) ((zr)))((2)
(11)
p(x)
11. p(a;) (z))((z)
(12)
p(x)
II
g ((z1)2)))(()) (2;;)
p(z)
(13)
Σx q (z:(5)p(())
p(z)
11.
p(z)
Il (
(14)
p(x))
II, ((2))
p(z)
II p(x:)
(15)
p(x)
((2))
Intuitively, this makes it easier for the model to tighten the gap because change in ((z|ær entails change (())a)). experiments, we observe the difference by a considerably faster optimization in the first steps during iteration. However, once the decoder is close to deterministic, both approaches with the Bayes formulation and the separate network will reach a similar variational gap as pl (a|z) will drop out of the Equation II with being 1.
Note that this variational gap is very similar to that from variational dequantization for integers, which is being used in most SOTA image modeling architectures. The difference to dequantization is that the decoder [2)) has been manually fixed, but yet represents the Bayes formulation of the here, i.e. the original integers with random encoder (22(z) (the latent variable z represents variable u between 0 and l).
A.1.2 TRAINING AND TESTING
Below, we lay out the specifics for training and testing a categorical normalizing flow with the logistic mixture encoding. Training and testing are almost identical to normalizing flows trained on image modeling, except for the loss calculation and encoding. Algorithm shows the training procedure. First of all. we determine the prior p(z;) over categories. which can be done by counting the number of occurrences of each category and divide by the sum of all. The difference between the prior probabilities in the training and testing is usually neglectable, while the training set commonly is
14

---

end for Published as a conference paper at ICLR 2021
larger and hence provides better overall data statistic. After this, the training can start by iterating over batches in the training set D. For encoding, we can make use of the reparameterization trick and simply shift and scale the samples of standard logistic distribution. The loss is the lower bound of Equation
Algorithm Training procedure for the logistic mixture encoding CNFs
1: Calculate prior probabilities p(r;) on the training dataset;
2: for æ E D do
3: Sample a random logistic variable z': x' ~ LogisticDist((.. 1);
4: Reparameterizatioo trick for encoding: 2; = 2 (z:) + 1xx3))
5: Negative log-likelihood calculation L log (2model II Eg (1)
((a))
(Eq.
6: Minimize loss L by updating parameters in Pmodel (z) and ((2,(22))
7: end for
During testing, we make use of importance sampling tighten the gap, given that log Eg [p(x)] 2
Ez log N Σm p(sr)) > Er [lg)(())] (Burda al. 2016). This again standard technique for evaluating normalizing flows on images, and can improve the bits per dimension score slightly. In our experiments however, we did experience significant difference between N and N 1000. The bits per dimension score is calculated by using the log with base 2 on the likelihood, and divide by the number of dimensions/elements in categorical space (denoted by S).
Algorithm 2 Test procedure for the logistic mixture encoding in CNFs. We use N 1000 in our experiments, although the difference between N = 1 and N = 1000 was marginal for most cases.
1: for E D do
2: for 1..... do
3: Encode r as continuous variable ((((x);
4: Determine likelihood L. - Pmodel (z) 11: Ee (3))(22)(3)
((z.)
5: end for
6:
Determine bits per dimension score: - -log2
N Σm Lm /5;
7:
A.1.3 EXAMPLE ENCODING ON MOLECULE GENERATION
An example of trained model encoding shown in Figure Here, we visualize the encoding the edge attributes in GraphCNF trained the molecule generation. In this setup. have 3 categories, representing the single, double and triple bond. While the single bond category is clearly the dominant one due the higher prior probability, did not observe any specific trends across trained models of the position or scale of the distributions.
Visualizations graph coloring show similar behavior as Figure 5 because all three categories have the same prior probability. Other encoding distributions such as the node types (atoms) cannot be so easily visualized because of their higher dimensionality than 2. We tried applying dimensionality reduction techniques for those but experienced that those do not capture the distribution shape well.
A.2 LINEAR FLOWS
The flexibility of the mixture model can be increased by applying normalizing flows each mixture that dependent on the discrete category. We refer to this approach as linear flows as the flows are applied for each categorical input variable independently. We visualize possible encoding distributions
15

---

Published as a conference paper at ICLR 2021
150
19.
Sigll
Singlo bend
Doubre bond
100
100
5.0
23
150 155 -50
-25
-10
-25
8.5
25
1.0
25
-7.5
+5.0
-25
0.0
2.5
5.0
7.5
10.0
125
19.
21
21
(a) Encoding distribution g(zi(ar))
(b) Decoder partitioning p(zx)2z))
Figure 6: Visualization the mixture model encoding of the edge attributes for a trained model on molecule generation. The majority of the space is assigned to category 1, the single bond, by far the most common edge type. Across multiple models however, we did not see consistent trend of position//sanddad deviation of each category's encoding.
with linear flows in Figure 7 Formally, write the distribution as:
N
(<(xx) = II (()))a)))
(16)
II det (54(2(0-1);;;; where = z(K)
a (2(8)(z²) = g( (2
(17)
(2-1-
where f1...., fK invertible, smooth mappings. In particular, coupling layers with activation normalization and invertible 1xl convolutions (Kingma and Dhariwal 2018). Both the activation normalization and coupling use the category z; additional external input to determine their transformation parameters by a neural network. The class-conditional transformations could also be implemented by storing K parameter sets for the coupling layer neural networks, which is however inefficient for a larger number of categories. Furthermore, in coupling layers, we apply channel mask that splits z; over latent dimensionality into two equally sized parts, of which one is transformed using the other as input.
here again sequence of
(a) Encoding distribution q(2, (x:))
(b) Decoder partitioning p(1;; (22)
Figure 7: Visualization of the linear flow encoding and decoding for 3 categories. Best viewed in color. (a) The distribution per category not restricted to a simple logistic and can be multi-modal, rotated transformed even more. (b) The posterior partitions the latent space which visualize by the background color. The borders show from when on we have an almost unique decoding of the corresponding category distribution (> 0.95 decoding probability).
16

---

Published as a conference paper at ICLR 2021
Similarly to the mixture model, we can calculate the true posterior pzz;(2;) using the Bayes rule. Thereby, we sample from the flow for 2; and need to inverse the flows for all other categories. Note that as the inverse of the flow also needs to be differentiable in this situation, we apply affine coupling layers instead of logistic mixture layers. However, this gets computationally expensive for more than 20 categories, and thus used single-layer linear network as posterior in these situations. The partitions of the latent space that can be learned by the encoding distribution are much more flexible, as illustrated in Figure 7
We experimented with increasing sizes of linear flows but noticed that the encoding distribution usually fell back rotated logistic distributions. The fact that the added complexity and flexibility by the flows are not being used further supports our observation that mixture models are indeed sufficient for representing categorical data well normalizing flows.
A.3 VARIATIONAL. ENCODING
The third encoding distribution we experimented with is inspired by variational dequantization (Ho et al. 2019) and models ((= by one flow across all categorical variables. Still, the posterior p(zx/22) is applied per categorical variable independently to maintain unique decoding and parti- tioning of the latent space. The normalizing flow again consists of sequence of logistic mixture coupling layers with activation normalization and invertible Ixl convolutions. The inner feature network of the coupling layers depends the task the normalizing flow is applied on. Hence, for sets, we used a transformer architecture, while for the graph experiments, we used a GNN. On the language modeling task, used Bi-LSTM model to generate the transformation parameters. All those networks use the discrete, categorical data x: as additional input.
As the true posterior cannot be found for this distribution, determine p();)2;). While the reconstruction error was again very low, we again experienced that the model mainly relied on a logistic mixture model, even if initialize it differently beforehand. Variational dequantization is presumably important for images as every pixel value has its own independent Gaussian noise signal. This noise can be nicely modeled by flexible dequantization distributions which need to be complex enough to capture the true mean and variance of this Gaussian noise. In categorical distributions, however, we do not have such noise signals and therefore seem not to benefit from variational encodings.
apply a two-layer linear network to
A.4 LATENT NORMALIZING FL.OW
A.4.1 ENCODING VISUALIZATION
C-0
C-4
C-6
C-8
C-10
C-12
C-14
C3
C-5
C-7
C-9
C-11
C-13
C-15
Figure Visualization of the encoding distribution q(2(==)) of latent NF on set summation.
In this section, we show visualizations of the learned encoding distribution of Latent NF ZZiegler and Rush 2019) on the task of set summation. As shown in Figure the encoder learned clear mixture distribution although not being restricted too. This shows that the possible complexity in the decoder is not being used. The figure was generated by sampling 5, 000 examples and merging them into one plot. Each latent variable z; is two dimensional, hence two-dimensional plot. The
17

---

Published as a conference paper at ICLR 2021
colors represent different categories. For readability, each color is normalized independently (i.e. highest value I, lowest 0) as otherwise, the colors of the stretched mixtures would be too low. For visualizations of the latent space in language modeling, refer the reader to Ziegler and Rush (0019).
A.4.2 LOSS COMPARISON
In the following, we compare LatentNF and CNF on their training behavior and loss fluctuation. Figure visualizes the loss curve for both models on the task of language modeling, trained on the text8 dataset (Mahoney 2011). Note that both models use the same hyperparameters and model except that the reconstruction loss is weighted 10 times higher in the first 5k iterations, and decays exponentially over consecutive iterations. This is why the initial loss LatentNF is higher, but the reconstruction loss becomes close to zero after 3k iterations in the example of Figure P Interestingly, we experience high fluctuations of the loss with LatentNF during the first iterations. Although Figure 9 shows only single run, similar fluctuations have occurred in all trained LatentNF models but at different training iterations. Hence, we decided to plot a single loss curve instead of the average of multiple. The fluctuations can most likely be explained by the need for strong loss scheduling. At the beginning of the training, the decoder loss weighted significantly higher than the prior component. Thus, the backpropagated gradients mostly focus on the decoder, which can peak for rare sntlegories/occurenee
in the dataset. These peaks cause the encoding distribution to change abruptly, and the prior has to adapt to these changes within the next iterations. However, this can again lead to a high loss when the transformations in the prior do not fit to the new encoding, and thus map them to points of low likelihood. Thus, we see a loss peak across a couple of iterations until the model balances itself out again.
Losa curve for compaing Liten?NF and CNF Text5
Losa curve for comparing LatendiF and CNF Texi
Latent
Latent NF
CNT
CNF
40
Sts per divessins
20
2.0
1000
2000
3000
4000
5000
1000
2000
2000
4000
5000
Namber of training Becetions
Number towining (lwerdons
(a) Loss plot on text8
(b) Zoomed loss plot on text8
Figure 9: Plotting the loss over batch iterations on the language modeling dataset text8 for LatentNF and CNF. The batch size is 128, and average the loss over the last 10 iterations. Sub-figure (b) shows a zoomed version of sub-figure (a) to show the larger fluctuation even on smaller scale, while CNF provides smooth optimization.
Similarly, when training on Wikitext103 (Merity al.] 2017), we experience even more frequent peaks in the loss, as Wikitext103 with 10k categories contains even more rare occurrences (see Figure TOp. Reducing the learning rate did not show to improve the stability while considerably increasing the training time. When reducing the weight of the decoder, the model was not able to optimize as well as before and usually reached bits per dimension scores of 10-20.
B IMPLEMENTATION DETAILS OF GRAPHCNF
In this section, we describe further implementation details of GraphCNF. We detail the implementation of the Edge-GNN model used in the coupling layers of GraphCNF, and discuss how we encode graphs of different sizes.
B.1 EDGE GRAPH NEURAL NETWORK
GraphCNF implements a three-step generation approach, for which the second and third step also models latent variables for edges. Hence, in the coupling layers, we need graph neural network
18

---

Published as a conference paper at ICLR 2021
Loss curve for comparing LanentuF and CNF Wikk lesd
Loos care for comparing LownthFF and CNF on Withwe
LF
Latert NF ONF
ONF
45
150
(5t per dineeeeeo
100
30
26
20
10020
15000
25000
10000
Nunber oftraiiing tocations
Number draiing
(a) Loss plot on Wikitext103
(b) Znomed loss plot on Wikitext103
Figure 10: Plotting the loss over batch iterations on the language modeling dataset Wikitext103 for LatentNF and CNF. The batch size is 128, and we average the loss over the last 10 iterations. Sub-figure (b) shows a zoomed version of sub-figure (a) show the details of the peaks in LatentNF's loss curve.
which supports both node and edge features. We implement this by alternating between updates of the edge and the node features. Specifically, given node features ot and edge features e' layer t, we update those as follows:
r'+1 Snode ee) c'41 fedge (e'g ott1)
(18)
(19)
The update functions, fnode and Sadges are both common GNN layers with slight adjustments to allow a communication between nodes and edges. Before detailing the update layers, it should be noted that we use Highway GNNs (Rahimi et al. (2018) which apply gating mechanism. Specifically, the updates for the nodes are determined by:
w+1
v' T (i+ + H (ū++)) (1- T (o*+!))
(20)
where 61+1 is the output of the GNN layer. H and T represent single linear layer networks where T has a consecutive sigmoid activation to limit the outputs between 0 and 1. The edge updates are applied in the similar manner. We experienced that such a gated update functions helps the gradient flow through the layers back to the input. This is important for normalizing flows as coupling layers or transformations in general strongly depend on previous transformations. Hence, we apply the same gating mechanism in the first step of GraphCNF, f1.
Next, we detail the GNN layers to obtain et+1 and öt+1, The edge update layer fedge resembles graph convolutional layer (Zhou et al. 2018) 8), and can be specified as follows:
2++2 g (Wi + W W!,))
(21)
where Gi represents the features of the edge between node and j. g stands for a GELU (Hendrycks)
and Gimpel 2016) non-linear activation. Using more complex transformations did not show to significantly improve the performance of GraphCNF.
To update the node representations, we took inspiration of the transformer architecture Vaswani jtt al. (017) and use modified multi-head attention layer. particular, a linear transformation maps each node to key, query and value vector:
K Qu, Vo, Wxe,, Wor Wvef
(22)
The attention value is usually computed based on the dot product between two nodes. However, as we explicitly have features for the edge between the two nodes, we use those to control the attention mechanism. Hence, we have an additional weight matrix u to map the edge features to an attention bias:
àg = Qu, KT
(23)
where d represents the hidden dimensionality of the features. Finally, we also add a edge-based value vector to allow full communication from edges to nodes. Overall, the updates node features
19

---

Published as a conference paper at ICLR 2021
calculated by:
exp (āis) Σm exp (āsm)*
aij
(24)
oft1
Σ aij [V, W
(25)
Alternatively to transformers, we also experimented with Graph Attention Networks (Veličković et al. 2018). However, those showed slightly worse results which is why used the transformer-basee layer.
In step 2, the (binary) adjacency matrix is given such that each node has a limited number of neighbors. A full transformer-based architecture as above is then not necessary anymore as every atom has usually between and neighbors. Especially the node-to-node dot product expensive to perform. Hence, we experimented with a node update layer where the attention is purely based on the edge features in step We found both work equally well while the second computationally more efficient.
B.2 ENCODING GRAPH SIZE
The number of nodes N varies across graphs in the dataset, and hence a generative model needs to be flexible regarding N. To encode the number of nodes, we use a similar approach as ZZegler and Rush (2019) for sequences and add simple prior over N. The prior is parameterized based on the graph size frequency in the training set. Alternatively, to integrate the number of nodes the latent space, we could add virtual nodes to the model, similar virtual edges. Every graph in the training dataset would be filled up to the maximum number of nodes (38 for Zinc250k (Irwin et al.) 2012)) by adding such virtual nodes. Meanwhile, during sampling, remove virtual nodes if the model generates such. GraphNVP (Madhawa et al., (019) uses such an encoding as their coupling layers did not support flexible graph sizes. However, in experiments, we obtained similar performance with both size encodings while the external prior is computationally more efficient and therefore used in this paper.
C ADDITIONAL RESULTS ON MOLECULE GENERATION
In this section, we present additional results on the molecule generation task. Table shows the results of our model on the Zine250k (Irwin et al.) 2012) dataset including the likelihood on the test set in bits per node. We calculate this metric by summing the log- likeehhoodd of all latent variables, both nodes, and edges, and divide by the number of nodes. Although the number of edges scales with O(N2), higher proportion of those are virtual and did not have a significant contribution to the likelihood. Thus, bits per node constitutes good metric for comparing the likelihood of molecules of varying size. Additionally, we also report the standard deviation for all metries over independent runs. For this, we initialized the random number generator with the seeds 42, 43, 44, and 45 before creating the model. The specific validity values we obtained are 80.74%, 81.16%, 85.3% and 86.44% (in no particular order). It should be noted that the standard deviation among those models is considerably high. This is because the models in molecule generation are trained on maximizing the likelihood of the training dataset and not explicitly on generating valid molecules. We experienced that among over seeds, models that perform better in terms of likelihood do not necessarily perform better in terms of validity.
Table 4: Performance molecule generation trained on Zinc250k (Irwin et al. 2012 with standard
deviation is calculated over 4 independent runs. See Table for baselines.
Method
Validity
Uniqueness
Novelty
Reconstruction
Bits per node
GraphCNF
83.41% 1%
99.99%
100%
100%
5. 5.17bpd (+:0.05)
(=2.88)
(+0.01)
(+0000)
(+0000)
+ Sub-graphs
96.35%
99.98%
99.98%
100%
2=22)
(10001)
(40002)
(1:0.00
20

---

Published as a conference paper at ICLR 2021
We also evaluated GraphCNN the Moses (Polykovskiy al. 2018) molecule dataset. Moses contains 1.9 million molecules with up to 30 heavy atoms different types. Again, we follow the preprocessing of Shi al. (2220) and represent molecules in kekulized form in which hydrogen removed. The results can be found in Table and show that we achieve very similar scores to the experiments Zinc250k. Compared the normalizing flow baseline GraphAF, GraphCNF generates considerably more valid atoms while being parallel in generation in contrast to GraphAF being autoregressive. JT-VAE uses manually encoded rules for generating valid molecules only such that the validity rate is 100%. Overall, the experiment on Moses validates that GraphCNF is not specialized on single dataset but can improve on current flow-based graph models across datasets.
Table 5: Performance on molecule generation Moses (Polykosskiy et al. 2018 calculated on 10k samples and averaged over 4 runs. Score for GraphAF taken from (hhi (2020) and JT-VAE from Polykovskiy et al. (2018).
Method
Validity
Uniqueness
Novelty
Bits per node
JT-VAE (Jin al. 2018
100%
99.92%
91.53%
GraphAF (Shi et al. 2020
71%
99.99%
100%
GraphCNF
82.56%
100.0%
100%
4.94bpd (+0.04)
(#2.34)
(+00000
(=00000)
+ Sub-graphs
95.66%
99.98%
100%
(+22588
(100001
(+00000
Finally, show 12 randomly sampled molecules from our model in Figure I1 In general, GraphCNF is able to generate very diverse set of molecules with variety of atom types. This qualitative analysis endorses the previous quantitative results of obtaining close to 100% uniqueness on 10k samples.
Hyb
Mil
20H
OH2
WH
OHz
Figure 11: Visualization of molecules generated by GraphCNF which has been trained on the Zinc250k (Irwin et al. (012) dataset. Nodes with black connections and no description represent carbon atoms. All of the presented molecules are valid. Best viewed in color and electronically for large molecules.
Besides generating multiple sub-graphs, the most common failure case we have found are single, invalid edges in large molecule, shown in four examples in Figure 12
21

---

Published as a conference paper at ICLR 2021
0
NH1
NI
NHH
OH2
hitt
Nily
Figure 12: Failure cases molecule generation besides sub-graph generation. The invalid edges have been indicates with green arrow. Changing the edges to single bonds in the examples above would constitute valid molecule.
D EXPERIMENTAL SETTINGS
In this section, we detail the hyperparameter settings and datasets for all experiments. All experiments have been implemented using the deep learning framework PyTorrh (Paszke al., 2019). The experiments for graph coloring and molecule generation have been executed on a single NVIDIA TitanRTX GPU. The average training time between 1 and 2 days. The and language experiments have been executed on a single NVIDIA GTX1080Ti in 4 to 16 hours. All experiments have been repeated with least 3 different random seeds.
D.1 SET MODELING
Dataset details We use two toy datasets, set shuffling and set summation, to simulate discrete distribution over sets our experiments. Note that we do have a classical split of train/val/test dataset, but instead train and test the models samples from the same discrete distribution. This is because we want to verify whether a categorical normalizing flow and other baselines can model an arbitrary discrete distribution. The special property of sets is that permuting the elements of set still represent the same set. However, a generative model still has to learn all possible permutations. While an autoregressive model considers those permutations different data points, cormutation-invariii model as Categorical Normalizing Flow contains an inductive bias to assign the same likelihood to any permutation.
In set shuffling, we only have one set to model which is the following (with categories C1 to C16):
(C1, C2, C3, CA, Cs, Ca. C7, Cs,, Cg, Cıo. Cn1, C12, C13, C14, C1s, C16)
This set has 16! possible permutations and therefore challenging to model. The optimal likelihood in bits per element is calculated by log2 (16!) /16 = 2.77.
The dataset set summing contains of 2200 valid sets for N = 16 and L = 42. An example for a valid set is:
For readability, the set is sorted by ascending values, although any permutation of the elements represent the exact same set. Taking into account all possible permutations of the sets in the dataset, we obtain optimal likelihood of log, (6.3 10a0) /16 2.24. The values for the sequence length N and sum L was chosen such that the task is challenging enough to show the differences between Categorical Normalizing Flows and its baselines, but also not too challenging to prevent unnecessarily long training times and model complexities.
Hyperparameter details Table an overview of the hyperparameters per model applied on set modeling. We use the notation "[val1, val2....) to show the different values we have tried during hyperparameter search. Thereby, the underlined value denotes hyperparameter value with the best performance and finally was being used to generate the results in Table
The number encoding coupling layers Categorical Normalizing Flows are sorted by the used encoding distribution. The mixture model uses no additional coupling layers, while for the linear flows, we apply 4 affine coupling layers using an external input for the discrete category. For the variational encoding distribution q [=), we 4 mixture coupling layers across the all latent variables with external input for æ. larger dimensionality of the latent space per element showed to be beneficial for all encoding distributions. Note that due to dimensionality larger than per
22

---

Published as a conference paper at ICLR 2021
element, can apply the channel mask instead of chess mask and maintain permutation invariance compared to the baselines.
In variational dequantization and Discrete NF, sort the categories randomly for set shuffling (the distribution is invariant to the category order/assignment) and in ascending order for set summation. In Discrete NE, we followed the published code from Tran et al. (2019) for their coupling layers and implemented it in Py TyTorhh (Paszke et al.] 2019). We use a discrete prior over the set elements which jointly optimized with the flow. However, we experienced significant optimization issues due the straight-through gradient estimator in the Gumbel Softmax.
Across this paper, experiment with the two optimizers Adam (Kingma and Ba 2013) and RAdam (Liu et al.] 2020), and experienced RAdam to work slightly better. The learning rate decay is applied every update and leads to exponential decay. However, we did not observe the choice of this hyperparameter to be crucial.
Table 6: Hyperparameter overview for the set modeling experiments presented in Table 3
Hyperparameters
Categorical NF
Var. dequant.
Discrete NF
Latent dimension
4, 6) 6)
1
16
#Encoding couplings
//4/4
4
#Coupling layers Coupling network
8
(4. 8)
Transformer
Transformer
Transformer
Number of layers
2
2
Hidden size
256
256
256
Mask
Channel mask
Chess mask
Chess mask
#mixtures
8 1024
Batch size
1024
1024
Training iterations
100k
100k
100k (SGD, Adam, RAdam)
Optimizer
( (Adam, RAdam)
RAdam
Learning rate
7.5e-4
7.5e-4
(1e-3, le-4, le-5)
Learning rate decay
0.999975
0.999975
0.999975 (0.1, 0.2, 0.5)
Temperature (GS)
D.2 GRAPH COLORING
Dataset details our experiments, we focus on the 3-color problem meaning that graph has to be colored using K = 3 colors. We generate the datasets by randomly sampling graph and using an SAT solver for finding one valid coloring assignment. In case no solution can be found, we discard the graph and sample new graph. We further ensure that every graph cannot be colored by less than colors to exclude too simple graphs. For creating the graphs, we take inspiration from Lemos et al. (2019) and first uniformly sample the number of nodes between 10 N |V| < 20 for the small dataset, and 25 S [V| < 50 for the large dataset. Next, we sample a value p between 0.1 and 0.3 which represents the probability of having an edge between random pair of nodes. Thus, p controls how dense a graph is, and we aim to have both dense and sparse graphs in our dataset. Finally, for each pair of nodes, we sample from Bernoulli distribution with probability p of adding an edge between the two nodes or not. Finally. we check whether each node has at least one connection and that all nodes can reached from any other node. This ensures that we have one connected graph and not multiple sub sub-graphs. Overall, we create train/val/test size of 192k/24k/24k for the small dataset, and 450k/20k/30k for the large graphs. We visualize examples the datasets Figure I3
During training, we randomly permute the colors of graph (e.g. red becomes blue, blue becomes green, green becomes red) as any permutation a valid color assignment. When we sample a color assignment from our models, we explicitly use temperature value of 1.0. For the autoregressive
We have used the following solver from the OR-Tools http://lleebbergggoo.ooooommmmmm
library
in python:
23

---

Published as a conference paper at ICLR 2021
model and the VAE, this means that sample from the softmax output. A common alternative is to take the argmax, which corresponds to a temperature value of 0.0. However, we stick to the original distribution because want to test whether the models capture the full discrete distribution of valid color assignments and not only the most likely solution. For the normalizing flow, a temperature of 1.0 corresponds sampling from the prior distribution as it used during training.
(a) V = 10
(b) | = 19
(c) V = 25
(d) V = 30
Figure 13: Examples of valid graph color assignments from the dataset (best viewed color). Due to the graph sizes and dense adjacency matrices, edges can be occluded or cluttered in (c) and (d).
Hyperparameter details 7 shows overview of the used hyperparameters. If "/" is used in the table, the first parameter refers to the hyperparameter value used on small dataset and the second for the larger dataset. The activation function used within the graph neural networks is GELU (Hendrycks and Gimpel| 2016). Interestingly we experience that a larger latent space dimensionality crucial for larger graphs despite having the same number of categories as the small dataset. This shows that having an encoding being flexible the number of dimensions can be further important for datasets where complex relations between categorical variables need be modeled. Increasing the number of dimensions on the small dataset did not show any significant differences in performance. The number of mixtures in the mixture coupling layers in general beneficial be large. However, this can also increase the sampling time. In the case of sampling time being crucial, the number of mixtures can be decreased for the tradeoff of slightly worse performance.
The input to the autoregressive model is the graph with the color assignment at time step T where each category including unassigned nodes is represented by an embedding vector. We experiment with an increasing number of hidden layers. While more layers are especially important for sub-optimal node ordering, the performance does not significantly improve for more than layers. As the sampling time also increases linearly with the number of layers, we use 5 hidden layers for the models.
For the variational autoencoder, we encode each node by a latent vector of size 4. As VAFs have shown to benefit from slowly adding the KL divergence between prior and posterior to the loss, experiment with a scheduler where the slope is based on a sigmoid and stretched over 10k iterations. We apply 5 layer graph attention network for both the encoder and decoder. Increasing the number of layers did not show a significant gain while making the loss scheduling more difficult, which is why stuck with 5 layers.
Detailed results Table shows the standard deviation of the results reported in Table with different seeds.
Each model
24

---

Published as a conference paper at ICLR 2021
Table 7: Hyperparameter overview for graph coloring experiments presented in Table
Hyperparameters
GraphCNF
Variational AE
Autoregressive
Latent dimension
(2, 4) (2, 4, 6,
4
#Coupling layers
[6, 8)
(Coupling) network
GAT
GAT
GAT [3, 5, 6, 7) 384
Number of layers
(3, 4, 5)
5
Hidden size
384
384
Number of heads
4
4
4
Mask
Channel mask
#mixtures
[4, 8, 16) [4, 8, 16)
Batch size
384 128
384 / 128
384 128
Training iterations
200k
200k
100k
Optimizer
RAdam
RAdam
RAdam
Learning rate
7.5c-4
7.5e-4
7.5e-4
KL scheduler
(1.0, 0.1- 0.1-+1.03
Table 8: Results on the graph coloring problem, including standard deviation over 3 seeds. The column time is excluded since the execution time per batch is constant seeds.
10 S [V| S
20
25 S V S 50
Method
Validity
Bits per node
Validity
Bits per node
VAE
44.95% +5.32%
0.84 +0.04
7.75% +1.5%
0.64 #0.02
ANN+-Smallest_irg
76.86% +0.84%
0.73 +0.02
32.27%
+1.41%
0.50 #0.01
Random
88.62% 10.65%
0.70 #0.01
49.28%
11.5%%
0.46 10.01
LaNgear first
93. 93.41% +0.42%
0.68 10.01
71.32% 10.77%
0.43 10.01
GraphCNF
94.56% +0.55%
0.67 +0.00
66.80% +1.14%
0.45 +0.01
Affine
93.90% +0.68%
0.69 #0.02
65.78%
#1.03%
0.47 #0.01
D.3 MOLECULE GENERATION
Dataset details The Zinc250k IIwin et al. 2012) dataset we use contains 239k molecules of which we use 214k molecules for training, 8k for validation, and 17k for testing. We follow the preprocessing Shi et al. (2020) and represent molecules in kekulized form in which hydrogen is removed. This leaves the molecules with up to 38 heavy atoms, with mean and median size of about 23. The smallest graph consists of nodes. Thereby, Zinc250k considers molecule with 8 different atom types where the distribution significantly imbalanced. The most common atom is carbon with 73% of all nodes in the dataset. Besides oxygen (10% and nitrogen (12%), the rest of the atoms occur in less than 2% of all nodes, with the rarest atom being Bromine (0.002%). Between those atoms, the dataset contains 3 different bonds or edge types, namely single, double and triple covalent bonds describing how many electrons are shared among the atoms. In over 90% of all node pairs there exist no bond. In 7% of the cases, the atoms are connected with a single connection, 2.4% with a double, and 0.02% with a triple connection. A similar imbalance is present in the Moses dataset and is based on the properties of molecules. Nevertheless, we experienced that GraphCNF was able to generate similar distribution, where adding the third stage (adding virtual edges later) considerably helped to stabilize the edge imbalance.
Hy Hypeeraaraeeeer details We summarize our hyperparameters in Table g Generally, a higher latent dimensionality is beneficial for representing nodes/atoms, similar to the graph coloring task. However, we experienced that a lower dimensionality for edges is slightly better, presumably because the flow already has a significant amount of latent variables for edges. Many edges, especially the virtual ones, do not contain much information. Besides, a deeper flow showed to gain better results offering more complex transformations. However, in contrast the graph coloring model, GraphCNF on
25

---

Published as a conference paper at ICLR 2021
molecule generation requires considerable amount of memory have to model feature vector per edge. Nevertheless, we did not experience any issues due to the limited batch size of 96, and during testing, could scale up the batch size easily to more than 128 NVIDIA GTX 1080Ti for both datasets.
Table 9: Hyperparameter overview for molecule generation experiments presented in Table I and
Hyperparameters
GraphCNF
Latent dimension (V/E)
(4, 64) (2. 3, 4)
#Coupling layers (fi/f2/fa)
4/ (4, 6) (4, 6)
Coupling network (f1/2,,)
Relational GCN/ Edge-GNN /3/3/3, 3/4/4, 4/4/4)
Number of layers (fa/f//fa)
Hidden size (V/E)
(256. 384) (128, 192) Channel mask
Mask
#mixtures (V/E)
{8, 16) / (4, 8, 16)
Batch size (Zinc250k/Moses)
64 / 96
Training iterations
150k
Optimizer
RAdam LLuu al. 2020
Learning rate
2c-4, 5e-4, 7.5e-4, le-3
D.4 LANGUAGE MODELING
Dataset details The three datasets we use for language modeling are the Penn Treebank (Marcus
al.. [994), text8 and Wikitext103 (Merity al. 2017). The Penn Treebank with a preprocessing of Mikolov et al. (2012) consists of approximately 5M characters and has a vocabulary size of K = 51. We follow the setup of Ziegler and Rush 2019) and split the dataset into sentences of maximum length of 288. Furthermore, instead of an end-of-sentence token, the length is passed to the model and encoded by an external discrete prior which is created based on the sentence lengths the training dataset.
Text8 contains about 100M characters and has a vocabulary size of K 27. We again follow the preprocessing Mikolov et al. (2012) and split the dataset into 90M characters for training, and 5M characters each for validation and testing. We train and test the models on a sequence length of 256.
In contrast to the previous two datasets, we use Wikitext103 as word-level language dataset. First, we create a vocabulary and limit it to the most frequent 10,000 words in the training corpus. We thereby pre-trained Glove (Pennington al., 2014) embeddings to represent the words in the baseline LST' networks and to determine the logistic mixture parameters in the encoding distribution of Categorical Normalizing Flows. Due to this calculation of the mixture parameters, linear network as a decoder. The linear network consists of three linear layers of hidden size 512 with GELU (Hendrycks and Gimpel| 2016) activation and output size of 10,000 (the vocabulary size). Similarly to text8, we train and test the models on an input sequence length of 256.
a small
Hyperparameter details The hyperparameters for the language modeling experiments are summa- rized in Table IO We apply the same hyperparameters for the flow and baseline if applicable. The best latent dimensionality for character-leved similar performance. For the word-level dataset, it is beneficial to increase the latent dimensionality to 10. However, note that is still significantly smaller than the Glove vector size of 300. As Penn Treebank has limited training dataset which LSTM networks casily overfit, we use dropout (Srivastava et al. 2014)) of 0.3 throughout the models and dropout an input token with a chance of 0. The other datasets seemed to benefit slightly by small input dropout to prevent overfitting at later stages of the training.
3, although larger dimensionality showed gain
Detailed results Table II shows the detailed numbers and standard deviations of the results reported in Figure 4 Each model was run with different seeds based on which the standard deviation has been calculated.
26

---

Published as a conference paper at ICLR 2021
Table 10: Hyperparameter overview for the language modeling experiments presented in Table I1
Hyperparameters
Penn Treebank
text8
Wikitext103
(Max) Sequence length
288
256
256
Latent dimension
[2, 3, 4]
12, 3, 4)
{8, 10, 12)
#Coupling layers
1
1
Coupling network
LSTM
LSTM
LSTM
Number layers
2
2
Hiddenn size
1024
1024
1024
Dropout
{0.0, 0.3)
0.0
0.0
Input dropout
[0.0, 0.05, 0.1, 0.2)
[0.0, 0.05, 0.1)
(0.0, 0.05, 0.1)
#mixtures
27
64
Batch size
128
128
128
Training iterations
100k
150k
150k
Optimizer
RAdam
RAdam
RAdam
Learning rate
7.5c-4
7.5e-4
7.5e-4
Table 11: Results on language modeling. The reconstruction error is shown in brackets.
Model
Penn Treebank
text8
Wikitext103
LSTM baseline
1.28 #0.01
1.44 #0.01
4.81 =0.05
Latent NF- layer
1.30+40.01 (0.01)
1.61 (0.03)
6.39-10.19 (1.78)
Categorical NF- layer
1.27 #0.01 (0.00)
1.45 #0.01 (0.00)
5.43 +0.09 (0.32)
D.5 REAL-WORLD EXPERIMENTS
To verify that Categorical Normalizing Flows can also be applied to real-world data, we have experi- mented on a credit-card risk record (Dua and Grafi 2019). The data contains different categorical attributes regarding potential credit risk, and we have used the following 9 attributes: "check- ina_status", "credit_history", history", "savings_stare status", "employment", "housing". "job", "own_ nown_telephone "foreign_ "ookeign and "class". The task is to model the joint density function of those 9 attributes over dataset of 1000 entries. We have used the first 750 entries for training. 100 for validation, and 150 for testing. The results are shown in Table 12 Both the LatentNF and CNF perform equally well, given the small size of the dataset and number of categories. The results have been averaged over three seeds. However, we experienced a relatively high standard deviation due to the small size of the training and test set.
Table 12: Results on credit card risk dataset (uua and Graff 2019). The reconstruction error is shown in brackets where applicable.
Model
Likelihood in bpd
Variational Dequantization
1.95-10.04 1.36=#0.03 (0.01) 1.37 +0.03 (0.00)
Latent NF
Categorical NF
27


Reviews:
[{"review_id": "mkVl0fF7CBq", "reviewer": "AnonReviewer4", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "Summary:\nThe 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.\n\nReasons for the score:\nI vote for accepting the paper, with minor improvements. The problem is well motivated, and the proposed algorithm improves significantly on previous approaches. I would encourage the authors to include generated samples/visualizations which may reflect some pathological failure modes of the algorithm. This would help with the readability of the paper, and improve understanding.\n\nStrengths:\n+ The CNF framework helps scaling normalizing-flow based approaches, and improves the stability and performance on standard benchmarks.\n+ GraphCNF outperforms baselines like GraphAF/GraphNVP on Zinc250 molecule generation. \n+ On language modelling tasks, CNF improves on Latent NFs and works competitively as an LSTM baseline.\n\nWeaknesses:\n- Plots such as comparing the training loss of different approaches would help understanding the stability of the learning algorithm. \n- Adding some visualizations/samples from the learned models would help with clarity of section5.\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "Y1o98Y3bgSs", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The paper investigates normalizing flows for the general case of categorical data. The authors propose continuous encodings in which different categories correspond to unique volumes in a continuous latent space. \n\nUsing the proposed Categorical Normalizing Flows (CNF) the authors present GraphCNF, a permutation-invariant normalizing flow on graph generation.\n\nIn particular the paper presents a new approach for normalizing flow on categorical data. Then a three step generation approach is presented for graph generation with CNF.\n\nThe paper is well written and the proposed solution is clearly presented in sections 2 and 3.\n\nThe first results on molecule generation are very good when compared to those obtained with other approaches. The second experiment on graph coloring shows the validity of the proposed CNF even on an NP-hard optimization problem. Even the results on language modeling are encouraging. A final experiment shows how CNF can model discrete distributions precisely.\n\nAn experiment could be done on real-world datasets usually used for density estimation in order to assess the validity of CNF in therms of likelihood. \n\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "4-bSS_H_j7", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "# Summary #\n\nThis paper proposes a new normalizing flow model for categorial data, where the typical dequantization is not applicable. \n\nWe assume a categorical sample x has S variables, and each attribute x_i is a categorical variable. \nWe want to model the probability mass of the S variable categorical data, and devise an invertible map that can convert x into the continuous latent variable z. \nFor simplicity, we assume that each attribute x_i has its own latent continuous probability distribution p(z_i). \nWe expect an encoder, q(z_i|x_i), map categorical x into a continuous space where all categories are well partitioned. For that purpose, the paper proposes to formulate q(z_i|x_i) by a mixture of logistic distributions. \n\nA graph generative flow model is proposed as an application. \nExisting flow models for graphs do not handle the categorical data in proper manners and are permutation-dependent. \nThe proposed categorical flow can develop a permutaion-invariant graph generative flow model. \n\nThe proposed model performs better than the existing graph flow models in the molecular graph generation experiments. Typical invalid generation examples include isolated nodes. If we only focus on the largest sub-graphs, the proposed model can almost perfect graph generations. \nThe permutation-invariant nature of the proposed model results in a stable performance on the graph coloring problem, while the baseline RNN models are deeply affected by the order choices. \n\n# Comment\n\nI found the mixture of logistic regression is a good idea. Figure 2 and 3 in appendix indicate that this formulation can pratitioin the latent space into categories. \n\nI have a few questions to confirm my understanding of the paper. \n\nQ1. The proposed categorical normalizing flow with K-logistic mixture provides an approximated invertible map for the true distribution of the categorical samples x. Is this correct? Namely, there is a non-zero KL divergence between the evidence P(X) and the marginalized ``likelihood'' q(X)?? \n\nQ2. The paper says \"we do not have a unknown KL divergence between approximate and true posterior constituting an ELBO\". Does this mean we can compute the KL(q||p) analytically for the categorical normalizing flow? \n\nConcerning the Q2. it is better if the final objective function to maximize/minimize, and the actual procedure for model training is clearly written in the main manuscript or the appendix. \n\nConcerning the molecular graph generation experiments, I'm interested in how the latent representations of the graphs are distributed in the space of Z. It is preferable if the paper can provide a visualization of the latent space for the actual molecular graph generation experiments, not the simulated ones of Figure 2 and 3. \n\nPresentaions of the experimental results in the main manuscript totally rely on the tables. However, current tables are not much effective to tell the significance of the proposed method. \nPlease consider visual presentations: we may use plots or bar graphs to compare several methods for example if the detailed numbers are not important. \nThe actual numbers can be moved to the appendix. \n\n# Evaluation points\n\n(+) A new approach to apply the normalizing flow. \n\n(+) Truly permutation-invariant NF for graph generation is great.\n\n(-) insufficient explanations for the optimization procedure\n\n(-) more visual results may improve impressions of the manuscript (especially for non-expert readers)\n", "extra_scores": {}, "extra_text": {}}, {"review_id": "0_y4dflcdsi", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "### Summary:\n\nThis 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.\n\n### Strength:\n\nI like the idea; it is simple and seems to be very useful in specific problems, especially those without a natural order on categorical variables. The authors also make an excellent effort to gather empirical evidence in favor of their method. The literature review is comprehensive, and the method is well-placed among the cited papers. Several sections are well written and have a nice natural flow. I find the application to graphs an excellent use of the multi-scale architectures in coupling-based flows and compliment the authors for the superb visualization in Figure 1. \n\n### Concerns:\n\nI found the most critical section in the paper least clearly explained--section 2.2. I appreciate the overall idea of this work and believe I understand the work well, yet the use of ill-defined terms and lack of clarity hamper my confidence. \n\nIn my current understanding of the paper, $p(z)$ is modeled using a large flow. Authors use the motivation that $KL[q||p]$ is negligible and approximate $p(x_i|z_i)$ using $q(z_i|x_i)$. With this, eq. 2 can be used to train; however, I am unsure if that is indeed the case; I will appreciate it if the rebuttal absolutely clarifies this.  To make it easier for the authors, I effort to pinpoint the specific instances where they lose me in section 2.2. \n\nIn the first paragraph, I find the last line: \"Specifically, as normalizing ..\" a bit hard to parse. In particular, the term \"one distribution\" is unclear. The next paragraph tries to explain why the decoder becomes almost deterministic; despite using words like \"therefore\" and \"as a consequence,\" I found the explanations insufficient. \n\nThe term \"overall encoding distribution\" is unspecified in the paper,  and I am skeptical of the use of \"true decoder likelihood.\" Maybe I am missing something obvious, but why does the expression for \"true decoder likelihood\" involve the variational approximation? Is it because we work under the assumption that the KL divergence is negligible? \n\nAlso, I am unsure what the objective is when you have \"removed\" the ELBO (see the first line, last paragraph, section 2.2.). \n\n \n\n#### Minor bits:\n\"... we do not have *an unknown ...\" --using an in place of a.\n\n### Updates after the rebuttal\n\nI find the revised manuscript to be clear and more transparent. After reading the reviews, the response, and the extended appendix, I am increasing my score and vote to accept this paper. \n\n\n\n", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Well-written paper that proposes a flow-based model for categorical data, and applies it to graph generation with good results. Extending flow-models to handle types of data that are not continuous is a useful contribution, and graph generation is an important application. Overall, the reviewers were positive about the paper, and only few negative points were raised, so I'm happy to recommend acceptance.
Decision: Accept (Poster)