Original: I NTERPRETABLE C LASSIFICATION VIA S UPERVISED V ARIATIONAL A UTOENCODERS AND D IFFERENTIABLE D ECISION T REES
Revised: I NTERPRETABLE C LASSIFICATION VIA S UPERVISED V ARIATIONAL A UTOENCODERS AND D IFFERENTIABLE D ECISION T REES

Original: Anonymous authors
Revised: Anonymous authors

Original: Paper under double-blind review
Revised: Paper under double-blind review

Original: A BSTRACT
Revised: A BSTRACT

Original: As deep learning-based classiﬁers are increasingly adopted in real-world applications, the importance of understanding how a particular label is chosen grows.
Revised: As deep learning-based classiﬁers are increasingly adopted in real-world applications, the importance of understanding how a particular label is chosen grows.

Original: Single decision trees are an example of a simple, interpretable classiﬁer, but are unsuitable for use with complex, high-dimensional data.
Revised: Single decision trees are an example of a simple, interpretable classiﬁer, but are unsuitable for use with complex, high-dimensional data.

Original: On the other hand, the variational autoencoder (VAE) is designed to learn a factored, low-dimensional representation of data, but typically encodes high-likelihood data in an intrinsically non-separable way.
Revised: On the other hand, the variational autoencoder (VAE) is designed to learn a factored, low-dimensional representation of data, but typically encodes high-likelihood data in an intrinsically non-separable way.

Original: We introduce the differentiable decision tree (DDT) as a modular component of deep networks and a simple, differentiable loss function that allows for end-to-end optimization of a deep network to compress high-dimensional data for classiﬁcation by a single decision tree.
Revised: We introduce the differentiable decision tree (DDT) as a modular component of deep networks and a simple, differentiable loss function that allows for end-to-end optimization of a deep network to compress high-dimensional data for classiﬁcation by a single decision tree.

Original: We also explore the power of labeled data in a supervised VAE (SVAE) with a Gaussian mixture prior, which leverages label information to produce a high-quality generative model with excellent values of negative log-likelihood.
Revised: We also explore the power of labeled data in a supervised VAE (SVAE) with a Gaussian mixture prior, which leverages label information to produce a high-quality generative model with improved bounds on log-likelihood.

Original: We combine the SVAE with the DDT to get our classiﬁer+VAE (C+VAE), which is competitive in both classiﬁcation error and NLL, despite optimizing both simultaneously and using a very simple encoder/decoder architecture.
Revised: We combine the SVAE with the DDT to get our classiﬁer+VAE (C+VAE), which is competitive in both classiﬁcation error and log-likelihood, despite optimizing both simultaneously and using a very simple encoder/decoder architecture.

Original:  I NTRODUCTION
Revised:  I NTRODUCTION

Original: While deep learning approaches are very effective in many classiﬁcation problems, interpretability of the classiﬁer (why a particular classiﬁcation was made) can be very difﬁcult, yet critical for many applications.
Revised: While deep learning approaches are very effective in many classiﬁcation problems, interpretability of the classiﬁer (why a particular classiﬁcation was made) can be very difﬁcult, yet critical for many applications.

Original: Decision trees are highly interpretable classiﬁers, so long as the data is encoded such that the classes can be easily separated.
Revised: Decision trees are highly interpretable classiﬁers, so long as the data is encoded such that the classes can be easily separated.

Original: We present a differentiable decision tree (DDT) that we connect to a variational autoencoder (VAE) to learn an embedding of the data that the tree can classify with low expected loss.
Revised: We present a differentiable decision tree (DDT) that we connect to a variational autoencoder (VAE) to learn an embedding of the data that the tree can classify with low expected loss.

Original: The expected loss of the DDT is differentiable, so standard gradient-based methods may be applied in training.
Revised: The expected loss of the DDT is differentiable, so standard gradient-based methods may be applied in training.

Original: Since we work in a supervised learning setting, it is natural to exploit the label information when training the VAE.
Revised: Since we work in a supervised learning setting, it is natural to exploit the label information when training the VAE.

Original: Thus, we employ a supervised VAE (SVAE) that uses a class-speciﬁc Gaussian mixture distribution as its prior distribution.
Revised: Thus, we employ a supervised VAE (SVAE) that uses a class-speciﬁc Gaussian mixture distribution as its prior distribution.

Original: We found that the SVAE was very effective in exploiting label information, resulting in a large improvement in negative log-likelihood (NLL) over approaches that ignore class information.
Revised: We found that the SVAE was very effective in exploiting label information, resulting in improved log-likelihood due to separation of classes in latent space.

Original: Further, when we combined SVAE with DDT (yielding our Classiﬁer+VAE, or C+VAE), we got a model that isvery competitive in both classiﬁcation error and NLL, despite optimizing both simultaneously and using a very simple encoder/decoder architecture.
Revised: Further, when we combined SVAE with DDT (yielding our Classiﬁer+VAE, or C+VAE), we got a model that is competitive in both classiﬁcation error and log-likelihood, despite optimizing both simultaneously and using a very simple encoder/decoder architecture.

Original: Further, the resultant decision tree revealed clear semantic meanings in its internal nodes.
Revised: Further, the resultant decision tree revealed clear semantic meanings in its internal nodes.

Original: 1. O UR C ONTRIBUTIONS
Revised: 1. O UR C ONTRIBUTIONS

Original: Our ﬁrst contribution is a demonstration of the power of labeled data in autoencoding, using a simple class-based Gaussian mixture model as a VAE’s prior, trained with fully labeled data.
Revised: Our ﬁrst contribution is a demonstration of the power of labeled data in autoencoding, using a simple class-based Gaussian mixture model as a VAE’s prior, trained with fully labeled data.

Original: Speciﬁcally, our VAE objective function regularizes w.r.t.
Revised: Speciﬁcally, our VAE objective function regularizes w.r.t.

Original: a class-speciﬁc Gaussian mixture model rather than to N ( 0 , I ) .
Revised: a class-speciﬁc Gaussian mixture model rather than to N ( 0 , I ) .

Original: Training this supervised VAE (SVAE) requires class labels, but results inconsiderably better negative log-likelihood (NLL) values and excellent image generation on MNIST data, very effectively taking advantage of the class information.
Revised: Training this supervised VAE (SVAE) requires class labels, but results in better loglikelihood bounds than an unmodiﬁed VAE and excellent image generation on MNIST data, very effectively taking advantage of the class information.

Original: Our second contribution is a differentiable decision tree (DDT), which allows us to differentiate the expected loss of a speciﬁc tree w.r.t.
Revised: Our second contribution is a differentiable decision tree (DDT), which allows us to differentiate the expected loss of a speciﬁc tree w.r.t.

Original: a speciﬁc distribution.
Revised: a speciﬁc distribution.

Original: This is applicable to learning embeddings, allowing us to compute the gradient of expected loss for weight updates.
Revised: This is applicable to learning embeddings, allowing us to compute the gradient of expected loss for weight updates.

Original: Our third contribution combines SVAE with DDT to get Classiﬁer+VAE (C+VAE), which learns a latent variable distribution suitable for both classiﬁcation and generation simultaneously.
Revised: Our third contribution combines SVAE with DDT to get Classiﬁer+VAE (C+VAE), which learns a latent variable distribution suitable for both classiﬁcation and generation simultaneously.

Original: Trained on MNIST data, C+VAE produced an encoder that is competitive in both classiﬁcation error and NLL, using very simple encoders and decoders (each with two hidden layers of 500 MLPs per layer).
Revised: Trained on MNIST data, C+VAE produced an encoder that is competitive in both classiﬁcation error and log-likelihood, using very simple encoders and decoders.

Original: Our ﬁnal contribution is an analysis of the interpretability of a DDT trained on MNIST.
Revised: Our ﬁnal contribution is an analysis of the interpretability of a DDT trained on MNIST.

Original: Each internal node of the tree tests one of the 50 encoded dimensions when making a classiﬁcation decision.
Revised: Each internal node of the tree tests one of the 50 encoded dimensions when making a classiﬁcation decision.

Original: We examine the values of the MNIST test data in the encoded dimensions used by the tree, analyzing the semantics of each dimension.
Revised: We examine the values of the MNIST test data in the encoded dimensions used by the tree, analyzing the semantics of each dimension.

Original: We found that the dimensions used by the tree to discriminate correspond to meaningful macro-features learned by the encoder, and that the tree itself effectively summarizes the classiﬁcation process.
Revised: We found that the dimensions used by the tree to discriminate correspond to meaningful macro-features learned by the encoder, and that the tree itself effectively summarizes the classiﬁcation process.

Original: The rest of this paper is organized as follows.
Revised: The rest of this paper is organized as follows.

Original: In Section 2 we give relevant background.
Revised: In Section 2 we give relevant background.

Original: Then in
Revised: Then in

Original: Section 3 we describe the SVAE and present our differentiable decision tree and our combined model.
Revised: Section 3 we describe the SVAE and present our differentiable decision tree and our combined model.

Original: Our experimental results appear in Section 4.
Revised: Our experimental results appear in Section 4.

Original: Finally, we present related work in Section 5, and conclude in Section 6 with a discussion of future work.
Revised: Finally, we present related work in Section 5, and conclude in Section 6 with a discussion of future work.

Original:  B ACKGROUND
Revised:  B ACKGROUND

Original: We begin with a dataset of pairs, { ( x 1 , y 1 ) , . . .
Revised: We begin with a dataset of pairs, { ( x 1 , y 1 ) , . . .

Original: , ( x n , y n ) } , where x i ∈ X ⊂ R m is the i th observation and the class label is y i .
Revised: , ( x n , y n ) } , where x i ∈ X ⊂ R m is the i th observation and the class label is y i .

Original: Kingma & Welling (2014) introduced the variational autoencoder (VAE) as a latent variable model for efﬁcient maximum marginal likelihood learning.
Revised: Kingma & Welling (2014) introduced the variational autoencoder (VAE) as a latent variable model for efﬁcient maximum marginal likelihood learning.

Original: The VAE performs density estimation on p ( x , z ) where z are latent variables, to maximize the likelihood of the observed training data x : log p ( X ) = 
Revised: The VAE performs density estimation on p ( x , z ) where z are latent variables, to maximize the likelihood of the observed training data x : [Equation] = [Equation]

Original: N (cid:88)
Revised: log

Original: [Equation]
Revised: [Equation]

Original: Since this marginal likelihood is difﬁcult to work with directly for non-trivial models, instead a parametric inference model q ( z | x ) is used to optimize the variational lower bound on the marginal log-likelihood:
Revised: Since this marginal likelihood is difﬁcult to work with directly for non-trivial models, instead a parametric inference model q ( z | x ) is used to optimize the variational lower bound on the marginal log-likelihood:

Original: [Equation]
Revised: [Equation]

Original: where θ indicates the parameters of the encoder p and decoder q models, and KL ( ·||· ) is
Revised: where θ indicates the parameters of the encoder p and decoder q models, and KL ( ·||· ) is

Original: the
Revised: the

Original: Kullback-Leibler divergence (Kullback & Leibler (1951)).
Revised: Kullback-Leibler divergence (Kullback & Leibler (1951)).

Original: The VAE optimizes the lower bound by reparameterizing q ( z | x ) (Kingma & Welling (2014); Rezende et al.
Revised: The VAE optimizes the lower bound by reparameterizing q ( z | x ) (Kingma & Welling (2014); Rezende et al.

Original: The ﬁrst term of L above corresponds to the reconstruction error of the decoder p ( x | z ) , and the second term regularizes the distribution parameterized by the encoder p ( z | x ) to minimize the K-L divergence from a chosen prior distribution, usually an isotropic, centered Gaussian.
Revised: The ﬁrst term of L above corresponds to the reconstruction error of the decoder p ( x | z ) , and the second term regularizes the distribution parameterized by the encoder p ( z | x ) to minimize the K-L divergence from a chosen prior distribution, usually an isotropic, centered Gaussian.

Original: The simplicity of this prior distribution has the downside of restricting the ﬂexibility of latent variable assignment, but allows the VAE to be easily used as a decoder-based generative model as
Revised: The simplicity of this prior distribution has the downside of restricting the ﬂexibility of latent variable assignment, but allows the VAE to be easily used as a decoder-based generative model as

Original: [Equation]
Revised: [Equation]

Original: While more sophisticated approximate posteriors have been used to improve variational inference (Salimans et al.
Revised: While more sophisticated approximate posteriors have been used to improve variational inference (Salimans et al.

Original: ; Rezende & Mohamed (2015); Tran et al. (2015); Kingma et al. (2016);
Revised: ; Rezende & Mohamed (2015); Tran et al. (2015); Kingma et al. (2016);

Original: Burda et al.
Revised: Burda et al.

Original: , for the sake of efﬁciency and simplicity, in our work we use only the simple approximate posterior q ( z | x ) stated above.
Revised: , for the sake of efﬁciency and simplicity, in our work we use only the simple approximate posterior q ( z | x ) stated above.

Original: It is straightforward to combine the more sophisticated approaches with our model.
Revised: It is straightforward to combine the more sophisticated approaches with our model.

Original:  O UR M ODEL
Revised:  O UR M ODEL

Original: Now we describe our model.
Revised: Now we describe our model.

Original: The main contributions that this model brings in are (1) A differentiable decision tree, where we describe how to compute the expected probability distribution over predicted labels and use this to differentiate the expected loss of the tree (which can be used to optimize learned embeddings to minimize expected loss); (2) an explicit concept of class-derived data (supervised
Revised: The main contributions that this model brings in are (1) A differentiable decision tree, where we describe how to compute the expected probability distribution over predicted labels and use this to differentiate the expected loss of the tree (which can be used to optimize learned embeddings to minimize expected loss); (2) an explicit concept of class-derived data (supervised

Original: VAE) by modifying the prior p ( z ) to a mixture of Guassians; and (3) A combined VAE model using modiﬁcations (1) and (2) designed to learn a latent variable distribution suitable simultaneously for classifying and generating data (C+VAE).
Revised: VAE) by modifying the prior p ( z ) to a mixture of Guassians; and (3) A combined VAE model using modiﬁcations (1) and (2) designed to learn a latent variable distribution suitable simultaneously for classifying and generating data (C+VAE).

Original: 3. S UPERVISED VAE
Revised: 3. S UPERVISED VAE

Original: Previous work has explored the potential of the VAE’s encoder to learn the data manifold for nonlinear dimensionality reduction and semi-supervised classiﬁcation (Kingma et al.
Revised: Previous work has explored the potential of the VAE’s encoder to learn the data manifold for nonlinear dimensionality reduction and semi-supervised classiﬁcation (Kingma et al.

Original: A difﬁculty arises from the VAE objective as the encoder learns to cluster the latent codes of high-probability data as close as possible to the mean of the prior distribution to minimize the K-L divergence term,
Revised: A difﬁculty arises from the VAE objective as the encoder learns to cluster the latent codes of high-probability data as close as possible to the mean of the prior distribution to minimize the K-L divergence term,

Original: KL ( q ( z | x ) (cid:107) p ( z )) , with the typical choice of unit Gaussian for the prior.
Revised: KL ( q ( z | x ) (cid:107) p ( z )) , with the typical choice of unit Gaussian for the prior.

Original: This mixes data of various classes, making them difﬁcult to separate for classiﬁcation.
Revised: This mixes data of various classes, making them difﬁcult to separate for classiﬁcation.

Original: We address this issue by changing the prior from an isotropic, unit Gaussian of the standard VAE to a mixture of unit Gaussians.
Revised: We address this issue by changing the prior from an isotropic, unit Gaussian of the standard VAE to a mixture of unit Gaussians.

Original: Speciﬁcally, we modify the generation procedure to be class-focused, rather than assuming otherwise undifferentiated data.
Revised: Speciﬁcally, we modify the generation procedure to be class-focused, rather than assuming otherwise undifferentiated data.

Original: To build the notion of class-distributed data into the VAE objective, we use a the prior to the following
Revised: To build the notion of class-distributed data into the VAE objective, we use a the prior to the following

Original: Gaussian mixture:
Revised: Gaussian mixture:

Original: [Equation]
Revised: [Equation]

Original: where y is a class label, and µ y is the posterior mean of class y , and π is a probability vector which may be pre-computed with the assumption that the labels in the training dataset are iid.
Revised: where y is a class label, and µ y is the posterior mean of class y , and π is a probability vector which may be pre-computed with the assumption that the labels in the training dataset are iid.

Original: As µ y is calculated empirically from the posterior, it can be initialized to small random values for all classes and updated regularly throughout training.
Revised: As µ y is calculated empirically from the posterior, it can be initialized to small random values for all classes and updated regularly throughout training.

Original: The objective for this VAE is
Revised: The objective for this VAE is

Original: [Equation]
Revised: [Equation]

Original: Since we train in a fully supervised model, each training instance is of the form ( x , y ) .
Revised: Since we train in a fully supervised model, each training instance is of the form ( x , y ) .

Original: Thus, when y is instantiated as the class variable, we get p ( z ) =
Revised: Thus, when y is instantiated as the class variable, we get p ( z ) =

Original: p ( z | y ) in Equation 1, so the KL term of
Revised: p ( z | y ) in Equation 1, so the KL term of

Original: Equation 2 becomes KL ( q ( z | x ) (cid:107)N ( z | µ y , I )) , yielding an objective of
Revised: Equation 2 becomes KL ( q ( z | x ) (cid:107)N ( z | µ y , I )) , yielding an objective of

Original: [Equation]
Revised: [Equation]

Original: In the remainder of this paper, we refer to a VAE trained with this objective function as the supervised
Revised: In the remainder of this paper, we refer to a VAE trained with this objective function as the supervised

Original: VAE (SVAE) .
Revised: VAE (SVAE) .

Original: Our approach to utilizing a Gaussian mixture as a prior distribution is similar to that of Dilokthanakul et al.
Revised: Our approach to utilizing a Gaussian mixture as a prior distribution is similar to that of Dilokthanakul et al.

Original: A key difference between their work and ours is that our use of class labels enhances training, obviating the need to marginalize over all classes to compute the K-L divergence.
Revised: A key difference between their work and ours is that our use of class labels enhances training, obviating the need to marginalize over all classes to compute the K-L divergence.

Original: This helps avoid the over-regularization problem that they discuss in their paper, while achieving high sample quality in our generated images.
Revised: This helps avoid the over-regularization problem that they discuss in their paper, while achieving high sample quality in our generated images.

Original: 3. D IFFERENTIABLE D ECISION T REE
Revised: 3. D IFFERENTIABLE D ECISION T REE

Original: The decision tree is a simple, interpretable model used for non-parametric classiﬁcation.
Revised: The decision tree is a simple, interpretable model used for non-parametric classiﬁcation.

Original: Typically, a tree is constructed by an algorithm like CART or C4.5 (Breiman et al.
Revised: Typically, a tree is constructed by an algorithm like CART or C4.5 (Breiman et al.

Original: (1984); Quinlan (1993)) that recursively divides the dataset at each node of the tree, greedily minimizing the weighted Gini coefﬁcient or entropy of two subsets by choosing a dividing line in one dimension.
Revised: (1984); Quinlan (1993)) that recursively divides the dataset at each node of the tree, greedily minimizing the weighted Gini coefﬁcient or entropy of two subsets by choosing a dividing line in one dimension.

Original: Inference is performed by walking an input from the root to a leaf according to this series of inequalities, and then assigning a class probability vector according to the leaf.
Revised: Inference is performed by walking an input from the root to a leaf according to this series of inequalities, and then assigning a class probability vector according to the leaf.

Original: Decision trees often classify well when using data with a few, richly descriptive features, and are very interpretable in their decision making processes.
Revised: Decision trees often classify well when using data with a few, richly descriptive features, and are very interpretable in their decision making processes.

Original: We are interested in learning a deep network for non-linear dimensionality reduction to allow the decision tree to classify a low-dimensional embedding of data that is normally highdimensional and very structured.
Revised: We are interested in learning a deep network for non-linear dimensionality reduction to allow the decision tree to classify a low-dimensional embedding of data that is normally highdimensional and very structured.

Original: Toward this end, we utilize a probabilistic generalization of decision trees, where each leaf returns a distribution over all classes.
Revised: Toward this end, we utilize a probabilistic generalization of decision trees, where each leaf returns a distribution over all classes.

Original: I.e., if instance z lands in leaf (cid:96) of tree T , then T returns a distribution P T ( y | (cid:96) ) .
Revised: I.e., if instance z lands in leaf (cid:96) of tree T , then T returns a distribution P T ( y | (cid:96) ) .

Original: Fix a decision tree T with leaves L that takes as input instance z and outputs a label y .
Revised: Fix a decision tree T with leaves L that takes as input instance z and outputs a label y .

Original: We observe that each leaf (cid:96) of T corresponds to a region in an axis-aligned rectilinear partitioning covering the data space whose bounds are deﬁned by the inequalities encoded in the path from the root.
Revised: We observe that each leaf (cid:96) of T corresponds to a region in an axis-aligned rectilinear partitioning covering the data space whose bounds are deﬁned by the inequalities encoded in the path from the root.

Original: To compute the expected loss of T on instance z drawn according to probability distribution D ( z ) , ﬁrst consider one leaf (cid:96) , and let R (cid:96)
Revised: To compute the expected loss of T on instance z drawn according to probability distribution D ( z ) , ﬁrst consider one leaf (cid:96) , and let R (cid:96)

Original: ⊂ R d be the region of T ’s input space that is covered by (cid:96) .
Revised: ⊂ R d be the region of T ’s input space that is covered by (cid:96) .

Original: Then the probability that randomly drawn instance z falls into leaf (cid:96) is
Revised: Then the probability that randomly drawn instance z falls into leaf (cid:96) is

Original: [Equation]
Revised: [Equation]

Original: Then the probability that randomly drawn instance z is predicted to be class y is
Revised: Then the probability that randomly drawn instance z is predicted to be class y is

Original: [Equation]
Revised: [Equation]

Original: This calculation is simpliﬁed if D ( z ) is restricted to have diagonal covariance, i.e., D ( z ) = (cid:81) di
Revised: This calculation is simpliﬁed if D ( z ) is restricted to have diagonal covariance, i.e., D ( z ) = (cid:81) di

Original: =1 D i ( z i ) .
Revised: =1 D i ( z i ) .

Original: Then, the integral of the PDF over R (cid:96) =
Revised: Then, the integral of the PDF over R (cid:96) =

Original: [ r − (cid:96), 1 , r + (cid:96), 1 ] × · · · × [ r − (cid:96),d , r + (cid:96),d ] may be replaced with a product of the integrals in each dimension.
Revised: [ r − (cid:96), 1 , r + (cid:96), 1 ] × · · · × [ r − (cid:96),d , r + (cid:96),d ] may be replaced with a product of the integrals in each dimension.

Original: Further, each integral may be calculated as the difference of the cumulative distribution function at the upper ( r + (cid:96),i ) and lower ( r − (cid:96),i ) bounds of the partition in each dimension i .
Revised: Further, each integral may be calculated as the difference of the cumulative distribution function at the upper ( r + (cid:96),i ) and lower ( r − (cid:96),i ) bounds of the partition in each dimension i .

Original: The full inference may be re-written as
Revised: The full inference may be re-written as

Original: [Equation]
Revised: [Equation]

Original: where CDF i is the cumulative distribution function of D i ( z i ) .
Revised: where CDF i is the cumulative distribution function of D i ( z i ) .

Original: Given a user-speciﬁed loss function loss T ( z , y ) , our goal is to minimize the expected loss L T =
Revised: Given a user-speciﬁed loss function loss T ( z , y ) , our goal is to minimize the expected loss L T =

Original: E z ∼ D
Revised: E z ∼ D

Original: [ loss T ( z , y )] .
Revised: [ loss T ( z , y )] .

Original: In our work, we use as D the distribution q ( z | x ) =
Revised: In our work, we use as D the distribution q ( z | x ) =

Original: N ( x ; µ x , σ x I ) , where µ x and σ x are the outputs of the encoder on input x .
Revised: N ( x ; µ x , σ x I ) , where µ x and σ x are the outputs of the encoder on input x .

Original: To perform gradient-based optimization of L T , the gradient w.r.t.
Revised: To perform gradient-based optimization of L T , the gradient w.r.t.

Original: each parameter is calculated as ∂
Revised: each parameter is calculated as ∂

Original: [Equation]
Revised: [Equation]

Original: and ∂
Revised: and ∂

Original: [Equation]
Revised: [Equation]

Original: [Table]
Revised: [Table]

Original: where PDF i ( r | µ x ,i , σ x ,i ) is the value of the Gaussian PDF of dimension i evaluated at r .
Revised: where PDF i ( r | µ x ,i , σ x ,i ) is the value of the Gaussian PDF of dimension i evaluated at r .

Original: This allows optimization of the distribution parameters for maximum likelihood w.r.t.
Revised: This allows optimization of the distribution parameters for maximum likelihood w.r.t.

Original: an existing decision tree T .
Revised: an existing decision tree T .

Original: Thus, an embedding of the data may be learned in an EM-style manner, alternately learning a tree on the embedding produced by the parameters of a deep encoder and optimizing the embedding parameters to better ﬁt the class-based partitioning induced by the learned decision tree.
Revised: Thus, an embedding of the data may be learned in an EM-style manner, alternately learning a tree on the embedding produced by the parameters of a deep encoder and optimizing the embedding parameters to better ﬁt the class-based partitioning induced by the learned decision tree.

Original: 3. O UR C OMBINED M ODEL
Revised: 3. O UR C OMBINED M ODEL

Original: The supervised, Gaussian-mixture-based VAE and decision tree inference can be used with a VAE model to both classify and reconstruct data from the encoded parameters of its latent variable distribution.
Revised: The supervised, Gaussian-mixture-based VAE and decision tree inference can be used with a VAE model to both classify and reconstruct data from the encoded parameters of its latent variable distribution.

Original: Although an embedding could be learned by only optimizing the classiﬁcation accuracy of the decision tree, the additional reconstruction objective ensures that the learned representation is non-arbitrary and contains more than just class information for downstream use.
Revised: Although an embedding could be learned by only optimizing the classiﬁcation accuracy of the decision tree, the additional reconstruction objective ensures that the learned representation is non-arbitrary and contains more than just class information for downstream use.

Original: Our new architecture C+VAE (Classiﬁer+VAE) uses a deep encoder network to parameterize a
Revised: Our new architecture C+VAE (Classiﬁer+VAE) uses a deep encoder network to parameterize a

Original: Gaussian distribution, which is then used as the input for classifying with the DDT and to reconstruct the encoded data with a deep decoder network.
Revised: Gaussian distribution, which is then used as the input for classifying with the DDT and to reconstruct the encoded data with a deep decoder network.

Original: Generally, the combined modiﬁcations can also be applied to existing VAE architectures when label information is available.
Revised: Generally, the combined modiﬁcations can also be applied to existing VAE architectures when label information is available.

Original: The C+VAE training procedure begins by randomly initializing the encoder/decoder parameters and encoding the training data to initialize the decision tree and aggregate posterior class means.
Revised: The C+VAE training procedure begins by randomly initializing the encoder/decoder parameters and encoding the training data to initialize the decision tree and aggregate posterior class means.

Original: Training then proceeds by running several epochs of gradient updates before re-training the decision tree and updating the aggregate posterior class means until the model converges.
Revised: Training then proceeds by running several epochs of gradient updates before re-training the decision tree and updating the aggregate posterior class means until the model converges.

Original: The optimization function of our combined model consists of a linear combination of the objective of the supervised VAE and the expected error of the current decision tree T .
Revised: The optimization function of our combined model consists of a linear combination of the objective of the supervised VAE and the expected error of the current decision tree T .

Original: However, since the effect of the DDT gradient is to separate the class means and the supervised VAE K-L divergence term measures w.r.t.
Revised: However, since the effect of the DDT gradient is to separate the class means and the supervised VAE K-L divergence term measures w.r.t.

Original: these movable class-based means (rather than the distance from the mean of p ( z ) ), the parameters learned by the encoder could diverge, driving the class means arbitrarily far from the origin.
Revised: these movable class-based means (rather than the distance from the mean of p ( z ) ), the parameters learned by the encoder could diverge, driving the class means arbitrarily far from the origin.

Original: Thus, an additional regularizing L 22 -loss is imposed on the encoded µ posterior value to keep the class means from “drifting” from the origin and encourage the model to learn common factors of variation between classes.
Revised: Thus, an additional regularizing L 22 -loss is imposed on the encoded µ posterior value to keep the class means from “drifting” from the origin and encourage the model to learn common factors of variation between classes.

Original: We observed experimentally that this additional drift loss term increases the training stability and classiﬁcation performance of the model.
Revised: We observed experimentally that this additional drift loss term increases the training stability and classiﬁcation performance of the model.

Original: The modiﬁed VAE objective of the
Revised: The modiﬁed VAE objective of the

Original: C+VAE to be minimized is
Revised: C+VAE to be minimized is

Original: [Equation]
Revised: [Equation]

Original:  E XPERIMENTS
Revised:  E XPERIMENTS

Original: Our experiments are designed to empirically study the following claims:
Revised: Our experiments are designed to empirically study the following claims:

Original: The supervised VAE very effectively takes advantage of class labels to improve generative performance.
Revised: The supervised VAE very effectively takes advantage of class labels to improve generative performance.

Original: The C+VAE classiﬁes competitively with other tree-based embedding methods while simultaneously maintaining a generative model competitive with the literature.
Revised: The C+VAE classiﬁes competitively with other tree-based embedding methods while simultaneously maintaining a generative model competitive with the literature.

Original: The differentiable decision tree is an interpretable classiﬁer that, when used in C+VAE, can learn the semantics of the macro-features learned by the underlying VAE.
Revised: The differentiable decision tree is an interpretable classiﬁer that, when used in C+VAE, can learn the semantics of the macro-features learned by the underlying VAE.

Original: The MNIST dataset was used for all experiments, as it is widely understood and commonly used for both classiﬁcation and generation tasks.
Revised: The MNIST dataset was used for all experiments, as it is widely understood and commonly used for both classiﬁcation and generation tasks.

Original: We applied the C+VAE modiﬁcations to a standard VAE (Kingma & Welling (2014)) with two-layer MLPs of 500 hidden units as encoder and decoder models and a 50-dimensional latent variable z without importance sampling or an autoregressive prior.
Revised: We applied the C+VAE modiﬁcations to a standard VAE (Kingma & Welling (2014)) with two-layer MLPs of 500 hidden units as encoder and decoder models and a 50-dimensional latent variable z without importance sampling or an autoregressive prior.

Original: The
Revised: The

Original: CART algorithm as made available in scikit-learn (Pedregosa et al.
Revised: CART algorithm as made available in scikit-learn (Pedregosa et al.

Original: (2011)) was used to train the decision tree.
Revised: (2011)) was used to train the decision tree.

Original: This was regularized by annealing the maximum depth of the decision tree from 1 to 8 as training proceeded, incrementing every 15 epochs, and by setting the minimum proportion of samples in a leaf to be 2% of the training set.
Revised: This was regularized by annealing the maximum depth of the decision tree from 1 to 8 as training proceeded, incrementing every 15 epochs, and by setting the minimum proportion of samples in a leaf to be 2% of the training set.

Original: Unless otherwise noted, we used γ = 1000 and λ = 0 .in
Revised: Unless otherwise noted, we used γ = 1000 and λ = 0 .in

Original: the objective function of the C+VAE (Equation (4)), and n = 5 epochs of gradient steps between each update of both the decision tree and the aggregate posterior class means.
Revised: the objective function of the C+VAE (Equation (4)), and n = 5 epochs of gradient steps between each update of both the decision tree and the aggregate posterior class means.

Original: Adam (Kingma & Ba (2014)) was used for optimization and the data was not pre-processed or augmented.
Revised: Adam (Kingma & Ba (2014)) was used for optimization and the data was not pre-processed or augmented.

Original: Table 1 lists the classiﬁcation performance of a number of tree-based and VAE-based models.
Revised: Table 1 lists the classiﬁcation performance of a number of tree-based and VAE-based models.

Original: The
Revised: The

Original: M1:
Revised: M1:

Original: SVAE+CART model trains the supervised VAE to convergence, and then trains a standard decision tree with CART to classify its latent code in the style of M1 (Kingma et al.
Revised: SVAE+CART model trains the supervised VAE to convergence, and then trains a standard decision tree with CART to classify its latent code in the style of M1 (Kingma et al.

Original: The intent is to highlight the effect of training without the backpropagated classiﬁcation loss from the
Revised: The intent is to highlight the effect of training without the backpropagated classiﬁcation loss from the

Original: DDT.
Revised: DDT.

Original: C+VAE sans reconstruction zeros the reconstruction loss term of the objective function to highlight the effect of training a model that only learns an embedding suitable for classiﬁcation with the DDT.
Revised: C+VAE sans reconstruction zeros the reconstruction loss term of the objective function to highlight the effect of training a model that only learns an embedding suitable for classiﬁcation with the DDT.

Original: The boundary tree (BT) with embedding is from Zoran et al.
Revised: The boundary tree (BT) with embedding is from Zoran et al.

Original: (2017), M1+M2 is from Kingma et al.
Revised: (2017), M1+M2 is from Kingma et al.

Original: (2014), and the Ladder Network is from Rasmus et al.
Revised: (2014), and the Ladder Network is from Rasmus et al.

Original: Table 2 lists the performance of a number of VAE-based generative models.

Original: The importance weighted autoencoder (IWAE) comes from Burda et al.

Original: (2015), Auxiliary VAE with L=1, IW=1 is from Maaløe et al.

Original: VAE with IAF posterior (IAF VAE) is from Kingma et al.

Original: and Variational Lossy Autoencoder (VLAE) is from Chen et al.

Original: The results presented in these tables are discussed below.

Original: [Table]

Original: [Table]

Original: 4. E VALUATING THE S UPERVISED VAE
Revised: 4. E VALUATING THE S UPERVISED VAE

Original: We ﬁrst evaluate the efﬁcacy of leveraging label data in a supervised VAE in generation.I
Revised: We ﬁrst evaluate the efﬁcacy of leveraging label data in a supervised VAE in generation.

Original: [Table] e., the effect on generation of making the prior distribution a Gaussian mixture and taking advantage of class label information.
Revised: I.e., the effect on generation of making the prior distribution a Gaussian mixture and taking advantage of class label information.

Original: This is equivalent to using C+VAE with γ = λ = 0 in Equation (4).
Revised: This is equivalent to using C+VAE with γ = λ = 0 in Equation (4).

Original: The ﬂexibility of a Gaussian mixture and the fact that the data is clearly multi-modal both contribute to  the SVAE log likelihood of − 71 . 86 , which isconsiderably better thanany (unsupervised) results from the literature.
Revised: The ﬂexibility of a Gaussian mixture and the fact that the data is clearly multi-modal both contribute to [Table] the SVAE log-likelihood of − 102 . 77 , which is better than the log-likelihood of − 109 .

Revised: 56 using our implementation of the VAE of Kingma & Welling (2014), which uses an unmodiﬁed Gaussian prior.

Original: We expect this  to be the result of using a ﬂexible prior that is more faithful to the true prior.
Revised: We expect this difference to be the result of using a ﬂexible prior that is more faithful to the true prior.

Original: This ﬂexibility is similar to that seen in techniques like normalizing ﬂows (Rezende & Mohamed (2015)) [Table] but modiﬁes the prior rather than the posterior and uses the additional information provided by label information, rather than adding additional computation.
Revised: This ﬂexibility is similar to that seen in techniques like normalizing ﬂows (Rezende & Mohamed (2015)), but modiﬁes the prior rather than the posterior and uses the additional information provided by label information, rather than adding additional computation.

Original: 4. E VALUATING THE C+VAE
Revised: 4. E VALUATING THE C+VAE

Original: We next empirically evaluate C+VAE for both classiﬁcation performance and generative ability.
Revised: We next empirically evaluate C+VAE for both classiﬁcation performance and generative ability.

Original: As a baseline, we ﬁrst examine how well a standard (non-differentiable) decision tree from CART can classify when the data is encoded by a supervised VAE (but without any error feedback: γ = λ = 0 ).
Revised: As a baseline, we ﬁrst examine how well a standard (non-differentiable) decision tree from CART can classify when the data is encoded by a supervised VAE (but without any error feedback: γ = λ = 0 ).

Original: This is similar to M1 from Kingma et al.
Revised: This is similar to M1 from Kingma et al.

Original: (2014) with a different VAE.
Revised: (2014) with a different VAE.

Original: In Table 1, row M1:
Revised: In Table 1, row M1:

Original: SVAE+CART shows that without the error feedback from the tree, it is unlikely that the embedding will be useful in classiﬁcation by a decision tree.
Revised: SVAE+CART shows that without the error feedback from the tree, it is unlikely that the embedding will be useful in classiﬁcation by a decision tree.

Original: This motivates our use of the DDT.
Revised: This motivates our use of the DDT.

Original: To test the beneﬁt of reconstruction in learning an embedding that can be classiﬁed well, we ran a test in which we switched off the reconstruction error feedback in learning.
Revised: To test the beneﬁt of reconstruction in learning an embedding that can be classiﬁed well, we ran a test in which we switched off the reconstruction error feedback in learning.

Original: I.e., we removed the ﬁrst term of Equation (3).
Revised: I.e., we removed the ﬁrst term of Equation (3).

Original: In Table 1, row C+VAE sans reconstruction shows a signiﬁcant improvement in classiﬁcation error over M1:
Revised: In Table 1, row C+VAE sans reconstruction shows a signiﬁcant improvement in classiﬁcation error over M1:

Original: SVAE+CART, but still quite high.
Revised: SVAE+CART, but still quite high.

Original: Row C+VAE in Tables 1and 2 shows our combined method’s performance with γ = 1000 and λ = 0 .
Revised: Row C+VAE in Table 1 shows our combined method’s performance with γ = 1000 and λ = 0 .

Original: We see a large improvement in classiﬁcation error over C+VAE sans reconstruction, demonstrating the importance of both types of feedback in training.
Revised: We see a large improvement in classiﬁcation error over C+VAE sans reconstruction, demonstrating the importance of both types of feedback in training.

Original: While C+VAE’s  performance is worse than  the state of the art in both error and NLL, it’s competitive in both, and to our knowledge, no other method in the literature does both simultaneously.
Revised: While C+VAE’s classiﬁcation performance is worse than results from the literature, it’s still competitive, despite simultaneously optimizing both classiﬁcation and log-likelihood.

Original: Further, we note that for the sake of comparison to the literature, our results are based on very small networks (the encoder and decoder each with two hidden layers of500 MLPs per layer) similar to theoriginal VAE  Kingma & Welling (2014).
Revised: Also, we note that C+VAE’s log-likelihood of − 110 . 12 is comparable to the − 109 . 56 from our implementation of the VAE of Kingma & Welling (2014), which uses the same encoder-decoder pair as C+VAE.

Original: A more powerful encoder  could conceivably improve both error and NLL even further.
Revised: A more powerful encoder or the use of more recent techniques (e.g., normalizing ﬂows, importance weighting, etc.) could conceivably improve both error and log-likelihood even further.

Original: Figure 7 in the appendix presents sample MNIST digits generated by C+VAE.
Revised: Figure 7 in the appendix presents sample MNIST digits generated by C+VAE.

Original: Each set of digits is generated from one of the empirical aggregated class means.
Revised: Each set of digits is generated from one of the empirical aggregated class means.

Original: 4. E VALUATING THE I NTERPRETABILITY OF THE DDT
Revised: 4. E VALUATING THE I NTERPRETABILITY OF THE DDT

Original: The ﬁnal decision tree learned by the C+VAE is shown in Figure 1 (a landscape version of the same tree is in Figure 8 in the appendix).
Revised: The ﬁnal decision tree learned by the C+VAE is shown in Figure 1 (a landscape version of the same tree is in Figure 8 in the appendix).

Original: This tree performs feature selection over the 50 available latent dimensions, using only 8 to classify with 98.02% accuracy with one dimension (21) used to split twice.
Revised: This tree performs feature selection over the 50 available latent dimensions, using only 8 to classify with 98.02% accuracy with one dimension (21) used to split twice.

Original: We were able to leverage the simplicity of the decision tree to assign meaning to the latent dimensions used by the tree to classify inputs.
Revised: We were able to leverage the simplicity of the decision tree to assign meaning to the latent dimensions used by the tree to classify inputs.

Original: Each node divides inputs according to a threshold value in a single dimension, which corresponds to detecting the most salient macro-feature that distinguishes the divided subsets.
Revised: Each node divides inputs according to a threshold value in a single dimension, which corresponds to detecting the most salient macro-feature that distinguishes the divided subsets.

Original: Dimension 21 of the latent code is the macro-feature used by the decision tree to discriminate between digits ‘6’ and ‘0’, as well as ‘4’ and ‘7’.
Revised: Dimension 21 of the latent code is the macro-feature used by the decision tree to discriminate between digits ‘6’ and ‘0’, as well as ‘4’ and ‘7’.

Original: Figure 2 visualizes the macro-feature corresponding to dimension 21.
Revised: Figure 2 visualizes the macro-feature corresponding to dimension 21.

Original: Speciﬁcally, the top image (starting with ‘6’) was generated by ﬁxing the otherdimensions to be the values of µ 6 and varying the value of dimension 21 in even steps from −to 2 .
Revised: Speciﬁcally, the top image (starting with ‘6’) was generated by ﬁxing the otherdimensions to be the values of µ 6 and varying the value of dimension 21 in even steps from −to 2 .

Original: The bottom image was generated the same way, with the other 49 dimensions initialized to values from µ 4 .
Revised: The bottom image was generated the same way, with the other 49 dimensions initialized to values from µ 4 .

Original: In both image sequences, we see that a high value of dimension 21 emphasizes the macro-feature of a ﬂat top bar of a digit whereas a low value removes it.
Revised: In both image sequences, we see that a high value of dimension 21 emphasizes the macro-feature of a ﬂat top bar of a digit whereas a low value removes it.

Original: To illustrate the effect of varying dimension 21 independently of a speciﬁc class mean, we generate
Revised: To illustrate the effect of varying dimension 21 independently of a speciﬁc class mean, we generate

Original: Figure 3 by an identical process, but ﬁx the other 49 dimensions to be the mean of all 10 class means
Revised: Figure 3 by an identical process, but ﬁx the other 49 dimensions to be the mean of all 10 class means.

Original: The center image is the mean of all 50 dimensions, and is provided for contrast.
Revised: The center image is the mean of all 50 dimensions, and is provided for contrast.

Original: The left and right images show the effect of a low or high value of dimension 21 on this ‘average digit’.
Revised: The left and right images show the effect of a low or high value of dimension 21 on this ‘average digit’.

Original: The clearest effect of this variation is that the ﬂat top macro-feature is present when this value is high, and absent when it is low, just as in Figure 2.
Revised: The clearest effect of this variation is that the ﬂat top macro-feature is present when this value is high, and absent when it is low, just as in Figure 2.

Original: We generate Figures 4, 5, and 6 by the same process, varying dimensions 10, 26, and 45, respectively.
Revised: We generate Figures 4, 5, and 6 by the same process, varying dimensions 10, 26, and 45, respectively.

Original: The decision tree uses dimension 0 to separate digits ‘3’, ‘5’, and ‘8’ from digits ‘2’, ‘4’, and ‘7’.
Revised: The decision tree uses dimension 0 to separate digits ‘3’, ‘5’, and ‘8’ from digits ‘2’, ‘4’, and ‘7’.

Original: Figure 4 shows that a low value of dimension 0 correlates strongly with rounded digits, while a high value creates an emphasized right side and a more angular appearance.
Revised: Figure 4 shows that a low value of dimension 0 correlates strongly with rounded digits, while a high value creates an emphasized right side and a more angular appearance.

Original: Figure 5 shows the effect of varying dimension 26, used by the tree to separate ‘2’ from ‘4’ and ‘7’.
Revised: Figure 5 shows the effect of varying dimension 26, used by the tree to separate ‘2’ from ‘4’ and ‘7’.

Original: The most notable impact of a low dimension 26 is the exaggerated lower-left corner, which is absent when that dimension is high.
Revised: The most notable impact of a low dimension 26 is the exaggerated lower-left corner, which is absent when that dimension is high.

Original: Figure 6 shows the effect of varying dimension 45.
Revised: Figure 6 shows the effect of varying dimension 45.

Original: This is the ﬁrst dimension used by the decision tree, to split ‘1’ from the other nine classes.
Revised: This is the ﬁrst dimension used by the decision tree, to split ‘1’ from the other nine classes.

Original: A clear vertical line near the center of the digit is emphasized by low values of dimension 1.
Revised: A clear vertical line near the center of the digit is emphasized by low values of dimension 1.

Original: Latent codes with high values of dimension 1, a feature common to the other nine classes, lack that central vertical line.
Revised: Latent codes with high values of dimension 1, a feature common to the other nine classes, lack that central vertical line.

Original:  R ELATED W ORK
Revised:  R ELATED W ORK

Original: Differentiable Trees
Revised: Differentiable Trees

Original: Previous work that uses deep networks for representation learning with decision trees is the Deep Neural Decision Forests of Kontschieder et al.
Revised: Previous work that uses deep networks for representation learning with decision trees is the Deep Neural Decision Forests of Kontschieder et al.

Original: (2015), which stochastically make routing decisions through a decision tree according to the outputs of a deep convolutional network.
Revised: (2015), which stochastically make routing decisions through a decision tree according to the outputs of a deep convolutional network.

Original: The setup achieved good performance at its classiﬁcation task, but it is not clear how to interpret the proposed classiﬁcation process, especially when more than one tree is combined into a forest classiﬁer.
Revised: The setup achieved good performance at its classiﬁcation task, but it is not clear how to interpret the proposed classiﬁcation process, especially when more than one tree is combined into a forest classiﬁer.

Original: As a method of making the decision tree differentiable, our proposed inference method of integrating a probability distribution over the decision regions of the tree is a novel approach.
Revised: As a method of making the decision tree differentiable, our proposed inference method of integrating a probability distribution over the decision regions of the tree is a novel approach.

Original: Another tree-based method uses differentiable boundary trees to learn an embedding suitable for k-nearest neighbor classiﬁcation (Zoran et al.
Revised: Another tree-based method uses differentiable boundary trees to learn an embedding suitable for k-nearest neighbor classiﬁcation (Zoran et al.

Original: The learned representation allows a small, interpretable boundary tree to classify effectively, similar to our technique.
Revised: The learned representation allows a small, interpretable boundary tree to classify effectively, similar to our technique.

Original: The classiﬁcation accuracy of the technique marginally outperforms our combined model, but the C+VAE also acts as a generative model and does not suffer from the signiﬁcant complexity of having to use dynamically constructed computation graphs.
Revised: The classiﬁcation accuracy of the technique marginally outperforms our combined model, but the C+VAE also acts as a generative model and does not suffer from the signiﬁcant complexity of having to use dynamically constructed computation graphs.

Original: VAE learning Other work in classifying the latent codes produced by a VAE includes Kingma et al.
Revised: VAE learning Other work in classifying the latent codes produced by a VAE includes Kingma et al.

Original: (2014), whose M1 semi-supervised model learns to classify from the latent embedding similarly to our combined classiﬁer.
Revised: (2014), whose M1 semi-supervised model learns to classify from the latent embedding similarly to our combined classiﬁer.

Original: However, M1 trains the discriminator separately from the VAE and lacks interpretability as the class separation is performed solely by a black-box discriminator.
Revised: However, M1 trains the discriminator separately from the VAE and lacks interpretability as the class separation is performed solely by a black-box discriminator.

Original: The M model is similar to the supervised VAE, but doesn’t change the VAE prior.
Revised: The M model is similar to the supervised VAE, but doesn’t change the VAE prior.

Original: Dilokthanakul et al.
Revised: Dilokthanakul et al.

Original: (2017) present a Gaussian Mixture Variational Autoencoder to learn a classfocused latent representation.
Revised: (2017) present a Gaussian Mixture Variational Autoencoder to learn a classfocused latent representation.

Original: Our work assumes a supervised, rather than the GMVAE’s unsupervised environment.
Revised: Our work assumes a supervised, rather than the GMVAE’s unsupervised environment.

Original: This allows the classifying modiﬁcation to the VAE framework to remain simpler and more interpretable, as well as more tractable optimization.
Revised: This allows the classifying modiﬁcation to the VAE framework to remain simpler and more interpretable, as well as more tractable optimization.

Original:  F UTURE W ORK
Revised:  F UTURE W ORK

Original: Future work includes applying our approach to other data sets such as CIFAR-10, and using more powerful encoders and decoders to see how performance is affected.
Revised: Future work includes applying our approach to other data sets such as CIFAR-10, and using more powerful encoders and decoders to see how performance is affected.

Original: We will also look into extending our approach to handle unlabeled data in applications such as semi-supervised learning and clustering. 
Revised: We will also look into extending our approach to handle unlabeled data in applications such as semi-supervised learning and clustering. A CKNOWLEDGEMENTS

Original: B L ANDSCAPE V ERSION OF D ECISION T REE
Revised: B L ANDSCAPE V ERSION OF D ECISION T REE