Title: GAN-BASED GAUSSIAN MIXTURE MODEL RESPONSI-BILITY LEARNING

Abstract: Mixture Model (MM) is a probabilistic framework which allows us to define a dataset containing K different modes. When each of the modes is associated with a Gaussian distribution, we refer it as Gaussian MM, or GMM. Given a data point x, GMM may assume the existence of a random index k ∈ {1, . . . , K} identifying which Gaussian the particular data is associated with. In a traditional GMM paradigm, it is straightforward to compute in closed-form, the conditional likelihood p(x|k, θ), as well as responsibility probability p(k|x, θ) which describes the distribution index corresponds to the data. Computing the responsibility allows us to retrieve many important statistics of the overall dataset, including the weights of each of the modes. Modern large datasets often contain multiple unlabelled modes, such as paintings dataset containing several styles; fashion images containing several unlabelled categories. In its raw representation, the Euclidean distances between the data do not allow them to form mixtures naturally, nor it's feasible to compute responsibility distribution, making GMM unable to apply. To this paper, we utilize the Generative Adversarial Network (GAN) framework to achieve an alternative plausible method to compute these probabilities at the data's latent space z instead of x. Instead of defining p(x|k, θ) explicitly, we devised a modified GAN to allow us to define the distribution using p(z|k, θ), where z is the corresponding latent representation of x, as well as p(k|x, θ) through an additional classification network which is trained with the GAN in an "end-toend" fashion. These techniques allow us to discover interesting properties of an unsupervised dataset, including dataset segments as well as generating new "outdistribution" data by smooth linear interpolation across any combinations of the modes in a completely unsupervised manner.

1. INTRODUCTION:

Gaussian Mixture Model (GMM) is one of the commonly used probabilistic models for dataset enjoy multiple modes. It assumes that all data points come from a mixture of a finite number of Gaussian distributions. The density function of the GMM is defined below:

EQUATION

)

where K is the total number of Gaussians in the mixture, and k th component is characterized by a Gaussian distribution with weight α k , mean µ k and covariance matrix Σ k .

Given x is the data and k is the (latent) index of the mixture density, GMM allows us to compute the conditional likelihood p(x|k, θ) as well as responsibility probability p(k|x, θ). In the Bayesian paradigm, one may refer p(k|x, θ) as the posterior density where the prior is p(k) ≡ U ({1, . . . , K}).

This will further allow us to retrieve many important statistics and properties about the dataset, for example the segment membership of the data, the overall weights of the modes, we may even synthesizing meaningful "out-distribution" data using a convex combination of any two modes.

However, the conditional likelihood is hard to compute when dealing with high dimensional data, such as images, as data in its raw form do not form mixtures naturally.

At the same time, Generative Adversarial Network (GAN) gives us a way to compute the latent representation z associated with the data x: GAN (Goodfellow et al. (2014) ), introduces a 2-player non-cooperative game by a generator G and a discriminator D. The generator produces samples from the random noise vector z. The Discriminator differentiates between true samples and fake samples. The objective function of the game is given as follows:

EQUATION

The problem of the original GAN methodology is that the association between the latent vector z from a low-dimension latent space to a data sample x (in high-dimension data space) is one way, i.e., one must generate z before generating x, making the backward association from x → z infeasible.

Since the generation of z is independent of x, any arbitrary distribution should theoretically achieve its take, making a uniform distribution or a standard Gaussian distribution a popular choice.

In our proposed work, we are to devise methods in which we are able to compute p(z|k, θ) and p(z|x, θ) to replace p(x|k, θ) and p(k|x, θ) respectively, through the use of GAN training. p(z|x, θ) is achieved by upgrading prior distribution p(z) to a GMM instead of a simple Gaussian. At the same time, the responsibility probability p(k|x, θ) is learned through a classification network which is trained from end-to-end with the GAN. The learned classifier can subsequently be used to classify new data points or perform segmentation on the test set. We also observe through experiments that a smooth linear interpolation can be performed across multiple distribution modes to create the desired effects of "out-distribution" data.

Below we introduce some previous research that are related to our work. GAN Eghbal-zadeh & Widmer (2017) integrate a GMM into the GAN framework. Both of the means and covariance matrices are trainable through the generator loss. Besides, instead of applying the classic adversarial loss as in Equation 1, the authors proposed to use GMM likelihood.

1.1. GAUSSIAN MIXTURE MODEL IN:

Ben-Yosef & Weinshall (2018) also proposed a method named GM-GAN which used Gaussian Mixture to model the distribution over the latent space, in addition to its variant for the conditional generation. The supervised GM-GAN modifies the Discriminator, so that instead of a single scalar, it returns a vector o ∈ R N . Each element of o represents the probability of the given sample being in each class. This will allow that images generated from the generator will be classified by the discriminator as a certain class.

In addition, in the absence of prior knowledge, the GMM is assumed to be uniform for both works above, i.e., ∀k ∈ {1, . . . , K} α k = 1 K . On the contrary, our work makes no such assumption, the segmentation of data can be inferred from the trained model.

1.2. LINEAR INTERPOLATION IN GAN:

Previously researchers have studied linear interpolation in the trained GAN models, where linear interpolations in the noise space lead to semantic interpolations in the generated images. Bojanowski et al. (2017) models the latent space as a unit sphere and learns the correspondence from the sphere space to the data space without the adversarial loss. Such a trained model is able to achieve smooth linear interpolation output between any two random vectors on the unit sphere. Chen et al. (2016) proposed InfoGAN which learns to maximize the mutual information between a subset of latent variables and the observation, i.e. I(c; G(z, c)), where c is the class of the real sample. A learned InfoGAN model can disentangle discrete and continuous latent factors. Thus, linear interpolation can be performed using the continuous latent code, for example, from a thin digit to a wide digit; but not possible on the discrete codes, such as across categories of images.

In the next section, we explain in details how each part of the proposed algorithm works.

1.3. PAPER ORGANIZATION:

In Section 2, we describe how each component of the proposed mechanism works. In Section 3, we demonstrate the performance of the proposed method by comparing it against several baseline models. Section 4 concludes the paper.

2. ARCHITECTURE:

The proposed architecture consists of three networks. First, a classifier C which outputs the possibility that a given image belonging to each Gaussian, and the possibility is used to construct a GMM distribution. Second, a Generator G which produces synthetic samples from GMM random vectors. Third, a Discriminator D which encodes samples to feature vectors and discriminates between real and synthetic samples. The overall architecture design is shown below in Figure 1 . In the following sections, we explain the details of the three networks. 

2.1. THE CLASSIFIER:

For an image x i , the classifier C outputs the probability α i = {α 1 i , . . . , α k i , . . . , α K i } of the image belonging to each Gaussian indexed from 1 to K. These probability values are considered to be the mixture weights of the GMM which corresponds to the specific image, i.e, the density function of the GMM model which corresponds to image The design of the classifier has two alternatives, the first is to share the feature encoding layers with the Discriminator. The shared feature encoding network encodes each image to a feature vector, and the Classifier will be a simple standard linear Soft-max classifier built on top of the features. The second is to build a stand-alone network which contains multiple CNN layers to classify images. Details of both network designs can be found below in Table 1 . In Table 1 , we use acronyms for operations in the table: "Conv" is the convolution operation, of which the kernel and stride size are in the bracket; "Batch norm" is short for batch normalization; "Flatten" refers to the operation that flattens a tensor to 1D array. The same acronyms are also used in Table 2 and Table 3 . D img , D h and D z are related to the datasets we use, the exact value of each are reported in Section 3.

x i is K k=1 α k i N (µ k , Σ k ). µ k and Σ k for k = {1, 2, . . . ,

In Section 3.3, we report the results of both classifier designs in terms of generation performance and computation costs. Parameters of the classifier are optimized by both the adversarial loss and the mutual information loss, the details of both loss functions are discussed in Section 2.2, and how the update is performed is introduced in Section 2.4.

The Classifier is not required for generating new images after the model is fully trained. During testing, a random vector is sampled directly from the GMM model for the generation. The Classifier can then used to assign an unseen image to Gaussians and perform segmentation on the testing set.

2.2. THE GENERATOR:

Given the classification result α i for the input image x i , ideally, multinomial sampling is performed to select one Gaussian out of K. The density function of this sampling is

K k=1 α k i * N (µ k , Σ k

), and a random vector z i is sampled from the chosen Gaussian distribution to perform the image generation.

However, our design expects the training to be performed in an "end-to-end" fashion, and the sampling process would break the chain of the back-propagation from the adversarial loss to the classifier parameters. Therefore, we instead sample one random vector from each Gaussian represented as z i = [z 1 i , . . . , z K i ]. These vectors are used to generate K synthetic images [x 1 i , . . . , xK i ]. α i is used to weigh the adversarial loss calculated from each pair of a generated sample xk i and the real image x i .

In addition, the reparameterization trick is applied to the z i sampling process so that the backpropagation can be used to update the parameters µ k and Σ k of each Gaussian. Instead of sampling

z k i ∼ N (µ k , Σ k ), we define z k i = A k + µ k ∀k ∈ [1, K],

where is sampled as ∼ N (0, I), and

Σ k = A k A k .

The adversarial loss for the proposed framework is thus calculated over K pairs of real and synthetic samples as:

EQUATION

In addition to the adversarial loss, we also encourage the classifier to restore the input classification output from the generated image. Therefore, we also employ mutual information loss between the classification result from the real and the generated image. Following InfoGAN by Chen et al. (2016) where the mutual information loss is formulated as

I(c, G(z, c)) = E c∼P (c),x∼G(z,c) [logQ(c|x)] + H(c).

where c is the class of real sample. We define the mutual information loss in our work as:

EQUATION

The Generator structure used in our experiments is as below in Table 2 . In the table, "LeakyReLU" which is short for "leaky rectified linear unit" and "Tanh" are activation functions.

Stage Name Input Tensors:

Output Tensors

G Linear + Batch norm + LeakyReLU + Reshape 1 × D z D img /4 × D img /4 × 256 Transposed Conv (kernel=5, stride=1) + Batch norm + LeakyReLU D img /4 × D img /4 × 256 D img /4 × D img /4 × 128 Transposed Conv (kernel=5, stride=2) + Batch norm + LeakyReLU D img /4 × D img /4 × 128 D img /2 × D img /2 × 64 Transposed Conv (kernel=5, stride=2) + Tanh D img /2 × D img /2 × 64 D img × D img × D h Table 2:

The network structure of the Generator G.

2.3. THE DISCRIMINATOR:

As we mentioned in Section 2.1, the design of the Discriminator has two options: whether or not the image encoding layers are shared with the classifier. The image encoding network is followed by a standard linear logistic regression to identify the given image to be real or fake. Details of the network are shown below in Table 3 . 

2.4. TRAINING PARAMETERS FOR THE PRIOR DISTRIBUTION:

In our setting, the prior distribution is a GMM model, the two trainable variables are means for K Gaussians µ ∈ R K×Dz and standard deviations for K Gaussians A ∈ R K×Dz×Dz . Both variables are updated by the adversarial loss in addition to the mutual information loss as the training is performed "end-to-end".

Below we give the pseudocode about how the updates are performed on each network in one iteration.

Algorithm 1 Training the proposed model for 1 iteration:

Require:

X = [x 1 , x 2 , . . . , x M ] -M training images in one batch 1: for i = 1 . . . M do 2: Classify x i into α i = [α 0 i , α 1 i , . . . , α k i , . . . , α K i ] 3: for k = 1 . . . K do 4:

∼ N (0, I) 5:

z k = A k + µ k 6: xk i ← G(z k ) 7: Classify xk into αi = [α 0 i , α 1 i , . . . , α k i , . . . , α K i ] 8:

end for 9:

Calculate L adversarial from x i and [x 1 i . . . xK i ] as in Equation 310:

Calculate L I from α i and αi as in Equation 411:

Update the Discriminator with L adversarial 12:

Update both the Generator and the Classifier with L adversarial 13:

Update the Classifier with L I 14: end for

3. EXPERIMENTS:

In this section, we evaluate the performance of the proposed method by comparing it with several baselines.

3.1. EXPERIMENT SETUP:

The datasets we use are the MNIST (LeCun et al. ( 2010)), Fashion-MNIST (Xiao et al. (2017) ) and Oxford-102 Flower (Krizhevsky (2009) ) datasets. The details are listed below in Table 4 . In Under review as a conference paper at ICLR 2020 particular, we only select a subset of Oxford-102 to perform the training, which is the images which belong to the first 10 classes. For experiments performed on each dataset, we used different hyperparameters, the details are listed in Table 5 

3.2. LINEAR INTERPOLATION ACROSS GAUSSIAN:

Below in Figure 2 , we show samples generated by the proposed model trained on several datasets in the completely unsupervised manner. We set the number of Gaussians equal to the total number of classes of the dataset in all experiments. When we are performing the linear interpolation as in the right panels, the random vector z for each image generation is calculated as z = A k + µ k ∀k ∈ {1, . . . , K}, where is sampled as ∼ N (0, I). is kept the same for all images for each dataset.

We can draw two conclusions from the results in Figure 2 . First, a fully trained proposed method can learn to "allocate" each class of image to a Gaussian. Second, the trained model can be used to perform smooth linear interpolation between Gaussians and even among more than two Gaussians.

In Figure 3 , we demonstrate the linear interpolation performed over 3 categories. The proportion of Gaussians of the synthetic images can be set manually.

3.3. IMAGE GENERATION QUALITY:

The generation performance is measured with two commonly used metrics: Inception score (Salimans et al. ( 2016)) and Frchet Inception Distance (FID) score (Heusel et al. (2017) ).

Inception score is calculated as I = exp(E x D KL (p(y|x)||p(y))), and a higher value generally indicates a better performance.

where x is a generated image and y is the label predicted by the Inception model (Szegedy et al. (2015) ).

FID score is another metrics that measures the image generation quality. A lower value shows a better image quality and diversity. It calculates the difference between real images x and generated images g as FID(x, g)

EQUATION

2 ).

Note that limitations of both Inception and FID score have been pointed out in several previous literature (Barratt & Sharma (2018) , Lucic et al. (2018) ), and there is currently no "perfect" metrics at this moment. These two metrics are used as an indication rather than a hard measure.

This evaluation is performed on the Oxford-102 dataset. As the Inception score is suggested to be evaluated on a large enough number of samples, we generate 5K synthetic samples to calculate both values. Below in Figure 4 , we show the plot of Inception scores and FID scores calculated over the training epochs. In Table 6 we report the number of trainable parameters, the best Inception and FID score.

While the four algorithms use the same network settings for the Generator and Discriminator, it is clear to see that the proposed method, whether or not the encoding layers are shared between the 

4. CONCLUSION:

In this paper we propose a novel framework which we incorporate the GMM to model the latent prior distribution. We also use a Classifier that learns to categorize an image to Gaussians. The Classifier is trained with the GAN model in an "end-to-end" fashion. We demonstrate through experiments that the proposed method surpass previous baselines in terms of the image generation performance with only minor growth on the size of the network. A trained model is also able to perform smooth linear interpolation across Gaussians, i.e. generate images with mixed styles from multiple categories.

:



Figures and Tables:

Figure 1: The overall architecture. The feed-forward logic of the Classifier C, the Generator G and the Discriminator D are marked with different colours.

Figure2: Samples generated by the proposed models trained on the MNIST (top panels), Fashion-MNIST (middle panels) and CIFAR-10 (bottom panels) datasets. In the left panels, the Gaussian mixture contains K = 10 Gaussians, and each row contains images sampled from a different Gaussian. In the right panels, each row shows the linear interpolation result from one Gaussian to another. Note that the index of Gaussian does not necessarily correspond to the actual digit, generated images are reordered as the order of the digit only for demonstration purpose.

Figure 3: Linear interpolation over 3 Gaussians on the MNIST, Fashion-MNIST and Oxford-102 dataset. Each row and each column contains linear interpolation performed between 2 Gaussians. Images along the diagonal shows the interpolation across 3 Gaussians.

Figure 4: Inception score and FID score over training epochs.

K} are trainable parameters. During test time, the learned Classifier can easily output p(k|x, θ) for any test image x.

Table content: <table/>

The network structure of the Classifier C. Each column from left to right reports the network, sub-stage name, name of operations and shapes of input and output tensors.

Table content: <table><tr><td colspan="2">Stage Sub-stage</td><td>Name</td><td>Input Tensors</td><td>Output Tensors</td></tr><tr><td>C</td><td>Encoding Network (if shared)</td><td>Conv (kernel=5, stride=2) + LeakyReLU</td><td>Dimg × Dimg × Dh</td><td>Dimg/2 × Dimg/2 × 64</td></tr><tr><td/><td/><td colspan="2">Conv (kernel=4, stride=2) + Batch norm + LeakyReLU + Flatten Dimg/2 × Dimg/2 × 64</td><td>Dimg/4 × Dimg/4 × 128</td></tr><tr><td/><td>Encoding Network (if not shared)</td><td>Conv (kernel=5, stride=1) + ReLU</td><td>Dimg × Dimg × Dh</td><td>Dimg × Dimg × 32</td></tr><tr><td/><td/><td>MaxPool (pool size=2, stride=2)</td><td>Dimg × Dimg × Dh</td><td>Dimg/2 × Dimg/2 × 32</td></tr><tr><td/><td/><td>Conv (kernel=5, stride=2) + ReLU + Flatten</td><td>Dimg/2 × Dimg/2 × 32</td><td>Dimg/4 × Dimg/4 × 16</td></tr><tr><td/><td/><td>Linear + ReLU</td><td>Dimg/4 × Dimg/4 × 16</td><td>1024</td></tr><tr><td/><td>Classification Network (if shared)</td><td>Linear</td><td colspan="2">Dimg/4 × Dimg/4 × 128 K</td></tr><tr><td/><td colspan="2">Classification network (if not shared) Linear</td><td>1024</td><td>K</td></tr></table>

The network structure of the Discriminator D.

Table content: <table/>

.

Table content: <table><tr><td>Dataset</td><td colspan="5">Number of Classes Data Dimension Train Samples Validation Samples Test Samples</td></tr><tr><td>MNIST</td><td>10</td><td>28 × 28 × 1</td><td>60,000</td><td>-</td><td>10,000</td></tr><tr><td>Fashion-MNIST</td><td>10</td><td>28 × 28 × 1</td><td>60,000</td><td>-</td><td>10,000</td></tr><tr><td colspan="2">Oxford-102 Flower 102</td><td>64 × 64 × 3</td><td>1,020</td><td>1,020</td><td>6,149</td></tr></table>

Statistics of the different datasets used in the empirical evaluation.

Table content: <table><tr><td>Dataset</td><td colspan="4">Number of Epochs Learning Rate γ D img D h</td></tr><tr><td>MNIST</td><td>200</td><td>0.0002</td><td>28</td><td>1</td></tr><tr><td colspan="2">Fashion-MNIST 200</td><td>0.0002</td><td>28</td><td>1</td></tr><tr><td>Oxford-102</td><td>10,000</td><td>0.0002</td><td>64</td><td>3</td></tr></table>

Hyper-parameters for the training performed on each dataset.

Table content: <table/>

Number of parameters, Inception scores and FID scores of the proposed method and the baselines.

Table content: <table/>

Peer review (iclr2020):
This paper proposes a modification of GANs where the latent space follows a distribution modelled by a Gaussian Mixture Model. While the idea of using GMMs in GANs is not novel, the main contribution of the paper is to add a classification models that enables posterior inference. The whole model is trained jointly end-to-end using both an adversarial loss and a mutual information loss. The procedure is then tested on MNIST, Fashion-MNIST and a subset of Oxford-102 Flower. While the problem of posterior inference is interesting, the novelty of the paper is quite limited. The overall structure is clear, although writing can be improved. However, a part from the form, I have other concerns about the paper: - The main purpose of the paper is to tackle large image datasets that are hard to address by classical GMMs. This being said, the used datasets are composed of small and modal enough images that it seems hard to validate the claim of the paper using only these data. It seems to me that for the claim of the paper to be verified, larger scale/more complex datasets are needed. - In all experiments, the authors suppose they have access to the number of classes/modes in the data, which is a huge assumption. It would be interesting to see if it would be possible to automatically accurately select the number of modes, e.g. on a held out validation set. - One problem of GANs that the authors do not seem to consider is mode collapse. It would be interesting to do experiments with unbalanced datasets (e.g. MNIST 1 vs all) to see if the proposed architecture will model the data correctly. - I am confused about the use of Mutual Information loss. The author claim that they would like to enforce each of the generated images to be from the same class as the input image. This would make sense if the authors used the multinomial sampling for the latent variable generation. However, the K generated samples are from K different Gaussians. It seems unreasonable to require of the classifier to render the same result. - On the same note, in order to both use the classical multinomial sampling in GMMs and not break the backpropagation, have the authors considered updating the classifier and the GAN in separately in expectation-maximization fashion? - Finally, I don*t see the point of weighting the adversarial loss by the weights of Gaussians. All the generated images are fake and should be equally detected as such. Although the idea is interesting, I think the paper, at its current status, not ready for publication. Minor: - P. 4: Generator: The sampling density from the multinomial distribution seems incorrect. However, as the authors skip the sampling step to be able to back propagate through the model, this is not significant. - P. 1: The paragraph before the last: may even synthesizing -> synthesize - P. 3: Architecture: Possibility -> Probability - P. 7: Figure title: CIFR10 -> Oxford-102 - Algorithm notation: *Indexes for alpha and alpha_hat from 1 and not 0 to be consistent with the rest of the text *Add hats to the entries of alpha_hat * LI from alpha_i and alpha_hat_i as in Equation 4 -> maybe change alpha_hat by x_hat to be consistent with the notation of eq.4 (although the meaning is clear here).