Title: SPATIAL VARIATIONAL AUTO-ENCODING VIA MATRIX-VARIATE NORMAL DISTRIBUTIONS

Abstract: The key idea of variational auto-encoders (VAEs) resembles that of traditional auto-encoder models in which spatial information is supposed to be explicitly encoded in the latent space. However, the latent variables in VAEs are vectors, which can be interpreted as multiple feature maps of size 1x1. Such representations can only convey spatial information implicitly when coupled with powerful decoders. In this work, we propose spatial VAEs that use feature maps of larger size as latent variables to explicitly capture spatial information. This is achieved by allowing the latent variables to be sampled from matrix-variate normal (MVN) distributions whose parameters are computed from the encoder network. To increase dependencies among locations on latent feature maps and reduce the number of parameters, we further propose spatial VAEs via low-rank MVN distributions. Experimental results show that the proposed spatial VAEs outperform original VAEs in capturing rich structural and spatial information.

1. INTRODUCTION:

The mathematical and computational modeling of probability distributions in high-dimensional space and generating samples from them are highly useful yet very challenging. With the development of deep learning methods, deep generative models have been shown to be effective and scalable (Kingma & Welling, 2013; Rezende et al., 2014; Burda et al., 2015; Gulrajani et al., 2016; Makhzani et al., 2015; Goodfellow et al., 2014; Radford et al., 2015) in capturing probability distributions over high-dimensional data spaces and generating samples from them. Among them, variational auto-encoders (VAEs) (Kingma & Welling, 2013; Rezende et al., 2014; Doersch, 2016) are one of the most promising approaches. In machine learning, the auto-encoder architecture is applied to train scalable models by learning latent representations. For image modeling tasks, it is preferred to encode spatial information into the latent space explicitly. However, the latent variables in VAEs are vectors, which can be interpreted as 1 × 1 feature maps with no explicit spatial information. While such lack of explicit spatial information does not lead to major performance problems on simple tasks such as digit generation from the MNIST dataset (LeCun et al., 1998b) , it greatly limits the model's abilities when images are more complicated (Krizhevsky & Hinton, 2009; Liu et al., 2015) .

To overcome this limitation, we propose spatial VAEs that employ d × d (d > 1) feature maps as latent representations. Such latent feature maps are generated from matrix-variate normal (MVN) distributions whose parameters are computed from the encoder network. Specifically, MVN distributions are able to generate feature maps with appropriate dependencies among locations. To increase dependencies among locations on latent feature maps and reduce the number of parameters, we further propose spatial VAEs via low-rank MVN distributions. In this low-rank formulation, the mean matrix of MVN distribution is computed as the outer product of two vectors computed from the encoder network. Experimental results on image modeling tasks demonstrate the capabilities of our spatial VAEs in complicated image generation tasks.

It is worth noting that the original VAEs can be considered as a special case of spatial VAEs via MVN distributions. That is, if we set the size of feature maps generated via MVN distributions to 1 × 1, spatial VAEs via MVN distributions reduce to the original VAEs. More importantly, when the size of feature maps is larger than 1 × 1, direct structural ties have been built into elements of the feature maps via MVN distributions. Thus, our proposed spatial VAEs are intrinsically different with the original VAEs when the size of feature maps is larger than 1 × 1. Specifically, our proposed spatial VAEs cannot be obtained by enlarging the size of the latent representations in the original VAEs.

2. BACKGROUND AND RELATED WORK:

2.1 AUTO-ENCODER ARCHITECTURES Auto-encoder (AE) is a model architecture used in tasks like image segmentation (Zeiler et al., 2010; Ronneberger et al., 2015; Long et al., 2015) , machine translation (Bahdanau et al., 2014; Sutskever et al., 2014) and denoising reconstruction (Vincent et al., 2008; 2010) . It consists of two parts: an encoder that encodes the input data into lower-dimensional latent representations and a decoder that generates outputs by decoding the representations. Depending on different tasks, the latent representations will focus on different properties of input data. Nevertheless, these tasks usually require outputs to have similar or exactly the same structure as inputs. Thus, structural information is expected to be preserved through the encoder-decoder process.

In computer vision tasks, structural information usually means spatial information of images. There are two main strategies to preserve spatial information in AE for image tasks. One is to apply very powerful decoders, like conditional pixel convolutional neural networks (PixelCNNs) (Oord et al., 2016; van den Oord et al., 2016; Salimans et al., 2017; Gulrajani et al., 2016) , that generate output images pixel-by-pixel. In this way, the decoders can recover spatial information in the form of dependencies among pixels. However, pixel-by-pixel generation is very slow, resulting in major speed problems in practice. The other method is to let the latent representations explicitly contain spatial information and apply decoders that can make use of such information. To apply this strategy for image tasks, usually the latent representations are feature maps of size between the size of a pixel (1 × 1) and that of the input image, while the decoders are deconvolutional neural networks (DCNNs) (Zeiler et al., 2010) . Since most computer vision tasks only require high-level spatial information like relative locations of objects instead of detailed relationships among pixels, preserving only rough spatial information is enough, and this strategy is proved effective and efficient.

2.2. VARIATIONAL AUTO-ENCODERS:

In unsupervised learning, generative models aim to modeling the underlying data distribution. Formally, for data space X , let p true (x) denote the probability density function (PDF) of the true data distribution for x ∈ X . Given a dataset D = {x (i) } N i=1 of i.i.d samples from X , generative models try to approximate p true (x) using a model distribution p θ (x) where θ represents model parameters. To train the model, maximum likelihood (ML) inference is performed on θ; that is, parameters are updated to optimize log p θ (D) = log p θ (x (1) , . . . , x (N ) ) = N i=1 log p θ (x (i) ). The approximation quality of p θ (x) relies on the generalization ability of the model. In machine learning, it highly depends on learning latent representations which can encode common features among data samples and disentangle abstract explanatory factors behind the data (Bengio et al., 2013) . In data generation tasks, we apply p θ (x) = p θ (x|z)p θ (z)dz for modeling, where p θ (z) is the PDF of the distribution of latent representations and p θ (x|z) represents a complex mapping from the latent space to the data space. A major advantage of using latent representations is dimensionality reduction of data since they are low-dimensional. The prior p θ (z) can be simple and easy to model while the mapping represented by p θ (x|z) can be learned through complicated deep learning models automatically.

Recently, Kingma & Welling (2013) point out that the above model has intractability problems and can only be trained by costly sampling-based methods. To tackle this, they propose variational auto-encoders (VAEs), which instead maximize a variational lower bound of the log-likelihood as

EQUATION

where q φ (z|x) is an approximation model to the intractable p θ (z|x), parameterized by φ, D KL [•] represents the Kullback-Leibler divergence. In VAEs, p θ (x|z) = N (x; f θ (z), σ 2 I), q φ (z|x) = N (z; µ φ (x), Σ φ (x)), and p θ (z) = N (z; 0, I) are modeled as multivariate Gaussian distributions with diagonal covariance matrices. Here, f θ (z), µ φ (x) and Σ φ (x) are computed with deep neural networks like CNNs. Figure 1 shows the architecture of VAEs. The model parameters θ and φ can be trained using the reparameterization trick (Rezende et al., 2014) , where the sampling process z ∼ q φ (z|x) = N (z; µ φ (x), Σ φ (x)) is decomposed into two steps as

∼ N ( ; 0, I), z = µ φ (x) + Σ 1 2 φ (x) * .

(2)

3. SPATIAL VARIATIONAL AUTO-ENCODERS:

In this section, we analyze a problem of the original VAEs and propose spatial VAEs in Section 3.1 to overcome it. Afterwards, several ways to implement spatial VAEs are discussed. A naïve implementation is introduced and analyzed in Section 3.2, followed by a method that incorporates the use of matrix-variate normal (MVN) distributions in Section 3.3. Finally, we propose our final model, spatial VAEs via low-rank MVN distributions, by applying a low-rank formulation of MVN distributions in Section 3.4.

3.1. OVERVIEW:

Note that p θ (x|z) and q φ (z|x) in VAEs resemble the encoder and decoder, respectively, in AE for image reconstruction tasks, where z represents the latent representations. However, in VAE, z is commonly a vector, which can be considered as multiple 1 × 1 feature maps. While z may implicitly preserve some spatial information of the input image x, it raises the requirement for a more complex decoder. Given a fixed architecture, the hypothesis space of decoder models is limited. As a result, the optimal decoder may not lie in the hypothesis space (Zhao et al., 2017) . This problem significantly hampers the performance of VAEs, especially when spatial information is important for images in X .

Based on the above analysis, it is beneficial to either have larger hypothesis space for decoders or let z explicitly contain spatial information. Note that these two methods correspond to the two strategies introduced in Section 2.1. Gulrajani et al. (2016) follow the first strategy and propose PixelVAEs whose decoders are conditional PixelCNNs (van den Oord et al., 2016) instead of simple DCNNs.

As conditional PixelCNNs themselves are also generative models, PixelVAEs can be considered as conditional PixelCNNs with the conditions replaced by z. In spite of their impressive results, the performance of PixelVAEs and conditional PixelCNNs is similar, which indicates that conditional PixelCNNs are responsible for capturing most properties of images in X . In this case, z contributes little to the performance. In addition, applying conditional PixelCNNs leads to very slow generation process in practice. In this work, the second strategy is explored by constructing spatial latent representations z in the form of feature maps of size larger than 1 × 1. Such feature maps can explicitly contain spatial information. We term VAEs with spatial latent representations as spatial VAEs.

The main distinction between spatial VAEs and the original VAEs is the size of latent feature maps. By having d × d (d > 1) feature maps instead of 1 × 1 ones, the total dimension of the latent representations z significantly increases. However, spatial VAEs are essentially different from the original VAEs with a higher-dimensional latent vector z. Suppose the vector z is extended by d 2 times in order to match the total dimension, the number of hidden nodes in each layer of decoders will explode correspondingly. This results in an explosion in the number of decoders' parameters, which slows down the generation process. Whereas in spatial VAEs, decoders becomes even simpler since d × d is closer to the required size of output images. From the other side, when using decoders of similar capacities, spatial VAEs must have higher-dimensional latent representations than the original VAEs. It is demonstrated that this only slightly influences the training process by requiring more outputs from encoders, while the generation process that only involves decoders remains unaffected.

Our experimental results show that with proper designs, spatial VAEs substantially outperform the original VAEs when applying similar decoders.

3.2. NAÏVE SPATIAL VAES:

To achieve spatial VAEs, a direct and naïve way is to simply reshape the original vector z into N feature maps of size d×d. But this naïve way is problematic since the sampling process does not change.

Note that in the original VAEs, the vector z is sampled from q φ (z|x) = N (z; µ φ (x), Σ φ (x)). The 

k diag  k  ( ) k diag  ( ) diag  ( ) diag    ( ) k diag  d d d d N N ~( , ) z MVN   N d d ~( , ) z N   C 1 1 ( ) diag   C C

Figure 1 : Illustration of the differences between the proposed spatial VAEs via low-rank MVN distributions and the original VAEs. At the top is the architecture of the original VAEs where the latent z is a vector sampled from a multivariate Gaussian distribution with a diagonal covariance matrix. Below is the proposed model which is explained in detail in Section 3.4. Briefly, it modifies the sampling process by incorporating a low-rank formulation of the MVN distributions and produces latent representations that explicitly retain spatial information. covariance matrix Σ φ (x) is diagonal, meaning each variable is uncorrelated. In particular, for multivariate Gaussian distributions, uncorrelation implies independence. Therefore, z's components are independent random variables and the variances of their distributions correspond to entries on the diagonal of Σ φ (x). Specifically, suppose z is a C-dimensional vector, the i th component is a random variable that follows the univariate normal distribution as z i ∼ N (z i ; µ φ (x) i , diag(Σ φ (x)) i ), i = 1, . . . , C, where diag(•) represents the vector consisting of a matrix's diagonal entries. After applying the reparameterization trick, we can rewrite Equation 2 as i ∼ N ( i ; 0, 1),

EQUATION

To sample N feature maps of size d × d in naïve spatial VAEs, the above process is followed by a reshape operation while setting C = d 2 N .

However, between two different components z i and z j , the only relationship is that their respective distribution parameters 

(µ φ (x) i , diag(Σ φ (x)) i ) and (µ φ (x) j , diag(Σ φ (x)) j )

∈ R m×m , det(Ω) > 0, Ψ ∈ R n×n , det(Ψ) > 0, if vec(A T ) follows the multivariate normal distribution N (vec(A T ); vec(M T ), Ω ⊗ Ψ).

Here, ⊗ denotes the Kronecker product and vec(•) denotes transforming a R m×n matrix into an mn-dimensional vector by concatenating the columns.

In MVN distributions, Ω and Ψ capture the relationships across rows and columns, respectively, of a matrix. By constructing the covariance matrix through the Kronecker product of these two matrices, dependencies among values in a matrix can be modeled. In spatial VAEs, a feature map F can be considered as a R d×d matrix that follows a MVN distribution N d,d (F ; M, Ω ⊗ Ψ), where Ω ∈ R d×d and Ψ ∈ R d×d are diagonal matrices. Although within F the random variables corresponding to each location are still independent since Ω ⊗ Ψ is diagonal, MVN distributions are able to add direct structural ties among locations through their variances. For example, for two locations (i 1 , j 1 ) and (i 2 , j 2 ) in F ,

EQUATION

Here, F (i1,j1) and F (i2,j2) are independently sampled from two univariate Gaussian distributions. However, the variances diag(Ω⊗Ψ) i1 * j1 and diag(Ω⊗Ψ) i2 * j2 have built direct interactions through the Kronecker product. Based on this, we propose spatial VAEs via MVN distributions, which samples N feature maps of size d × d from N independent MVN distributions as

EQUATION

where M kφ (x), Ω kφ (x) and Ψ kφ (x) are computed through the encoder. Here, compared to the original VAEs, q φ (z|x) is replaced but p θ (z) remains the same. Since MVN distributions are defined based on multivariate Gaussian distributions, the term D KL [q φ (z|x)|p θ (z)] in Equation 1 can be calculated in a similar way.

To demonstrate the differences with naïve spatial VAEs, we reexamine the original VAEs. Note that naïve spatial VAEs have the same sampling process as the original VAEs. The original VAE samples a C = d 2 N -dimensional vector z from q φ (z|x) = N (z; µ φ (x), Σ φ (x)) where µ φ (x) is a C-dimensional vector and Σ φ (x) is a R C×C diagonal matrix. Because Σ φ (x) is diagonal, it can be represented by the C-dimensional vector diag(Σ φ (x)). To summarize, the encoder of the original VAEs outputs 2C = 2d 2 N values which are interpreted as µ φ (x) and diag(Σ φ (x)).

In spatial VAEs via MVN distributions, according to Equation 5, M kφ (x) is a R d×d matrix while Ω kφ (x) and Ψ kφ (x) are R d×d diagonal matrices that can be represented by ddimensional vectors. In this case, the required number of outputs from the encoder is changed to

(d 2 + 2d)N , corresponding to [M 1φ (x), . . . , M N φ (x)], [diag(Ω 1φ (x)), . . . , diag(Ω N φ (x))] and [diag(Ψ 1φ (x)), . . . , diag(Ψ N φ (x))].

As has been explained in Section 3.2, since Ω kφ (x) ⊗ Ψ kφ (x) is diagonal, sampling the matrix F k is equivalent to sampling d × d scalar numbers from d × d independent univariate normal distributions. So the modified sampling process with the reparameterization trick is ,j) . Here, we take advantage of the fact that for diagonal matrices, the Kronecker product is equivalent to the out-product of vectors. To be specific, suppose D 1 and D 2 are two R d×d diagonal matrices, then d 1 = diag(D 1 ) and d 2 = diag(D 2 ) are two d-dimensional vectors and satisfy

(i,j,k) ∼ N ( (i,j,k) ; 0, 1), z (i,j,k) = µ kφ (x) (i,j) + diag(Ω kφ (x) ⊗ Ψ kφ (x)) 1 2 i * j * (i,j,k) , i, j = 0, . . . , d, k = 1, . . . , N, (6) where diag(Ω kφ (x) ⊗ Ψ kφ (x)) i * j = [diag(Ω kφ (x))diag T (Ψ kφ (x))] (i

EQUATION

It is worth noting that, compared to naïve spatial VAEs, the required number of outputs from the encoder decreases from 2d 2 N to (d 2 + 2d)N . As a result, spatial VAEs via MVN distributions leads to a simpler model while adding structural ties among locations. Note that the original VAEs can be considered as a special case of the spatial VAEs via MVN distributions. That is, if we set d = 1, spatial VAEs via MVN distributions reduce to the original VAEs.

3.4. A LOW-RANK FORMULATION:

The use of MVN distributions makes locations directly related to each other within a feature map by adding restrictions on variances. However, in probability theory, variance only measures the expected distance from the mean. To have more direct relationships, it is preferred to have restricted means. In this section, we introduce a low-rank formulation of MVN distributions (Allen & Tibshirani, 2010) for spatial VAEs.

The low-rank formulation of a MVN distribution N m,n (M, Ω ⊗ Ψ) is denoted as N m,n (µ, ν, Ω ⊗ Ψ) where the mean matrix M is computed by the out-product µν T instead. Here, µ and ν are mdimensional and n-dimensional vectors, respectively. Similar to computing the covariance matrix through the Kronecker product of two separate matrices, it explicitly forces structural interactions among entries of the mean matrix. Applying this low-rank formulation leads to our final model, spatial VAEs via low-rank MVN distributions, which is illustrated in Figure 1 . By using two distinct d-dimensional vectors to construct M iφ (x) ∈ R d×d , Equation 5 is modified as

EQUATION

where µ kφ (x) and ν kφ (x) are d-dimensional vectors. For the encoder, the number of outputs is further reduced to 4dN from (d 2 + 2d)N , replacing d 2 N outputs for (M 1φ (x), . . . , M N φ (x)) with dN outputs for (µ 1φ (x), . . . , µ N φ (x)) and another dN outputs for (ν 1φ (x), . . . , ν N φ (x)). In contrast to Equation 6, the two-step sampling process can be expressed as

EQUATION

where diag(

Ω kφ (x) ⊗ Ψ kφ (x)) i * j = [diag(Ω kφ (x))diag T (Ψ kφ (x))] (i,j) .

As has been demonstrated in Section 3.1, spatial VAEs require more outputs from encoders than the original VAEs, which slows down the training process. Spatial VAEs via low-rank MVN distributions properly address the problem while achieving appropriate spatial latent representations. According to the experimental results, they outperform the original VAEs in several image generation tasks when similar decoders are used.

4. EXPERIMENTAL STUDIES:

We use the original VAEs as the baseline models in our experiments, as most recent improvements on VAEs are derived from the vector latent representations and can be easily incorporated into our matrix-based models. To elucidate the performance differences of various spatial VAEs, we compare the results of three different spatial VAEs as introduced in Section 3; namely naïve spatial VAEs, spatial VAEs via MVN distributions and spatial VAEs via low-rank MVN distributions. We train the models on the CelebA, CIFAR-10 and MNIST datasets, and analyze sample images generated from the models to evaluate the performance. For the same task, the encoders of all compared models are composed of the same convolutional neural networks (CNNs) and a fully-connected output layer (LeCun et al., 1998a; Krizhevsky et al., 2012) . While the fully-connected layer may differ as required by different numbers of output units, it only slightly affects the training process. As discussed in Section 3.1, it is reasonable to compare spatial VAEs with the original VAEs in the case that their decoders have similar architectures and model capabilities. Therefore, following the original VAEs, deconvolutional neural networks (DCNNs) are used as decoders in spatial VAEs. Meanwhile, the total number of trainable parameters in the decoders of all compared models are set to be as similar as possible while accommodating different input sizes.

4.1. CELEBA:

The CelebA dataset contains 202, 599 colored face images of size 64 × 64. The generative models are supposed to generate faces that are similar but not exactly the same to those in the dataset. For this task, the CNNs in the encoders have 3 layers while the decoders are 5 or 6-layer DCNNs corresponding to spatial VAEs and the original VAEs, respectively. This difference is caused by the Figure 2 shows sample face images generated by the original VAEs and three different variants of spatial VAEs. It is clear that spatial VAEs can generate images with more details than the original VAEs. Due to the lack of explicit spatial information, the original VAEs produce face images with little details like hair near the borders. While naïve spatial VAEs seem to address this problem, most faces have only incomplete hairs as naïve spatial VAEs cannot capture the relationships among different locations. Theoretically, spatial VAEs via MVN distributions are able to incorporate interactions among locations. However, the results are strange faces with some distortions. We believe the reason is that adding dependencies among locations through restrictions on distribution variances is not effective and sufficient. Spatial VAEs via low-rank MVN distributions that have restricted means tackle this well and generate faces with appealing visual appearances.

4.2. CIFAR-10:

The CIFAR-10 dataset consists of 60, 000 color images of 32 × 32 in 10 classes. VAEs usually perform poorly in generating photo-realistic images since there are significant differences among images in different classes, indicating that the underlying true distribution of the data is a multimodel. In this case, VAEs tend to output very blurry images (Theis et al., 2015; Goodfellow et al., 2014; Goodfellow, 2016) . However, comparison among different models can still demonstrate the differences in terms of generative capabilities. In this experiment, we set d = 3 and N = 128, and the dimension of z in the original VAEs is 150. The encoders have 4 layers while the decoders have 4 or 5 layers. Some sample images are provided in Figure 3 . The original VAEs only produce images composed of several colored areas, which is consistent to the results of a similar model reported in Rezende et al. (2014) . It is obvious that all three implementations of spatial VAEs generate images with more details. However, naïve spatial VAEs still produce meaningless images as there is no relationship among different parts. The images generated by spatial VAEs via MVN distributions look like some distorted objects, which have similar problems to the results of the CelebA dataset. Again, spatial VAEs via low-rank MVN distributions outperform the other models, producing blurry but object-like images. We perform quantitative analysis on real-valued MNIST dataset by employing the Parzen window log-likelihood estimates (Breuleux et al., 2011) . This evaluation method is used for several generative models where the exact likelihood is not tractable (Goodfellow et al., 2014; Makhzani et al., 2015) . The results are reported in Table 2 . Despite of the difference in visual quality of generated images, spatial VAE via low-rank MVN distributions shares similar quantitative results with the original VAE. Note that generative models for images are supposed to capture the underlying data distribution by maximizing log-likelihood and generate images that are similar to real ones. However, it has been pointed in Theis et al. (2015) that these two objectives are not consistent, and generative models need to be evaluated directly with respect to the applications for which they were intended. A model that can generates samples with good visual appearances may have poor average log-likelihood on test dataset and vice versa. Common examples of deep generative models are VAEs and generative adversarial networks (GANs) (Goodfellow et al., 2014) . VAEs usually have higher average log-likelihood while GANs can generate more photo-realistic images. This is basically caused by the different training objectives of these two models (Goodfellow, 2016) . Currently there is no commonly accepted standard for evaluating generative models.

4.4. TIMING COMPARISON:

To show the influence of different spatial VAEs to the training process, we compare the training time on the CelebA dataset. Theoretically, spatial VAEs slow down training due to the larger numbers of outputs from encoders. To keep the number of trainable parameters in decoders roughly equal, we set the dimension of z in the original VAEs to be 81 while d = 3 and N = 64 for spatial VAEs.

According to Section 3, the numbers of outputs from their encoders are 162, 1152, 960, and 768 for the original VAE, naïve spatial VAE, spatial VAE via MVN distributions and spatial VAE via lowrank MVN distributions, respectively. We train our models on a Nvidia Tesla K40C GPU and report the average time for training one epoch in Table 2 . Comparisons of the time for generating 10, 000 images are also provided to show that the increase in the total dimension of latent representations does not affect the generation process.

The results show consistent relationships between the training time and the number of outputs from encoders; that is, spatial VAEs cost more time than the original VAE but spatial VAEs via low-rank MVN distributions can alleviate this problem. Moreover, spatial VAEs only slightly slow down the training process since they only affect one single layer in the models. 

:



Figures and Tables:

Figure 2: Sample face images generated by different VAEs when trained on the CelebA dataset. The first and second rows shows training images and images generated by the original VAEs. The remaining three rows are the results of naïve spatial VAEs, spatial VAEs via MVN distributions and spatial VAEs via low-rank MVN distributions, respectively.

Figure 3: Sample images generated by different VAEs when trained on the CIFAR-10 dataset. From top to bottom, the five rows are training images and images generated by the original VAEs, naïve spatial VAEs, spatial VAEs via MVN distributions, spatial VAEs via low-rank MVN distributions, respectively.

are both computed from x. Such dependencies are implicit and weak. It is obvious that after reshaping, there is no direct relationship among locations within each feature map, while spatial latent representations should contain spatial information like dependencies among locations. To overcome this limitation, we propose spatial VAEs via matrix-variate normal distributions.3.3 SPATIAL VAES VIA MATRIX-VARIATE NORMAL DISTRIBUTIONSInstead of obtaining N feature maps of size d × d by first sampling a d 2 N -dimensional vector from multivariate normal distributions and then reshaping, we propose to directly sample d×d matrices as Definition: A random matrix A ∈ R m×n is said to follow a matrix-variate normal distribution N m,n (A; M, Ω⊗Ψ) with mean matrix M ∈ R m×n and covariance matrix Ω⊗Ψ, where Ω

Table content: <table/>

Parzen window log-likelihood estimates of test data on the MNIST dataset. We follow the same procedure as inGoodfellow et al. (2014).

Table content: <table><tr><td>Model</td><td>Log-Likelihood</td></tr><tr><td>Original VAE</td><td>297</td></tr><tr><td>Naïve spatial VAE</td><td>275</td></tr><tr><td>Spatial VAE via MVN distributions</td><td>267</td></tr><tr><td>Spatial VAE via low-rank MVN distributions</td><td>296</td></tr><tr><td>4.3 MNIST</td><td/></tr></table>

Training and generation time of different models when trained on the CelebA dataset using a Nvidia Tesla K40C GPU. The average time for training one epoch and the time for generating 10, 000 images are reported and compared. this work, we propose spatial VAEs for image generation tasks, which improve VAEs by requiring the latent representations to explicitly contain spatial information of images. Specifically, in spatial VAEs, d × d (d > 1) feature maps are sampled to serve as spatial latent representations in contrast to a vector. This is achieved by sampling the latent feature maps from MVN distributions, which can model dependencies between the rows and columns in a matrix. We further propose to employ a lowrank formulation of MVN distributions to establish stronger dependencies. Qualitative results on different datasets show that spatial VAEs via low-rank MVN distributions substantially outperform the original VAEs.

Table content: <table><tr><td>Model</td><td colspan="2">Training time Generation time</td></tr><tr><td>Original VAE</td><td>167.0309s</td><td>1.3892s</td></tr><tr><td>Naïve spatial VAE</td><td>178.8601s</td><td>1.3676s</td></tr><tr><td>Spatial VAE via MVN distributions</td><td>177.4387s</td><td>1.3767s</td></tr><tr><td>Spatial VAE via low-rank MVN distributions</td><td>172.9639s</td><td>1.3686s</td></tr><tr><td>5 CONCLUSION</td><td/><td/></tr></table>

Peer review (iclr2018):
 - Summary The paper proposes a VAE implementation that allows the latent space to carry useful spatial information. Specifically, it proposes using a convolutional encoder that outputs a spatial feature map for the VAE encoder where the feature map is used to infer the means and variances of the latent representation. The sampled latent code is then passed to a deconvolutional decoder for reconstructing the input image. In order to properly sampling from the spatial feature maps in a way that respects the interaction between the neurons at different locations, the paper proposes modeling the latent space using a matrix variate normal distribution. Furthermore, it proposed a low-rank modeling for the pixel mean. In the experiment section, improved log-likelihood scores on the MNIST dataset are shown across different variants of the proposed method as well as the VAE baseline. Visualization comparison with different variants of the proposed method as well as the VAE baseline was reported using the CIFA10 and CelebA dataset. - Missing comparison to closely related works As also reviewed in the paper, PixelVAEs are closely related. However, the paper fails to include a comparison with the PixelVAE work. In addition, several papers improved upon the VAE work including VAE with IAF (NIPS*16). To further establish the benefit of the proposed method, proper comparison with the prior works should be included. Also, VAE methods are often benchmarked using the CIFAR10 datasets using the log-likelihood score. The paper should also report performances using the same setting. This way the reader can better understand the performance of the proposed method. - Poor justification It is not clear from the paper why using the matrix variate normal distribution or the low rank approximation is a better approach, other than reducing the number of parameters in the model. If the goal is reduce number of parameters, isn*t it can be achieved by reducing the feature dimension?