Generative Image Modeling Using Spatial LSTMs

Introduction

The last few years have seen tremendous progress in learning useful image representations . While early successes were often achieved through the use of generative models [e.g., 13, 23, 30], recent breakthroughs were mainly driven by improvements in supervised techniques [e.g., 20, 34]. Yet unsupervised learning has the potential to tap into the much larger source of unlabeled data, which may be important for training bigger systems capable of a more general scene understanding. For example, multimodal data is abundant but often unlabeled, yet can still greatly benefit unsupervised approaches .

Generative models provide a principled approach to unsupervised learning. A perfect model of natural images would be able to optimally predict parts of an image given other parts of an image and thereby clearly demonstrate a form of scene understanding. When extended by labels, the Bayesian framework can be used to perform semi-supervised learning in the generative model while it is less clear how to combine other unsupervised approaches with discriminative learning. Generative image models are also useful in more traditional applications such as image reconstruction or compression .

Recently there has been a renewed strong interest in the development of generative image models [e.g., 47, 18, 45, 10, 8, 31, 11, 35, 24, 4]. Most of this work has tried to bring to bear the flexibility of deep neural networks on the problem of modeling the distribution of natural images. One challenge in this endeavor is to find the right balance between tractability and flexibility. The present article contributes to this line of research by introducing a fully tractable yet highly flexible image model.

Our model combines multi-dimensional recurrent neural networks with mixtures of experts. More specifically, the backbone of our model is formed by a spatial variant of long short-term memory (LSTM) . One-dimensional LSTMs have been particularly successful in modeling text and speech [e.g., 38, 39], but have also been used to model the progression of frames in video and very recently to model single images . In contrast to earlier work on modeling images, here we use multi-dimensional LSTMs which naturally lend themselves to the task of generative image modeling due to their spatial structure and ability to capture long-range correlations.

To model the distribution of pixels conditioned on the hidden states of the neural network, we use mixtures of conditional Gaussian scale mixtures (MCGSMs) . This class of models can be viewed as a generalization of Gaussian mixture models, but their parametrization makes them much more suitable for natural images. By treating images as instances of a stationary stochastic process, this model allows us to sample and capture the correlations of arbitrarily large images.

A recurrent model of natural images

In the following, we first review and extend the MCGSM and multi-dimensional LSTMs before explaining how to combine them into a recurrent image model. Section 3 will demonstrate the validity of our approach by evaluating and comparing the model on a number of image datasets.

One successful approach to building flexible yet tractable generative models has been to use fully-visible belief networks . To apply such a model to images, we have to give the pixels an ordering and specify the distribution of each pixel conditioned on its parent pixels. Several parametrizations have been suggested for the conditional distributions in the context of natural images . We here review and extend the work of Theis et al. who proposed to use mixtures of conditional Gaussian scale mixtures (MCGSMs).

Let $\mathbf{x}$ be a grayscale image patch and $x_{ij}$ be the intensity of the pixel at location $ij$. Further, let $\mathbf{x}_{<ij}$ designate the set of pixels $x_{mn}$ such that $m<i$ or $m=i$ and $n<j$ (Figure 1A). Then

for the distribution of any parametric model with parameters $\bm{\theta}$. Note that this factorization does not make any independence assumptions but is simply an application of the probability chain rule. Further note that the conditional distributions all share the same set of parameters. One way to improve the representational power of a model is thus to endow each conditional distribution with its own set of parameters,

Applying this trick to mixtures of Gaussian scale mixtures (MoGSMs) yields the MCGSM . Untying shared parameters can drastically increase the number of parameters. For images, it can easily be reduced again by adding assumptions. For example, we can limit $\mathbf{x}_{<ij}$ to a smaller neighborhood surrounding the pixel by making a Markov assumption. We will refer to the resulting set of parents as the pixel’s causal neighborhood (Figure 1B). Another reasonable assumption is stationarity or shift invariance, in which case we only have to learn one set of parameters $\bm{\theta}_{ij}$ which can then be used at every pixel location. Similar to convolutions in neural networks, this allows the model to easily scale to images of arbitrary size. While this assumption reintroduces parameter sharing constraints into the model, the constraints are different from the ones induced by the joint mixture model.

The conditional distribution in an MCGSM takes the form of a mixture of experts,

where the sum is over mixture component indices $c$ corresponding to different covariances and scales $s$ corresponding to different variances. The gates and experts in an MCGSM are given by

where $\mathbf{K}_{c}$ is positive definite. The number of parameters of an MCGSM still grows quadratically with the dimensionality of the causal neighborhood. To further reduce the number of parameters, we introduce a factorized form of the MCGSM with additional parameter sharing by replacing $\mathbf{K}_{c}$ with $\sum_{n}\beta_{cn}^{2}\mathbf{b}_{n}\mathbf{b}_{n}^{\top}$. This factorized MCGSM allows us to use larger neighborhoods and more mixture components. A detailed derivation of a more general version which also allows for multivariate pixels is given in Supplementary Section 1.

2 Spatial long short-term memory

In the following we briefly describe the spatial LSTM (SLSTM), a special case of the multi-dimensional LSTM first described by Graves & Schmidhuber . At the core of the model are memory units $\mathbf{c}_{ij}$ and hidden units $\mathbf{h}_{ij}$. For each location $ij$ on a two-dimensional grid, the operations performed by the spatial LSTM are given by $\displaystyle\mathbf{c}_{ij}$ $\displaystyle=\mathbf{g}_{ij}\odot\mathbf{i}_{ij}+\mathbf{c}_{i,j-1}\odot\mathbf{f}_{ij}^{c}+\mathbf{c}_{i-1,j}\odot\mathbf{f}_{ij}^{r},$ $\displaystyle\mathbf{h}_{ij}$ $\displaystyle=\tanh\left(\mathbf{c}_{ij}\odot\mathbf{o}_{ij}\right),$ $\displaystyle\begin{pmatrix}\mathbf{g}_{ij}\\ \mathbf{o}_{ij}\\ \mathbf{i}_{ij}\\ \mathbf{f}_{ij}^{r}\\ \mathbf{f}_{ij}^{c}\end{pmatrix}=\begin{pmatrix}\tanh\\ \sigma\\ \sigma\\ \sigma\\ \sigma\end{pmatrix}T_{\mathbf{A},\mathbf{b}}\begin{pmatrix}\mathbf{x}_{<ij}\\ \mathbf{h}_{i,j-1}\\ \mathbf{h}_{i-1,j}\end{pmatrix},$ (6) where $\sigma$ is the logistic sigmoid function, $\odot$ indicates a pointwise product, and $T_{\mathbf{A},\mathbf{b}}$ is an affine transformation which depends on the only parameters of the network $\mathbf{A}$ and $\mathbf{b}$. The gating units $\mathbf{i}_{ij}$ and $\mathbf{o}_{ij}$ determine which memory units are affected by the inputs through $\mathbf{g}_{ij}$, and which memory states are written to the hidden units $\mathbf{h}_{ij}$. In contrast to a regular LSTM defined over time, each memory unit of a spatial LSTM has two preceding states $\mathbf{c}_{i,j-1}$ and $\mathbf{c}_{i-1,j}$ and two corresponding forget gates $\mathbf{f}_{ij}^{c}$ and $\mathbf{f}_{ij}^{r}$.

3 Recurrent image density estimator

We use a grid of SLSTM units to sequentially read relatively small neighborhoods of pixels from the image, producing a hidden vector at every pixel. The hidden states are then fed into a factorized MCGSM to predict the state of the corresponding pixel, that is, $p(x_{ij}\mid\mathbf{x}_{<ij})=p(x_{ij}\mid\mathbf{h}_{ij})$. Importantly, the state of the hidden vector only depends on pixels in $\mathbf{x}_{<ij}$ and does not violate the factorization given in Equation 1. Nevertheless, the recurrent network allows this recurrent image density estimator (RIDE) to use pixels of a much larger region for prediction, and to nonlinearly transform the pixels before applying the MCGSM. We can further increase the representational power of the model by stacking spatial LSTMs to obtain a deep yet still completely tractable recurrent image model (Figure 1C).

4 Related work

Larochelle & Murray derived a tractable density estimator (NADE) in a manner similar to how the MCGSM was derived , but using restricted Boltzmann machines (RBM) instead of mixture models as a starting point. In contrast to the MCGSM, NADE tries to keep the weight sharing constraints induced by the RBM (Equation 1). Uria et al. extended NADE to real values and introduced hidden layers to the model . Gregor et al. describe a related autoregressive network for binary data which additionally allows for stochastic hidden units.

Gregor et al. used one-dimensional LSTMs to generate images in a sequential manner (DRAW). Because the model was defined over Bernoulli variables, normalized RGB values had to be treated as probabilities, making a direct comparison with other image models difficult. In contrast to our model, the presence of stochastic latent variables in DRAW means that its likelihood cannot be evaluated but has to be approximated.

Ranzato et al. and Srivastava et al. use one-dimensional recurrent neural networks to model videos, but recurrency is not used to describe the distribution over individual frames. Srivastava et al. optimize a squared error corresponding to a Gaussian assumption, while Ranzato et al. try to side-step having to model pixel intensities by quantizing image patches. In contrast, here we also try to solve the problem of modeling pixel intensities by using an MCGSM, which is equipped to model heavy-tailed as well as multi-modal distributions.

Experiments

RIDE was trained using stochastic gradient descent with a batch size of 50, momentum of 0.9, and a decreasing learning rate varying between 1 and $10^{-4}$. After each pass through the training set, the MCGSM of RIDE was finetuned using L-BFGS for up to 500 iterations before decreasing the learning rate. No regularization was used except for early stopping based on a validation set. Except where indicated otherwise, the recurrent model used a 5 pixel wide neighborhood and an MCGSM with 32 components and 32 quadratic features ($\mathbf{b}_{n}$ in Section 2.1). Spatial LSTMs were implemented using the Caffe framework . Where appropriate, we augmented the data by horizontal or vertical flipping of images.

We found that conditionally whitening the data greatly sped up the training process of both models. Letting $\mathbf{y}$ represent a pixel and $\mathbf{x}$ its causal neighborhood, conditional whitening replaces these with

where $\mathbf{C}_{\mathbf{yx}}$ is the covariance of $\mathbf{y}$ and $\mathbf{x}$, and $\mathbf{m}_{\mathbf{x}}$ is the mean of $\mathbf{x}$. In addition to speeding up training, this variance normalization step helps to make the learning rates less dependent on the training data. When evaluating the conditional log-likelihood, we compensate for the change in variance by adding the log-Jacobian $\log|\det\mathbf{W}|$. Note that this preconditioning introduces a shortcut connection from the pixel neighborhood to the predicted pixel which is not shown in Figure 1C.

Uria et al. found that forming ensembles of their autoregressive model over different pixel orderings significantly improved performance. We here consider a simple trick to produce an ensemble without the need for training different models or to change training procedures. If $\mathbf{T}_{k}$ are linear transformations leaving the targeted image distribution invariant (or approximately invariant) and if $p$ is the distribution of a pretrained model, then we form the ensemble $\frac{1}{K}\sum_{k}p(\mathbf{T}_{k}\mathbf{x})|\det\mathbf{T}_{k}|$. Note that this is simply a mixture model over images $\mathbf{x}$. We considered rotating as well as flipping images along the horizontal and vertical axes (yielding an ensemble over 8 transformations). While it could be argued that most of these transformations do not leave the distribution over natural images invariant, we nevertheless observed a noticeable boost in performance.

2 Natural images

Several recent image models have been evaluated on small image patches sampled from the Berkeley segmentation dataset (BSDS300) . Although our model’s strength lies in its ability to scale to large images and to capture long-range correlations, we include results on BSDS300 to make a connection to this part of the literature. We followed the protocol of Uria et al. . The RGB images were turned to grayscale, uniform noise was added to account for the integer discretization, and the resulting values were divided by 256. The training set of 200 images was split into 180 images for training and 20 images for validation, while the test set contained 100 images. We extracted 8 by 8 image patches from each set and subtracted the average pixel intensity such that each patch’s DC component was zero. Because the resulting image patches live on a 63 dimensional subspace, the bottom-right pixel was discarded. We used $1.6\cdot 10^{6}$ patches for training, $1.8\cdot 10^{5}$ patches for validation, and $10^{6}$ test patches for evaluation.

MCGSMs have not been evaluated on this dataset and so we first tested MCGSMs by training a single factorized MCGSM for each pixel conditioned on all previous pixels in a fixed ordering. We find that already an MCGSM (with 128 components and 48 quadratic features) outperforms all single models including a deep Gaussian mixture model (Table 3). Our ensemble of MCGSMsDetails on how the ensemble of transformations can be applied despite the missing bottom-right pixel are given in Supplementary Section 2.1. outperforms an ensemble of RNADEs with 6 hidden layers, which to our knowledge is currently the best result reported on this dataset.

Training the recurrent image density estimator (RIDE) on the 63 dimensional dataset is more cumbersome. We tried padding image patches with zeros, which was necessary to be able to compute a hidden state at every pixel. The bottom-right pixel was ignored during training and evaluation. This simple approach led to a reduction in performance relative to the MCGSM (Table 3). A possible explanation is that the model cannot distinguish between pixel intensities which are zero and zeros in the padded region. Supplying the model with additional binary indicators as inputs (one for each neighborhood pixel) did not solve the problem.

However, we found that RIDE outperforms the MCGSM by a large margin when images were treated as instances of a stochastic process (that is, using infinitely large images). MCGSMs were trained for up to 3000 iterations of L-BFGS on $10^{6}$ pixels and corresponding causal neighborhoods extracted from the training images. Causal neighborhoods were 9 pixels wide and 5 pixels high. RIDE was trained for 8 epochs on image patches of increasing size ranging from 8 by 8 to 22 by 22 pixels (that is, gradients were approximated as in backpropagation through time ). The right column in Table 3 shows average log-likelihood rates for both models. Analogously to the entropy rate , we have for the expected log-likelihood rate:

where $\mathbf{x}$ is an $N$ by $N$ image patch. An average log-likelihood rate can be directly computed for the MCGSM, while for RIDE and ensembles we approximated it by splitting the test images into 64 by 64 patches and evaluating on those.

A comparison of the log-likelihood rates reveals that an MCGSM with 16 components applied to large images already captures more correlations than any model applied to small image patches. The difference is particularly striking given that the factorized MCGSM has approximately 3,000 parameters while a GMM with 200 components has approximately 400,000 parameters. Using an ensemble of RIDEs, we are able to further improve this number significantly (Table 3).

Another dataset frequently used to test generative image models is the dataset published by van Hateren and van der Schaaf . Details of the preprocessing used in this paper are given in Supplementary Section 3. We reevaluated several models for which the likelihood has been reported on this dataset . Likelihood rates as well as results on 16 by 16 patches are given in Table 3. Because of the larger patch size, RIDE here already outperforms the MCGSM on patches.

3 Dead leaves

Dead leaf images are generated by superimposing disks of random intensity and size on top of each other . This simple procedure leads to images which already share many of the statistical properties and challenges of natural images, such as occlusions and long-range correlations, while leaving out others such as non-stationary statistics. They therefore provide an interesting test case for natural image models.

We used a set of 1,000 images, where each image is 256 by 256 pixels in size. We compare the performance of RIDE to the MCGSM and a very recently introduced deep multiscale model based on a diffusion process . The same 100 images as in previous literature were used for evaluation and we used the remaining images for training. We find that the introduction of an SLSTM with 64 hidden units greatly improves the performance of the MCGSM. We also tried an extended version of the SLSTM which included memory units as additional inputs (right-hand side of Equation 6). This yielded a small improvement in performance (5th row in Table 5) while adding layers or using more hidden units led to more drastic improvements. Using 3 layers with 128 hidden units in each layer, we find that our recurrent image model is on a par with the deep diffusion model. By using ensembles, we are able to beat all previously published results for this dataset (Table 5).

Figure 5 shows that the improved performance of RIDE is not simply due to an effectively larger causal neighborhood but that the nonlinear transformations performed by the SLSTM units matter. Simply increasing the neighborhood size of an MCGSM does not yield the same improvement. Instead, the performance quickly saturates. We also find that the performance of RIDE slightly deteriorates with larger neighborhoods, which is likely caused by optimization difficulties.

4 Texture synthesis and inpainting

To get an intuition for the kinds of correlations which RIDE can capture or fails to capture, we tried to use it to synthesize textures. We used several 640 by 640 pixel textures published by Brodatz . The textures were split into sixteen 160 by 160 pixel regions of which 15 were used for training and one randomly selected region was kept for testing purposes. RIDE was trained for up to 6 epochs on patches of increasing size ranging from 20 by 20 to 40 by 40 pixels.

Samples generated by an MCGSM and RIDE are shown in Figure 6. Both models are able to capture a wide range of correlation structures. However, the MCGSM seems to struggle with textures having bimodal marginal distributions and periodic patterns (D104, D34, and D110). RIDE clearly improves on these textures, although it also struggles to faithfully reproduce periodic structure. Possible explanations include that LSTMs are not well suited to capture periodicities, or that these failures are not penalized strong enough by the likelihood. For some textures, RIDE produces samples which are nearly indistinguishable from the real textures (D106 and D110).

One application of generative image models is inpainting [e.g., 33, 12, 35]. As a proof of concept, we used our model to inpaint a large (here, 71 by 71 pixels) region in textures (Figure 7). Missing pixels were replaced by sampling from the posterior of RIDE. Unlike the joint distribution, the posterior distribution cannot be sampled directly and we had to resort to Markov chain Monte Carlo methods. We found the following Metropolis within Gibbs procedure to be efficient enough. The missing pixels were initialized via ancestral sampling. Since ancestral sampling is cheap, we generated 5 candidates and used the one with the largest posterior density. Following initialization, we sequentially updated overlapping 5 by 5 pixel regions via Metropolis sampling. Proposals were generated via ancestral sampling and accepted using the acceptance probability

where here $\mathbf{x}_{ij}$ represents a 5 by 5 pixel patch and $\mathbf{x}_{ij}^{\prime}$ its proposed replacement. Since evaluating the joint and conditional densities on the entire image is costly, we approximated $p$ using RIDE applied to a 19 by 19 pixel patch surrounding $ij$. Randomly flipping images vertically or horizontally in between the sampling further helped. Figure 7 shows results after 100 Gibbs sampling sweeps.

Conclusion

We have introduced RIDE, a deep but tractable recurrent image model based on spatial LSTMs. The model exemplifies how recent insights in deep learning can be exploited for generative image modeling and shows superior performance in quantitative comparisons. RIDE is able to capture many different statistical patterns, as demonstrated through its application to textures. This is an important property considering that on an intermediate level of abstraction natural images can be viewed as collections of textures.

We have furthermore introduced a factorized version of the MCGSM which allowed us to use more experts and larger causal neighborhoods. This model has few parameters, is easy to train and already on its own performs very well as an image model. It is therefore an ideal building block and may be used to extend other models such as DRAW or video models .

Deep generative image models have come a long way since deep belief networks have first been applied to natural images . Unlike convolutional neural networks in object recognition, however, no approach has as of yet proven to be a likely solution to the problem of generative image modeling. Further conceptual work will be necessary to come up with a model which can handle both the more abstract high-level as well as the low-level statistics of natural images.

The authors would like to thank Aäron van den Oord for insightful discussions and Wieland Brendel, Christian Behrens, and Matthias Kümmerer for helpful input on this paper. This study was financially supported by the German Research Foundation (DFG; priority program 1527, BE 3848/2-1).

References