Discrete Tree Flows via Tree-Structured Permutations
Abstract
While normalizing ﬂows for continuous data have been extensively researched, ﬂows for discrete data have only recently been explored.
These prior models, however, suffer from limitations that are distinct from those of continuous ﬂows.
Most notably, discrete ﬂow-based models cannot be straightforwardly optimized with conventional deep learning methods because gradients of discrete functions are undeﬁned or zero, and backpropagation can be computationally burdensome compared to alternative discrete algorithms such as decision tree algorithms.
Previous works approximate pseudogradients of the discrete functions but do not solve the problem on a fundamental level.
Our approach seeks to reduce computational burden and remove the need for pseudo-gradients by developing a discrete ﬂow based on decision trees—building upon the success of efﬁcient treebased methods for classiﬁcation and regression for discrete data.
We ﬁrst deﬁne a tree-structured permutation (TSP) that compactly encodes a permutation of discrete data where the inverse is easy to compute; thus, we can efﬁciently compute the density value and sample new data.
We then propose a decision tree algorithm to learn
TSPs that estimates the tree structure and simple permutations at each node via a novel criteria.
We empirically demonstrate the feasibility of our method on multiple datasets.
1. Introduction
Discrete data is abundant in numerous applications and domains, from DNA sequences and medical records to text
Third workshop on Invertible Neural Networks, Normalizing Flows, and Explicit Likelihood Models (ICML 2021).
Copyright 2021 by author data and many forms of tabular data.
Learning a density model for discrete data is a fundamental task useful for understanding or leveraging such data.
However, discrete data can be hard to model especially when the number of categories and the dimension of the data is large.
Prior methods include probabilistic graphical models (Inouye et al., 2014; 2017; Wainwright & Jordan, 2008), VAEs (Kingma & Welling, 2019), or autoregressive models (Germain et al., 2015).
However, these works either lack exact likelihood computation or require expensive sampling procedures.
As an alternative, discrete ﬂows provide exact likelihood computation and efﬁcient sampling.
The two key components of discrete ﬂows are a base distribution (similar to continuous ﬂows) and a permutation of the discrete conﬁguration values (a discrete version of invertible functions)
(Tran et al., 2019).
In its full generality, this requires optimizing over the set of all permutations, whose size is doubly exponential in terms of the feature dimension.
Thus, parameterizing and optimizing over a set of possible permutations is critical for practical algorithms and generalizability.
Tran et al.
(2019) propose to parameterize these permutations by neural networks and optimize them using standard backpropagation.
Tran et al.
(2019) estimate pseudo-gradients of discrete functions using the straightthrough gradient estimator (Bengio et al., 2013) along with a Gumbel-softmax distribution for backpropagation, which is an approximation of the forward pass equation (i.e., onehot of argmax function).
van den Berg et al.
(2020) later show that the architecture of the coupling layers is significantly more important than the gradient bias issue.
However, none of these discrete ﬂow methods address the discrete nature of the problem on a fundamental level and may be computationally expensive compared to alternative discrete-oriented algorithms such as those based on decision trees—which have seen wide success in classiﬁcation and regression for discrete data (e.g., XGBoost (Chen &
Guestrin, 2016)).
Thus, we seek to answer the following research question: Can we design a more computationally efﬁcient discrete ﬂow algorithm using decision trees that handles discrete data on a fundamental level?
To answer this, we propose a novel tree-structured permutation (TSP) model that can compactly represent permutations and a novel decision tree algorithm that optimizes over the space of these permutations.
Moreover, for more powerful permutations, we can iteratively build up a sequence of TSPs to form a deep permutation.
We also propose a novel decision tree algorithm for learning TSPs that includes a novel splitting criteria and an efﬁcient algorithm for determining the best permutation at each tree node.
Background: Discrete Flows The Discrete ﬂow change of variables formula is given by: P x ( x ) =
Q z ( f ( x )) , where Q z is some known base distribution that is usually a simple distribution (e.g. an independent categorical distribution), and f is an invertible function.
Importantly, in the discrete case, the invertible function can only be a permutation of the possible discrete conﬁgurations of x so there is no change in volume, hence no need to incorporate the
Jacobian determinant term as done in continuous ﬂow.
In its full generality, discrete ﬂows would require optimizing over the set of all possible permutations, whose size is doubly exponential in terms of the dimension—which emphasizes the computational intractability of the general problem.
Thus, two key ingredients for practical discrete ﬂows are a simple and compact parameterization of permutations and an algorithm to optimize over these permutations.
Tran et al.
(2019) propose two different architectures (autoregressive ﬂows and bipartite ﬂows) to parameterize the discrete ﬂow permutations.
Hoogeboom et al.
and van den Berg et al.
(2020) propose integer discrete ﬂows for compression of image-based data (rather than categorical tabular data as in this paper).
2. Tree-Structured Permutations
Our goal is to deﬁne a set of permutations that is both computationally tractable (in terms of evaluating the permutation and its inverse and optimizing over the set) and generalizable (i.e., the class of permutations is more likely to generalize well to new test data).
To achieve this goal, we introduce tree-structured permutations (TSP) as new model for discrete ﬂows that utilizes trees for compactly parameterizing a set a of permutations.
Notation We will denote a discrete dataset as X ∈ Z n × d where n is the number of samples, d is the number of dimensions, and Z is a set of discrete values, and where k is the maximum number of possible discrete values per feature (i.e., the number of categories).
We will denote tree nodes by N .
The split information at each node N will be encoded by a feature index s ∈ { 0 , 1 ...., d − 1 } and a split value v such that data with the s -th feature equal to v will go to the left node and all other data will go to the right node.
We deﬁne the node domain , denoted by D ( N ) , to be the set of all discrete conﬁgurations that could reach this node when traversing the decision tree.
We will denote permutations by π ( · ) ∈ Π , where Π will denote a set of permutations.
Given that the number of discrete conﬁgurations of d features with k possible discrete values is k d , the number of all possible permutations is ( k d )! , which is
O (exp( k d )) , i.e., doubly exponential in d .
We will denote the parameters of the permutation by θ π , where the exact parametrization may depend on context.
2.1. Deﬁnition of TSP
A tree-structured permutation is a binary decision tree where each node N contains both a permutation π and the usual split information (i.e., a split feature s and split value v ).
To evaluate a TSP, an input vector traverses the tree from the root to a leaf node based on the split information and applies node permutations as soon as it reaches the node (i.e., the node permutation will be applied before determining the split).
More formally, we can deﬁne the evaluation of a TSP recursively as the evaluation of a node f N where the initial N is the root node:
[Equation]
where f N is the evaluation of a node, f left ( N ) denotes the evaluation of the left child node (and similarly for the right node), π N is the permutation associated with the node, and [ π N ( x )]
j = v denotes the condition that the j -th feature after the permutation has value v .
We illustrate the idea of the forward traversal with an example in Figure 1.
Note that even though each node’s permutation is invertible, it is unclear whether the entire forward evaluation of a TSP is invertible (i.e., if the path and sequence of permutations implied by the tree structure can be inverted)—which is critical for their use in discrete ﬂows.
Thus, we explore the conditions for the invertibility of TSPs.
2.2. Invertibility Constraint and Calculation for TSPs
To ensure our TSPs are invertible (and thus applicable to discrete ﬂows), we prove that a simple and intuitive con- straint on the node permutations is sufﬁcient for invertibility (in the appendix, we derive a necessary and sufﬁcient condition for invertibility but it is challenging to check or enforce).
Theorem 1 (TSP Invertibility Constraint) .
A TSP is invertible if all node permutations π N do not permute conﬁgurations that are outside of the node’s domain, i.e., π N ( x ) = x, ∀ x (cid:54)∈ D ( N ) .
We will denote the set of permutations that satisfy this constraint for a node N as Π( N ) .
The proof of the if direction is constructive and relies on the following lemma.
Lemma 2.
If this invertibility constraint is satisﬁed, the
TSP tree traversal path for any input can be recovered from the output.
Leveraging Lemma 2, the inverse can be computed by traversing the tree in the forward direction (i.e., from root to leaves) without applying the node permutations but keeping track of the node path.
Once we reach a leaf, we apply the permutations along the path in reverse order to compute the inverse.
We illustrate this inverse computation using the example output point z
= [2 , 0] from Figure 1.
Starting at the root, we see that because z 0 (cid:54) = 0 , z will move right to node 2.
Then, at node 2, the node permutation will be ignored (for the inverse) and the split condition again will be checked on z such that z will then move to node 4 because z 1 (cid:54) = 1 , etc.
Finally, the inverses of permutations along the path will be applied in reverse order.
2.3. Na¨ıve TSP: Restricting to the class of independent permutations
For computational tractability and generalizability, we will further restrict the possible permutations at each node to be feature-wise independent (i.e., the permutation for one feature does not depend on values of other features) and, for k > 2 , value-pair swaps (i.e., permute by swapping one pair of categorical values).
We will denote these restrictions as π ∈ Π IndValPair , where the size of the permutation space is O ( k 2 d ) .
Because of independence of the permutation and the base distribution, however, the computational cost can be reduced to O ( dk 2 ) (more details in the appendix).
Additionally, while we could permute all d features independently at each tree node, we choose to only permute the single best feature at each node to avoid overﬁtting the dataset, denoted by Π TSP ( N , s ) , where s is the feature index of the single best feature to permute.
An example of this na¨ıve TSP is shown in Figure 1.
To increase model expressiveness despite these restrictions, we propose to compose a sequence of na¨ıve TSPs into a deep
TSP permutation that we call Discrete Tree Flows (DTF).
3. Learning TSP Discrete Flows
We will present a greedy approach for building up the tree similar to other decision tree methods where the goal is to minimize the negative log likelihood (NLL) over na¨ıve TSP permutations and independent base distributions Q z given our dataset X :
[Equation]
where L ( π, Q z ) denotes the negative log-likelihood, Q Ind is the set of independent distributions over categorical data, and Π TSP is the set of na¨ıve TSP permutations.
Assuming that the base distribution is independent enables efﬁcient methods for determining the best permutations and the best features using only feature-wise category counts at each node.
Additionally, an independent base distribution enables fast sampling, unlike autoregressive base distributions or other complex base distributions.
Our proposed learning algorithm can be decomposed into node-wise subproblems where two main steps are required.
First, we need to determine the best permutation for each node.
Second, as with all decision tree algorithms, we need to determine the best node to split among all current leaf nodes (i.e., a splitting criteria).
We describe our objectives for both of these steps in the next section.
3.1. Permutation Criteria for Na¨ıve TSP construction
Our method of na ¨ ıve TSP construction revolves around building a tree that minimizes the negative log-likelihood (NLL) of our training data.
We select the best singlefeature node permutation π that minimizes the change in the NLL:
[Equation]
where Q (cid:48)
z is a reﬁtted base distribution after applying permutation π , id denotes that no permutation is applied, and Π TSP ( N , s ) denotes the class of independent permutations that act only on the domain of the node (so that the TSP will be invertible) and restricted to the s -th feature.
Note that because Q (cid:48) z is an independent base distribution, the optimal Q (cid:48)
z can be computed efﬁciently only using the empirical (and smoothed) frequencies of each feature for the training data at the current node.
Given that each feature is independent (both in the permutation and the base distributions), we can greatly simplify the computation of ∆ L (details in appendix).
3.2. Splitting Criteria for Na¨ıve TSP construction
To determine the best node split, we consider using a splitting criteria that maximizes the difference between the fac- torized distribution on the left and the factorized distribution on the right of the proposed split.
Intuitively, if the distributions on the right and left are different, then a permutation will be able to align them better (and thus help reduce NLL).
Let us denote Q left ( s, v ) and Q right ( s, v ) to be the best independent distributions on the left and right of a proposed split parametrized by s and v .
We want to maximize the divergence between these two distributions, where we will use generalized Jensen-Shannon Divergence (JSD).
The full derivation is given in the appendix, and the splitting criteria is:
[Equation]
where w 1 and w 2 are the relative empirical probability of data going left versus right.
As is standard in decision tree algorithms, we also take the minimum over all current leaf nodes N to select the next best leaf node to split.
3.3. TSP construction: Training the model
Training our model entails greedily constructing the TSP based on our training data by determining the node permutation, split feature s and split value v at each node N .
The tree construction takes in a maximum height of the tree given by the user ( M ), and constructs the TSP nodes node by node, until the maximum height is reached or until no further splits can be made (which can be due to reaching a very small number of samples at each node, or running out of viable split features or values).
At a low-level, the algorithm mainly requires keeping track of feature-wise counts at each node and a global count vector to optimize each criteria.
The computational complexity for the algorithm is
O (2 2 M nd 2 k 2 ) , where M is the max depth of the tree, but the quadratic d 2 term can be reduced to d using a Monte
Carlo approximation (details can be found in the appendix).
4. Experiments
We compared our model against the two models introduced in (Tran et al., 2019) and implemented in (Bricken, 2021).
All the models including ours use an independent base distribution.
We use the following abbreviations in our comparisons, AF for autoregressive ﬂow, BF for bipartite ﬂow.
The BF model implemented in (Bricken, 2021), however, had some issues and so we implemented some modiﬁcations to it (that we mention in more details in the appendix).
For the ﬁrst set of experiments (labelled exp=1 through 5), we generated synthetic data X ∈ Z n × d with k categories and split the data into training and validation datasets using cross validation with 5 folds.
For d = 2 and k = 2 , we chose to generate data that follows a known probability distribution.
Speciﬁcally, for exp = 1, d = 2 , k = 2 , we use the probabilities Pr([0 , 0])
= 1 / 3 , Pr([0 , 1])
= 1 / 6 , Pr([1 , 0])
= 1 / 6 , Pr([1 , 1]) = 1 / 3 , and for exp = 2, d = 2 , k = 2 , we use the probabilities Pr([0 , 0])
= 1 / 8 , Pr([0 , 1]) = 3 / 8 , Pr([1 , 0])
= 3 / 8 , Pr([1 , 1])
= 1 / 8 .
For the other combinations of d and k , we simulated the probability distribution over all possible conﬁgurations as a Dirichlet distribution (with α = 1 ) and we simulate n = 10000 data points.
We also ran an experiment for a real discrete dataset 1 , that has n = 8124 datapoints which is denoted “exp=6”.
For all exps we report the mean and standard deviation of NLL across the 5 test folds.
For the
DTF, the results are for a ﬂow of 1 TSP where the maximum depth M of the TSP is 2 for experiments 1 and 2, and
M = 4 for the rest.
For the AF ﬂows, we used a single ﬂow, that utilizes 64 hidden units (in the MADE layers).
We ran this model for a total of 1500 epochs, sampling 1024 data points in each epoch.
Corresponding to the conﬁguration in (Tran et al., 2019), 4 ﬂow layers were used for the BF model.
As can be observed from Table 1, our method gives comparable results to that of AF and BF models.
[Table]
For exp = 6, our model took 5.4372 sec for training (with an Std of 2.463 sec), the AF model took 39.5838 sec (with an Std of 0.2519 sec), and the BF model took 143.8929 sec (with an Std of 2.2670 sec).
For exp 1-5, timing results are available in the appendix.
Thus, our approach seems to be computationally less expensive than prior methods as we had hoped though further experimentation is likely needed.
5. Discussion
We presented a novel framework for discrete normalizing ﬂows that relies on tree-structured permutations (TSPs), which we deﬁne and develop.
Our model isn’t without limitations though, as our implementation relies on a greedy algorithm and thus may end up not achieving a global optimum when learning deep nai¨ve TSPs.
This is a problem that is similar in nature to problems faced by most decision tree algorithms.
Moreover, handling very high dimensional data may be computationally challenging.
However, previous decision tree algorithms were able to overcome these obstacles, and we expect that similar techniques could be used in our method.
Ultimately, we hope that our paper lays the groundwork for developing practical and effective discrete ﬂows using decision tree algorithms.
6. Acknowledgement
The authors acknowledge support from the Army Research
Lab through Contract number W911NF-2020-221.
A. Invertibility Proofs
A.1.
Proof of Lemma 2 in Section 2.
Lemma 2.
If the invertibility constraint is satisﬁed, the
TSP tree traversal path for any input can be recovered from the output.
Proof.
First, note that because π is a permutation (i.e., oneto-one mapping) and π N ( x ) = x, ∀ x (cid:54)∈ D ( N ) , then ∀ x ∈
D ( N ) , π N ( x ) ∈ D ( N ) , i.e., all node permutations π N do not permute conﬁgurations from in the domain to outside the domain.
Without loss of generality, we consider trees where all leaf nodes are at the max depth of M .
We will denote the tree traversal path of an input x up to tree level i to be P
i ( x ) (cid:44) ( N L 0 ( x ) , N L 1 ( x ) , · · · , N L i ( x ) ) , where N L j ( x ) is the tree node that x reaches at j -th level of the tree (i.e., L j ( x ) gives the index of the node for the particular input x at the j -th level), N L (0) is the root node.
Let y (cid:44) π TSP ( x ) denote the output of our TSP forward evaluation.
For all x ∈ Z d , let y (cid:44) π TSP ( x ) ≡ π N LM ( x ) ◦· · ·◦ π N L 1( x ) ◦ π N L 0( x ) ( x ) denote the output of our TSP forward evaluation.
We want to prove that P M ( x ) can be recovered from y and π TSP .
Let P (cid:48) i ( y ) (cid:44) ( N L (cid:48)
( y ) , N L (cid:48)
( y ) , · · · , N L (cid:48)
i ( y ) ) where we traverse the decision tree of π TSP without applying node permutations and where L (cid:48)
j ( y ) is the index of that y reaches at the j -th layer.
If we can prove P (cid:48)
M ( y ) =
P M ( x ) or equivalently ∀ j, L (cid:48)
j ( y ) =
L j ( x ) , then we are done.
Inductive hypothesis For all i ∈ { 0 , 1 , · · · , M } ,
P (cid:48) i ( y ) =
P i ( x ) .
Base case ( i = 0 )
Since both y and x start at the root node, then the path up to i = 0 is the same.
Induction step We need to prove that if P (cid:48) i ( y ) =
P i ( x ) , then P (cid:48)
i +1 ( y ) =
P i +1 ( x ) .
From assumption of inductive hypothesis, we know that N L (cid:48)
i ( y ) =
N L i ( x ) (i.e., x and the corresponding y are the same node for level i ).
Let x ( i ) =
π N Li ( x ) ◦· · ·◦ π N L 1( x ) ◦
π N L 0( x ) ( x ) where x ( M ) ≡ y .
Note that only the value v of the split feature s for both x ( i ) and y is relevant for determining whether to go left or right.
We prove that the chosen nodes are the same using contradiction.
Suppose x ( i ) s
= v such that x goes left and suppose y s (cid:54) = v such that y would go right.
We know that the s -th part of the domain of the left child has only v in it, i.e., D s ( N left ) = { v } .
Also, we know that domains of children are always smaller disjoint subsets of the parent, i.e., D s ( N L N ( x ) ) ⊆
D s ( N left ) for all N > i .
Thus, x ( N ) s = v for all N > i because the invertibility constraint ensures that we cannot permute a value inside the domain to a value outside the domain.
However, this is a contradiction to our assumption that x ( M ) s ≡ y s (cid:54) = v .
Therefore, if x goes left, then y will also go left.
In a similar way, now suppose x ( i ) s (cid:54) = v such that x goes right and suppose y s = v such that y would go left.
The sth part of the domain of the right child has D s ( N right ) =
{ a : a (cid:54) = v, a ∈ D s ( N ) } .
Again, D s ( N L N ( x ) ) ⊆
D s ( N right ) for all N > i because every child is a subset of the parent domain.
Therefore, x ( N ) s ∈ D s ( N right ) for all N > i , and thus in particular x ( M ) s ∈ D s ( N right ) , where x ( M ) s ≡ y s by deﬁnition.
However, this contradicts our assumption that y s = v (i.e., goes left) because v (cid:54)∈ D s ( N right ) .
Hence, if x goes right, y will also go right.
Combining these two we get that L (cid:48)
i +1 ( y ) =
L i +1 ( x ) (i.e., they will both go left or both go right), and thus we can recover the path for i +1 by adding the child node to the path for i , i.e., P (cid:48)
i +1 ( y ) =
P i +1 ( x ) .
This proves our inductive step and concludes the proof of the lemma.
A.2. Proof of Theorem 1 in Section 2.
Theorem 1 (TSP Invertibility Constraint (sufﬁcient only)) .
A TSP is invertible if all node permutations π N do not permute conﬁgurations that are outside of the node’s domain, i.e., π N ( x ) = x, ∀ x (cid:54)∈ D ( N ) .
Proof.
Lemma 2 (proven above) states that the TSP traversal path for any input can be recovered from the output.
Thus, the output path is identical to the input path
P (cid:48) i ( y ) =
P i ( x ) .
Therefore, x = π − 1 TSP ( y ) ≡ π − 1 N L (cid:48)
( y ) ◦· · ·◦ π − 1 N L (cid:48)
M − 1( y ) ◦ π − 1 N L (cid:48)
M
( y ) ( y ) because each node permutation is itself invertible by the deﬁnition of a permutation.
This can be seen as traversing the tree from the corresponding leaf node to the root node and applying the inverse node permutations along the path.
A.3. Necessary and Sufﬁcient Condition for Invertibility with Proof
We now present a corrected condition (i.e., disjoint ranges of leaf nodes) that is both necessary and sufﬁcient for invertibility.
Note that this new theorem can be easily used to prove Theorem 1 as a corollary because the original invertibility constraint set is a subset of the disjoint range of leaf nodes constraint.
However, the proof here does not provide an efﬁcient algorithm for determining the leaf node while the proof of Lemma 2 does provide an efﬁcient algorithm (i.e., merely traverse the tree as described in the lemma proof to determine the path).
First, we deﬁne the range of a node to simplify the deﬁnition of the theorem and proof.
Deﬁnition 1 (Range of a Node) .
We deﬁne the range of a node , denoted R ( N ) , as the image of D ( N ) under π N , i.e., R ( N ) (cid:44) { π N ( x ) : x ∈ D ( N ) } .
Now we present our theorem with a corrected condition that is both necessary and sufﬁcient.
Theorem 3 (TSP Necessary and Sufﬁcient Invertibility
Constraint) .
A TSP is invertible if and only if the range of each leaf node is disjoint from all other leaf nodes, i.e.,
[Equation]
T leaf such that N (cid:54)
[Equation]
N
[Equation]
, where T leaf are the set of leaves in the TSP tree.
Proof.
We use a constructive proof for the if direction (sufﬁciency).
Because each of the permutations themselves are invertible, the primary challenge is showing that we can ﬁnd the right path through the tree (as there could be multiple paths without constraints).
If the disjoint range condition is satisﬁed, then each possible output can be mapped to one of the leaves, i.e., N L M ( y ) is the leaf node such that y ∈ R ( N ) .
Given the leaf node, there is only one possible path through the decision tree back to the root.
Thus, the inverse can be computed by traversing from the leaf node to the root node and applying the inverse of each node’s permutation.
To prove the only-if direction (necessity), we will use a proof by contradiction.
Suppose a TSP is invertible but the disjoint range condition is not satisﬁed, then R ( N ) ∩
R ( N (cid:48) )
(cid:54) = ∅ .
Therefore, there exists an output y that is in the range of two leaf nodes, i.e., ∃ y such that y ∈ R ( N ) and y ∈ R ( N (cid:48) )
where N (cid:54) =
N (cid:48) .
Yet, each input traverses the TSP tree in a deterministic way and thus each unique input will always arrive at the same leaf node.
Therefore, there must exist two distinct inputs x (cid:54) =
x (cid:48)
such that π TSP ( x ) =
π TSP ( x (cid:48) )
=
y .
This means that two distinct inputs map to the same output (i.e., not one-to-one mapping) and violates invertibility.
However, this contradicts our assumption that the TSP is invertible.
B. A discussion of classes of permutations
Since the number of discrete conﬁgurations of d features with k possible discrete values is k d , the number of all possible permutations is ( k d )! , which is O (exp( k d )) , i.e., doubly exponential in d .
We chose to make some restrictions to our general permutations class as described below.
Independent feature-wise permutations We choose to restrict to the natural and computationally tractable class of independent feature-wise permutations, denoted Π Ind , which allows each feature to be permuted independently of the other features (i.e., the permutation of one feature cannot depend on the permutations of other features).
This class signiﬁcantly reduces the number of permutations compared to all possible permutations, i.e., | Π Ind | = ( k !)
d ≈ exp( kd ) (cid:28) exp
( k d ) ≈ ( k d )!
= | Π | .
For example, for binary data ( k = 3 ) and three features ( d = 5 ), we can have 3!
permutations for each of the ﬁve features, leading to a total of (3!)
= 7776 joint permutations—which is considerably smaller than the total number of permutations (3 5 )!
= 243!
≈ 10 474 .
This makes the independent class of permutations significantly more computationally tractable and more likely to generalize to new data points.
Additionally, if we assume that the base distribution Q z is independent, it is possible to separate the optimization problem into subproblems for each feature that can be solved independently.
Thus, each feature subproblem only needs to evaluate k ! permutations, and thus the computational complexity for our problem can be reduced to O ( d · k !) .
Permutations that swap a single pair of values for k >
While restricting to independent permutations Π Ind ( N ) can signiﬁcantly reduce our computational complexity to
O ( dk !) , the computationally complexity would still grow exponentially in terms of k , the number of features.
Thus, we further restrict the class of node permutations to be the class of permutations that swaps a single pair of the k possible values (or categories) while holding all other feature values constant.
We will denote this class by Π IndValPair and this reduces number of possible permutations per feature from k ! to (cid:0) k 2 (cid:1) + 1 which is O ( k 2 ) (where the +1 is to include the identity permutation).
Thus, with this restriction and our assumption about Q z being independent, we can reduce the computational complexity to O ( dk 2 ) , which is computationally in all parameters including k .
As an example, if k = 3 , then the possible permutations in this class are (012 , 102 , 210 , 021) which correspond to the original permutation, swapping 0 and 1, swapping 0 and 2, and swapping 1 and 2.
As mentioned in the main section, in our Na¨ıve TSP implementation, we restrict the node permutations to π ∈ Π IndValPair .
Additionally, while we could permute all d features independently at each tree node, we choose to only permute the single best feature at each node to avoid overﬁtting the dataset; we leave exploration of permuting multiple features to future work.
We denote this single-best permutation class as Π TSP , where s is the feature index of the single best feature to permute.
An example of this na¨ıve TSP is shown in Figure 1.
C. Derivation for Calculating the difference in NLL
We can greatly simplify the computation of ∆ L as follows:
[Equation]
Where id denotes that no permutations applied are applied (the identity), and Q (cid:48)
z is the independent base distribuition estimate after θ π is applied.
Thus, while constructing the tree, for every node N we seek to minimize the following term.
[Equation]
Where s ∗ is the best permutation feature and π ∗ is the best permutation to apply.
To simplify Eqn9, we notice that it can be rewritten if we applied the following properties, for a particular feature (s):
[Equation]
Where n c is the number of samples that has the category c in the s dimension index.
This property can be used in Eqn.9 to re-write it as:
[Equation]
Where n (cid:48) c is the number of samples with category c in the s dimension after applying the permutation
π s
Furthermore, since for nai ¨ ve TSPs, we have restricted our permutations to Π TSP which not only allow us to be computationally more efﬁcient but also allow us to further simplify the calculation of ∆ .
Let’s say that the permutation π s ﬂips values k = v 1 (that appears on the left child) with k = v 2 (that appears on the right child) then the number of counts only in the categories that are equal to v 1 and v 2 will change before and after applying the permutation, while the counts of the remaining categories will remain constant.
This means that Eqn(15) can be re-written as
[Equation]
where n (cid:48) v 1 and n (cid:48) v 2 are number of counts of categories vand v 2 respectively after applying the permutation π s and n v 1 and n v 2 are number of counts of categories v 1 and vrespectively before applying the permutation π s .
For all the other categories where k (cid:54) = v 1 and k (cid:54) = v 2 , n (cid:48)
c will be equal to n c , resulting in the summation in Eqn(15) being zero for those categories.
D. Derivation for calculating the splitting criteria
As mentioned in the main section, to determine the best node split, we consider using a splitting criteria that maximizes the difference between the factorized distribution on the left and the factorized distribution on the right of the proposed split.
Intuitively, if the distributions on the right and left are different, then a permutation will be able to align them better (and thus help reduce NLL).
Let us denote Q left ( s, v ) and Q right ( s, v ) to be the best independent distributions on the left and right of a proposed split parametrized by s and v .
We want to maximize the divergence between these two distributions, where we will use generalized Jensen-Shannon Divergence (JSD).
The generalized JSD can be written either as a weighted sum of KL divergences or a difference in entropy terms:
[Equation]
This second version can be seen as the entropy of the mixture minus the mixture of the entropies.
Given this deﬁnition, we can now deﬁne our splitting problem as:
[Equation]
where Q parent = w 1 Q ( s,v ) left
[Equation]
Q
[Equation]
s
[Equation] [Equation]
smoothed frequencies), and n left s,v are the number of samples where the s -th feature equals the value v on the left side— and similarly for the right side—, and where w ( s,v ) =
[Pr( x s = v ) , Pr( x s !
= v )] is the probability vector encoding the probability that samples go to the left or the right side of the split.
Notice that the entropy of the parent node is not required to optimize the JSD.
E. Discussion of Algorithmic complexity
The algorithmic complexity for our algorithm will depend on the number of leaves at any point during the tree construction which at its worse would be 2 M where M is the maximum depth of the tree.
The computational complexity of determining the best permutation for a node is O ( dk 2 ) , which is due to being able to solve each feature independently and considering only a single swap of feature values.
Thus the computational complexity of ﬁnding the node permutations for the whole tree is O (2 M dk 2 ) .
The complexity of ﬁnding the best split at a node is O ( nd 2 k ) which is due to the entropy calculation that must be computed over all features for every possible split value.
However, we can choose to make a Monte Carlo approximation to the entropy calculation by only sampling a relatively small constant number of dimensions (e.g., ten or twenty) and take the average entropy over these sampled dimensions rather than over all dimensions.
Thus, we would get an unbiased approximation to the entropy—similar in spirit to SGD.
This approximation allows for the complexity of ﬁnding the best split to be O ( ndk 2 M ) and since this must be done for all current leafs a the complexity is O ( ndk 2 2 M ) , which makes the overall complexity of constructing a na¨ıve TSP construction O (2 M dk 2 + 2 2 M ndk 2 ) .
Once a Na¨ıve TSP is constructed, passing a data in the forward direction is just a matter of the traversing the tree from root to leaf while applying the permutations stored in the tree node, which is dependant in complexity on the depth of the tree, and calculating the inverse of the data utilizes the same idea with the exception of needing additional memory to keep track of the permutations that we encounter, so both operations are O ( M ) in time complexity.
F. Modiﬁcation of the Bipartite ﬂow code
The BF model implemented in (Bricken, 2021), had some bugs in the code (e.g., running the bipartite model for data dimensions d > 2 would yield a runtime error).Thus, modiﬁcations were made to ﬁx the bipartite model code.
The model’s initial embedding ﬂow layers (commonly used for
NLP sequence data) were replaced by a single hidden layer network with an activation function.
A ReLU activation function was incorporated into the hidden layer to add nonlinearity.
The size of the hidden layer was proportional to the feature size times half the dimension ( 12 ∗ k ∗ d ) , since only half of the dimension would be mapped after the split.
Each ﬂow layer does a transformation on one of the splits, therefore, at least a paired ﬂow layers (even number of ﬂow layers) is required for a balanced dimensional transformation.
G. Timing Results for exp = 1 through
[Table]