summaries: The goal of this paper is to use deep generative models for missing data imputation. This paper proposes learning a latent variable deep generative model over every randomly sampled subset of observed features. First, a masking variable is sampled from a chosen prior distribution. The mask determines which features are observed. Then, the likelihood of the observed features is maximized via a lower bound. Inference in this latent variable model is achieved through the use of an inference network which conditions on the set of "missing" (to the generative model) features.

summaries: The paper presents a model for learning conditional distribution when arbitrary partitioning the input to observed and masked parts. The idea is to extend the conditional VAE framework such that the posterior is a function of an arbitrary subset of observed variables. Accordingly, reconstruction loss only penalizes the error in the reconstruction of masked (unobserved) variables. The method is compared against 1) classical approaches in missing data imputation on UCI benchmarks; 2) image inpainting against recently proposed GANS for the similar task, as well as; 3) against universal marginalizer, which learns conditional densities using a feedforward / autoregressive architecture.

abstractText: We propose a single neural probabilistic model based on variational autoencoder that can be conditioned on an arbitrary subset of observed features and then sample the remaining features in “one shot”. The features may be both real-valued and categorical. Training of the model is performed by stochastic variational Bayes. The experimental evaluation on synthetic data, as well as feature imputation and image inpainting problems, shows the effectiveness of the proposed approach and diversity of the generated samples.

authors:
Philip Bachman
Doina Precup

title: Data generation as sequential decision making

venue: In Advances in Neural Information Processing Systems,

authors:
Marcelo Bertalmio
Guillermo Sapiro
Vincent Caselles
Coloma Ballester

title: Image inpainting

venue: In Proceedings of the 27th Annual Conference on Computer Graphics and Interactive Techniques,

authors:
Florian Bordes
Sina Honari
Pascal Vincent

title: Learning to generate samples from noise through infusion training

venue: In International Conference on Learning Representations,

authors:
Yuri Burda
Roger Grosse
Ruslan Salakhutdinov

title: Importance weighted autoencoders

venue: arXiv preprint arXiv:1509.00519,

authors:
S van Buuren
Karin Groothuis-Oudshoorn

title: mice: Multivariate imputation by chained equations in r

venue: Journal of statistical software,

authors:
Laura Douglas
Iliyan Zarov
Konstantinos Gourgoulias
Chris Lucas
Chris Hart
Adam Baker
Maneesh Sahani
Yura Perov
Saurabh Johri

title: A universal marginalizer for amortized inference in generative models

venue: arXiv preprint arXiv:1711.00695,

authors:
Mathieu Germain
Karol Gregor
Iain Murray
Hugo Larochelle

title: Made: Masked autoencoder for distribution estimation

venue: Proceedings of the 32nd International Conference on Machine Learning,

authors:
Ian Goodfellow
Jean Pouget-Abadie
Mehdi Mirza
Bing Xu
David Warde-Farley
Sherjil Ozair
Aaron Courville
Yoshua Bengio

title: Generative adversarial nets

venue: In Advances in neural information processing systems,

authors:
Anirudh Goyal Alias Parth Goyal
Nan Ke
Surya Ganguli
Yoshua Bengio

title: Variational walkback: Learning a transition operator as a stochastic recurrent net

authors:
Kaiming He
Xiangyu Zhang
Shaoqing Ren
Jian Sun

title: Deep residual learning for image recognition

venue: In Proceedings of the IEEE conference on computer vision and pattern recognition,

authors:
Diederik P Kingma
Jimmy Ba

title: Adam: A method for stochastic optimization

venue: arXiv preprint arXiv:1412.6980,

authors:
Diederik P. Kingma
Max Welling

title: Auto-encoding variational bayes

authors:
Brenden M Lake
Ruslan Salakhutdinov
Joshua B Tenenbaum

title: Human-level concept learning through probabilistic program induction

authors:
Yann LeCun
Corinna Cortes
Christopher JC Burges

title: The mnist database of handwritten digits

authors:
Yijun Li
Sifei Liu
Jimei Yang
Ming-Hsuan Yang

title: Generative face completion

venue: In The IEEE Conference on Computer Vision and Pattern Recognition (CVPR),

authors:
Ziwei Liu
Ping Luo
Xiaogang Wang
Xiaoou Tang

title: Deep learning face attributes in the wild

venue: In Proceedings of International Conference on Computer Vision (ICCV),

authors:
Alireza Makhzani
Jonathon Shlens
Navdeep Jaitly
Ian Goodfellow

title: Adversarial autoencoders

authors:
Xiaojiao Mao
Chunhua Shen
Yu-Bin Yang

title: Image restoration using very deep convolutional encoderdecoder networks with symmetric skip connections

authors:
Mehdi Mirza
Simon Osindero

title: Conditional generative adversarial nets

authors:
Deepak Pathak
Philipp Krähenbühl
Jeff Donahue
Trevor Darrell
Alexei Efros

title: Context encoders: Feature learning by inpainting

venue: In Computer Vision and Pattern Recognition

authors:
Danilo Jimenez Rezende
Shakir Mohamed
Daan Wierstra

title: Stochastic backpropagation and approximate inference in deep generative models

venue: In International Conference on Machine Learning,

authors:
Olaf Ronneberger
Philipp Fischer
Thomas Brox

title: U-net: Convolutional networks for biomedical image segmentation

venue: In International Conference on Medical image computing and computer-assisted intervention,

authors:
Jascha Sohl-Dickstein
Eric Weiss
Niru Maheswaranathan
Surya Ganguli

title: Deep unsupervised learning using nonequilibrium thermodynamics

authors:
Kihyuk Sohn
Honglak Lee
Xinchen Yan

title: Learning structured output representation using deep conditional generative models

venue: Advances in Neural Information Processing Systems

authors:
Casper Kaae Sønderby
Tapani Raiko
Lars Maaløe
Søren Kaae Sønderby
Ole Winther

title: Ladder variational autoencoders

authors:
Daniel J Stekhoven
Peter Bühlmann

title: Missforest - non-parametric missing value imputation for mixedtype data

authors:
Benigno Uria
Marc-Alexandre Côté
Karol Gregor
Iain Murray
Hugo Larochelle

title: Neural autoregressive distribution estimation

venue: The Journal of Machine Learning Research,

authors:
C. Yang
X. Lu
Z. Lin
E. Shechtman
O. Wang
H. Li

title: High-resolution image inpainting using multiscale neural patch synthesis

authors:
Raymond Yeh
Chen Chen
Teck-Yian Lim
Mark Hasegawa-Johnson
Minh N. Do

title: Semantic image inpainting with perceptual and contextual losses

authors:
Raymond A Yeh
Chen Chen
Teck Yian Lim
Alexander G Schwing
Mark Hasegawa-Johnson
Minh N Do

title: Semantic image inpainting with deep generative models

venue: In Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition,

authors:
Jinsung Yoon
James Jordon
Mihaela van der Schaar

title: GAIN: Missing data imputation using generative adversarial nets

venue: Proceedings of the 35th International Conference on Machine Learning,

heading: 1 INTRODUCTION

text:
In past years, a number of generative probabilistic models based on neural networks have been proposed. The most popular approaches include variational autoencoder (Kingma & Welling, 2013) (VAE) and generative adversarial net (Goodfellow et al., 2014) (GANs). They learn a distribution over objects p(x) and allow sampling from this distribution.
In many cases, we are interested in learning a conditional distribution p(x|y). For instance, if x is an image of a face, y could be the characteristics describing the face (are glasses present or not; length of hair, etc.) Conditional variational autoencoder (Sohn et al., 2015) and conditional generative adversarial nets (Mirza & Osindero, 2014) are popular methods for this problem.
In this paper, we consider the problem of learning all conditional distributions of the form p(xI |xU\I), where U is the set of all features and I is its arbitrary subset. This problem generalizes both learning the joint distribution p(x) and learning the conditional distribution p(x|y). To tackle this problem, we propose a Variational Autoencoder with Arbitrary Conditioning (VAEAC) model. It is a latent variable model similar to VAE, but allows conditioning on an arbitrary subset of the features. The conditioning features affect the prior on the latent Gaussian variables which are used to generate unobserved features. The model is trained using stochastic gradient variational Bayes (Kingma & Welling, 2013).
We consider two most natural applications of the proposed model. The first one is feature imputation where the goal is to restore the missing features given the observed ones. The imputed values may be valuable by themselves or may improve the performance of other machine learning algorithms which process the dataset. Another application is image inpainting in which the goal is to fill in an unobserved part of an image with an artificial content in a realistic way. This can be used for removing unnecessary objects from the images or, vice versa, for complementing the partially closed or corrupted object.
∗Author is in DeepMind now.
The experimental evaluation shows that the proposed model successfully samples from the conditional distributions. The distribution over samples is close to the true conditional distribution. This property is very important when the true distribution has several modes. The model is shown to be effective in feature imputation problem which helps to increase the quality of subsequent discriminative models on different problems from UCI datasets collection (Lichman, 2013). We demonstrate that model can generate diverse and realistic image inpaintings on MNIST (LeCun et al., 1998), Omniglot (Lake et al., 2015) and CelebA (Liu et al., 2015) datasets, and works even better than the current state of the art inpainting techniques in terms of peak signal to noise ratio (PSNR).
The paper is organized as follows. In section 2 we review the related works. In section 3 we briefly describe variational autoencoders and conditional variational autoencoders. In section 4 we define the problem, describe the VAEAC model and its training procedure. In section 5 we evaluate VAEAC. Section 6 concludes the paper. Appendix contains additional explanations, theoretical analysis, and experiments for VAEAC.

heading: 2 RELATED WORK

text:
Universal Marginalizer (Douglas et al., 2017) is a model based on a feed-forward neural network which approximates marginals of unobserved features conditioned on observable values. A related idea of an autoregressive model of joint probability was previously proposed in Germain et al. (2015) and Uria et al. (2016). The description of the model and comparison with VAEAC are available in section 5.3.
Yoon et al. (2018) propose a GANs-based model called GAIN which solves the same problem as VAEAC. In contrast to VAEAC, GAIN does not use unobserved data during training, which makes it easier to apply to the missing features imputation problem. Nevertheless, it turns into a disadvantage when the fully-observed training data is available but the missingness rate at the testing stage is high. For example, in inpainting setting GAIN cannot learn the conditional distribution over MNIST digits given one horizontal line of the image while VAEAC can (see appendix D.4). The comparison of VAEAC and GAIN on the missing feature imputation problem is given in section 5.1 and appendix D.2.
Rezende et al. (2014) [Appendix F], Sohl-Dickstein et al. (2015), Goyal et al. (2017), and Bordes et al. (2017) propose to fill missing data with noise and run Markov chain with a learned transition operator. The stationary distribution of such chains approximates the true conditional distribution of the unobserved features. Bachman & Precup (2015) consider missing feature imputation in terms of Markov decision process and propose LSTM-based sequential decision making model to solve it. Nevertheless, these methods are computationally expensive at the test time and require fully-observed training data.
Image inpainting is a classic computer vision problem. Most of the earlier methods rely on local and texture information or hand-crafted problem-specific features (Bertalmio et al., 2000). In past years multiple neural network based approaches have been proposed.
Pathak et al. (2016), Yeh et al. (2016) and Yang et al. (2017) use different kinds and combinations of adversarial, reconstruction, texture and other losses. Li et al. (2017) focuses on face inpainting and uses two adversarial losses and one semantic parsing loss to train the generative model. In Yeh et al. (2017) GANs are first trained on the whole training dataset. The inpainting is an optimization procedure that finds the latent variables that explain the observed features best. Then, the obtained latents are passed through the generative model to restore the unobserved portion of the image. We can say that VAEAC is a similar model which uses prior network to find a proper latents instead of solving the optimization problem.
All described methods aim to produce a single realistic inpainting, while VAEAC is capable of sampling diverse inpaintings. Additionally, Yeh et al. (2016), Yang et al. (2017) and Yeh et al. (2017) have high testtime computational complexity of inpainting, because they require an optimization problem to be solved. On the other hand, VAEAC is a “single-shot” method with a low computational cost.

heading: 3 BACKGROUND

heading: 3.1 VARIATIONAL AUTOENCODER

text:
Variational autoencoder (Kingma & Welling, 2013) (VAE) is a directed generative model with latent variables. The generative process in variational autoencoder is as follows: first, a latent variable z is generated from the prior distribution p(z), and then the data x is generated from the generative distribution pθ(x|z), where θ are the generative model’s parameters. This process induces the distribution pθ(x) = Ep(z)pθ(x|z). The distribution pθ(x|z) is modeled by a neural network with parameters θ. p(z) is a standard Gaussian distribution.
The parameters θ are tuned by maximizing the likelihood of the training data points {xi}Ni=1 from the true data distribution pd(x). In general, this optimization problem is challenging due to intractable posterior inference. However, a variational lower bound can be optimized efficiently using backpropagation and stochastic gradient descent:
log pθ(x) = Eqφ(z|x) log pθ(x, z)
qφ(z|x) +DKL(qφ(z|x)‖p(z|x, θ))
≥ Eqφ(z|x) log pθ(x|z)−DKL(qφ(z|x)‖p(z)) = LV AE(x; θ, φ) (1)
Here qφ(z|x) is a proposal distribution parameterized by neural network with parameters φ that approximates the posterior p(z|x, θ). Usually this distribution is Gaussian with a diagonal covariance matrix. The closer qφ(z|x) to p(z|x, θ), the tighter variational lower bound LV AE(θ, φ). To compute the gradient of the variational lower bound with respect to φ, reparameterization trick is used: z = µφ(x) + εσφ(x) where ε ∼ N (0, I) and µφ and σφ are deterministic functions parameterized by neural networks. So the gradient can be estimated using Monte-Carlo method for the first term and computing the second term analytically:
∂LV AE(x; θ, φ)
∂φ = Eε∼N (0,I)
∂
∂φ log pθ(x|µφ(x) + εσφ(x))−
∂
∂φ DKL(qφ(z|x)‖p(z)) (2)
So LV AE(θ, φ) can be optimized using stochastic gradient ascent with respect to φ and θ.

heading: 3.2 CONDITIONAL VARIATIONAL AUTOENCODER

text:
Conditional variational autoencoder (Sohn et al., 2015) (CVAE) approximates the conditional distribution pd(x|y). It outperforms deterministic models when the distribution pd(x|y) is multi-modal (diverse xs are probable for the given y). For example, assume that x is a real-valued image. Then, a deterministic regression model with mean squared error loss would predict the average blurry value for x. On the other hand, CVAE learns the distribution of x, from which one can sample diverse and realistic objects.
Variational lower bound for CVAE can be derived similarly to VAE by conditioning all considered distributions on y:
LCV AE(x, y; θ, ψ, φ) = Eqφ(z|x,y) log pθ(x|z, y)−DKL(qφ(z|x, y)‖pψ(z|y)) ≤ log pθ,ψ(x|y) (3)
Similarly to VAE, this objective is optimized using the reparameterization trick. Note that the prior distribution pψ(z|y) is conditioned on y and is modeled by a neural network with parameters ψ. Thus, CVAE uses three trainable neural networks, while VAE only uses two.
Also authors propose such modifications of CVAE as Gaussian stochastic neural network and hybrid model. These modifications can be applied to our model as well. Nevertheless, we don’t use them, because of their disadvantage which is described in appendix C.

heading: 4 VARIATIONAL AUTOENCODER WITH ARBITRARY CONDITIONING

heading: 4.1 PROBLEM STATEMENT

text:
Consider a distribution pd(x) over a D-dimensional vector x with real or categorical components. The components of the vector are called features.
Let binary vector b ∈ {0, 1}D be the binary mask of unobserved features of the object. Then we describe the vector of unobserved features as xb = {xi:bi=1}. For example, x(0,1,1,0,1) = (x2, x3, x5). Using this notation we denote x1−b as a vector of observed features.
Our goal is to build a model of the conditional distribution pψ,θ(xb|x1−b, b) ≈ pd(xb|x1−b, b) for an arbitrary b, where ψ and θ are parameters that are used in our model at the testing stage.
However, the true distribution pd(xb|x1−b, b) is intractable without strong assumptions about pd(x). Therefore, our model pψ,θ(xb|x1−b, b) has to be more precise for some b and less precise for others. To formalize our requirements about the accuracy of our model we introduce the distribution p(b) over different unobserved feature masks. The distribution p(b) is arbitrary and may be defined by the user depending on the problem. Generally it should have full support over {0, 1}D so that pψ,θ(xb|x1−b, b) can evaluate arbitrary conditioning. Nevertheless, it is not necessary if the model is used for specific kinds of conditioning (as we do in section 5.2).
Using p(b) we can introduce the following log-likelihood objective function for the model:
max ψ,θ Epd(x)Ep(b) log pψ,θ(xb|x1−b, b) (4)
The special cases of the objective (4) are variational autoencoder (bi = 1 ∀i ∈ {1, . . . , D}) and conditional variational autoencoder (b is constant).

heading: 4.2 MODEL DESCRIPTION

text:
The generative process of our model is similar to the generative process of CVAE: for each object firstly we generate z ∼ pψ(z|x1−b, b) using prior network, and then sample unobserved features xb ∼ pθ(xb|z, x1−b, b) using generative network. This process induces the following model distribution over unobserved features:
pψ,θ(xb|x1−b, b) = Ez∼pψ(z|x1−b,b)pθ(xb|z, x1−b, b) (5)
We use z ∈ Rd, and Gaussian distribution pψ over z, with parameters from a neural network with weights ψ: pψ(z|x1−b, b, ψ) = N (z|µψ(x1−b, b), σ2ψ(x1−b, b)I). The real-valued components of distribution pθ(xb|z, x1−b, b) are defined likewise. Each categorical component i of distribution pθ(xi|z, x1−b, b) is parameterized by a function wi,θ(z, x1−b, b), whose outputs are logits of probabilities for each category: xi ∼ Cat[Softmax(wi,θ(z, x1−b, b))]. Therefore the components of the latent vector z are conditionally independent given x1−b and b, and the components of xb are conditionally independent given z, x1−b and b.
The variables xb and x1−b have variable length that depends on b. So in order to use architectures such as multi-layer perceptron and convolutional neural network we consider x1−b = x ◦ (1 − b) where ◦ is an element-wise product. So in implementation x1−b has fixed length. The output of the generative network also has a fixed length, but we use only unobserved components to compute likelihood.
The theoretical analysis of the model is available in appendix B.1.

heading: 4.3 LEARNING VARIATIONAL AUTOENCODER WITH ARBITRARY CONDITIONING

heading: 4.3.1 VARIATIONAL LOWER BOUND

text:
We can derive a lower bound for log pψ,θ(xb|x1−b, b) as for variational autoencoder:
log pψ,θ(xb|x1−b, b) = Eqφ(z|x,b) log pψ,θ(xb, z|x1−b, b)
qφ(z|x, b) +DKL(qφ(z|x, b)‖pψ,θ(z|x, b))
≥ Eqφ(z|x,b) log pθ(xb|z, x1−b, b)−DKL(qφ(z|x, b)‖pψ(z|x1−b, b)) = LV AEAC(x, b; θ, ψ, φ) (6)
Therefore we have the following variational lower bound optimization problem:
max θ,ψ,φ Epd(x)Ep(b)LV AEAC(x, b; θ, ψ, φ) (7)
We use fully-factorized Gaussian proposal distribution qφ which allows us to perform reparameterization trick and compute KL divergence analytically in order to optimize (7).

heading: 4.3.2 PRIOR IN LATENT SPACE

text:
During the optimization of objective (7) the parameters µψ and σψ of the prior distribution of z may tend to infinity, since there is no penalty for large values of those parameters. We usually observe the growth of ‖z‖2 during training, though it is slow enough. To prevent potential numerical instabilities, we put a Normal-Gamma prior on the parameters of the prior distribution to prevent the divergence. Formally, we redefine pψ(z|x1−b, b) as follows:
pψ(z, µψ, σψ|x1−b, b) = N (z|µψ, σ2ψ)N (µψ|0, σµ)Gamma(σψ|2, σσ) (8)
As a result, the regularizers − µ 2 ψ
2σ2µ and σσ(log(σψ) − σψ) are added to the model log-likelihood. Hyperpa-
rameter σµ is chosen to be large (104) and σσ is taken to be a small positive number (10−4). This distribution is close to uniform near zero, so it doesn’t affect the learning process significantly.

heading: 4.3.3 MISSING FEATURES

text:
The optimization objective (7) requires all features of each object at the training stage: some of the features will be observed variables at the input of the model and other will be unobserved features used to evaluate the model. Nevertheless, in some problem settings the training data contains missing features too. We propose the following slight modification of the problem (7) in order to cover such problems as well.
The missing values cannot be observed so xi = ω ⇒ bi = 1, where ω describes the missing value in the data. In order to meet this requirement, we redefine mask distribution as conditioned on x: p(b) turns into p(b|x) in (4) and (7). In the reconstruction loss (5) we simply omit the missing features, i. e. marginalize them out:
log pθ(xb|z, x1−b, b) = ∑
i:bi=1,xi 6=ω
log pθ(xi|z, x1−b, b) (9)
The proposal network must be able to determine which features came from real object and which are just missing. So we use additional missing features mask which is fed to proposal network together with unobserved features mask b and object x.
The proposed modifications are evaluated in section 5.1.

heading: 5 EXPERIMENTS

text: In this section we validate the performance of VAEAC using several real-world datasets. In the first set of experiments we evaluate VAEAC missing features imputation performance using various UCI datasets (Lichman, 2013). We compare imputations from our model with imputations from such classical methods as MICE (Buuren & Groothuis-Oudshoorn, 2010) and MissForest (Stekhoven & Bühlmann, 2011) and recently proposed GANs-based method GAIN (Yoon et al., 2018). In the second set of experiments we use VAEAC to solve image inpainting problem. We show inpainitngs generated by VAEAC and compare our model with models from papers Pathak et al. (2016), Yeh et al. (2017) and Li et al. (2017) in terms of peak signal-to-noise ratio (PSNR) of obtained inpaintings on CelebA dataset (Liu et al., 2015) . And finally, we evaluate VAEAC against the competing method called Universal Marginalizer (Douglas et al., 2017). Additional experiments can be found in appendices C and D. The code is available at https://github.com/tigvarts/ vaeac.

heading: 5.1 MISSING FEATURES IMPUTATION

text:
The datasets with missing features are widespread. Consider a dataset with D-dimensional objects x where each feature may be missing (which we denote by xi = ω) and their target values y. The majority of discriminative methods do not support missing values in the objects. The procedure of filling in the missing features values is called missing features imputation.
In this section we evaluate the quality of imputations produced by VAEAC. For evaluation we use datasets from UCI repository (Lichman, 2013). Before training we drop randomly 50% of values both in train and test set. After that we impute missing features using MICE (Buuren & Groothuis-Oudshoorn, 2010), MissForest (Stekhoven & Bühlmann, 2011), GAIN (Yoon et al., 2018) and VAEAC trained on the observed data. The details of GAIN implementation are described in appendix A.4.
Our model learns the distribution of the imputations, so it is able to sample from this distribution. We replace each object with missing features by n = 10 objects with sampled imputations, so the size of the dataset increases by n times. This procedure is called missing features multiple imputation. MICE and GAIN are also capable of multiple imputation (we use n = 10 for them in experiments as well), but MissForest is not.
For more details about the experimental setup see appendices A.1, A.2, and A.4.
In table 1 we report NRMSE (i.e. RMSE normalized by the standard deviation of each feature and then averaged over all features) of imputations for continuous datasets and proportion of falsely classified (PFC) for categorical ones. For multiple imputation methods we average imputations of continuous variables and take most frequent imputation for categorical ones for each object.
We also learn linear or logistic regression and report the regression or classification performance after applying imputations of different methods in table 2. For multiple imputation methods we average predictions for continuous targets and take most frequent prediction for categorical ones for each object in test set.
As can be seen from the tables 1 and 2, VAEAC can learn joint data distribution and use it for missing feature imputation. The imputations are competitive with current state of the art imputation methods in terms of RMSE, PFC, post-imputation regression R2-score and classification accuracy. Nevertheless, we don’t claim that our method is state of the art in missing features imputation; for some datasets MICE or MissForest outperform it. The additional experiments can be found in appendix D.2.

heading: 5.2 IMAGE INPAINTING

text:
The image inpainting problem has a number of different formulations. The formulation of our interest is as follows: some of the pixels of an image are unobserved and we want to restore them in a natural way. Unlike the majority of papers, we want to restore not just one most probable inpainting, but the distribution over all possible inpaintings from which we can sample. This distribution is extremely multi-modal because often there is a lot of different possible ways to inpaint the image.
Unlike the previous subsection, here we have uncorrupted images without missing features in the training set, so p(b|x) = p(b). As we show in section 2, state of the art results use different adversarial losses to achieve more sharp and realistic samples. VAEAC can be adapted to the image inpainting problem by using a combination of those adversarial losses as a part of reconstruction loss pθ(xb|z, x1−b, b). Nevertheless, such construction is out of scope for this research, so we leave it for the future work. In the current work we show that the model can generate both diverse and realistic inpaintings.
In figures 1, 2, 3 and 4 we visualize image inpaintings produced by VAEAC on binarized MNIST (LeCun et al., 1998), Omniglot (Lake et al., 2015) and CelebA (Liu et al., 2015). The details of learning procedure and description of datasets are available in appendixes A.1 and A.3.
To the best of our knowledge, the most modern inpainting papers don’t consider the diverse inpainting problem, where the goal is to build diverse image inpaintings, so there is no straightforward way to compare with these models. Nevertheless, we compute peak signal-to-noise ratio (PSNR) for one random inpainting from VAEAC and the best PSNR among 10 random inpaintings from VAEAC. One inpainting might not be similar to the original image, so we also measure how good the inpainting which is most similar to the original image reconstructs it. We compare these two metrics computed for certain masks with the PSNRs for the same masks on CelebA from papers Yeh et al. (2017) and Li et al. (2017). The results are available in tables 3 and 4.
We observe that for the majority of proposed masks our model outperforms the competing methods in terms of PSNR even with one sample, and for the rest (where the inpaintings are significantly diverse) the best PSNR over 10 inpaintings is larger than the same PSNR of the competing models. Even if PSNR does not reflect completely the visual quality of images and tends to encourage blurry VAE samples instead of realistic GANs samples, the results show that VAEAC is able to solve inpainting problem comparably to the state of the art methods. The disadvantage of VAEAC compared to Yeh et al. (2017) and Li et al. (2017) (but
not Pathak et al. (2016)) is that it needs the distribution over masks at the training stage to be similar to the distribution over them at the test stage. However, it is not a very strict limitation for the practical usage.

heading: 5.3 UNIVERSAL MARGINALIZER

text:
Universal Marginalizer (Douglas et al., 2017) (UM) is a model which uses a single neural network to estimate the marginal distributions over the unobserved features. So it optimizes the following objective:
max θ Ex∼pd(x)Eb∼p(b) D∑ i=1 bi log pθ(xi|x1−b, b) (10)
For given mask b we fix a permutation of its unobserved components: (i1, i2, . . . , i|b|), where |b| is a number of unobserved components. Using the learned model and the permutation we can generate objects from joint distribution and estimate their probability using chain rule.
log pθ(xb|x1−b, b) = |b|∑ j=1 log pθ(xij |x1−(b−∑j−1k=1 eik ), b− j−1∑ k=1 eik) (11)
For example, pθ(x1, x4, x5|x2, x3) = pθ(x4|x2, x3)pθ(x1|x2, x3, x4)pθ(x5|x1, x2, x3, x4). Conditional sampling or conditional likelihood estimation for one object requires |b| requests to UM to compute pθ(xi|x1−b, b). Each request is a forward pass through the neural network. In the case of conditional sampling those requests even cannot be paralleled because the input of the next request contains the output of the previous one.
We propose a slight modification of the original UM training procedure which allows learning UM efficiently for any kind of masks including those considered in this paper. The details of the modification are described in appendix B.3.
1The results are from the paper (Yeh et al., 2017) 2The results are from the paper (Li et al., 2017)
Left: input. The gray pixels are unobserved. Middle: samples from VAEAC. Right: ground truth.
The results of using this modification of UM are provided in table 5. We can say that the relation between VAEAC and UM is similar to the relation between VAE and PixelCNN. The second one is much slower at the testing stage, but it easily takes into account local dependencies in data while the first one is faster but assumes conditional independence of the outputs. Nevertheless, there are a number of cases where UM cannot learn the distribution well while VAEAC can. For example, when the data is real-valued and marginal distributions have many local optima, there is no straightforward parametrization which allows UM to approximate them, and, therefore also the conditioned joint distribution. An example of such distribution and more illustrations for comparison of VAEAC and UM are available in appendix D.5.

heading: 6 CONCLUSION

text: In this paper we consider the problem of simultaneous learning of all conditional distributions for a vector. This problem has a number of different special cases with practical applications. We propose neural network based probabilistic model for distribution conditioning learning with Gaussian latent variables. This model is scalable and efficient in inference and learning. We propose several tricks to improve optimization and give recommendations about hyperparameters choice. The model is successfully applied to feature imputation and inpainting tasks. The experimental results show that the model is competitive with state of the art methods for both missing features imputation and image inpainting problems.

heading: APPENDIX

heading: A EXPERIMENTAL DETAILS

heading: A.1 NEURAL NETWORK ARCHITECTURES

text:
In all experiments we use optimization method Adam (Kingma & Ba, 2014), skip-connections between prior network and generative network inspired by (Mao et al., 2016), (Sønderby et al., 2016) and (Ronneberger et al., 2015), and convolutional neural networks based on ResNet blocks (He et al., 2016).
Without skip-connections all information for decoder goes through the latent variables. In image inpainting we found skip-connections very useful in both terms of log-likelihood improvement and the image realism, because latent variables are responsible for the global information only while the local information passes through skip-connections. Therefore the border between image and inpainting becomes less conspicuous.
The main idea of neural networks architecture is reflected in figure 5.
The number of hidden layers, their widths and structure may be different.
The neural networks we used for image inpainting have He-Uniform initialization of convolutional ResNet blocks, and the skip-connections are implemented using concatenation, not addition. The proposal network structure is exactly the same as the prior network except skip-connections.
Also one could use much simpler fully-connected networks with one hidden layer as a proposal, prior and generative networks in VAEAC and still obtain nice inpaintings on MNIST.

heading: A.2 MISSING FEATURES IMPUTATION

text:
We split the dataset into train and test set with size ratio 3:1. Before training we drop randomly 50% of values both in train and test set. We repeat each experiment 5 times with different train-test splits and dropped features and then average results and compute their standard deviation.
As we show in appendix B.2, the better results can be achieved when the model learns the concatenation of objects features x and targets y. So we treat y as an additional feature that is always unobserved during the testing time.
To train our model we use distribution p(bi|x) in which p(bi|xi = ω) = 1 and p(bi|x) = 0.2 otherwise. Also for VAEAC trainig we normalize real-valued features, fix σθ = 1 in the generative model of VAEAC in order to optimize RMSE, and use 25% of training data as validation set to select the best model among all epochs of training.
For the test set, the classifier or regressor is applied to each of the n imputed objects and the predictions are combined. For regression problems we report R2-score of combined predictions, so we use averaging as a combination method. For classification problem we report accuracy, and therefore choose the mode. We consider the workflow where the imputed values of y are not fed to the classifier or regressor to make a fair comparison of feature imputation quality.
NRMSE or PFC for dataset is computed as an average of NRMSE or PFC of all features of this dataset. NRMSE of a feature is just RMSE of imputations divided by the standard deviation of this feature. PFC of a feature is a proportion of imputations which are incorrect.
A.3 IMAGE INPAINTING DATASETS AND MASKS
MNIST is a dataset of 60000 train and 10000 test grayscale images of digits from 0 to 9 of size 28x28. We binarize all images in the dataset. For MNIST we consider Bernoulli log-likelihood as the reconstruction loss: log pθ(xb|z, x1−b, b) = ∑ i:bi=1
log Bernoulli(xi|pθ,i(z, x1−b, b)) where pθ,i(z, x1−b, b) is an output of the generative neural network. We use 16 latent variables. In the mask for this dataset the observed pixels form a three pixels wide horizontal line which position is distributed uniformly.
Omniglot is a dataset of 19280 train and 13180 test black-and-white images of different alphabets symbols of size 105x105. As in previous section, the brightness of each pixel is treated as a Bernoulli probability of it to be 1. The mask we use is a random rectangular which is described below. We use 64 latent variables. We train model for 50 epochs and choose best model according to IWAE log-likelihood estimation on the validation set after each epoch.
CelebA is a dataset of 162770 train, 19867 validation and 19962 test color images of faces of celebrities of size 178x218. Before learning we normalize the channels in dataset. We use logarithm of fully-factorized Gaussian distribution as reconstruction loss. The mask we use is a random rectangular which is describe below. We use 32 latent variables.
Rectangular mask is the common shape of unobserved region in image inpainting. We use such mask for Omniglot and Celeba. We sample the corner points of rectangles uniprobably on the image, but reject those rectangles which area is less than a quarter of the image area.
In Li et al. (2017) six different masks O1–O6 are used on the testing stage. We reconstruct the positions of masks from the illustrations in the paper and give their coordinates in table 6. The visualizations of the masks are available in figure 10.
At the training stage we used a rectangle mask with uniprobable random corners. We reject masks with width or height less than 16pt. We use 64 latent variables and take the best model over 50 epochs based on the validation IWAE log-likelihood estimation. We can obtain slightly higher PSNR values than reported in table 4 if use only masks O1–O6 at the training stage.
In Yeh et al. (2017) four types of masks are used. Center mask is just an unobserved 32x32 square in the center of 64x64 image. Half mask mean that one of upper, lower, left or right half of the image is unobserved. All these types of a half are equiprobable. Random mask means that we use pixelwise-independent Bernoulli
distribution with probability 0.8 to form a mask of unobserved pixels. Pattern mask is proposed in Pathak et al. (2016). As we deduced from the code 3, the generation process is follows: firstly we generate 600x600 one-channel image with uniform distribution over pixels, then bicubically interpolate it to image of size 10000x10000, and then apply Heaviside step function H(x− 0.25) (i. e. all points with value less than 0.25 are considered as unobserved). To sample a mask we sample a random position in this 10000x10000 binary image and crop 64x64 mask. If less than 20% or more than 30% of pixel are unobserved, than the mask is rejected and the position is sampled again. In comparison with this paper in section 5.2 we use the same distribution over masks at training and testing stages. We use VAEAC with 64 latent variables and take the best model over 50 epochs based on the validation IWAE log-likelihood estimation.

heading: A.4 GAIN IMPLEMENTATION DETAILS

text:
For missing feature imputation we reimplemented GAIN in PyTorch based on the paper (Yoon et al., 2018) and the available TensorFlow source code for image inpainting 4.
For categorical features we use one-hot encoding. We observe in experiments that it works better in terms of NRMSE and PFC than processing categorical features in GAIN as continuous ones and then rounding them to the nearest category.
For categorical features we also use reconstruction loss LM (xi, x′i) = − 1|Xi| ∑|Xi| j=1 xi,j log(x ′ i,j). |Xi| is the number of categories of the i-th feature, and xi,j is the j-th component of one-hot encoding of the feature xi. Such LM enforces equal contribution of each categorical feature into the whole reconstruction loss.
We use one more modification of LM (x, x′) for binary and categorical features. Cross-entropy loss in LM penalizes incorrect reconstructions of categorical and binary features much more than incorrect reconstructions for continuous ones. To avoid such imbalance we mixed L2 and cross-entropy reconstruction losses for binary and categorical features with weights 0.8 and 0.2 respectively:
L′M (xi, x ′ i) = 0.2 · LM (xi, x′i) + 0.8 ·
{ 1 |Xi| ∑|Xi| j=1(xi,j − x′i,j)2, if xi is categorical
(xi − x′i)2, if xi is binary (12)
We observe in experiments that this modification also works better in terms of NRMSE and PFC than the original model.
We use validation set which contains 5% of the observed features for the best model selection (hyperparameter is the number of iterations).
In the original GAIN paper authors propose to use cross-validation for hyper-parameter α ∈ {0.1, 0.5, 1, 2, 10}. We observe that using α = 10 and a hint h = b ◦ m + 0.5(1 − b) where vector b is sampled from Bernoulli distribution with p = 0.01 provides better results in terms of NRMSE and PFC than the original model with every α ∈ {0.1, 0.5, 1, 2, 10}. Such hint distribution makes model theoretically inconsistent but works well in practice (see table 7).
Table 7 shows that our modifications provide consistently not worse or even better imputations than the original GAIN (in terms of NRMSE and PFC, on the considered datasets). So in this paper for the missing feature imputation problem we report the results of our modification of GAIN.
3https://github.com/pathak22/context-encoder/blob/master/train_random.lua#

heading: L273

heading: B THEORY

text:
B.1 VAEAC UNIVERSALITY
The theoretical guarantees that VAEAC can model arbitrary distribution are based on the same guarantees for Condtitional Variational Autoencoder (CVAE). We prove below that if CVAE can model each of the conditional distributions p(xb|x1−b), then VAEAC can model all of them. We can imagine 2D CVAEs learned each for the certain mask. Because neural networks are universal approximators, VAEAC networks could model the union of CVAE networks, so that VAEAC network performs transformation defined by the same network of the corresponding to the given mask CVAE.
pψ,V AEAC(z|x1−b, b) = pψ,CV AE,1−b(z|x1−b) ∀x, b
pθ,V AEAC(xb|z, x1−b, b) = pθ,CV AE,1−b(xb|z, x1−b) ∀z, x, b So if CVAE models any distribution p(x|y), VAEAC also do. The guarantees for CVAE in the case of continuous variables are based on the point that every smooth distribution can be approximated with a large enough mixture of Gaussians, which is a special case of CVAE’s generative model. These guarantees can be extended on the case of categorical-continuous variables also. Actually, there are distributions over categorical variables which CVAE with Gaussian prior and proposal distributions cannot learn. Nevertheless, this kind of limitation is not fundamental and is caused by poor proposal distribution family.

heading: B.2 WHY VAEAC NEEDS TARGET VALUES FOR MISSING FEATURES IMPUTATION?

text:
Consider a dataset with D-dimensional objects x where each feature may be missing (which we denote by xi = ω) and their target values y. In this section we show that the better results are achieved when our model learns the concatenation of objects features x and targets y. The example that shows the necessity of it is following. Consider a dataset where x1 = 1, x2 ∼ N (x2|y, 1), pd(y = 0) = p(y = 5) = 0.5. In this
case pd(x2|x1 = 1) = 0.5N (x2|0, 1) + 0.5N (x2|5, 1). We can see that generating data from pd(x2|x1) may only confuse the classifier, because with probability 0.5 it generates x2 ∼ N (0, 1) for y = 5 and x2 ∼ N (5, 1) for y = 0. On the other hand, pd(x2|x1, y) = N (x2|y, 1). Filling gaps using pd(x2|x1, y) may only improve classifier or regressor by giving it some information from the joint distribution pd(x, y) and thus simplifying the dependence to be learned at the training time. So we treat y as an additional feature that is always unobserved during the testing time.

heading: B.3 UNIVERSAL MARGINALIZER: TRAINING PROCEDURE MODIFICATION

text:
The problem authors did not address in the original paper is the relation between the distribution of unobserved components p(b) at the testing stage and the distribution of masks in the requests to UM p̂(b). The distribution over masks p(b) induces the distribution p̂(b), and in the most cases p(b) 6= p̂(b). The distribution p̂(b) also depends on the permutations (i1, i2, . . . , i|b|) that we use to generate objects.
We observed in experiments, that UM must be trained using unobserved mask distribution p̂(b). For example, if all masks from p(b) have a fixed number of unobserved components (e. g., D2 ), then UM will never see an example of mask with 1, 2, . . . , D2 − 1 unobserved components, which is necessary to generate a sample conditioned on D2 components. That leads to drastically low likelihood estimate for the test set and unrealistic samples.
We developed an easy generative process for p̂(b) for arbitrary p(b) if the permutation of unobserved components (i1, i2, . . . , i|b|) is chosen randomly and equiprobably: firstly we generate b0 ∼ p(b), u ∼ U [0, 1], then b1 ∼ (Bernoulli(u))D and b = b0 ◦ b1. More complicated generative process exists for a sorted permutation where ij−1 < ij ∀j : 2 ≤ j ≤ |b|. In experiments we use uniform distribution over the permutations.

heading: C GAUSSIAN STOCHASTIC NEURAL NETWORK

text:
Gaussian stochastic neural network (13) and hybrid model (14) are originally proposed in the paper on Conditional VAE (Sohn et al., 2015). The motivation authors mention in the paper is as follows. During training the proposal distribution qφ(z|x, y) is used to generate the latent variables z, while during the testing stage the prior pψ(z|y) is used. KL divergence tries to close the gap between two distributions but, according to authors, it is not enough. To overcome the issue authors propose to use a hybrid model (14), a weighted mixture of variational lower bound (3) and a single-sample Monte-Carlo estimation of log-likelihood (13). The model corresponding to the second term is called Gaussian Stochastic Neural Network (13), because it is a feed-forward neural network with a single Gaussian stochastic layer in the middle. Also GSNN is a special case of CVAE where qφ(z|x, y) = pψ(z|y).
LGSNN (x, y; θ, ψ) = Epψ(z|y) log pθ(x|z, y) (13) L(x, y; θ, ψ, φ) = αLCV AE(x, y; θ, ψ, φ) + (1− α)LGSNN (x, y; θ, ψ), α ∈ [0, 1] (14)
Authors report that hybrid model and GSNN outperform CVAE in terms of segmentation accuracy on the majority of datasets.
We can also add that this technique seems to soften the “holes problem” (Makhzani et al., 2016). In Makhzani et al. (2016) authors observe that vectors z from prior distribution may be different enough from all vectors z from the proposal distribution at the training stage, so the generator network may be confused at the testing stage. Due to this problem CVAE can have good reconstructions of y given z ∼ qφ(z|x, y), while samples of y given z ∼ pψ(z|x) are not realistic.
The same trick is applicable to our model as well:
LGSNN (x, b; θ, ψ) = Epψ(z|x1−b,b) log pθ(xb|z, x1−b, b) (15) L(x, b; θ, ψ, φ) = αLV AEAC(x, b; θ, ψ, φ) + (1− α)LGSNN (x, b; θ, ψ), α ∈ [0, 1] (16)
In order to reflect the difference between sampling z from prior and proposal distributions, authors of CVAE use two methods of log-likelihood estimation:
log pθ,ψ(x|y) ≈ log 1
S S∑ i=1 pθ(x|zi, y), zi ∼ pψ(z|y) (17)
log pθ,ψ(x|y) ≈ log 1
S S∑ i=1 pθ(x|zi, y)pψ(zi|y) qφ(zi|x, y) , zi ∼ qφ(z|x, y) (18)
The first estimator is called Monte-Carlo estimator and the second one is called Importance Sampling estimator (also known as IWAE). They are asymptotically equivalent, but in practice the Monte-Carlo estimator requires much more samples to obtain the same accuracy of estimation. Small S leads to underestimation of the log-likelihood for both Monte-Carlo and Importance Sampling (Burda et al., 2015), but for Monte-Carlo the underestimation is expressed much stronger.
We perform an additional study of GSNN and hybrid model and show that they have drawbacks when the target distribution p(x|y) is has multiple different local maximums.

heading: C.1 THEORETICAL STUDY

text:
In this section we show why GSNN cannot learn distributions with several different modes and leads to a blurry image samples.
For the simplicity of the notation we consider hybrid model for a standard VAE:
L(x;φ, ψ, θ) = αEz∼qφ(z|x) log pθ(x|z)pψ(z) qφ(z|x) + (1− α)Ez∼pψ(z) log pθ(x|z) (19)
The hybrid model (16) for VAEAC can be obtained from (19) by replacing x with xb and conditioning all distributions on x1−b and b. The validity of the further equations and conclusions remains for VAEAC after this replacement.
Consider now a categorical latent variable z which can take one of K values. Let x be a random variable with true distribution pd(x) to be modeled. Consider the following true data distribution: pd(x = xi) = 1K for i ∈ {1, 2, . . . ,K} and some values x1, x2, . . . , xK . So the true distribution has K different equiprobable modes. Suppose the generator network NNθ which models mapping from z to some vector of parameters vz = NNθ(z). Thus, we define generative distribution as some function of these parameters: pθ(x|z) = f(x, vz). Therefore, the parameters θ are just the set of v1, v2, . . . , vK . For the simplicity of the model we assume pψ(z) = 1K . Taking into account pψ(z) = 1 K , we obtain optimal q(z = i|x) = f(x,vi)∑K j=1 f(x,vj) . Using (19) and the above formulas for qφ, pψ and pθ we obtain the following optimization problem:
max v1,v2,...,vK
1
K K∑ i=1 α K∑ j=1 f(xi, vj)∑K k=1 f(xi, vk) log f(xi, vj) 1 K f(xi,vj)∑K k=1 f(xi,vk) + (1− α) K∑ j=1 1 K log f(xi, vj)  (20)
It is easy to show that (20) is equivalent to the following optimization problem:
max v1,v2,...,vK K∑ i=1 α log ∑Kj=1 f(xi, vj) K + (1− α) K∑ j=1 1 K log f(xi, vj)  (21) It is clear from (21) that when α = 1 the log-likelihood of the initial model is optimized. On the other hand, when α = 0 the optimal point is v1 = v2 = · · · = vK = argmaxv ∑K i=1 log f(xi, v), i. e. z doesn’t influence the generative process, and for each z generator produces the same v which maximizes likelihood estimation of the generative model f(x, v) for the given dataset of x’s. For Bernoulli and Gaussian generative distributions f such v is just average of all modes x1, x2, . . . , xK . That explains why further we observe blurry images when using GSNN model.
The same conclusion holds for for continuous latent variables instead of categorical. Given K different modes in true data distribution, VAE uses proposal network to separate prior distribution intoK components (i. e. regions in the latent space), so that each region corresponds to one mode. On the other hand, in GSNN z is sampled independently on the mode which is to be reconstructed from it, so for each z the generator have to produce parameters suitable for all modes.
From this point of view, there is no difference between VAE and VAEAC. If the true conditional distribution has several different modes, then VAEAC can fit them all, while GSNN learns their average. If true conditional distribution has one mode, GSNN and VAEAC are equal, and GSNN may even learn faster because it has less parameters.
Hybrid model is a trade-off between VAEAC and GSNN: the closer α to zero, the more blurry and closer to the average is the distribution of the model. The exact dependence of the model distribution on α can be derived analytically for the simple data distributions or evaluated experimentally. We perform such experimental evaluation in the next sections.

heading: C.2 SYNTHETIC DATA

text:
In this section we show that VAEAC is capable of learning a complex multimodal distribution of synthetic data while GSNN and hybrid model are not. Let x ∈ R2 and p(b1 = 1) = p(b2 = 1) = 0.5. pd(x) = 1 8 ∑8 i=1N (x|µi, 1 10I) where µi ∼ N (µi|0, I). The distribution p(x) is plotted in figure 6. The dataset contains 100000 points sampled from pd(x). We use multi-layer perceptron with four ReLU layers of size 400-200-100-50, 25-dimensional Gaussian latent variables.
For different mixture coefficients α we visualize samples from the learned distributions pψ,θ(x1, x2), pψ,θ(x1|x2), and pψ,θ(x2|x1). The observed features for the conditional distributions are generated from the marginal distributions p(x2) and p(x1) respectively.
We see in table 8 and in figure 7, that even with very small weight GSNN prevents model from learning distributions with several local optimas. GSNN also increases Monte-Carlo log-likelihood estimation with
1 0 1 x1
1
0
1
x2
0.00 0.25 0.50 0.75 1.00 1.25 1.50 1.75 2.00
Figure 6: Probability density function of synthetic data distribution.
1 0 1
1
0
1
= 1
x1 unknown: x2 p(x2)
x1 p , (x1|x2)
1 0 1
1
0
1
x2 unknown: x1 p(x1)
x2 p , (x2|x1)
1 0 1
1
0
1
x2
x1, x2 unknown: x1, x2 p , (x1, x2)
1 0 1
1
0
1
= 0.99
1 0 1
1
0
1
1 0 1
1
0
1
x2
1 0 1 x1
1
0
1
= 0.9
1 0 1 x1
1
0
1
1 0 1 x1
1
0
1
x2
Figure 7: VAEAC for synthetic data.
(a) VAEAC (b) GSNN
Figure 8: MNIST inpaintings.
Left: input. The gray pixels are unobserved. Middle: samples from the model. Right: ground truth.
a few samples and decreases much more precise Importance Sampling log-likelihood estimation. When α = 0.9 the whole distribution structure is lost.
We see that using α 6= 1 ruins multimodality of the restored distribution, so we highly recommend to use α = 1 or at least α ≈ 1.

heading: C.3 COMPARISON ON THE IMAGE INPAINTING PROBLEM

text:
In figure 8 we can see that the inpaintings produced by GSNN are smooth, blurry and not diverse compared with VAEAC.
Table 9 shows that VAEAC learns distribution over inpaintings better than GSNN in terms of test loglikelihood. Nevertheless, Monte-Carlo estimations with a small number of samples sometimes are better for GSNN, which means less local modes in the learned distribution and more blurriness in the samples.

heading: D ADDITIONAL EXPERIMENTS

heading: D.1 CONVERGENCE SPEED

text: In figure 9 one can see that VAEAC has similar convergence speed to VAE in terms of iterations on MNIST dataset. In our experiments we observed the same behaviour for other datasets. Each iteration of VAEAC is about 1.5 times slower than VAE due to usage of three networks instead of two.

heading: D.2 MISSING FEATURES IMPUTATION

text:
We evaluate the quality of imputations on different datasets (mostly from UCI (Lichman, 2013)). The evaluation is performed for VAEAC, GSNN (15) and NN (neural network; can be considered as a special case of GSNN where pθ(z|x1−b, b) is delta-function; produces single imputation). We compare these methods with MICE (Buuren & Groothuis-Oudshoorn, 2010), MissForest (Stekhoven & Bühlmann, 2011), and GAIN (Yoon et al., 2018).
We see that for some datasets MICE and MissForest outperform VAEAC, GSNN and NN. The reason is that for some datasets random forest is more natural structure than neural network.
The results also show that VAEAC, GSNN and NN show similar imputation performance in terms of NRMSE, PFC, post-imputation R2-score and accuracy. Given the result from appendix C we can take this as a weak evidence that the distribution of imputations has only one local maximum for datasets from (Lichman, 2013).
Left: input. The gray pixels are unobserved. Middle: samples from VAEAC. Right: ground truth.

heading: D.3 FACE INPAINTINGS

text: In figure 10 we provide samples of VAEAC on the CelebA dataset for the masks from (Li et al., 2017).

heading: D.4 GAIN FOR IMAGE INPAINTING

text:
GAIN (Yoon et al., 2018) doesnt use unobserved data during training, which makes it easier to apply to the missing features imputation problem. Nevertheless, it turns into a disadvantage when the fully-observed training data is available but the missingness rate at the testing stage is high.
We consider the horizontal line mask for MNIST which is described in appendix A.3. We use the released GAIN code 5 with a different mask generator. The inpaintings from VAEAC which uses the unobserved pixels during training are available in figure 1. The inpaintings from GAIN which ignores unobserved pixels are provided in figure 11. As can be seen in figure 11, GAIN fails to learn conditional distribution for given mask distribution p(b).
Nevertheless, we don’t claim that GAIN is not suitable for image inpainting. As it was shown in the supplementary of (Yoon et al., 2018) and in the corresponding code, GAIN is able to learn conditional distributions when p(b) is pixel-wise independent Bernoulli distribution with probability 0.5.
5https://github.com/jsyoon0823/GAIN
Left: input. The gray pixels are unobserved. Middle: samples from the model. Right: ground truth.
Left: input. The gray pixels are unobserved. Middle: samples from the model. Right: ground truth.

heading: D.5 UNIVERSAL MARGINALIZER: ILLUSTRATIONS

text:
In figure 12 we provide samples of Universal Marginalizer (UM) and VAEAC for the same inputs.
Consider the case when UM marginal distributions are parametrized with Gaussians. The most simple example of a distribution, which UM cannot learn but VAEAC can, is given in figure 13.


summaries: This paper proposes to use 8/4-bit approximation of activations to save the memory cost during gradient computation.  The proposed technique is simple and straightforward. On the other hand, the proposed method only saves up to a constant cost of the storage. With the constant factor (4x, 8x) depending on whether fp16 or fp32 is used during computation. Notably, there is a small but noticeable accuracy drop in the final trained model using this mechanism.

summaries: The authors detail a procedure to reduce the memory footprint of deep networks by quantization of the activations only on back propagation. While this scheme does not benefit from computational speedups of activation quantization on both passes (and indeed has a slight computational overhead), the authors demonstrate that for common convolutional architectures it nicely preserves the accuracy of computation by computing the forward pass at full accuracy and limiting propagation of errors in the backward pass. This is possible because the majority of errors are introduced in gradient calculation of the weights and not the inputs each layer. The authors also wisely perform quantization after batch normalization and use the known mean and variance of the activations to scale the quantization and reduce errors. They demonstrate very slight drops in performance accuracy for ResNets on Cifar10, Cifar100, and ImageNet with memory compression factors up to 8. They also point to natural future directions such as using vector quantization to better leverage the activation statistics. The paper is also very clearly written with appropriate references to the relevant literature.

abstractText: With innovations in architecture design, deeper and wider neural network models deliver improved performance on a diverse variety of tasks. But the increased memory footprint of these models presents a challenge during training, when all intermediate layer activations need to be stored for back-propagation. Limited GPU memory forces practitioners to make sub-optimal choices: either train inefficiently with smaller batches of examples; or limit the architecture to have lower depth and width, and fewer layers at higher spatial resolutions. This work introduces an approximation strategy that significantly reduces a network’s memory footprint during training, but has negligible effect on training performance and computational expense. During the forward pass, we replace activations with lower-precision approximations immediately after they have been used by subsequent layers, thus freeing up memory. The approximate activations are then used during the backward pass. This approach limits the accumulation of errors across the forward and backward pass—because the forward computation across the network still happens at full precision, and the approximation has a limited effect when computing gradients to a layer’s input. Experiments, on CIFAR and ImageNet, show that using our approach with 8and even 4-bit fixed-point approximations of 32-bit floating-point activations has only a minor effect on training and validation performance, while affording significant savings in memory usage.

authors:
Martı́n Abadi
Paul Barham
Jianmin Chen
Zhifeng Chen
Andy Davis
Jeffrey Dean
Matthieu Devin
Sanjay Ghemawat
Geoffrey Irving
Michael Isard

title: Tensorflow: A system for largescale machine learning

venue: arXiv preprint arXiv:1605.08695,

authors:
Tianqi Chen
Bing Xu
Chiyuan Zhang
Carlos Guestrin

title: Training deep nets with sublinear memory

venue: cost. CoRR,

authors:
Jeffrey Dean
Greg Corrado
Rajat Monga
Kai Chen
Matthieu Devin
Mark Mao
Andrew Senior
Paul Tucker
Ke Yang
Quoc V Le

title: Large scale distributed deep networks

authors:
Aidan N. Gomez
Mengye Ren
Raquel Urtasun
Roger B. Grosse

title: The reversible residual network: Backpropagation without storing activations

venue: In Advances in Neural Information Processing Systems 30: Annual Conference on Neural Information Processing Systems

authors:
Audrunas Gruslys
Rémi Munos
Ivo Danihelka
Marc Lanctot
Alex Graves

title: Memory-efficient backpropagation through time

venue: In Advances in Neural Information Processing Systems 29: Annual Conference on Neural Information Processing Systems

authors:
Suyog Gupta
Ankur Agrawal
Kailash Gopalakrishnan
Pritish Narayanan

title: Deep learning with limited numerical precision

authors:
Song Han
Huizi Mao
William J Dally

title: Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding

venue: arXiv preprint arXiv:1510.00149,

authors:
Kaiming He
Xiangyu Zhang
Shaoqing Ren
Jian Sun

title: Deep residual learning for image recognition

venue: IEEE Conference on Computer Vision and Pattern Recognition,

title: Identity mappings in deep residual networks

venue: In Proceedings of the European conference on computer vision,

authors:
Gao Huang
Zhuang Liu
Laurens Van Der Maaten
Kilian Q Weinberger

title: Densely connected convolutional networks

venue: In Proc. CVPR,

authors:
Itay Hubara
Matthieu Courbariaux
Daniel Soudry
Ran El-Yaniv
Yoshua Bengio

title: Quantized neural networks: Training neural networks with low precision weights and activations

venue: Journal of Machine Learning Research (JMLR),

authors:
Sergey Ioffe
Christian Szegedy

title: Batch normalization: Accelerating deep network training by reducing internal covariate shift

venue: In Proceedings of the 32nd International Conference on Machine Learning, ICML 2015,

authors:
Alex Krizhevsky
Geoffrey Hinton

title: Learning multiple layers of features from tiny images

venue: Technical report,

authors:
James Martens
Ilya Sutskever

title: Training deep and recurrent networks with hessian-free optimization

venue: In Neural Networks: Tricks of the Trade - Second Edition,

authors:
Paulius Micikevicius
Sharan Narang
Jonah Alben
Gregory Diamos
Erich Elsen
David Garcia
Boris Ginsburg
Michael Houston
Oleksii Kuchaev
Ganesh Venkatesh

title: Mixed precision training

venue: arXiv preprint arXiv:1710.03740,

authors:
Benjamin Recht
Christopher Re
Stephen Wright
Feng Niu

title: Hogwild: A lock-free approach to parallelizing stochastic gradient descent

authors:
Herbert Robbins
Sutton Monro

title: A stochastic approximation method

venue: In Herbert Robbins Selected Papers,

heading: 1 INTRODUCTION

text:
Deeper neural network models are able to express more complex functions, and recent results have shown that with the use of residual (He et al., 2016a) and skip (Huang et al., 2017) connections to address vanishing gradients, such networks can be trained effectively to leverage this additional capacity. As a result, the use of deeper network architectures has become prevalent, especially for visual inference tasks (He et al., 2016b). The shift to larger architectures has delivered significant improvements in performance, but also increased demand on computational resources. In particular, deeper network architectures require significantly more on-device memory during training—much more so than for inference. This is because training requires retaining the computed activations of all intermediate layers since they are needed to compute gradients during the backward pass.
The increased memory footprint means fewer training samples can fit in memory and be processed as a batch on a single GPU. This is inefficient: smaller batches are not able to saturate all available parallel cores, especially because computation in “deeper” architectures is distributed to be more sequential. Moreover, smaller batches also complicate the use of batch-normalization (Ioffe & Szegedy, 2015), since batch statistics are now computed over fewer samples making training less stable. These considerations often force the choice of architecture to be based not just on optimality for inference, but also practical feasibility for training—for instance, deep residual networks for large images drop resolution early, so that most layers have smaller sized outputs.
While prior work to address this has traded-off memory for computation (Martens & Sutskever, 2012; Chen et al., 2016; Gruslys et al., 2016; Gomez et al., 2017), their focus has been on enabling exact gradient computation. However, since stochastic gradient descent (SGD) inherently works with noisy gradients at each iteration, we propose an algorithm that computes reasonably approximate gradients, while significantly reducing a network’s memory footprint and with virtually no additional computational cost. Our work is motivated by distributed training algorithms
that succeed despite working with approximate and noisy gradients aggregated across multiple devices (Recht et al., 2011; Dean et al., 2012; Seide et al., 2014; Wen et al., 2017). We propose using low-precision approximate activations—that require less memory—to compute approximate gradients during back-propagation (backprop) on a single device. Note that training with a lowerprecision of 16- instead of 32-bit floating-point representations is not un-common. But this lower precision is used for all computation, and thus allows only for a modest lowering of precision, since the approximation error builds up across the forward and then backward pass through all layers.
In this work, we propose a new backprop implementation that performs the forward pass through the network at full-precision, and incurs limited approximation error during the backward pass. We use the full-precision version of a layer’s activations to compute the activations of subsequent layers. However, once these activations have been used in the forward pass, our method discards them and stores a low-precision approximation instead. During the backward pass, gradients are propagated back through all the layers at full precision, but instead of using the original activations, we use their low-precision approximations. As a result, we incur an approximation error at each layer when computing the gradients to the weights from multiplying the incoming gradient with the approximate activations, but ensure the error in gradients going back to the previous layer is minimal.
Our experimental results show that even using only 4-bit fixed-point approximations, for the original 32-bit floating-point activations, causes only minor degradation in training quality. This significantly lowers the memory required for training, which comes essentially for “free”—incurring only the negligible additional computational cost of converting activations to and from low precision representations. Our memory-efficient version of backprop is thus able to use larger batch sizes at each iteration—to fully use available parallelism and compute stable batch statistics—and makes it practical for researchers to explore the use of much larger and deeper architectures than before.

heading: 2 RELATED WORK

text:
A number of works focus on reducing the memory footprint of a model during inference, e.g., by compression Han et al. (2015) and quantization Hubara et al. (2017), to ensure that it can be deployed on resource-limited mobile devices, while still delivering reasonable accuracy. These methods still require storing full versions of model weights and activations during network training, and assume there is sufficient memory to do so. However, training requires significantly more memory than inference because of the need to store all intermediate activations. And so, memory can be a bottleneck during training, especially with the growing preference for larger and deeper network architectures. A common recourse to this has been to simply use multiple GPUs during training. But, this is inefficient due to the overhead of intra-device communication, and often under-utilizes the available parallelism on each device—computation in deeper architectures is distributed more sequentially and, without sufficient data parallelism, often does not saturate all GPU cores.
A popular strategy to reduce training memory requirements is “checkpointing”. Activations for only a subset of layers are stored at a time, and the rest recovered by repeating forward computations (Martens & Sutskever, 2012; Chen et al., 2016; Gruslys et al., 2016). This affords memory savings with the trade-off of additional computational cost—e.g., Chen et al. (2016) propose a strategy that requires memory proportional to the square-root of the number of layers, while requiring up to the computational cost of an additional forward pass. In a similar vein, Gomez et al. (2017) considered network architectures with “reversible” or invertible layers. This allows re-computing intermediate input activations of reversible layers from their outputs during the backward pass.
These methods likely represent the best possible solutions if the goal is restricted to computing exact gradients. But SGD is fundamentally a noisy process, and the exact gradients computed over a batch at each iteration are already an approximation—of gradients of the model over the entire training set (Robbins & Monro, 1985). Researchers have posited that further approximations are possible without degrading training ability, and used this to realize gains in efficiency. For distributed training, asynchronous methods (Recht et al., 2011; Dean et al., 2012) delay synchronizing models across devices to mitigate communication latency. Despite each device now working with stale models, there is no major degradation in training performance. Other methods quantize gradients to two (Seide et al., 2014) or three levels (Wen et al., 2017) so as to reduce communication overhead, and again find that training remains robust to such approximation. Our work also adopts an approximation strategy to gradient computation, but targets the problem of memory usage on a each device. We approximate activations, rather than gradients, with lower-precision representations, and by doing so, we are able to achieve considerable reductions in a model’s memory footprint during training. Note that since our method achieves a constant factor saving in memory for back-propagation across any group of layers, it can also be employed within checkpointing to further improve memory cost.
It is worth differentiating our work from those that carry out all training computations at lowerprecision Micikevicius et al. (2017); Gupta et al. (2015). This strategy allows for a modest loss in precision: from 32- to 16-bit representations. In contrast, our approach allows for much greater reduction in precision. This is because we carry out the forward pass in full-precision, and approximate activations only after they have been used by subsequent layers. Our strategy limits accumulation of errors across layers, and we are able to replace 32-bit floats with 8- and even 4-bit fixed-point approximations, with little to no effect on training performance. Note that performing all computation at lower-precision also has a computational advantage: due to reduction in-device memory bandwidth usage (transferring data from global device memory to registers) in Micikevicius et al. (2017), and due to the use of specialized hardware in Gupta et al. (2015). While the goal of our work is different, our strategy can be easily combined with these ideas: compressing intermediate activations to a greater degree, while also using 16- instead of 32-bit precision for computation.

heading: 3 PROPOSED METHOD

heading: 3.1 BACKGROUND

text:
A neural network is composition of linear and non-linear functions that map the input to the final desired output. These functions are often organized into “layers”, where each layer consists of a single linear transformation—typically a convolution or a matrix multiply—and a sequence of nonlinearities. We use the “pre-activation” definition of a layer, where we group the linear operation
with the non-linearities that immediately preceed it. Consider a typical network whose lth layer applies batch-normalization and ReLU activation to its input Al:i followed by a linear transform:
[Batch Norm.] Al:1 = (σ2l + ) −1/2 ◦ (Al:i − µl), µl = Mean(Al:i), σ2 = Var(Al:i); (1) [Scale & Bias] Al:2 = γl ◦Al:1 + βl; (2) [ReLU] Al:3 = max(0, Al:2); (3) [Linear] Al:o = Al:3 ×Wl; (4)
to yield the output activations Al:o that are fed into subsequent layers. Here, each activation is a tensor with two or four dimensions: the first indexing different training examples, the last corresponding to “channels”, and others to spatial location. Mean(·) and Var(·) aggregate statistics over batch and spatial dimensions, to yield vectors µl and σ2l with per-channel means and variances. Element-wise addition and multiplication (denoted by ◦) are carried out by “broadcasting” when the tensors are not of the same size. The final operation represents the linear transformation, with × denoting matrix multiplication. This linear transform can also correspond to a convolution.
Note that (1)-(4) are defined with respect to learnable parameters γl, βl, and Wl, where γl, βl are both vectors of the same length as the number of channels in Al, and Wl denotes a matrix (for fullyconnected layers) or elements of a convolution kernel. These parameters are learned iteratively using SGD, where at each iteration, they are updated based on gradients—∇γl, ∇βl, and ∇Wl—of some loss function with respect to these parameters, computed on a batch of training samples.
To compute gradients with respect to all parameters for all layers in the network, the training algorithm first computes activations for all layers in sequence, ordered such that each layer in the sequence takes as input the output from a previous layer. The loss is computed with respect to activations of the final layer, and then the training algorithm goes through all layers again in reverse sequence, using the chain rule to back-propagate gradients of this loss. For the lth layer, given the gradients ∇Al:o of the loss with respect to the output, this involves computing gradients ∇γl,∇βl, and∇Wl with respect to the layer’s learnable parameters, as well as gradients∇Al:i with respect to its input for further propagation. These gradients are given by:
[Linear] ∇W = ATl:3 × (∇Al:o), ∇Al:3 = (∇Al:o)×WTl ; (5) [ReLU] ∇Al:2 = δ(Al:2 > 0) ◦ (∇Al:3); (6)
[Scale & Bias] ∇βl = Sum(∇Al:2), ∇γl = Sum (Al:1 ◦ (∇Al:2)) , ∇Al:1 = γl ◦ ∇Al:2; (7) [Batch Norm.] ∇Al:i = (σ2l + )−1/2 ◦ [∇Al:1 −Mean(∇Al:1)−Al:1 ◦Mean(Al:1 ◦ ∇Al:1)] ;
(8)
where Sum(·) and Mean(·) involve aggregation over all but the last dimension, and δ(A > 0) a tensor the same size as A that is one where the values in A are positive, and zero otherwise.
When the goal is to just compute the final output of the network, the activations of an intermediate layer can be discarded during the forward pass as soon as we finish processing the subsequent layer or layers that use it as input. However, we need to store all these intermediate activations during training because they are needed to compute gradients during back-propagation: (5)-(8) involve not just the values of the incoming gradient, but also the values of the activations themselves. Thus, training requires enough available memory to hold the activations of all layers in the network.

heading: 3.2 BACK-PROPAGATION WITH APPROXIMATE ACTIVATIONS

text:
We begin by observing we do not necessarily need to store all intermediate activations Al:1, Al:2, and Al:3 within a layer. For example, it is sufficient to store the activation values Al:2 right before the ReLU, along with the variance vector σ2l (which is typically much smaller than the activations themselves). Given Al:2, we can reconstruct the other activations Al:3 and Al:3 needed in (5)-(8) using element-wise operations, which typically have negligible computational cost compared to the linear transform itself. Some deep learning frameworks already use such “fused” layers to conserve memory, and we consider this to be our “baseline” for memory usage.
However, storing one activation tensor at full-precision for every layer still requires a considerable amount of memory. We therefore propose retaining an approximate low-precision version Ãl:2 of Al:2, that requires much less memory for storage, for use in (5)-(8) during back-propagation.
As shown in Fig. 1, we use full-precision versions of all activations during the forward pass to compute Al:o from Al:i as per (1)-(4), and use Al:2 to compute its approximation Ãl:2. The full precision approximations are discarded as soon they have been used—the intermediate activations Al:1, Al:2, Al:3 are discarded as soon as the approximation Ãl:2 and outputAl:o have been computed, and Al:o is discarded after it has been used by a subsequent layer. Thus, only the approximate activations Ãl:2 and (full-precision) variance vector σ2l are retained in memory for back-propagation.
We use a simple, computationally inexpensive approach to approximate Al:2 via a K-bit fixed-point representation for some desired value of K. Since Al:1 is normalized to be zero-mean and unitvariance, Al:2 has mean βl and variance γ2l . We compute an integer tensor Ã ∗ l:2 from Al:2 as:
Ã∗l:2 = ClipK( bAl:2 ◦ 2K(6 ∗ γl)−1c+ 2K−1 − bβl ◦ 2K(6 ∗ γl)−1c), (9)
where b·c indicates the “floor” operator, and ClipK(x) = max(0,min(2K − 1, x)). The resulting integers (between 0 and 2K − 1) can be directly stored with K-bits. When needed during backpropagation, we recover a floating-point tensor holding the approximate activations Ãl:2 as:
Ãl:2 = 2 −K(6 ∗ γl) ◦ (Ã∗l:2 + 0.5− 2K−1 + bβl ◦ 2K(6 ∗ γl)−1c). (10)
This simply has the effect of clipping Al:2 to the range βl ± 3γl (approximately, the range may be slightly asymmetric around βl because of rounding), and quantizing values in 2K fixed-size intervals (to the median of each interval) within that range. However, crucially, this approximation ensures that the sign of each value is preserved, i.e., δ(Al:2 > 0) = δ(Ãl:2 > 0).

heading: 3.3 APPROXIMATION ERROR IN TRAINING

text:
Since the forward computations happen in full-precision, there is no error introduced in any of the activations Al prior to approximation. To analyze the error introduced by our approach, we then consider the effect of using Ãl:2 instead of Al:2 (and equivalently, Ãl:1 and Ãl:3 derived from Ãl:2) to compute gradients in (5)-(8). We begin by noting that for all values of Al:2 that fall within the range βl ± 3γl (and are therefore not clipped), the worst-case approximation error in the activations themselves is bounded by half the width of the quantization intervals:
|Al:2 − Ãl:2| ≤ 3/2Kγ = 3/2K √ Var(Al:2), (11)
where Var(·) denotes per-channel variance (and the RHS is interpreted as applying to all channels). Hence, the approximation error is a fraction of the variance in the activations themselves, and is lower for higher values of K. It is easy to see that |Al:3 − Ãl:2| ≤ |Al:2 − Ãl:3| since Al:3 and Ãl:3 are derived from Al:2 and Ãl:3 by clipping negative values of both to 0, which only decreases the error. Further, since Al:2 is related to Al:1 by simply scaling, the error in Ãl:1 is also bounded as a fraction of its variance, which is one, i.e: |Al:1 − Ãl:1| ≤ 3/2K . We next examine how these errors in the activations affect the accuracy of gradient computations in (5)-(8). During the first back-propagation step in (5) through the linear transform, the gradient∇W to the learnable transform weights will be affected by the approximation error in Ãl:3. However, the gradient ∇Al:2 can be computed exactly (as a function of the incoming gradient to the layer ∇Al:o), because it does not depend on the activations. Back-propagation through the ReLU in (7) is also not affected, because it depends only on the sign of the activations, which is preserved by our approximation. When back-propagating through the scale and bias in (6), only the gradient ∇γ to the scale depends on the activations, but gradients to the bias βl and toAl:1 can be computed exactly.
And so, although our approximation introduces some error in the computations of ∇W and ∇γ, there is no error introduced in the gradient flowing towards the input of the layer, until it reaches the batch-normalization operation in (8). Here, we do incur an error, but note that this is only in one of the three terms of the expression for ∇Al:i—which accounts for back-propagating through the variance computation, and is the only term that depends on the activations. Hence, while our activation approximation does introduce some errors in the gradients for the learnable weights, we limit the accumulation of these errors across layers because a majority of the computations for backpropagation to the input of each layer are exact. This is illustrated in Fig. 1, with the use of green arrows to show computations that are exact, and red arrows for those affected by the approximation.

heading: 3.4 NETWORK ARCHITECTURES AND MEMORY USAGE

text:
Our full training algorithm applies our approximation strategy to every layer (defined by grouping linear transforms with preceding non-linear activations) during the forward and backward pass. Skip and residual connections are handled easily, since back-propagation through these connections involves simply copying to and adding gradients from both paths, and doesn’t involve the activations themselves. (Although we do not consider this in our implementation, older residual connections that are added after batch-normalization but before the ReLU can also be handled, but would require saving activations both before and after addition—in the traditional case, well as our approach).
Our method is predicated on the use of ReLU activations since its gradient depends only on the sign of the activations, and can be used for other such non-linearities such as “leaky”-ReLUs. Other activations (like sigmoid) may incur additional errors—in particular, we do not approximate the activations of the final output layer in classifier networks that go through a Soft-Max. However, since this is typically at the final layer, and computing these activations is immediately followed by back-propagating through that layer, approximating these activations offers no savings in memory. Our approach also handles average pooling by simply folding it in with the linear transform. For max-pooling, exact back-propagation through the pooling operation would require storing the argmax indices (the number of bits required to store these would depend on the max-pool receptive field size). However, since max-pool layers are used less often in recent architectures in favor of learned downsampling (ResNet architectures for image classification use max-pooling only in one layer), we instead choose not to approximate layers with max-pooling for simplicity.
Given a network with L layers, our memory usage depends on connectivity for these layers. Our approach requires storing the approximate activations for each layer, each occupying reduced memory rate at a fractional rate of α < 1. During the forward pass, we also need to store, at full-precision, those activations that are yet to be used by subsequent layers. This is one layer’s activations for feedforward networks, and two layers’ for standard residual architectures. More generally, we will need to store activations for upto W layers, where W is the “width” of the architecture—which we define as the maximum number of outstanding layer activations that remain to be used as process layers in sequence. During back-propagation, the same amount of space is required for storing gradients till they are used by previous layers. We also need space to re-create a layer’s approximate activations as full-precision tensors from the low-bit stored representation, for use in computation.
Thus, assuming that all activations of layers are the same size, our algorithm requiresO(W+1+αL) memory, compared to the standard requirement of O(L). This leads to substantial savings for deep networks with large L (note α = 1/4, 1/8 when approximating 32-bit floats with K = 8, 4 bits).

heading: 4 EXPERIMENTS

text:
We developed a library that implements the proposed method for approximate memory-efficient training, given a network architecture specification which can include residual layers (i.e., W = 2). As illustrated in Fig. 1, the method allocates a pair of global buffers for the direct and residual paths that is common to all layers. At any point during the forward pass, these buffers hold the fullprecision activations that are needed for computation of subsequent layers. The same buffers are used to store gradients during the back-ward pass. Beyond these common buffers, the library only stores the low-precision approximate activations for each layer for use during the backward-pass. Further details of our implementation are provided in the appendix. We compare our approximate training approach, with 8- and 4-bit activations, to exact training with full-precision activations as a baseline. For a fair comparison, we again only store one set of activations like our method for a group of batch-normalization, ReLU, and linear (convolution) operations. This is achieved with our library by storing Al:2 without approximation (Ãl:2 = Al:2).
CIFAR-10 and CIFAR-100. We begin with comparisons on 164-layer pre-activation residual networks (He et al., 2016b) on CIFAR-10 and CIFAR-100 (Krizhevsky & Hinton, 2009), using threelayer “bottlneck” residual units and parameter-free shortcuts for all residual connections. We train the network for 64k iterations with a batch size of 128, momentum of 0.9, and weight decay of 2 × 10−4. Following He et al. (2016b), the learning rate is set to 10−2 for the first 400 iterations, then increased to 10−1, and dropped by a factor of 10 at 32k and 48k iterations. We use standard data-augmentation with random translation and horizontal flips. We train these networks with our
approach using K = 8 and K = 4 bit approximations, and measure degradation in accuracy with respect to the baseline—repeating training for all cases with four random seeds. We visualize the evolution of training and test set error in Fig. 2, and report statistics of the final test error in Table 1.
We find that both training and test errors when using our low-memory approximation strategy closely follow those of exact back-propagation, throughout the training process. Moreover, the final median test errors of models trained with even 4-bit approximations (i.e., corresponding to α = 1/8) are higher only by 0.07% compared to those trained with exact computations. To examine the reason behind this robustness, Fig. 2 also visualizes the error in the final parameter gradients used to update the model. Specifically, we take two models for CIFAR-100—at the start and end of training—and then compute gradients for a 100 batches with respect to the convolution kernels of all layers exactly, and using our approximate strategy. We plot the average squared error between these gradients. We compare this approximation error to the “noise” inherent in SGD, due to the fact that each iteration considers a random batch of training examples. This is measured by average variance between the (exact) gradients computed in the different batches. We see that our approximation error is between one and two orders of magnitude below the SGD noise for all layers, both at the start and end of training. So while we do incur an error due to approximation, this is added to the much higher error that already exists due to SGD even in exact training, and hence further degradation is limited.
ImageNet. We also report results on training models for ImageNet (Russakovsky et al., 2015). Here, we consider two residual architectures: 34-layer (with two-layer units without bottlenecks) and 152- layer (with three-layer bottleneck units)—again using pre-activation parameter-free shortcuts. We
train with a batch size of 256 for a total of 640k iterations with a momentum of 0.9, weight decay of 10−4, and standard scale, color, flip, and translation augmentation. The initial learning rate is set to 10−1 with drops by factor of 10 every 160k iterations. Table. 1 reports top-5 validation accuracy (using 10 crops at a scale of 256) for models trained using exact computation, and our approach with K = 8 and K = 4 bit approximations. Again, the drop in accuracy is relatively small: at 0.7% and 0.5% for the 34- and 152-layer models respectively, for a memory savings factor of α = 1/8.
Memory and Computational Efficiency. For the CIFAR experiments, the full 128-size batch fit on a single 1080Ti GPU for both the baseline and our method. For ImageNet with ResNet-34, our method could fit the 256-sized batch, but not the baseline—for which we used two passes with 128-size batches and averaged the gradients for each iteration. For ResNet-152, we parallelized computation across two GPUs, and again, our method could fit half a batch (size 128) on each GPU, the baseline required two passes with 64−sized batches per-GPU per-pass. In the CIFAR experiments and ImageNet with ResNet-34, the batches were large enough to saturate all GPU cores for both our method and the baseline. In this case, the running times per iteration were almost identical—with a very slight increase in our case due to the cost computing approximations: exact vs approximate (4-bit) training took 0.66 seconds vs 0.72 seconds for CIFAR-100, and 1.68 seconds vs 1.71 seconds for ImageNet ResNet-34. But for ResNet-152 on ImageNet, the 64-sized batch for exact training underutilized the available parallelism, and the time per-iteration (across two GPUs) was 2s vs 1.7s for exact vs approximate (4-bit) training.
However, these represent comparisons restricted to have the same total batch size (needed to evaluate relative accuracy). For a more precise evaluation of memory usage, and the resulting computational efficiency from parallelism, we considered residual networks for CIFAR-10 of various depths up to 1001 layers—and additionally for the deepest network, a version with four times as many feature channels in each layer. For each network, we measured the largest batch size that could be fit in memory with our method (with K = 4) vs the baseline, i.e., b such that a batch of b + 1 caused an out-of-memory error on a 1080Ti GPU. We also measured the corresponding wall-clock training time per sample, computed as the training time per-iteration divided by this batch size. These results are summarized in Table 2. We find that in all cases, our method allows significantly larger batches to be fit in memory. Moreover, for larger networks, our method also yields a notable computational advantage since the larger batches permit full exploitation of available cores on the GPU.

heading: 5 CONCLUSION

text:
We introduced a new algorithm for approximate gradient computation in neural network training, that significantly reduces the amount of required on-device memory. Our experiments show that this comes at a minimal cost in terms of both quality of the learned models, and computational expense. With a lower memory footprint, our method allows training with larger batches in each iteration— improving efficiency and stability—and exploration of deeper architectures that were previously impractical to train. We will release our reference implementation on publication.
Our method shows that SGD is reasonably robust to working with approximate activations. While we used an extremely simple approximation strategy—uniform quantization—in this work, we are interested in exploring whether more sophisticated techniques—e.g., based on random projections or vector quantization—can provide better trade-offs, especially if informed by statistics of gradients and errors from prior iterations. We are also interested in investigating whether our approach to partial approximation can be utilized in other settings, especially to reduce inter-device communication for distributed training with data or model parallelism.