Title: SEMI-SUPERVISED LEARNING WITH NORMALIZING FLOWS

Abstract: We propose Flow Gaussian Mixture Model (FlowGMM), a general-purpose method for semi-supervised learning based on a simple and principled probabilistic framework. We approximate the joint distribution of the labeled and unlabeled data with a flexible mixture model implemented as a Gaussian mixture transformed by a normalizing flow. We train the model by maximizing the exact joint likelihood of the labeled and unlabeled data. We evaluate FlowGMM on a wide range of semi-supervised classification problems across different data types: AG-News and Yahoo Answers text data, MNIST, SVHN and CIFAR-10 image classification problems as well as tabular UCI datasets. FlowGMM achieves promising results on image classification problems and outperforms the competing methods on other types of data. FlowGMM learns an interpretable latent representation space and allows hyper-parameter free feature visualization at real time rates. Finally, we show that FlowGMM can be calibrated to produce meaningful uncertainty estimates for its predictions.

1. INTRODUCTION:

In many domains unlabeled data is plentiful, while labeled data may be scarce. Semi-supervised learning framework leverages both labeled and unlabeled data reducing the need for expensive manual annotation. Recently, consistency-based methods have shown outstanding performance in semisupervised image classification (Laine & Aila, 2016; Miyato et al., 2018; Tarvainen & Valpola, 2017; Athiwaratkun et al., 2019; Verma et al., 2019; Berthelot et al., 2019) and are currently state-of-the-art on challenging datasets like CIFAR-10, CIFAR-100, and ImageNet. However, these methods have not seen much application on domains other than images, where a suitable set of data perturbations to which the classifier should be invariant is not known a priori.

In this paper, we employ a simple and principled probabilistic framework for semi-supervised learning. We introduce a mixture distribution for modeling the data, where different components correspond to different classes. We can maximize the joint likelihood of unlabeled (using the mixture) and labeled (using individual mixture components) data. At test time, we classify the input to belong to the class corresponding to mixture component with the highest likelihood.

In order to apply this framework to complex data, we need to choose a sufficiently flexible family of distributions in the mixture, preserving the ability to compute exact likelihoods. We propose Flow Gaussian Mixture Model (FlowGMM), a mixture model based on normalizing flows (Dinh et al., 2014) . Each component of the mixture is modeled as a trainable invertible transformation (normalizing flow) of the corresponding component of a Gaussian mixture in the latent space. The transformation is shared among all the mixture components. Due to invertibility, we can compute exact likelihood of the data using the change of variable formula.

We illustrate FlowGMM on a toy problem in Figure 1 . We are solving a binary semi-supervised classification problem on the dataset shown in panel (a): the labeled data is shown with triangles colored according to their class, and unlabeled data is shown with blue circles. We introduce a Gaussian mixture with two components corresponding to each of the classes, shown in panel (c) in the latent space Z and an invertible transformation f . The transformation f is then trained to map the data distribution in the data space X to the latent Gaussian mixture in the Z space, mapping the labeled data to the corresponding mixture component. We visualize the learned transformation in panel (b), showing the positions of the images f (x) for all of the training data points. The inverse Samples from the model in the data space.

f -1 of this mapping serves as a class-conditional generative model, that we visualize in panel (d).

To classify a data point x in the input space we compute its image f (x) in the latent space, and pick the class corresponding to the Gaussian that is closest to f (x). We visualize the decision boundary of the learned classifier with a dashed line in panel (a).

Many of the previous attempts of using generative models for classification (Salimans et al., 2016; Nalisnick et al., 2019; Chen et al., 2019) have relied upon multitask learning, where a shared latent representation is learned for the generative model and the classifier. With the method of Chen et al. ( 2019), hybrid modeling is observed to reduce performance for both tasks in the supervised case.

For GANs, Dai et al. (2017) have pointed out that classification performance and generative performance are in direct conflict: a perfect generator yields no benefit to classification performance. FlowGMM takes a different approach, where the generative model is used directly as a Bayes classifier, and in the limit of a perfect generative model the Bayes classifier achieves provably optimal misclassification rate (see e.g. Mohri et al., 2018) .

FlowGMM naturally encodes the clustering principle: the decision boundary between classes must lie in the low-density region in the data space. Indeed, in the latent space the decision boundary between two classes coincides with the hyperplane perpendicular to the line segment connecting means of the corresponding mixture components and passing through the midpoint of this line segment (assuming the components are normal distributions with identity covariance matrices); in panel (b) of Figure 1 we show the decision boundary in the latent space with a dashed line. The density of the latent distribution near the decision boundary is low. As the flow is trained to represent data as a transformation of this latent distribution, the density near the decision boundary should also be low.

In panel (a) of Figure 1 the decision boundary indeed lies in the low-density region.

The contributions of this work are as follows:

• We propose FlowGMM, a new probabilistic classification model based on normalizing flows, that can be naturally applied to semi-supervised learning. We evaluate FlowGMM on a range of semi-supervised classification benchmarks including text, tabular and image data, and it works well. FlowGMM outperforms alternative approaches on non-image data and shows promise for image data. We propose modified consistency regularization for FlowGMM and empirically demonstrate that it substantially improves performance of the method on image classification problems.

• We conduct a thorough empirical analysis of FlowGMM for supervised and semisupervised classification. One of the important features of FlowGMM is its interpretability.

To demonstrate it, we visualize the learned latent space representations for the proposed semi-supervised model and show that interpolations between data points from different classes pass through low-density regions. We show how our classification model can be used for optimization free feature visualization. We also study the predictive uncertainties produced by the method and show that they can be calibrated by scaling the variances of mixture components.

2. BACKGROUND: NORMALIZING FLOWS:

The normalizing flow (Dinh et al., 2016 ) is an unsupervised model for density estimation defined as an invertible mapping f : X → Z from the data space X to the latent space Z. We can model the data distribution as a transformation f -1 : Z → X applied to a random variable from the latent distribution z ∼ p Z often chosen to be Gaussian. The density of the transformed random variable x = f -1 (z) is given by the change of variables formula

EQUATION

The mapping f is implemented as a sequence of invertible functions, parametrized by a neural network with architecture that is designed to ensure invertibility and efficient computation of logdeterminants, and a set of parameters θ that can be optimized. The model can be trained by maximizing the likelihood (equation 1) of the training data with respect to the parameters θ.

3. FLOW GAUSSIAN MIXTURE MODEL (FLOWGMM):

In FlowGMM, we introduce a discrete latent variable y for the class label, y ∈ {1 . . . C}. Our latent space distribution, conditioned on a given label k, is Gaussian with mean µ k and covariance Σ k :

EQUATION

The marginal distribution of z is then a Gaussian mixture. When the classes are balanced, this distribution is

EQUATION

Thus, the likelihood for labeled data is

p X (x|y = k) = N (f (x)|µ k , Σ k ) • det ∂f ∂x ,

and the likelihood for data with unknown label is p X (x) = k p X (x|y = k)p(y = k). If we have access to both a labeled dataset D and an unlabeled dataset D u , then we can train our model in a semi-supervised way to maximize the joint likelihood of the labeled and unlabeled data

EQUATION

over the parameters θ of the bijective function f , which learns a density model with a Bayes classifier: given a test point x, the model predictive distribution is given by

EQUATION

)

We can then make predictions for a test point x with the Bayes decision rule

y = arg max i∈{1,...,C} p X (y = i|x).

Alternatively to direct likelihood maximization, we can adapt Expectation Maximization algorithm for model training (discussed in Appendix A).

3.1. CONSISTENCY REGULARIZATION:

Most of the existing state-of-the-art approaches to semi-supervised learning on image data are based on consistency regularization (Laine & Aila, 2016; Miyato et al., 2018; Tarvainen & Valpola, 2017; Athiwaratkun et al., 2019; Verma et al., 2019) . These methods penalize changes in network predictions with respect to input perturbations, such as random translations and horizontal flips, with an additional loss term that can be computed on unlabeled data,

EQUATION

where x , x are random perturbations of x, and g is the vector of probabilities over the classes.

Motivated by these methods, we introduce a simple consistency regularization term for FlowGMM. Let y be the label predicted on image x by FlowGMM according to equation 5. We then define the consistency loss term as the negative log likelihood of the input x given the label y :

EQUATION

This loss term encourages the model to map small perturbations of the same unlabeled inputs to the same components of the Gaussian mixture distribution in the latent space. Unlike the standard consistency loss of equation 6, the proposed loss in equation 7 takes values on the same scale as the data log likelihood (equation 4), and empirically we found it to perform better. We refer to FlowGMM with the consistency term as FlowGMM-cons. The final loss for FlowGMM-cons is then the weighted sum of the consistency loss (equation 7) and the negative log likelihood of both labeled and unlabeled data (equation 4).

4. RELATED WORK:

Generative models have been used extensively in semi-supervised learning. Along with GANs and VAEs, Normalizing Flows (NF) (Dinh et al., 2014) present another major class of deep generative models. Unlike GANs and VAEs, normalizing flows can be trained using exact likelihood. NFs admit controllable latent representations and can be sampled efficiently unlike auto-regressive models (Papamakarios et al., 2017; Oord et al., 2016) . Recent work (Dinh et al., 2016; Kingma & Dhariwal, 2018; Behrmann et al., 2018) demonstrated that normalizing flows can produce high-fidelity samples for natural image datasets. Some normalizing flow papers (such as RealNVP (Dinh et al., 2016) ) have used class-conditional sampling, where the transformation is conditioned on the class label. To do so, they pass the class label as an input to coupling layers, conditioning the output of the flow on the class.

Deep Invertible Generalized Linear Model (DIGLM, Nalisnick et al., 2019) , most closely related to our work, trains a classifier on the latent representation of a normalizing flow to perform supervised or semi-supervised image classification. Our approach is principally different, as we use a mixture of Gaussians in the latent space Z and perform classification based on class-conditional likelihoods (see equation 5), rather than training a separate classifier. One of the key advantages of our approach is the explicit encoding of clustering principle in the method and a more natural probabilistic interpretation.

5. EXPERIMENTS:

In all experiments, we use the RealNVP normalizing flow architecture. Throughout training, Gaussian mixture parameters are fixed: the means are initialized randomly from the standard normal distribution and the covariances are set to I. See Appendix B for further discussion on GMM initialization and training.

5.1. SYNTHETIC DATA:

We first apply FlowGMM to a range of two-dimensional synthetic datasets, in order to get a better visual intuition for the method. We use RealNVP architecture with 5 coupling layers, defined by fully-connected shift and scale networks, each with 1 hidden layer of size 512. In addition to the semi-supervised setting, we also trained the method only using the labeled data. In Figure 2 we visualize the decision boundaries of the classifier corresponding to FlowGMM for both of these We compare the classifier decision boundaries when only using labeled data (panels b, d) and when using both labeled and unlabeled data (panels a, c) on two circles (panels a, b) and pinwheel (panels c, d) datasets. FlowGMM leverages unlabeled data to push the decision boundary to low-density regions of the space. settings on the two circles and pinwheel datasets. On both datasets FlowGMM is able to leverage the unlabeled data to push the decision boundary to a low-density region, as expected. On the two circles dataset the method is unable to fit the data perfectly, as it is impossible to represent this dataset as an invertible continuous mapping of two Gaussians, because they are topologically distinct. FlowGMM still produces a reasonable decision boundary and improves over the case when only labeled data is available. We provide additional visualizations in Appendix C, Figure 4 .

5.2. TEXT AND TABULAR DATA:

We believe that FlowGMM can be especially useful for semi-supervised learning on tabular data. Consistency-based semi-supervised methods have mostly been developed for image classification, where the predictions of the method are regularized to be invariant to random flips and translations of the image. On tabular data, where the structure is less prominent, finding suitable transformations to apply for consistency-based methods is not-trivial. Similarly, approaches based on GANs have mostly been developed for images. We evaluate FlowGMM on the Hepmass and Miniboone UCI classification datasets (previously used in Papamakarios et al. (2017) for density estimation).

Along with standard tabular UCI datasets, we also consider text classification on AG-News and Yahoo Answers datasets. Using the recent advances in transfer learning for NLP, we construct embeddings for input texts using the BERT transformer model (Devlin et al., 2018) trained on a corpus of Wikipedia articles, and then train FlowGMM and other baselines on the embeddings.

We compare FlowGMM to the graph based label spreading method from Zhou et al. (2004) , a Π-Model (Laine & Aila, 2016 ) that uses dropout perturbations, as well as supervised logistic regression, k-nearest neighbors, and neural network that were trained on the labeled data only. We report the results in Table 5 .1, where FlowGMM beats the competing semi-supervised learning methods on each of the considered datasets. Implementation details for FlowGMM, the baselines, and the dataset preprocessing details are listed in Appendix D.

5.3. IMAGE CLASSIFICATION:

We next evaluate the proposed method on semi-supervised image classification benchmarks on CIFAR-10, MNIST and SVHN datasets. For all the datasets, we use RealNVP (Dinh et al., 2016) architecture. Exact implementation details are listed in the appendix E. The supervised model is trained using the same loss (equation 4), where all the data points are labeled (n u = 0).

We present the results for FlowGMM and FlowGMM-cons in Table 2 . We also report results from DIGLM (Nalisnick et al., 2019) (which only report semi-supervised performance on MNIST and supervised performance on MNIST and SVHN) and the M1+M2 VAE model (Kingma et al., 2014) . FlowGMM outperforms M1+M2 model and performs better or on par with DIGLM. Furthermore, FlowGMM-cons improves over FlowGMM on all three datasets, suggesting that consistency regularization is very beneficial for the proposed model when useful perturbations are available.

Following Oliver et al. (2018) , we evaluate FlowGMM-cons varying the number of labeled data points. Specifically, we follow the setup of Kingma et al. (2014) and train FlowGMM-cons on MNIST with 100, 600, 1000 and 3000 labeled data points. We present the results in Table 3 . FlowGMM-cons outperforms the M1+M2 model of Kingma et al. (2014) in all the considered settings.

We note that the results presented in this section are not directly comparable with the state-of-theart methods using GANs or consistency regularization (see e.g. Laine & Aila, 2016; Dai et al., 2017; Athiwaratkun et al., 2019; Berthelot et al., 2019) , as the architecture we employ is much less powerful for classification than the ConvNet and ResNet architectures that have been designed for classification without the constraint of invertibility. We believe that invertible architectures with better inductive biases for classification (possibly like iResNet (Behrmann et al., 2018) ) may help bridge this gap. The space of images is challenging to model and while there has been substantial progress using normalizing flows, there is still much ground to cover. In many machine learning applications, it is crucial to understand how confident a model is in its predictions. In classification problems, well-calibrated models are expected to output meaningful probabilities of belonging to a particular class. In Guo et al. (2017) , it was observed that modern deep learning models are highly overconfident, however, simple temperature scaling can substantially improve model's calibration. In this section, we analyze the predictive uncertainties produced by the FlowGMM. In Appendix Section F, we additionally study the ability of FlowGMM to detect out-of-domain data.

When using FlowGMM for classification, the class predictive probabilities are p(y

= k|x) = N (x|µ k , Σ k )/ m N (x|µ m , Σ m ).

Since we initialize Gaussian mixture means randomly from the standard normal distribution and do not train them along with the flow parameters (see section B), FlowGMM predictions become inherently overconfident due to the curse of dimensionality. Indeed, consider two Gaussians with means sampled independently from the standard normal µ 1 , µ 2 ∼ N (0, I) in D-dimensional space. If s 1 ∼ N (µ 1 , I) is a sample from the first Gaussian, then its expected squared distances to both mixture means are E s 1 -µ 1 2 = D and E s 1 -µ 2 2 = 3D (for a detailed derivation see Appendix Section G.). In high dimensional spaces, such logits would lead to hard label assignment in FlowGMM (p(y|x) = 1 for exactly one class). In fact, in the experiments we observe that FlowGMM is overconfident and performs hard label assignment: predicted class probabilities are all close to either 1 or 0.

We address this problem by learning a single scalar parameter σ 2 for all components in the Gaussian mixture (the component k will be N (µ k , σ 2 I)) by minimizing the negative log likelihood on a validation set. This way we can re-calibrate the variance of the latent GMM in a natural way. This procedure is also equivalent to applying temperature scaling (Guo et al., 2017) to logits log N (x|µ k , Σ k ). We test FlowGMM calibration on MNIST and CIFAR datasets in the supervised setting. On MNIST we restricted the training set size to 1000 objects, since on the full dataset the model makes too few mistakes which makes evaluating calibration harder. In Table 4 , we report negative log likelihood and expected calibration error (ECE, see Guo et al. (2017) for a description of this metric). We can see that re-calibrating variances of the Gaussians in the mixture significantly improves both metrics and mitigates overconfidence. The effectiveness of this simple rescaling procedure suggests that the latent space distances learned by the flow model are correlated with the probabilities of belonging to a particular class: the closer a datapoint is to the mean of a Gaussian in the latent space, the more likely it belongs to the corresponding class.

Under review as a conference paper at ICLR 2020 

6.2. LEARNED LATENT REPRESENTATIONS:

We next analyze the latent representation space learned by FlowGMM. We examine latent interpolations between members of the same class in Figure 3 (a) and between different classes in Figure 3 (b) for our MNIST FlowGMM-cons model trained with n = 1k labels. As expected, inter-class interpolations pass through regions of low-density, leading to low quality samples but intra-class interpolations do not. These observations suggest that, as expected, the model learns to put the decision boundary in the low-density region of the data space.

In Appendix section H, we present images corresponding to the means of the Gaussian mixture and class-conditional samples from FlowGMM.

Distance to Decision Boundary:

To explicitly test this conclusion, we compute the distribution of distances from unlabeled data to the decision boundary for FlowGMM-cons and FlowGMM Sup trained on labeled data only. In order to compute this distance exactly for an image x, we find the two closest means µ , µ to the corresponding latent variable z = f (x), and evaluate the

expression d(x) = µ -f (x) 2 -µ -f (x) 2 2 µ -µ

. We visualize the distributions of the distances for the supervised and semi-supervised method in Figure 3 (c) . While most of the unlabeled data are far from the decision boundary for both methods, the supervised method puts a substantially larger fraction of data close to the decision boundary. For example, the distance to the decision boundary is smaller than 5 for 1089 unlabeled data points with supervised model, but only 143 data points with FlowGMM-cons. This increased separation suggests that FlowGMM-cons indeed pushes the decision boundary away from the data distribution in agreement with the clustering principle.

6.3. FEATURE VISUALIZATION:

Feature visualization has become an important tool for increasing the interpretability of neural networks. The majority of methods rely on maximizing the activations of a given neuron, channel, or layer over a parametrization of an input image with different kinds of image regularization (Szegedy et al., 2013; Olah et al., 2017; Mahendran & Vedaldi, 2015) . These methods, while effective, require optimization and regularization hyper-parameters and iterative optimization too costly for real time interactive exploration.

Since our classification model uses a flow which is a sequence of invertible transformations:

f (x) = f :L (x) := f L • f L-1 • . . . f 1 (x)

, intermediate activations can be inverted directly. This means that we can combine the methods of feature inversion and feature maximization directly by feeding in a set of input images, modifying intermediate activations arbitrarily, and inverting the representation.

Given a set of activations in the th layer a [c, i, j] = f : (x) cij with channels c and spatial extent i, j, we may perturb a single neuron with

EQUATION

where δ c is a one hot vector at channel c; and σ c is the standard deviation of the activations in channel c over the the training set and spatial locations. This can be performed at real time rates to explore the activation parametrized by α and the location cij without any optimization or hyper-parameters.

The feature visualization of intermediate layers on CIFAR10 test images are shown in panel (d) of Figure 3 . The given channel being visualized appears to activate on the zeroed pixels from random translations as well as the green channel, giving us insight into the workings of the model.

7. DISCUSSION:

We proposed FlowGMM, a natural and interpretable model for semi-supervised learning with normalizing flows. FlowGMM outperforms graph-based and consistency-based baselines on tabular data including semi-supervised text classification with BERT embeddings. On image classification, FlowGMM is not yet competitive with the state-of-the-art approaches (Athiwaratkun et al., 2019; Verma et al., 2019) , but we expect that the results could be greatly improved by using a better flow architecture that has stronger inductive biases for classification.

FlowGMM is greatly extensible and can be easily adapted to different settings. For example, consider an extreme case of semi-supervised learning where some classes may be completely unlabeled, and the model has to identify those extra classes and cluster them separately. An effective model that can operate in this setting could, for example, be used for automated scientific discovery: when applied to classify certain experimental data, it can identify new separate classes that were not considered by experts. FlowGMM can be applied in this scenario: instead of the standard Gaussian mixture model (GMM) in the latent space we can use the Chinese Restaurant Process GMM (CRP-GMM), which can infer the number of mixture components from data automatically (Rasmussen, 2000) .

Further, we can easily adapt FlowGMM to few-shot learning setting where the model has to generalize to classes not seen at train time by only using a few examples of that new class. In particular, we can fix the parameters of the flow model and add a new Gaussian to the mixture whenever we add a new class. We can estimate the parameters of this Gaussian from the few available examples.

We view interpretability as another strong advantage of FlowGMM. The access to latent space representations and feature visualization technique discussed in Section 6 as well as the ability to sample from the model can be used to get insights into the performance of the model in practical applications.

A EXPECTATION MAXIMIZATION:

As an alternative to direct optimization of the likelihood (equation 4), we consider Expectation-Maximization algorithm (EM). EM is a popular approach for finding maximum likelihood estimates in mixture models. Suppose X = {x i } n i=1 is the observed dataset, T = {t i } n i=1 are corresponding unobserved latent variables (often denoting the component in mixture model) and θ is a vector of model parameters. EM algorithm consists of the two alternating steps: on E-step, we compute posterior probabilities of latent variables for each data point q(t i |x i ) = P (t i |x i , θ); and on M-step, we fix q and maximize the expected log likelihood of the data and latent variables with respect to θ: E q log P (X, T |θ) → max θ . The algorithm can be easily adapted to the semi-supervised setting where a subset of data is labeled with {y l i } n l i=1 : then, on E-step we have hard assignment to the true mixture component q(t i |x i ) = I[t i = y l i ] for labeled data points. EM algorithm is applicable in our setting which is fitting the transformed mixture of Gaussians. We can perform the exact E-step for unlabeled data in the model since

q(t|x) = p(x|t, θ) p(x|θ) = N (f (x)|µ t , Σ t ) • det ∂f ∂x C k=1 N (f (x)|µ k , Σ k ) • det ∂f ∂x = N (f (x)|µ t , Σ t ) C k=1 N (f (x)|µ k , Σ k )

which coincides with the E-step of EM algorithm on Gaussian mixture model. On M-step, the objective has the following form:

n l i=1 log N (f θ (x l i )|µ y l i , Σ y l i ) ∂f θ ∂x l i + nu i=1 E q(ti|x u i ,θ) log N (f θ (x u i )|µ ti , Σ ti ) ∂f θ ∂x u i .

Since the exact solution is not tractable due to complexity of the flow model, we perform a stochastic gradient step to optimize the expected log likelihood with respect to flow parameters θ.

Note that unlike regular EM algorithm for mixture models, we have Gaussian mixture parameters {(µ k , Σ k )} C k=1 fixed in our experiments, and on M-step the update of θ induces the change of z i = f θ (x i ) latent space representations.

Using EM algorithm for optimization in the semi-supervised setting on MNIST dataset with 1000 labeled images, we obtain 98.97% accuracy which is comparable to the result for FlowGMM with regular SGD training. However, in our experiments, we observed that on E-step, hard label assignment happens for unlabeled points (q(t|x) ≈ 1 for one of the classes) because of the high dimensionality of the problem (see section 6.1) which affects the M-step objective and hinders training.

B LATENT DISTRIBUTION MEAN AND COVARIANCE CHOICES:

Initialization In our experiments, we draw the mean vectors µ i of Gaussian mixture model randomly from the standard normal distribution µ i ∼ N (0, I), and set the covariance matrices to identity Σ i = I for all classes; we fixed GMM parameters throughout training. However, one could potentially benefit from data-dependent placing of means in the latent space. We experimented with different initialization methods, in particular, initializing means using the mean point of latent representations of labeled data in each class:

µ i = (1/n i l ) n i l m=1 f (x i m )

where x i m represents labeled data points from class i and n i l is the total number of labeled points in that class. In addition, we can scale all means by a scalar value μi = rµ i to increase or decrease distances between them. We observed that such initialization leads to much faster convergence of FlowGMM on semi-supervised classification on MNIST dataset, however, the final performance of the model was worse compared to the one with random mean placing. We hypothesize that it becomes easier for the flow model to warm up faster with data-dependent initialization because Gaussian means are closer to the initial latent representations, but afterwards the model gets stuck in a suboptimal solution.

GMM training FlowGMM would become even more flexible and expressive if we could learn Gaussian mixture parameters in a principled way. In the current setup where means are sampled from the standard normal distribution, the distances between mixture components are about

√:

2D where D is the dimensionality of the data (see Appendix G). Thus, classes are quite far apart from each Two Circles other in the latent space, which, as observed in Section 6.1, leads to model miscalibration. Training GMM parameters can further increase interpretability of the learned latent space representations: we can imagine a scenario in which some of the classes are very similar or even intersecting, and it would be useful to represent it in the latent space. We could train GMM by directly optimizing likelihood (equation 4), or using expectation maximization (see Section A), either jointly with the flow parameters or iteratively switching between training flow parameters with the fixed GMM and training GMM with the fixed flow. In our initial experiments on semi-supervised classification on MNIST, training GMM jointly with the flow parameters did not improve performance or lead to substantial change of the latent representations. Further improvements require careful hyperparameter choice which we leave for future work.

EQUATION

C SYNTHETIC EXPERIMENTS:

In Figure 4 we visualize the classification decision boundaries of FlowGMM as well as the learned mapping to the latent space and generated samples for three different synthetic datasets.

D TABULAR DATA PREPARATION AND HYPERPARAMETERS:

The AG-News and Yahoo Answers were constructed by applying BERT embeddings to the text input, yielding a 768 dimensional vector for each data point. AG-News has 4 classes while Ya- hoo Answers has 10. The UCI datasets Hepmass and Miniboone were constructed using the data preprocessing from Papamakarios et al. (2017) , but with the inclusion of the removed background process class so that the two problems can be used for binary classification. We then subsample the fraction of background class examples so that the dataset is balanced. For each of the datasets, a separate validation set of size 5k was used to tune hyperparameters. All neural network models use the ADAM optimizer (Kingma & Ba, 2014).

k-Nearest Neighbors: We tested both using both L2 distance and L2 with inputs normalized to unit norm, (sin 2 distance), and the latter performed the best. The value k chosen in the method was found sweeping over 1 -20, and the optimal values for each of the datasets are shown in 5.

3 Layer NN + Dropout: The 3-Layer NN + Dropout baseline network has three fully connected hidden layers with inner dimension k = 512, ReLU nonlinearities, and dropout with p = 0.5. We use the learning rate 3e-4 for training the supervised baseline across all datasets.

Π-Model: The Π-Model uses the same network architecture, and dropout for the perturbations. The additional consistency loss per unlabeled data point is computed as

L Unlab = ||g(x ) -g(x )|| 2 ,

where g is are the output probabilities after the softmax layer of the neural network and the consistency weight λ = 30 which worked the best across the datasets. The model was trained for 50 epochs with labeled and unlabeled batch size n for AG-News and Yahoo Answers, and labeled and unlabeled batch sizes n and 2000 for Hepmass and Miniboone.

Label Spreading: We use the local and global consistency method from Zhou et al. (2004) , Y * = (I -αS) -1 Y where in our case Y is the matrix of labels for the labeled, unlabeled, and test data but filled with zeros for unlabeled and test. S = D -1/2 W D -1/2 computed from the affinity matrix W ij = exp (-γ sin 2 (x i , x j )) where sin 2 (x i , x j ) := 1 -xi,xj xi xj . This is equivalent to L2 distance on the inputs normalized to unit magnitude. Because the algorithm scales poorly with number of unlabeled points for dense affinity matrices, O(n 3 u ), we we subsampled the number of unlabeled data points to 10k and test data points to 5k for this graph method. However, we also evaluate the label spreading algorithm with a sparse kNN affinity matrix on using a larger subset 20k of unlabeled data. The two hyperparameters for label spreading (γ/k and α) were tuned by separate grid search for each of the datasets. In both cases, we use the inductive variant of the algorithm where the test data is not included in the unlabeled data.

FlowGMM::

We train our FlowGMM model with a RealNVP normalizing flow, similar to the architectures used in Papamakarios et al. (2017) . Specifically, the model uses 7 coupling layers, with 1 hidden layer each and 256 hidden units for the UCI datasets but 512 for text classification. UCI models were trained for 50 epochs of unlabeled data and the text datasets were trained for 30 epochs of unlabeled data. The labeled and unlabeled batch sizes are the same as in the Π-Model.

The tuned learning rates for each of the models that we used for these experiments are shown in table 5 .

E IMAGE DATA PREPARATION AND HYPERPARAMETERS:

We use the CIFAR-10 multi-scale architecture with 2 scales, each containing 3 coupling layers defined by 8 residual blocks with 64 feature maps. We use Adam optimizer (Kingma & Ba, 2014) with learning rate 10 -3 for CIFAR-10 and SVHN and 10 -4 for MNIST. We train the supervised model for 100 epochs, and semi-supervised models for 1000 passes through the labeled data for CIFAR-10 and SVHN and 3000 passes for MNIST. We use a batch size of 64 and sample 32 labeled and 32 unlabeled data points in each mini-batch. For the consistency loss term (equation 7), we linearly increase the weight from 0 to 1 for the first 100 epochs following Athiwaratkun et al. (2019) .

For FlowGMM and FlowGMM-cons, we re-weight the loss on labeled data by λ = 3 (value tuned on validation (Kingma et al., 2014) on CIFAR-10), as otherwise, we observed that the method underfits the labeled data.

F OUT-OF-DOMAIN DATA DETECTION Density models have held promise for being able to detect out-of-domain data, an especially important task for robust machine learning systems (Nalisnick et al., 2019) . Recently, it has been shown that existing flow and autoregressive density models are not as apt at this task as previously thought, yielding high likelihood on images coming from other (simpler) distributions. The conclusion put forward is that datasets like SVHN are encompassed by, or have roughly the same mean but lower variance than, more complex datasets like CIFAR10 (Nalisnick et al., 2018) . We examine this hypothesis in the context of our flow model which has a multi-modal latent space distribution unlike methods considered in Nalisnick et al. (2018) . Using a fully supervised model trained on MNIST, we evaluate the log likelihood for data points coming from the NotMNIST dataset, consisting of letters instead of digits, and the FashionM-NIST dataset. We then train a supervised model on the more complex dataset FashionMNIST and evaluate on MNIST and NotMNIST. The distribution of the log likelihood log p X (•) = log p Z (f (•)) + log det ∂f ∂x on these datasets is shown in Figure 5 . For the model trained on MNIST we see that the data from Fashion MNIST and NotMNIST is assigned lower likelihood, as expected. However, the model trained on FashionMNIST predicts higher likelihoods for MNIST images. The majority (≈ 75%) of the MNIST data points get mapped into the mode of the Fashion-MNIST model corresponding to sandals, which is the class with the largest fraction of pixels that are zero. Similarly, for the model trained on MNIST the image of all zeros has very high likelihood and gets mapped to the mode corresponding to the digit 1 which has the largest fraction of empty space.

G EXPECTED DISTANCES BETWEEN GAUSSIAN SAMPLES:

Consider two Gaussians with means sampled independently from the standard normal µ 1 , µ 2 ∼ N (0, I) in D-dimensional space. If s 1 ∼ N (µ 1 , I) is a sample from the first Gaussian, then its expected squared distances to both mixture means are:

E s 1 -µ 1 2 = E E s 1 -µ 1 2 |µ 1 = E D i=1 E (s 1,i -µ 1,i ) 2 |µ 1,i = E D i=1 E[s 2 1,i ] -2µ 2 1,i + µ 2 1,i = E D i=1 1 + µ 2 1,i -µ 2 1,i = D

:



Figures and Tables:

Figure 1: Illustration of semi-supervised learning with FlowGMM on a binary classification problem. Colors represent the two classes or the corresponding Gaussian mixture components. Labeled data is shown with triangles, colored by the corresponding class label, and blue dots represent unlabeled data. (a): Data distribution and the classifier decision boundary. (b): The learned mapping of the data to the latent space. (c): Samples from the Gaussian mixture in the latent space. (d): Samples from the model in the data space.

Figure 2: Illustration of FlowGMM performance on synthetic datasets. Labeled data is shown with colored triangles, and unlabeled data is shown with blue circles. Colors represent different classes.We compare the classifier decision boundaries when only using labeled data (panels b, d) and when using both labeled and unlabeled data (panels a, c) on two circles (panels a, b) and pinwheel (panels c, d) datasets. FlowGMM leverages unlabeled data to push the decision boundary to low-density regions of the space.

Figure 3: Visualizations of the latent space representations learned by supervised FlowGMM on MNIST. (a): Latent space interpolations between test images from the same class and (b): from different classes. Observe that interpolations between objects from different classes pass through low-density regions. (c): Histogram of distances from unlabeled data to the decision boundary for FlowGMM-cons trained on 1k labeled and 59k unlabeled data and FlowGMM Sup trained on 1k labeled data only. FlowGMM-cons is able to push the decision boundary away from the data distribution using unlabeled data. (d): Feature visualization for CIFAR10: four test reconstructions are shown as an intermediate feature is perturbed. The value of the perturbation α is shown in red vs the distribution of the channel activations. Observe that the channel visualized activates on zeroed out pixels to the left of the image mimicking the random translations applied to the training data.

Figure 4: Illustration of FlowGMM on synthetic datasets: two circles (top row), eight Gaussians (middle row) and pinwheel (bottom row). (a): Data distribution and classification decision boundaries. Unlabeled data is shown with blue circles and labeled data is shown with colored triangles, where color represents the class. Background color visualizes the classification decision boundaries of FlowGMM. (b): Mapping of the data to the latent space. (c): Gaussian mixture in the latent space. (d): Samples from the learned generative model corresponding to different classes, as shown by their color.

Figure 5: Left: Log likelihoods on in-and out-of-domain data for our model trained on MNIST. Center: Log likelihoods on in-and out-of-domain data for our model trained on FashionMNIST. Right: MNIST digits get mapped onto the sandal mode of the FashionMNIST model 75% of the time, often being assigned higher likelihood than elements of the original sandal class. Representative elements are shown above.



Table content: <table><tr><td>Method</td><td>AG-News</td><td>Yahoo Answers</td><td>Hepmass</td><td>Miniboone</td></tr><tr><td/><td>(200 / 200k, 4)</td><td>(800 / 50k, 10)</td><td colspan="2">(20 / 140k, 2) (20 / 65k, 2)</td></tr><tr><td>kNN</td><td>51.3</td><td>28.4</td><td>84.6</td><td>77.7</td></tr><tr><td>Logistic Regression</td><td>78.9</td><td>54.9</td><td>84.9</td><td>75.9</td></tr><tr><td>3-Layer NN + Dropout</td><td>78.1</td><td>55.6</td><td>84.4</td><td>77.3</td></tr><tr><td>RBF Label Spreading</td><td>54.6</td><td>30.4</td><td>87.1</td><td>78.8</td></tr><tr><td>kNN Label Spreading</td><td>56.7</td><td>25.6</td><td>87.2</td><td>78.1</td></tr><tr><td>Π-model</td><td>80.6</td><td>56.6</td><td>87.9</td><td>78.3</td></tr><tr><td>FlowGMM</td><td>84.8</td><td>57.4</td><td>88.8</td><td>80.6</td></tr></table>

Accuracy of the FlowGMM, VAE model (M1+M2 VAE, Kingma et al., 2014), DIGLM(Nalisnick et al., 2019) in supervised and semi-supervised settings on MNIST, SVHN, and CIFAR-10. FlowGMM Sup (All labels) as well as DIGLM Sup (All labels) were trained on full train datasets with all labels to demonstrate general capacity of these models. FlowGMM Sup (n l labels) was trained on n l labeled examples (and no unlabeled data). For reference, at the bottom we list the performance of the Π-Model(Laine & Aila, 2016) and BadGAN(Dai et al., 2017) as representative consistency-based and GAN-based state-of-the-art methods. Both of these methods use non-invertible architectures with substantially higher base performance and, thus, are not directly comparable.

Table content: <table><tr><td/><td/><td colspan="2">Dataset (n l / nu)</td></tr><tr><td>Method</td><td>MNIST</td><td>SVHN</td><td>CIFAR-10</td></tr><tr><td/><td colspan="2">(1k/59k) (1k/72k)</td><td>(4k/46k)</td></tr><tr><td>DIGLM Sup (All labels)</td><td>99.27</td><td>95.74</td><td>-</td></tr><tr><td>FlowGMM Sup (All labels)</td><td>99.63</td><td>95.81</td><td>88.44</td></tr><tr><td>M1+M2 VAE SSL</td><td>97.60</td><td>63.98</td><td>-</td></tr><tr><td>DIGLM SSL</td><td>99.0</td><td>-</td><td>-</td></tr><tr><td>FlowGMM Sup (n l labels)</td><td>97.36</td><td>78.26</td><td>73.13</td></tr><tr><td>FlowGMM</td><td>98.94</td><td>82.42</td><td>78.24</td></tr><tr><td>FlowGMM-cons</td><td>99.0</td><td>86.44</td><td>80.9</td></tr><tr><td>BadGAN</td><td>-</td><td>95.75</td><td>85.59</td></tr><tr><td>Π-Model</td><td>-</td><td>94.57</td><td>87.64</td></tr></table>

Semi-supervised classification accuracy for FlowGMM-cons and VAE M1 + M2 model(Kingma et al., 2014)  on MNIST for different number of labeled data points n l .

Table content: <table><tr><td>Method</td><td>n l = 100</td><td>n l = 600</td><td>n l = 1000</td><td>n l = 3000</td></tr><tr><td>M1+M2 VAE SSL (n l labels)</td><td>96.67</td><td colspan="3">97.41 ± 0.05 97.60 ± 0.02 97.82 ± 0.04</td></tr><tr><td>FlowGMM-cons (n l labels)</td><td>98.2</td><td>98.7</td><td>99</td><td>99.2</td></tr></table>

Negative log-likelihood and Expected Calibration Error for supervised FlowGMM trained on MNIST (1k train, 1k validation, 10k test) and CIFAR-10 (50k train, 1k validation, 9k test). FlowGMM-temp stands for tempered FlowGMM where a single scalar parameter σ 2 was learned on a validation set for variances in all components.

Table content: <table><tr><td/><td colspan="2">MNIST (test acc 97.3%)</td><td colspan="2">CIFAR-10 (test acc 89.3%)</td></tr><tr><td/><td colspan="4">FlowGMM FlowGMM-temp FlowGMM FlowGMM-temp</td></tr><tr><td>NLL ↓</td><td>0.295</td><td>0.094</td><td>2.98</td><td>0.444</td></tr><tr><td>ECE ↓</td><td>0.024</td><td>0.004</td><td>0.108</td><td>0.038</td></tr><tr><td colspan="2">6 MODEL ANALYSIS</td><td/><td/><td/></tr><tr><td colspan="3">6.1 UNCERTAINTY AND CALIBRATION</td><td/><td/></tr></table>



Table content: <table><tr><td>Method Learning Rate</td><td colspan="4">AG-News Yahoo Answers Hepmass Miniboone</td></tr><tr><td>3-Layer NN + Dropout</td><td>3e-4</td><td>3e-4</td><td>3e-4</td><td>3e-4</td></tr><tr><td>Π-model</td><td>1e-3</td><td>1e-4</td><td>3e-3</td><td>1e-4</td></tr><tr><td>FlowGMM</td><td>1e-4</td><td>1e-4</td><td>3e-3</td><td>3e-4</td></tr><tr><td>kNN</td><td>k = 4</td><td>k = 18</td><td>k = 9</td><td>k = 3</td></tr></table>

Peer review (iclr2020):
The paper describes a normalising flow with the prior distribution represented by a Gaussian mixture model (GMM). The method, FlowGMM, maps each class of the dataset to a Gaussian distribution in the latent space by optimising the joint likelihood of both labelled and unlabelled data, thus making the method useful for semi-supervised problems. Predictions are made using the maximum a posteriori estimate of the class label. To make the method robust to small perturbations of the inputs, the authors introduce a novel consistency regularisation term to the total loss function, which maximises the likelihood of predicting the same class after a perturbation. The authors further examine the learnt latent space by considering two simple, synthetic datasets that can be easily visualised, showing that the latent space behaves in a way one would intuitively expect. The method is evaluated on both tabular and image data, showing promising results in terms of accuracy (presumably, see below). As the model is found to be overconfident in its predictions, the authors introduce a calibration scheme and empirically verify that it improves the uncertainty estimates. Lastly, the authors introduce a feature visualisation scheme and use it to illustrate the effect of perturbing the activations of the invertible transformations. I generally like the proposed method, which seems useful and intuitive. I am particularly happy with the discussion on uncertainty calibration, where the authors suggest an elegant addition to the model to increase the variance of the mixture components. I do, however, have significant concerns about the novelty of the paper as well as its structure and clarity, as detailed below. I do, therefore, not recommend it for acceptance. The paper reads well, although I feel that it lacks some details and explanations. For example, in table 2, it is never mentioned what *FlowGMM Sup* refers to and if it is different from *FlowGMM Supervised*. It is also not clear what *(All labels)* refers to - does it mean that labels were provided for the entire dataset or that the models were trained only on the small subset with labels? Or something else? Which performance metric is used in the tables? The accuracy, presumably, but this is never specifically stated. Similarly, the number of datapoints and ratio of labelled to unlabelled data for the synthetic datasets are not reported. They are not crucial to know but should be included for completeness. While the first half of the paper is informative and well-structured, the second half appears a bit less so. From experimentally verifying that the method works, the paper goes on to discuss uncertainty calibration, examine the latent space representations, and visualising the effect of feature perturbations. While I greatly appreciate the focus on interpreting the trained model, I think it appears somewhat chaotic, as if the authors tried to squeeze too much into the paper. For example, the feature visualisation technique is quite neat, but it works for any flow and is not really used for anything in the paper. I would suggest saving it for a dedicated paper. I am not convinced by the novelty of this paper. The authors list two contributions: 1) the model itself, 2) an empirical analysis with much focus on the interpretability of the model. While the model is, to my knowledge, indeed novel, the analysis is quite standard, and the interpretability even appears to be oversold. GMMs are nice and intuitive, but not novel in any way, yet the authors seem to be describing the properties of GMMs as specific to their method. In particular, the authors go to great lengths to show that the latent space representations cluster around the means of the mixture components and that the decision boundary lies in low-density regions of the latent space. I do not see why these properties should be so surprising since the method directly optimises the likelihood of the data under the mixture distribution. That this is also empirically observed is, of course, reassuring, but these observations are better suited for the appendix, in particular given that the paper went over the recommended page limit. I think that much of the claimed second contribution follows directly from the GMM aspect of the model. Instead of claiming the standard GMM properties as contributions, I think the proposed consistency loss term should be highlighted as a contribution on its own. I find it elegant, and I guess it would be particularly useful for NLP tasks where sentences can be phrased in different ways but still mean the same. Instead of discussing the latent space, I would have preferred to see extra evaluations of the method, like convergence rates of both FlowGMM and FlowGMM-cons compared to the competing models. Furthermore, a major limitation of the model is that knowledge of the correct number of classes in the data - even in the unsupervised setting. The authors hint at extensions to mitigate this in the discussion (using a Chinese Restaurant Process GMM or by adding extra Gaussians to the mixture during training), but these should have been investigated in the current paper. In conclusion, I think that the paper lacks novelty and that it spends far too much space on *trivial* properties of the model instead of addressing shortcomings, like the prior specification of the number of classes, which the authors even point out in the discussion. Minor comments: - p 4, bottom: *each with 1 hidden layers* -> *each with 1 hidden layer* - p 5, middle: *FlowGMM is able to leveraged* -> *FlowGMM is able to leverage* - p 5, bottom: *Table 5.1* -> *Table 1*