AMORTISED MAP INFERENCE FOR IMAGE SUPER-RESOLUTION
Image super-resolution (SR) is an underdetermined inverse problem, where a large number of plausible high resolution images can explain the same downsampled image. Most current single image SR methods use empirical risk minimisation, often with a pixel-wise mean squared error (MSE) loss. However, the outputs from such methods tend to be blurry, over-smoothed and generally appear implausible. A more desirable approach would employ Maximum a Posteriori (MAP) inference, preferring solutions that always have a high probability under the image prior, and thus appear more plausible. Direct MAP estimation for SR is nontrivial, as it requires us to build a model for the image prior from samples. Here we introduce new methods for amortised MAP inference whereby we calculate the MAP estimate directly using a convolutional neural network. We first introduce a novel neural network architecture that performs a projection to the affine subspace of valid SR solutions ensuring that the high resolution output of the network is always consistent with the low resolution input. Using this architecture, the amortised MAP inference problem reduces to minimising the cross-entropy between two distributions, similar to training generative models. We propose three methods to solve this optimisation problem: (1) Generative Adversarial Networks (GAN) (2) denoiser-guided SR which backpropagates gradient-estimates from denoising to train the network, and (3) a baseline method using a maximum-likelihoodtrained image prior. Our experiments show that the GAN based approach performs best on real image data. Lastly, we establish a connection between GANs and amortised variational inference as in e. g. variational autoencoders.
1 INTRODUCTION
Image super-resolution (SR) is the underdetermined inverse problem of estimating a high resolution (HR) image given the corresponding low resolution (LR) input. This problem has recently attracted significant research interest due to the potential of enhancing the visual experience in many applications while limiting the amount of raw pixel data that needs to be stored or transmitted. While SR has many applications in for example medical diagnostics or forensics (Nasrollahi & Moeslund, 2014, and references therein), here we are primarily motivated to improve the perceptual quality when applied to natural images. Most current single image SR methods use empirical risk minimisation, often with a pixel-wise mean squared error (MSE) loss (Dong et al., 2016; Shi et al., 2016). However, MSE, and convex loss functions in general, are known to have limitations when presented with uncertainty in multimodal and nontrivial distributions such as distributions over natural images. In SR, a large number of plausible images can explain the LR input and the Bayes-optimal behaviour for any MSE trained model is to output the mean of the plausible solutions weighted according to their posterior probability. For natural images this averaging behaviour leads to blurry and over-smoothed outputs that generally appear implausible, i.e. the produced estimates have low probability under the natural image prior.
An idealised method for our applications would use a full-reference perceptual loss function that describes the sensitivity of the human visual perception system to different distortions. However the ∗Work done while CKS was an intern at Twitter
shown ( ). b) Trained model outputs for x ∈ [−8, 8] and estimated gradients from a denoising function trained on pY . Note the AffGAN( ) and AffDG( ) models fit the posterior mode well whereas the MSE ( ) and MAE ( ) model outputs generally fall in low probability regions.
H[qθ, pY ] `MSE(x,Ay)
MAP 3.15 - MSE 9.10 1.25 · 10−2 MAE 6.30 4.04 · 10−2 AffGAN 4.10 0.0 SoftGAN 4.25 8.87 · 10−2 AffDG 3.81 0.0 SoftDG 4.19 1.01 · 10−1
Table 1: Directly estimated cross-entropy H[qθ, pY ] values. The AffGAN and AffDG achieves cross-entropy values close to the MAP solution confirming that they minimize the desired quantity. The MSE and MAE models performs worse since they do not minimize the cross-entropy. Further the models using affine projections (Aff) performs better than the soft constrained models.
most widely used loss functions MSE and the related peak-signal-to-noise-ratio (PSNR) metric have been shown to correlate poorly with human perception of image quality (Laparra et al., 2016; Wang et al., 2004). Improved perceptual quality metrics have been proposed, the most popular being structural similarity (SSIM) (Wang et al., 2004) and its multi-scale variants (Wang et al., 2003). Although the correlation of these metrics with human perception has improved, they still do not provide a fully satisfactory alternative to MSE for training of neural networks (NN) for SR.
In lieu of a satisfactory perceptual loss function, we leave the empirical risk minimisation framework and present methods based only on natural image statistics. In this paper we argue that a desirable approach is to employ amortised Maximum a Posteriori (MAP) inference, preferring solutions that have a high posterior probability and thus high probability under the image prior while keeping the computational benefits of amortised inference. To motivate why MAP inference is desirable consider the toy problem in Figure 1a, where the HR data is two-dimensional y = [y1, y2] and distributed according to the Swiss-roll density. The LR observation is defined as the average of the two pixels x = y1+y22 . Consider observing a LR data point x = 0.5: the set of possible HR solutions is the line y1 = 2x − y2, more generally an affine subspace, which is shown by the dashed line in Figure 1a. The posterior distribution p(y|x) is thus degenerate, and corresponds to a slice of the prior along this line, as shown by the red shading. If one minimise MSE or Mean Absolute Error (MAE), the Bayes-optimal solution will lie at the mean or the median along the line, respectively. This example illustrates that MSE and MAE can produce output with very low probability under that data prior whereas MAP inference would always find the mode which by definition is in a high-probability region. See Section 5.6 for a discussion of possible limitations of the MAP inference approach.
Our first contribution is a convolutional neural networks (CNN) architecture designed to exploit the structure of the SR problem. Image downsampling is a linear transformation, and can be modelled as a strided convolution. As Figure 1a illustrates, the set of HR images y that are compatible with any LR image x span an affine subspace. We show that by using specifically chosen linear convolution and deconvolution layers we can implement a projection to this affine subspace. This ensures that our CNNs always output estimates that are consistent with the inputs. The affine projection layer can be added to any CNN, or indeed, any other trainable SR algorithm. Using this architecture we show that training the model for MAP inference reduces to minimising the cross-entropy H[qG, pY ] between the HR data distribution pY and the implied distribution qG of the model’s output when evaluated at random LR images. As a result, we don’t need corresponding HR and LR image pairs any more, and training becomes more akin to training generative models. However direct minimisation of the cross-entropy is not possible and instead we develop three approaches, all depending on projecting the model output to the affine subspace of valid solution, to approximate it directly from data:
1. We present a variant of the Generative Adversarial Networks (GAN) (Goodfellow et al., 2014) which approximately minimises the Kullback–Leibler divergence (KL) and cross-entropy between qG and pY . Our analysis provides theoretical grounding for using GANs in image SR (Ledig et al., 2016). We also introduce a trick that we call instance noise that can be generally applied to address the instability of training GANs. 2. We employ denoising as a way to capture natural image statistics. Bayes-optimal denoising approximately learn to take a gradient step along the log-probability of the data distribution (Alain & Bengio, 2014). These gradient estimates from denoising can be directly backpropagated through the network to minimise cross-entropy between qG and pY via gradient descent. 3. We present an approach where the probability density of data is directly modelled via a generative model trained by maximum likelihood. We use a differentiable generative model based on PixelCNNs (Oord et al., 2016) and Mixture of Conditional Gaussian Scale Mixtures (MCGSM, Theis et al., 2012) whose performance we believe is very close to the-state-of-the-art in this category.
In section 5 we empirically demonstrate the behaviour of the proposed methods on both the two dimensional toy dataset and on real image datasets. Lastly, in Appendix F we show that a stochastic version of AffGAN performs amortised variational inference, which for the first time establishes a connection between GANs and variational inference as in e. g. variational autoencoders (Kingma & Welling, 2014).
2 RELATED WORK
The GAN framework was introduced by Goodfellow et al. (2014) which also showed that these models minimise the Shannon-Jensen Divergence between qG and pY under certain conditions. In Section 3.2, we present an update rule that corresponds to minising KL[qG‖pY ]. Recently, Nowozin et al. (2016) presented a more general treatment that connects GANs to f -divergence minimisation. In parallel to our contributions, theoretical work by Mohamed & Lakshminarayanan (2016) presented a unifying view on learning in GAN-style algorithms, of which our variant can be regarded a special case. The focus of several recent papers on GANs were algorithmic tricks to improve their stability (Radford et al., 2015; Salimans et al., 2016). In Section 3.2.1 we introduce another such trick we call instance noise. We discuss theoretical motivations for this and compare it to one-sided label smoothing proposed by Salimans et al. (2016). We also refer to parallel work by Arjovsky & Bottou (2017) proposing a similar method. Recently, several attempts have been made to improve perceptual quality of SR using deep representations of natural images. Bruna et al. (2016) and Li & Wand (2016) measure the Euclidean distance in the nonlinear feature space of a deep NN pre-trained to perform object classification. Dosovitskiy & Brox (2016) and Ledig et al. (2016) use a similar approach and also add an adversarial loss term. Unpublished work by Garcia (2016) explored combining GANs with an L1 penalty between the LR input and the down-sampled output. We note that the soft L2 or L1 penalties used in these methods can be interpreted as assuming Gaussian and Laplace observation noise. In contrast, our approach assumes no observation noise and satisfies the consistency of inputs and outputs exactly by using an affine projection as explained in Section 3.1. In other work, Larsen et al. (2015) proposed to replace the pixel-wise MSE used for training of variational autoencoders with a learned metric from the GAN discriminator. Our denoiser based method exploits a fundamental connection between probabilistic modelling and learning to denoise (see e. g. Vincent et al., 2008; Alain & Bengio, 2014; Särelä & Valpola, 2005; Rasmus et al., 2015; Greff et al., 2016): a Bayes-optimal denoiser can be used to estimate the gradient of the log probability of data. To our knowledge this work is the first time that the output of a denoiser is explicitly back-propagated to train another network. Lastly, we note that denoising has been used to solve inverse problems in compressed sensing as in approximate message passing (Metzler et al., 2015).
3 THEORY
Consider a function fθ(x) parametrised by θ which maps a LR observation x to a HR estimate ŷ. Most current SR methods optimise model parameters via empirical risk minimization:
argmin θ
Ey,x[`(y, fθ(x))] (1)
Where y is the true target and ` is some loss function. The loss function is typically a simple convex function most often MSE `MSE(y, ŷ) = ‖y − ŷ‖22 as in (Dong et al., 2016; Shi et al., 2016). Here,
we seek to perform MAP inference instead. For a single LR observation the MAP estimate is
ŷ(x) = argmax y
log pY |X(y|x) (2)
Instead of calculating ŷ for each x separately we perform amortised inference, i. e. we would like to train the SR function fθ(x) to calculate the MAP estimate. A natural loss function for learning the parameters θ is the average log-posterior:
argmax θ
Ex log pY |X(fθ(x)|x), (3)
where the expectation is taken over the distribution of LR observations x. This loss depends on the unknown posterior distribution pY |X . We proceed by decomposing the log-posterior using Bayes’ rule as follows.
argmax θ Ex log pX|Y (x|fθ(x))︸ ︷︷ ︸ Likelihood +Ex log pY (fθ(x))︸ ︷︷ ︸ Prior − Ex log pX(x)︸ ︷︷ ︸ Marginal Likelihood  . (4)
3.1 HANDLING THE LIKELIHOOD TERM
Notice that the last term of Eqn. (4), the marginal likelihood, does not depend on θ, so we only have to deal with the likelihood and image prior. The observation model in SR can be described as follows. x = Aŷ, (5) where A is a linear transformation used for image downsampling. In general, A can be modelled as a strided two-dimensional convolution. Therefore, the likelihood term in Eqn. (4) is degenerate p(x|fθ(x)) = δ(x−Afθ(x)), and Eqn. (4) can be rewritten as constrained optimisation:
argmax θ
∀x:Afθ(x)=x
Ex[log pY (fθ(x))] (6)
To satisfy the constraints, we introduce a parametric function class that always guaranteesAfθ(x) = x. Specifically, we propose to use functions of the form
gθ(x) = Π A x fθ(x) = (I −A+A)fθ(x) +A+x (7)
where fθ is an arbitrary mapping from LR to HR space, ΠAx a projection to the affine subspace {y : yA = x}, and A+ is the Moore-Penrose pseudoinverse of A, which satisfies AA+A = A and A+AA+ = A+. Conveniently, if A is a strided two-dimensional convolution, then A+ becomes a deconvolution or up-convolution, which is a standard operation used in deep learning (e. g. Shi et al., 2016). It is important to stress that the optimal deconvolution A+ is not simply the transpose of A, Figure 2 illustrates the upsampling kernel (A+) that corresponds to a Gaussian downsampling kernel (A). For any A the deconvolution A+ can be easily found, here we used numerical methods as detailed in Appendix B. Intuitively, A+x can be thought of as a baseline SR solution, while (I − A+A)fθ is the residual. The operation (I − A+A) is a projection to the null-space of A, therefore when we downsample the residual (I − A+A)fθ we are guaranteed to get 0 no matter what fθ is. By using functions of this form we can turn Eqn. (6) into an unconstrained optimization problem.
argmax θ
Ex log pY (ΠAx fθ(x)) (8)
Interestingly, the objective above can be expressed in terms of the probability distribution of the model output qθ(y) := ∫ δ ( y −ΠAx fθ(x) ) pX(x)dx as follows.
argmax θ Ex log pY (ΠAx fθ(x)) = argmax θ Eŷ∼qθ log pY (ŷ) = argmin θ H[qθ, pY ], (9)
where H[q, p] denotes the cross-entropy between q and p and we used H[qθ, pY ] = Eŷ∼qθ [− log pY (ŷ)]. To minimise this objective, we do not need matched input-output pairs as in empirical risk minimisation. Instead we need to match the marginal distribution of reconstructed images qθ to that of the distribution of HR images. In this respect, the problem becomes more akin to unsupervised learning or generative modelling. In the following sections we present three approaches to finding the optimal θ utilising the properties of the affine projection.
3.2 AFFINE PROJECTED GENERATIVE ADVERSARIAL NETWORKS
Generative Adversarial Networks (Goodfellow et al., 2014) consist of a generator G that turns noise sampled from some distribution z ∼ pZ into images G(z) via a parametric mapping, and a discriminator D that learns to distinguish between real and synthetic images. The generator and discriminator are updated in tandem resulting in the generative distribution qG moving closer to the distribution of real data pY . The behaviour of GANs depends on the specifics of how the generator and the discriminator are trained. We use the following objective functions for D and G:
L(D;G) = −Ey∼pY logD(y)− Ez∼pZ log(1−D(G(z)), (10)
L(G;D) = −Ez∼pZ log D(G(z))
1−D(G(z)) .
The algorithm iterates two steps: first, it updates D by lowering L(D;G) keeping G fixed, then it updates G by lowering L(G;D) keeping D fixed. It can be shown that this amounts to minimising KL[qG‖pY ], where qG is the distribution of samples generated by G. See Appendix A for a proof1 In the context of SR, the affine projected SR function ΠAx fθ takes the role of the generator. Instead of noise, the generator is now fed low-resolution images x ∼ pX . Leaving everything else unchanged, we can deploy the GAN algorithm to minimise KL[qθ‖pY ]. We call this algorithm affine projected GAN or AffGAN for short. Similarly, we introduce notation SoftGAN to denote the GAN algorithm without the affine projection, which instead uses an additional soft-constraint `LR = MAE(x,Aŷ) as in (Garcia, 2016). Note that the difference between the cross-entropy and the KL divergence is the entropy of qθ: H[qθ, pY ] − KL[qθ‖pY ] = H[qθ]. Hence, we can expect AffGAN to favour approximate MAP solutions that lead to higher entropy and thus more diverse solutions overall.
3.2.1 INSTANCE NOISE
The theory suggests that GANs should be a convergent algorithm. If a unique optimal discriminator exists and it is reached by optimising D to perfection at each step, technically the whole algorithm corresponds to gradient descent on an estimate of KL[qθ‖pY ] with respect to θ. In practice, however, GANs tend to be highly unstable. So where does the theory go wrong? We think the main reason for the instability of GANs stems from qθ and pY being concentrated distributions whose support does not overlap. The distribution of natural images pY is often assumed to concentrate on or around a low-dimensional manifold. In most cases, qθ is degenerate and manifold-like by construction, such as in AffGAN. Therefore, odds are that especially before convergence is reached, qθ and pY can be perfectly separated by several Ds violating a condition for the convergence proof. We try to remedy this problem by adding instance noise to both SR and true image samples. This amounts to minimising the divergence dσ(qθ, pY ) = KL [pσ ∗ qθ‖pσ ∗ pY ], where pσ ∗ qθ denotes convolution of qθ with the noise distribution pσ . The noise level σ can be annealed during training, and the noise allows us to safely optimise D until convergence in each iteration. The trick is related to one-sided label noise introduced by Salimans et al. (2016), however without introducing a bias in the optimal discriminator, and we believe it is a promising technique for stabilising GAN training in general. For more details please see Appendix C
3.3 DENOISER GUIDED SUPER-RESOLUTION
To optimise the criterion Eqn. (6) via gradient descent we need its gradient with respect to θ:
∂
∂θ Ex[log p(ΠAx fθ(x))] = Ex
[ ∂
∂y log p(y) ∣∣∣∣ y=ΠAx fθ(x) ·ΠAx ∂ ∂θ fθ(x) ] (11)
Here ∂∂θfθ are the gradients of the SR function which can be calculated via back-propagation whereas ∂∂y log pY (y) requires estimation since pY is unknown. We use results from (Alain & Bengio, 2014; Särelä & Valpola, 2005) showing that in the limit of infinitesimal Gaussian noise, optimal denoising functions can be used to estimate this gradient:
f∗σ = argmin f
Ey∼pY `MSE(f(y + σ ), y) =⇒ f∗(y)− y σ2 ≈ ∂ ∂y log pY (y), (12)
1First shown in (Huszár, 2016).
where ∼ N (0, I) is Gaussian white noise, f∗σ is the Bayes-optimal denoising function for noise level σ. Using these results we can maximise Eqn. (9) by first training a neural network to denoise samples from pY and then backpropagate the gradient estimates from Eqn. (12) via the chain rule in Eqn. (11) to update θ. Well call this method AffDG, as it uses the affine subspace projection and is guided by the gradient from the DAE. Similar to above we’ll call the similar algorithm soft-enforcing Eqn. (5) SoftDG.
3.4 DENSITY GUIDED SUPER-RESOLUTION
As a more direct baseline model for amortised MAP inference we fit a tractable, yet powerful density model to pY using maximum likelihood, and then use cross entropy with respect to the generative model to approximate Eqn. (9). We use a deep generative model similar to the pixelCNN (Oord et al., 2016) but with a continuous (and differentiable) MCGSM (Theis et al., 2012) likelihood. These type of models are state-of-the-art in density estimation, are relatively fast to evaluate and produce visually interesting samples (Oord et al., 2016). We call this method AffLL, as it uses the affine projection and is guided by the log-likelihood of a density model.
4 EXPERIMENTS
We designed our experiments to address the following questions:
• Are the methods proposed in Section 3 successful at minimising cross-entropy?→Section 5.1 • Does the affine projection layer hurt the performance of CNNs for image SR?→Section 5.2 • Do the proposed methods produce perceptually superior SR results? →Sections 5.3-5.5 We initially illustrate the behaviour of the proposed algorithms on data where exact MAP inference is computationally tractable. Here the HR data y = [y1, y2] is drawn from a two-dimensional noisy Swiss-roll distribution and the one-dimensional LR data x is simply the average of the two HR pixels. Next we tested the proposed algorithm in a series of experiments on natural images using 4× downsampling.. For the first dataset, we took random crops from HR images containing grass texture. SR of random textures is known to be very hard using MSE or MAE loss functions. Finally, we tested the proposed models on real image data of faces (Celeb-A) and natural images (ImageNet). All models were convolution neural networks implemented using Theano (Team et al., 2016) and Lasagne (Dieleman et al., 2015). We refer to Appendix D for full experimental details.
5 RESULTS AND DISCUSSION
5.1 2D MAP INFERENCE: SWISS-ROLL
In this experiment we wanted to demonstrate that AffGAN and AffDG are indeed minimising the MAP objective in Eqn. (9). For this we used the two-dimensional toy problem where pY can be evaluated using brute-force Monte Carlo. Figure 1b) shows the outputs for x = [−8, 8] for models trained with different criterion. The AffGAN and AffDG solutions largely fit the dominant mode similar to MAP inference. For the MSE and MAE models the output generally falls in regions with low prior density. Table 1 shows the cross-entropy H[qθ, pY ] achieved by different methods, averaged over 10 independent trials with random initialisation. The cross-entropy values for the GAN and DAE based models are relatively close to the optimal MAP solution, which in this case we can find in a brute-force way. As expected the MSE and MAE models perform worse as these models do not minimize H[qθ, pY ]. We also calculated the average MSE between the network input and the downsampled network output. For the affine projected models, this error is exactly 0. The soft constrained models only approximately satisfy this constraint, even after extensive training (Table 1 second column). Further, we observe that the affine projected models generally found a lower cross-entropy H[qθ, pY ] when compared to soft-constrained versions.
5.2 AFFINE PROJECTED NETWORKS: PROOF OF CONCEPT USING MSE CRITERION
Adding the affine projection ΠAx restricts the class of functions that the SR network can model, so it is important to verify that the network is still capable of achieving the same performance in
SR as unconstrained CNN architectures. To test this, we trained CNNs with and without affine projections to perform SR on the CelebA dataset using MSE as the objective function. Results are shown in Figure 2. First note that when using affine projections, a randomly initialised network starts learning from a lower initial loss as the low-frequency components of the network output already match those of the target image. We observed that the affine projected networks generally train faster than unconstrained ones. Furthermore, the affine projected networks tend to find a better solution as measured by MSE and SSIM (Figure 2a-b). To investigate which aspects of the network architecture are responsible for the improved performance, we evaluated two further models: In one variant, we initialise the affine projected CNN to implement the correct projection, but then treat A+ as a trainable parameter. In the final variant, we keep the architecture the same, but initialise the final deconvolution layer A+ randomly and allow it to be trained. We found that initialising A+ to the correct Moore-Penrose inverse is important, and we get the similar results irrespective of whether or not it is fixed during training. Figure 2c shows the error between the network input and the downsampled network output. We can see that the exact affine projected network keeps this error at virtually 0.0 (up to numerical precision), whereas any other network will violate this consistency. In Figure 2d we show the downsampling kernel A and the corresponding optimal kernel for A+.
5.3 GRASS TEXTURES
Random textures are known to be hard model using MSE loss function. Figure 3 shows 4× SR of grass texture patches using identical affine projected CNNs trained with different loss functions. When randomly initialised, affine projected CNNs always produce an output with the correct lowfrequency components,as illustrated by the third panel labelled Affinit in Figure 3. The AffGAN model produces clearly the sharpest images, and we found the images to be plausible given the LR inputs. Notice that the reconstruction is not perfect pixel-by-pixel, but it has the correct statistical properties for the human visual system to recognise it as grass texture. The AffDG and AffLL models both produced blurry results which we where unable to improve upon using various optimization methods. Due to these findings we choose not to perform any further experiments with these models and concentrate on AffGAN instead. We refer to Appendix E for discussion of the results of these models.
5.4 CELEBA FACES
In Figure 4 the SR results are seen for several models trained using different loss functions. The MSE trained models outputs somewhat generic and over-smoothed images as expected. For the GAN models the global content is correct for both the affine projected and soft constrained models. Comparing the AffGAN and SoftGAN outputs the AffGAN model produces slightly sharper images
which however also seem to contain slightly more high frequency noise. We observed some colour drifting for the soft constrained models. Table 2 shows quantitative results for the same four models where, in terms of PSNR and SSIM, the MSE model achieves the best scores as expected. The consistency between input and output clearly shows that the models using the affine projections satisfy Eqn. (5) better than the soft constrained versions for both MSE and GAN losses.
5.5 NATURAL IMAGES
In Figure 5 we show the results for 4× SR from 32× 32 to 128× 128 pixels for AffGAN trained on natural images from ImageNET. For most of the images the results are sharp and corresponds well with the LR input. However we still see the high-frequency noise present in most GAN results in some of the images. Interestingly the snake depicted in the third column is super resolved into water which is obviously wrong but still a very plausible image considering the LR input image. Further, water will likely have a higher density under the image prior than snakes which suggests that the GAN model dreams up reasonable data.
5.6 CRITICISM AND FUTURE DIRECTIONS
One argument against MAP inference is that the mode of a distribution is dependent on the representation: transforming a variable through an invertible transformation and performing MAP inference in the transformed space may lead to different answers depending on the transformation. As an extreme example, consider transforming a continuous random scalar Y with its cumulative distribution function F = P(Y ≤ ·). The resulting variable F (Y ) is uniformly distributed, so any value in the interval (0, 1] can be the mode. Thus, the MAP estimate is not unique if one allows for alternative representations, and there is no guarantee that the MAP estimate in 24-bit RGB pixel representation which we seek in this paper is in any way special. One may arrive at a different solution when performing MAP estimation in the feature space of a convolutional neural network, or even if merely an alternative colour space is used. Interestingly, AffGAN is more resilient to coordinate transformations: Eqn. (10) includes the extra term H[qθ] which is effected by transformations the same way as H[qθ, pY ]. The second argument relates to the assumption that MAP estimates appear plausible. Although by definition the mode lies in a high-probability region, it does not guarantee that its appearance is anything like that of a random sample. Consider for example data drawn from a d-dimensional standard Normal distribution. Due to concentration of measure, as d increases the norm of a typical sample will be approximately √ d with very high probability. The mode, however, has a norm of 0. In this sense, the mode of the distribution is highly atypical. Indeed human observers can easily tell apart a typical sample from the noise distribution and the mode, but would have a hard time noticing the difference between two random samples. This argument suggests that sampling from the posterior pY |X may be a good or even preferable way to obtain plausible reconstructions. In Appendix F we establish a connection between variational inference, such as in varational autoencoders (Kingma & Welling, 2014), and a stochastic version of AffGAN, however leaving emperical studies as further.
6 CONCLUSION
In this work we developed methods for approximate MAP inference in SR. We first introduced an architectural restriction to neural networks projecting the model output to the affine subspace of valid solutions. We then proposed three methods, based on GANs, denoising or density models, for amortised MAP inference in SR using this affine projection. In high dimensions we empirically found that the GAN based approach, AffGAN produced the most visually appealing results. Our work follows successful demonstrations of GAN-based algorithms for image SR (Ledig et al., 2016), and we provide additional theoretical motivation for why this approach makes sense. In future work we plan to focus on a stochastic extension of AffGAN which can be seen as performing amortised variational inference.
A GENERATIVE ADVERSARIAL NETWORKS FOR MINIMISING KL-DIVERGENCE
First note that for a fixed generator G the discriminator D maximises:
Ey∼pY logDψ(y) + Ez∼N log(1−Dψ(Gθ(z)) = (13) Ey∼pY logDψ(y) + Ey∼qG [log(1−Dψ(y))] = (14)∫ y pY (y) logDψ(y) + qG(y) log(1−Dψ(y))dy (15)
where qG is the generative distribution. A function of the form a log(x) + b log(1− x) always has maximum at a
a+b and we find the Bayes-optimal discriminator to be (assuming equal prior class probabilities)
D∗(y) = pY (y)
pY (y) + qG(y) (16)
Let’s assume that this Bayes-optimal discriminator is unique and can be approximated closely by our neural network (see Appendix C for more discussion on this assumption).
Using the modified update rule proposed here the combined optimization problem for the discriminator and generator is
V (ψ, θ) = max ψ,θ Ey∼pY logDψ(y) + Ez∼N [log(Dψ(Gθ(z))− log(1−Dψ(Gθ(z))] (17)
Starting from the definition of KL[qG|pY ]
KL[qG|pY ] = Ey∼qG log qG(y)
pY (y) (18)
= Ey∼qG log 1−D∗(y) D∗(y) (insert Bayes-optimal classifier) (19)
≈ Ey∼qG log 1−Dψ(y) Dψ(y) = −Ey∼qG log Dψ(y) 1−Dψ(y) (20)
Which is equal to the terms affecting the generator in Eqn. (17).
B AFFINE PROJECTION
B.1 NUMERICAL ESTIMATION OF THE PSEUDO-INVERSE
In practice we implement the down-sampling projection A as a strided convolution with a fixed Gaussian smoothing kernel where the stride corresponds to the down-sampling factor. A+ is implemented as a transposed convolution operation with parameters optimised numerically via stochastic gradient descent on the following objective function:
`1(B) = Ey∼Nrd ‖Ay −ABAy‖ 2 2 (21) `2(B) = Ex∼Nr ‖Bx−BABx‖ 2 2 (22)
A+ = argmin B {`1(B) + `2(B)} (23)
Where Nd is the d-dimensional standard normal distribution, and d is the dimensionality of LR data x. `1 and `2 can be thought of as a Monte Carlo estimate of the spectral norm of the transformations A − ABA and B − BAB, respectively. The Monte Carlo formulation above has the advantage that it can be optimised via stochastic gradient descent. The operation ABA can be thought of as a three-layer fully linear convolutional neural network, where A corresponds to a strided convolution with fixed kernels, while B is a trainable deconvolution. We note that for certain downsampling kernels A the exact A+ would have an infinitely large kernel, although it can always be approximated with a local kernel. At convergence we found `1 + `2 to be between 10−12 and 10−8 depending on the down-sampling factor, width of the Gaussian kernel used forA and the filter sizes of A and B.
B.2 GRADIENTS
The gradients of the affine projected SR models is derived by applying the chain rule
fθ(x) = (I −A+A)gθ(x) +A+x (24)
∂fθ(x) ∂θ = ∂fθ(x) ∂gθ(x) ∂gθ(x) ∂θ = (I −A+A)∂gθ(x) ∂θ (25)
Which is essentially the high-pass filtered version of the gradient of gθ(x).
C INSTANCE NOISE
GANs are notoriously unstable to train, and several papers exist that try to improve their convergence properties (Salimans et al., 2016; Radford et al., 2015) via various tricks. Consider the following idealised GAN algorithm, each iteration consisting of the following steps:
1. we train the discriminator D via logistic regression between qθ vs pY , until convergence 2. we extract from D an estimate of the logarithmic likelihood ratio s(y) = log qθ(y)
p(y)
3. we update θ by taking a stochastic gradient step with objective function Ey∼qθs(y)
If qθ and pY are well-conditioned distributions in a low-dimensional space, this algorithm performs gradient descent on an approximation to the KL divergence, so it should converge. So why is it highly unstable in practical situations?
Crucially, the convergence of this algorithm relies on a few assumptions that don’t always hold: (1) that the log-likelihood-ratio log qθ(y)
p(y) is finite, (2) that the Jensen-Shannon divergence JS[qθ‖p] is a well-behaved func-
tion of θ and (3) that the Bayes-optimal solution to the logistic regression problem is unique. We stipulate that in real-world situations neither of these holds, mainly because qθ and pY are concentrated distributions whose support may not overlap. In image modelling, distribution of natural images pY is often assumed to be concentrated on or around a lower-dimensional manifold. Similarly, qθ is often degenerate by construction. The odds that the two distributions share support in high-dimensional space, especially early in training, are very small. If qθ and pY have non-overlapping support (1) the log-likelihood-ratio and therefore KL divergence is infinite (2) the Jensen-Shannon divergence is saturated so its maximum value and is locally constant in θ and (3) there may be a large set of near-optimal discriminators whose logistic regression loss is very close to the Bayes optimum, but each of these possibly provides very different gradients to the generator. Thus, training the discriminator D might find a different near-optimal solution each time depending on initialisation, even for a fixed qθ and pY .
The main ways to avoid these pathologies involve making the discriminator’s job harder. For example, in most GAN implementations the discriminator is only partially updated in each iteration, rather than trained until convergence. Another way to cripple the discriminator is adding label noise, or equivalently, one-sided label smoothing as introduced by Salimans et al. (2016). In this technique the labels in the discriminator’s training data are randomly flipped. However we do not believe these techniques adequately address all of the concerns described above.
In Figure 6a we illustrate two almost perfectly separable distributions. Notice how the large gap between the distributions means that there are large number of possible classifiers that tell the two distributions apart and achieve similar logistic loss. The Bayes-optimal classifier may not be unique, and the set of near-optimal classifiers is very large and diverse. In Figure 6b we show the effect of one sided label smoothing or equivalently, adding label noise. In this technique, the labels of some real data samples y ∼ pY are flipped so the discriminator is trained thinking they were samples from qθ . The discriminator indeed has a harder task now, but all classifiers are penalised almost equally. As a result, there is still a large set of discriminators which achieve near-optimal loss, it’s just that the near-optimal loss is now larger. Label smoothing does not help if the Bayes-optimal classifier is not unique.
Instead we propose to add noise to the samples, rather than labels, which we denote instance noise. Using instance noise the support of the two distributions is broadened and they are no longer perfectly separable as illustrated in Figure 6c. Adding noise, the Bayes-optimal discriminator becomes unique, the discriminator is less prone to overfitting because it has a wider training distribution, and the log-likelihood-ratio becomes better behaved. The Jensen-Shannon divergence between the noisy distributions is now a non-constant function of θ. Using instance noise, is easy to construct an algorithm that minimises the following divergence:
dσ(qθ, pY ) = KL [pσ ∗ qθ‖pσ ∗ pY ] , (26)
where σ is the parameter of the noise distribution. Logistic regression on the noisy samples provides an estimate of sσ(x) = log pσ∗qθpσ∗pY . When updating the generator we have to minimise the mean of s(σ) on noisy samples from qθ . We know that, if pσ is Gaussian, dσ is a Bregman-divergence, and that it is 0 if and only if the two distributions are equal. Because of the added noise, dσ is less sensitive to local features of the distribution. We found that in our experiments instance noise helped the convergence of AffGAN. We have not tested the instance noise in the generative modelling application. Because we don’t have to worry about over-training the discriminator, we can train it until convergence, or take more gradient steps between subsequent updates to the generator. One critical hyper-parameter of this method is the noise distribution. We used additive Gaussian noise, whose variance we annealed during training. We propose a heuristic annealing schedule where the noise is adapted so as to keep the optimal discriminator’s loss constant during training. It is possible that other noise distributions such as heavy-tailed or spike-and-slab would work better but we have not investigated these options.
D EXPERIMENTAL DETAILS
LOSS FUNCTIONS
For the GAN models the generative and discriminative parameters were updated using Eqn. (10). For the models enforcing Eqn. (5) using a soft-constraint we added an extra MAE loss term to the generative parameters `MAE =
1 N ∑ i ‖xi −Aŷi‖, where i i runs over the number of data samples N .
The denoiser guided models were trained in a two step procedure. Initially we pre-trained a DAE to denoise samples from the data distribution by minimising
`DAE = 1
N
∑ ‖y, fσDAE(ỹ)‖2 (27)
ỹ = y + , ∼ N (0, σI)
During training we anneal the noise level σ and continuously save the model parameters of the DAE fσ(DAE) trained at increasingly smaller noise levels. We then learn the parameters of the generator by following the gradient in Eqn. (11) using the DAE to estimate ∂
∂y log p(y)
∂
∂θ Ex[log p(ŷ)] = Ex
[ ∂
∂y log p(y) · ∂ ∂θ ŷ
] (28)
= Ex [ fσDAE(ỹ)− y
σ2 · ∂ ∂θ ŷ
] (29)
θi+1 ← θi + α ∂
∂θ Ex[log p(ŷ)] (30)
Where α is the learning rate. During training we continuously load parameters of the DAE trained at increasingly low noise levels to get gradients pointing in the approximate correct direction in beginning of training while covering a large data space and precise gradients close to the data manifold in the end of the training.
For the density guided models we first pre-train a density model by maximising the tractable log-likelihood
L(y) = ∑ j log p(yj |y<j) (31)
Where the joint density have been decomposed using the chain rule and j runs over the pixels. Similar to the DAE we continuously save the parameters of the density model during training. We then learn the parameters of the generator by directly minimising the negative log-likelihood of the generated samples under the learned density model.
` = −L(ŷ) = −L(fθ(x)) (32)
2D SWISS-ROLL
The 2D target data y = [y1, y2] was sampled from the 2D Swiss-Roll defined as:
ν1 ∼ N (µ1, σ1), ν2 ∼ N (µ2, σ2) (33) r = 0.4ν1 + ν2 (34) y = [cos(ν1) ∗ r, sin(ν1) ∗ r] (35)
Where µ1 = 10, σ1 = 3, µ2 = 0 and σ2 = 0.2. The LR input was defined as x = y1+y22 . The cross-entropy H[qθ, pY ] were calculated by estimating the probability density function using a Gaussian kernel density estimator fitted to 50.000 samples from a noiseless Swiss Roll density i.e. σ2 = 0, and setting the bandwidth of each kernel to σ2 = 0.2. All generator and discriminators were 2-layered fully connected NNs with 64 units in each layer. For the AffDG model the DAE was a two layered NN with 256 units in each layer trained while annealing the standard deviation of the Gaussian noise from 0.5 to 0.25.
IMAGE DATA
For all image experiments we set A to a convolution using a Gaussian smoothing kernel of size 9 × 9 using a stride of 4 corresponding to 4× down-sampling. A+ were set to a convolution operation with 42 kernels of size 5 × 5 followed by a reordering of the pixel with the output corresponding to 4× up-sampling convolution as described in (Shi et al., 2016). The parameters of the A+ was optimised numerically as described in Appendix B. All down-sampling were done using the A projection. For all image models we used convolutional models using ReLU nonlinearities and batch normalization in all layers except the output. All generators used skip connections similar to (Huang et al., 2016) and a final sigmoid non-linearity was applied to output of the model which were either used directly or feed through the affine transformation layers parameterised by A and A+. The discriminators were standard convolutional networks followed by a final sigmoid layer.
For the grass texture experiments we used randomly extracted patches of data from high resolution grass texture images. The generators used 6 layers of convolutions with 32, 32, 64, 64, 128 and filter maps and skip connections after every second layer. The discriminators had four layers of strided convolutions with 32, 64, 128 and 256 filter maps. For the AffDG model the DAE was a four layer convolutional network with 128 filter maps in each layer trained while annealing the standard deviation of the Gaussian noise from 0.5 to 0.01. The density model was implemented as a pixelCNN similar to Oord et al. (2016) with four layers of convolution
with 64 filter map with kernel sizes of 5, except for the first layers which used 7. The original PixelCNN uses a non-differentiable categorical distribution as the likelihood model why it can not be used for gradient based optimization. Instead we used a MCGSM as the likelihood model (Theis & Bethge, 2015), which have been shown to be a good density model for images (Theis et al., 2012), using 32 mixture components and 32 quadratic features to approximate the covariance matrices.
For the CelebA experiments the datasets were split into train, validation and test set using the standard splitting. All images were center cropped and resized to 64×64 before down-sampling to 16×16 usingA. All generators were 12 layer convolution networks with four layers of 128, 256 and 512 filter maps and skip connections between every fourth layer. The discriminators were 8 layer convolution nets with two layers of 128, 256, 512 and 1024 filter maps using a stride of 2 for every second layer.
For the ImageNET experiments the 2012 dataset were randomly split into train, validation and test set with 104 samples in the test and validation sets. All images below 20kB were then discarded to remove images with to low resolution. The images were center cropped and resized to 128 × 128 before down-sampling to 32 × 32 using A. The generator were a 8 layer convolutional network with 4 layers of 128 and 256 filter maps and skip connections between every second layer. The discriminators were 8 layer convolution nets with two layers of 128, 256, 512 and 1024 filter maps using a stride of 2 for every second layer. To stabilise training we used Gaussian instance noise linearly annealed from an initial standard deviation of 0.1 to 0. We were unable to stable train models without this extra regularization.
E ADDITIONAL RESULTS FOR DENOISER AND DENSITY GUIDED SUPER-RESOLUTION
Figure 7 show the PSNR and SSIM scores during training for the AffDG and AffLL models trained on the grass textures. Note that the models are converging, but as seen in Figure 3 the images are very blurry. For both models we had problems with diverging training. For the DAE models with high noise levels the gradients are only approximately correct but covers a large space around the data manifold whereas for small noise levels the gradients are more accurate in a small space around the data manifold. For the density model we believe a similar phenomenon is making the training diverge since for accurate density models the estimated density is likely very peaked around the data manifold making learning in the beginning of training difficult. To resolve these issue we started training using models with high noise levels or low log-likelihood values and then loaded model parameters during training with continuously smaller noise levels or better log-likelihood values. The effect of this can be clearly seen during training as the step like behavior of the AffDG in Figure 7. We note that the density model used for training the AffLL achieved a log-likelihood of −4.10 bits per dimension which is comparable to values obtained in Theis & Bethge (2015) on a texture dataset. Further the AffLL model achieved high log-likelihood values > −3.5 under this model suggesting that the density model is simply not providing an accurate enough representation of pY to provide precises scores for training the AffLL model.
F AMORTISED VARIATIONAL INFERENCE USING AFFGAN
Here we’ll show that a stochastic extension of the AffGAN model approximately minimises an amortised variational inference criterion as in e. g. variational autoencoders, which for the first time establishes a connection between adversarial methods of inferences and and variational inference. We introduce a variant of AffGAN where, in addition to the LR data x, the generator function also takes as input some independent noise variables z: we establish a connection between GANs and amortised variational
z ∼ pZ (36)
ŷ = ΠAx fθ(x, z) (37)
Similarly to how we defined qθ in Section 3.1 we introduce the following notation: qY ;θ := Ex∼pXEz∼pZ δ ( y −ΠAx fθ(x, z) ) (38)
qY |X;θ := Ez∼pZ δ ( y −ΠAx fθ(x, z) ) (39)
qX,Y ;θ := pX · qY |X;θ (40)
Here the affine projection ensures that under qX,Y ;θ , x and y are always consistent. Therefore, under qX,Y ;θ , the conditional of x given y is the same as the likelihood pX|Y = δ(x−Ay) by construction and the following equality holds:
qX,Y ;θ = qY ;θ · pX|Y = pX · qY |X;θ (41)
Applying Bayes’ rule to pX|Y = pY |XpX pY and substituting into the above equality we get:
qY ;θ · pY |X=Ay = pY ;θ · qY |X=Ay;θ (42)
The KL divergence that the AffGAN objective minimises can now be rewritten as.
KL [qY ;θ‖pY ] = EqY ;θ log qY ;θ(y)
pY (y) (43)
= EqX,Y ;θ log qY ;θ(y)
pY (y) (44)
= EqX,Y ;θ log qY |X;θ(y|x) pY |X(y|x)
(45)
= EpX KL [ qY |X;θ ∥∥pY |X] (46) Therefore we can conclude that the AffGAN algorithm described in Section 3.2 approximately minimizes the following amortised variational inference criterion:
argmin θ KL [qY ;θ‖pY ] = argmin θ
Ex∼pX KL [ qY |X;θ ∥∥pY |X] , (47) and in doing so it only requires samples from PY and PX .

The paper presents a new framework to solve the SR problem - amortized MAP inference and adopts a pre-learned affine projection layer to ensure the output is consistent with LR. Also, it proposes three different methods to solve the problem of minimizing cross-entropy. Generally, it is a great paper. However, I still have several comments:
1) The proposed amortized MAP inference is novel and different from the previous SR methods. Combined with GAN, this framework can obtain plausible and good results. Compared with another GAN-based SR methods - Photo-Realistic Single Image Super-Resolution Using a Generative Adversarial Network, question may arise as to what this new formulation adds to the latest state-of-the-art.
2) Using an affine projection architecture as a constraint, the model do not need any corresponding {HR, LR} image pairs for training. However, when training the affine projection layer, we still need the {HR, LR} image pairs. Does it mean that we merely transfer this training procedure to the training of affine projection?
3) The paper presents many results of the framework, including the results of natural images from ImageNet. Can the author also provide the results of Set5, Set14 or BSD100, which are conventional test dataset for SR, so that we can perform a fair comparison with previous work.
4) I see that the size of the results of nature images presented in this paper are limited to 128*128. Can this framework perform well on images with larger size? Because SR will encounter input with arbitrary size.
5) A normal GAN will have a noise term as a latent space, so that it can be better illustrated as learning a distribution. Do the author try the noise vector?
Overall, this paper provides a new framework for SR with solid theoretical analysis. The idea is novel and the author explore many methods. Though there still exist questions like the necessity and more experiments are needed. I think this work will will provide good inspiration to the community.


Deep neural network models, though very powerful and highly successful, are computationally expensive in terms of space and time. Recently, there have been a number of attempts on binarizing the network weights and activations. This greatly reduces the network size, and replaces the underlying multiplications to additions or even XNOR bit operations. However, existing binarization schemes are based on simple matrix approximations and ignore the effect of binarization on the loss. In this paper, we propose a proximal Newton algorithm with diagonal Hessian approximation that directly minimizes the loss w.r.t. the binarized weights. The underlying proximal step has an efficient closed-form solution, and the second-order information can be efficiently obtained from the second moments already computed by the Adam optimizer. Experiments on both feedforward and recurrent networks show that the proposed loss-aware binarization algorithm outperforms existing binarization schemes, and is also more robust for wide and deep networks.
1 INTRODUCTION
Recently, deep neural networks have achieved state-of-the-art performance in various tasks such as speech recognition, visual object recognition, and image classification (LeCun et al., 2015). Though powerful, the large number of network weights leads to space and time inefficiencies in both training and storage. For instance, the popular AlexNet, VGG-16 and Resnet-18 all require hundred of megabytes to store, and billions of high-precision operations on classification. This limits its use in embedded systems, smart phones and other portable devices that are now everywhere.
To alleviate this problem, a number of approaches have been recently proposed. One attempt first trains a neural network and then compresses it (Han et al., 2016; Kim et al., 2016). Instead of this two-step approach, it is more desirable to train and compress the network simultaneously. Example approaches include tensorizing (Novikov et al., 2015), parameter quantization (Gong et al., 2014), and binarization (Courbariaux et al., 2015; Hubara et al., 2016; Rastegari et al., 2016). In particular, binarization only requires one bit for each weight value. This can significantly reduce storage, and also eliminates most multiplications during the forward pass.
Courbariaux et al. (2015) pioneered neural network binarization with the BinaryConnect algorithm, which achieves state-of-the-art results on many classification tasks. Besides binarizing the weights, Hubara et al. (2016) further binarized the activations. Rastegari et al. (2016) also learned to scale the binarized weights, and obtained better results. Besides, they proposed the XNOR-network with both weights and activations binarized as in (Hubara et al., 2016). Instead of binarization, ternary-connect quantizes each weight to {−1, 0, 1} (Lin et al., 2016). Similarly, the ternary weight network (Li & Liu, 2016) and DoReFa-net (Zhou et al., 2016) quantize weights to three levels or more. However, though using more bits allows more accurate weight approximations, specialized hardwares are needed for the underlying non-binary operations.
Besides the huge amount of computation and storage involved, deep networks are difficult to train because of the highly nonconvex objective and inhomogeneous curvature. To alleviate this problem, Hessian-free methods (Martens & Sutskever, 2012) use the second-order information by conjugate gradient. A related method is natural gradient descent (Pascanu & Bengio, 2014), which utilizes ge-
ometry of the underlying parameter manifold. Another approach uses element-wise adaptive learning rate, as in Adagrad (Duchi et al., 2011), Adadelta (Zeiler, 2012), RMSprop (Tieleman & Hinton, 2012), and Adam Kingma & Ba (2015). This can also be considered as preconditioning that rescales the gradient so that all dimensions have similar curvatures.
In this paper, instead of directly approximating the weights, we propose to consider the effect of binarization on the loss during binarization. We formulate this as an optimization problem using the proximal Newton algorithm (Lee et al., 2014) with a diagonal Hessian. The crux of proximal algorithms is the proximal step. We show that this step has a closed-form solution, whose form is similar to the use of element-wise adaptive learning rate. The proposed method also reduces to BinaryConnect (Courbariaux et al., 2015) and the Binary-Weight-Network (Hubara et al., 2016) when curvature information is dropped. Experiments on both feedforward and recurrent neural network models show that it outperforms existing binarization algorithms. In particular, BinaryConnect fails on deep recurrent networks because of the exploding gradient problem, while the proposed method still demonstrates robust performance. Notations: For a vector x, √ x denotes the element-wise square root, |x| denotes the element-wise
absolute value, ‖x‖p = ( ∑ i |xi|p) 1 p is the p-norm of x, x 0 denotes that all entries of x are positive, sign(x) is the vector with [sign(x)]i = 1 if xi ≥ 0 and −1 otherwise, and Diag(x) returns a diagonal matrix with x on the diagonal. For two vectors x and y, x y denotes the elementwise multiplication and x y denotes the element-wise division. For a matrix X, vec(X) returns the vector obtained by stacking the columns of X, and diag(X) returns a diagonal matrix whose diagonal elements are extracted from diagonal of X.
2 RELATED WORK
2.1 WEIGHT BINARIZATION IN DEEP NETWORKS
In a feedforward neural network with L layers, let the weight matrix (or tensor in the case of a convolutional layer) at layer l be Wl. We combine the (full-precision) weights from all layers as w = [w>1 ,w > 2 , . . . ,w > L ] >, where wl = vec(Wl). Analogously, the binarized weights are denoted as ŵ = [ŵ>1 , ŵ > 2 , . . . , ŵ > L ] >. As it is essential to use full-precision weights during updates (Courbariaux et al., 2015), typically binarized weights are only used during the forward and backward propagations, but not on parameter update. At the tth iteration, the (full-precision) weight wtl is updated by using the backpropagated gradient∇l`(ŵt−1) (where ` is the loss and∇l`(ŵt−1) is the partial derivative of ` w.r.t. the weights of the lth layer). In the next forward propagation, it is then binarized as ŵtl = Binarize(w t l), where Binarize(·) is some binarization scheme.
The two most popular binarization schemes are BinaryConnect (Courbariaux et al., 2015) and Binary-Weight-Network (BWN) (Rastegari et al., 2016). In BinaryConnect, binarization is performed by transforming each element of wtl to −1 or +1 using the sign function:1
Binarize(wtl) = sign(w t l). (1)
Besides the binarized weight matrix, a scaling parameter is also learned in BWN. In other words, Binarize(wtl) = α t lb t l , where α t l > 0 and b t l is binary. They are obtained by minimizing the difference between wtl and α t lb t l , and have a simple closed-form solution:
αtl = ‖wtl‖1 nl , btl = sign(w t l), (2)
where nl is the number of weights in layer l. Hubara et al. (2016) further binarized the activations as x̂tl = sign(x t l), where x t l is the activation of the lth layer at iteration t.
2.2 PROXIMAL NEWTON ALGORITHM
The proximal Newton algorithm (Lee et al., 2014) has been popularly used for solving composite optimization problems of the form
min x f(x) + g(x),
1A stochastic binarization scheme is also proposed in (Courbariaux et al., 2015). However, it is much more computational expensive than (1) and so will not be considered here.
where f is convex and smooth, and g is convex but possibly nonsmooth. At iteration t, it generates the next iterate as
xt+1 = argmin x ∇f(xt)>(x− xt) + (x− xt)>H(x− xt) + g(x),
where H is an approximate Hessian matrix of f at xt. With the use of second-order information, the proximal Newton algorithm converges faster than the proximal gradient algorithm (Lee et al., 2014). Recently, by assuming that f and g have difference-of-convex decompositions (Yuille & Rangarajan, 2002), the proximal Newton algorithm is also extended to the case where g is nonconvex (Rakotomamonjy et al., 2016).
3 LOSS-AWARE BINARIZATION
As can be seen, existing weight binarization methods (Courbariaux et al., 2015; Rastegari et al., 2016) simply find the closest binary approximation of w, and ignore its effects to the loss. In this paper, we consider the loss directly during binarization. As in (Rastegari et al., 2016), we also binarize the weight wl in each layer as ŵl = αlbl, where αl > 0 and bl is binary.
In the following, we make the following assumptions on `. (A1) ` is continuously differentiable with Lipschitz-continuous gradient, i.e., there exists β > 0 such that ‖∇`(u)−∇`(v)‖2 ≤ β ‖u− v‖2 for any u,v; (A2) ` is bounded from below.
3.1 BINARIZATION USING PROXIMAL NEWTON ALGORITHM
We formulate weight binarization as the following optimization problem: minŵ `(ŵ) (3) s.t. ŵl = αlbl, αl > 0, bl ∈ {±1}nl , l = 1, . . . , L, (4) where ` is the loss. Let C be the feasible region in (4), and define its indicator function: IC(ŵ) = 0 if ŵ ∈ C, and∞ otherwise. Problem (3) can then be rewritten as
min ŵ `(ŵ) + IC(ŵ). (5)
We solve (5) using the proximal Newton method (Section 2.2). At iteration t, the smooth term `(ŵt) is replaced by the second-order expansion
`(ŵt−1) +∇`(ŵt−1)>(ŵt − ŵt−1) + 1 2 (ŵt − ŵt−1)>Ht−1(ŵt − ŵt−1),
where Ht−1 is an estimate of the Hessian of ` at ŵt−1. Note that using the Hessian to capture second-order information is essential for efficient neural network training, as ` is often flat in some directions but highly curved in others. By rescaling the gradient, the loss has similar curvatures along all directions. This is also called preconditioning in the literature (Dauphin et al., 2015a).
For neural networks, the exact Hessian is rarely positive semi-definite. This can be problematic as the nonconvex objective leads to indefinite quadratic optimization. Moreover, computing the exact Hessian is both time- and space-inefficient on large networks. To alleviate these problems, a popular approach is to approximate the Hessian by a diagonal positive definite matrix D. One popular choice is the efficient Jacobi preconditioner. Though an efficient approximation of the Hessian under certain conditions, it is not competitive for indefinite matrices (Dauphin et al., 2015a). More recently, it is shown that equilibration provides a more robust preconditioner in the presence of saddle points (Dauphin et al., 2015a). This is also adopted by popular stochastic optimization algorithms such as RMSprop (Tieleman & Hinton, 2012) and Adam (Kingma & Ba, 2015). Specifically, the second moment v in these algorithms is an estimator of diag(H2) (Dauphin et al., 2015b). Here, we use the square root of this v, which is readily available in Adam, to construct D = Diag([diag(D1)>, . . . , diag(DL)>]>), where Dl is the approximate diagonal Hessian at layer l. In general, other estimators of diag(H) can also be used.
At the tth iteration of the proximal Newton algorithm, the following subproblem is solved:
minŵt ∇`(ŵt−1)>(ŵt − ŵt−1) + 1
2 (ŵt − ŵt−1)>Dt−1(ŵt − ŵt−1) (6)
s.t. ŵtl = α t lb t l , α t l > 0, b t l ∈ {±1}nl , l = 1, . . . , L.
Proposition 3.1 Let dt−1l ≡ diag(D t−1 l ), and
wtl ≡ ŵt−1l −∇l`(ŵ t−1) dt−1l . (7)
The optimal solution of (6) can be obtained in closed-form as
αtl = ‖dt−1l wtl‖1 ‖dt−1l ‖1 , btl = sign(w t l). (8)
Theorem 3.1 Assume that [dtl ]k > β ∀l, k, t, the objective of (5) produced by the proximal Newton algorithm (with closed-form update of ŵt in Proposition 3.1) converges.
Note that both the loss ` and indicator function IC(·) in (5) are not convex. Hence, convergence analysis of the proximal Newton algorithm in (Lee et al., 2014), which is only for convex problems, cannot be applied. Recently, Rakotomamonjy et al. (2016) proposed a nonconvex proximal Newton extension. However, it assumes a difference-of-convex decomposition which does not hold here.
Remark 3.1 When Dt−1l = λI, i.e., the curvature is the same for all dimensions in the lth layer, (8) then reduces to the BWN solution in (2) In other words, BWN corresponds to using the proximal gradient algorithm, while the proposed method corresponds to the proximal Newton algorithm with diagonal Hessian. In composite optimization, it is known that the proximal Newton method is more efficient than the proximal gradient algorithm (Lee et al., 2014; Rakotomamonjy et al., 2016).
Remark 3.2 When αtl = 1, (8) reduces to sign(wtl), which is the BinaryConnect solution in (1).
From (7) and (8), each iteration first performs gradient descent along ∇l`(ŵt−1) with an adaptive learning rate 1 dt−1l , and then projects it to a binary solution. As discussed in (Courbariaux et al., 2015), it is important to keep a full-precision weight during training. Hence, we replace (7) by wtl ← w t−1 l − ∇l`(ŵt−1) d t−1 l . The whole procedure, which will be called Loss-Aware Binarization (LAB), is shown in Algorithm 1. In steps 5 and 6, following (Li & Liu, 2016), we first rescale input xt−1l to the lth layer with αl, so that multiplications in dot products and convolutions become additions.
While binarizing weights changes most multiplications to additions, binarizing both weights and activations saves even more computations as additions are further changed to XNOR bit operations (Hubara et al., 2016). Our Algorithm 1 can also be easily extended by binarizing the activations with the simple sign function.
3.2 EXTENSION TO RECURRENT NEURAL NETWORKS
The proposed method can be easily extended to recurrent neural networks. Let xl and hl be the input and hidden states, respectively, at time step (or depth) l. A typical recurrent neural network has a recurrence of the form hl = Wxxl +Whσ(hl−1) + b (equivalent to the more widely known hl = σ(Wxxl+Whhl−1+b) (Pascanu et al., 2013) ). We binarize both the input-to-hidden weight Wx and hidden-to-hidden weight Wh. Since weights are shared across time in a recurrent network, we only need to binarize Wx and Wh once in each forward propagation. Besides weights, one can also binarize the activations (of the inputs and hidden states) as in the previous section.
In deep networks, the backpropagated gradient takes the form of a product of Jacobian matrices (Pascanu et al., 2013). In a vanilla recurrent neural network,2 for activations hp and hq at depths p and q, respectively (where p > q), ∂hp∂hq = ∏ q<l≤p ∂hl ∂hl−1 = ∏ q<l≤pW > h diag(σ
′(hl−1)). The necessary condition for exploding gradients is that the largest singular value λ1(Wh) of Wh is larger than some given constant (Pascanu et al., 2013). The following Proposition shows that for any binary Wh, its largest singular value is lower-bounded by the square root of its dimension.
Proposition 3.2 For any W ∈ {−1,+1}m×n (m ≤ n), λ1(W) ≥ √ n.
2Here, we consider the vanilla recurrent neural network for simplicity. It can be shown that a similar behavior holds for the more commonly used LSTM.
Algorithm 1 Loss-Aware Binarization (LAB) for training a feedforward neural network. Input: Minibatch {(xt0,yt)}, current full-precision weights {wtl}, first moment {m t−1 l }, second moment {vt−1l }, and learning rate ηt. 1: Forward Propagation 2: for l = 1 to L do 3: αtl = ‖dt−1l w t l‖1
‖dt−1l ‖1 ;
4: btl = sign(w t l); 5: rescale the layer-l input: x̃tl−1 = α t lx t l−1; 6: compute ztl with input x̃ t l−1 and binary weight b t l ; 7: apply batch-normalization and nonlinear activation to ztl to obtain x t l ; 8: end for 9: compute the loss ` using xtL and y
t; 10: Backward Propagation 11: initialize output layer’s activation’s gradient ∂`
∂xtL ;
12: for l = L to 2 do 13: compute ∂`
∂xtl−1 using ∂` ∂xtl , αtl and b t l ;
14: end for 15: Update parameters using Adam 16: for l = 1 to L do 17: compute gradients∇l`(ŵt) using ∂`∂xtl and x t l−1; 18: update first moment mtl = β1m t−1 l + (1− β1)∇l`(ŵt); 19: update second moment vtl = β2v t−1 l + (1− β2)(∇l`(ŵt) ∇l`(ŵt)); 20: compute unbiased first moment m̂tl = m t l/(1− βt1); 21: compute unbiased second moment v̂tl = v t l/(1− βt2); 22: compute current curvature matrix dtl = 1 ηt ( 1+ √ v̂tl ) ; 23: update full-precision weights wt+1l = w t l − m̂tl dtl ; 24: update learning rate ηt+1 = UpdateRule(ηt, t+ 1); 25: end for
Thus, with weight binarization as in BinaryConnect, the exploding gradient problem becomes more severe as the weight matrices are often large. On the other hand, recall that λ1(cŴh) = cλ1(Ŵh) for any non-negative c. The proposed method alleviates this exploding gradient problem by adaptively learning the scaling parameter αh.
4 EXPERIMENTS
In this section, we perform experiments on the proposed binarization scheme with both feedforward networks (Sections 4.1 and 4.2) and recurrent neural networks (Sections 4.3 and 4.4).
4.1 FEEDFORWARD NEURAL NETWORKS
We compare the original full-precision network (without binarization) with the following weightbinarized networks: (i) BinaryConnect; (ii) Binary-Weight-Network (BWN); and (iii) the proposed Loss-Aware Binarized network (LAB). We also compare with networks having both weights and activations binarized:3 (i) BinaryNeuralNetwork (BNN) (Hubara et al., 2016), the weight-andactivation binarized counterpart of BinaryConnect; (ii) XNOR-Network (XNOR) (Rastegari et al., 2016), the counterpart of BWN; (iii) LAB2, the counterpart of the proposed method, which binarizes weights using proximal Newton method and binarizes activations using a simple sign function.
The setup is similar to that in Courbariaux et al. (2015). We do not perform data augmentation or unsupervised pretraining. Experiments are performed on three commonly used data sets:
3We use the straight-through-estimator (Hubara et al., 2016) to compute the gradient involving the sign function.
1. MNIST: This contains 28 × 28 gray images from ten digit classes. We use 50000 images for training, another 10000 for validation, and the remaining 10000 for testing. We use the 4-layer model:
784FC − 2048FC − 2048FC − 2048FC − 10SVM, where FC is a fully-connected layer, and SVM is a L2-SVM output layer using the square hinge loss. Batch normalization, with a minibatch size 100, is used to accelerate learning. The maximum number of epochs is 50. The learning rate for the weight-binarized (resp. weight-and-activation-binarized) network starts at 0.01 (resp. 0.005), and decays by a factor of 0.1 at epochs 15 and 25.
2. CIFAR-10: This contains 32 × 32 color images from ten object classes. We use 45000 images for training, another 5000 for validation, and the remaining 10000 for testing. The images are preprocessed with global contrast normalization and ZCA whitening. We use the VGG-like architecture:
(2×128C3)−MP2−(2×256C3)−MP2−(2×512C3)−MP2−(2×1024FC)−10SVM, where C3 is a 3×3 ReLU convolution layer, andMP2 is a 2×2 max-pooling layer. Batch normalization, with a minibatch size of 50, is used. The maximum number of epochs is 200. The learning rate for the weight-binarized (resp. weight-and-activation-binarized) network starts at 0.03 (resp. 0.02), and decays by a factor of 0.5 after every 15 epochs.
3. SVHN: This contains 32 × 32 color images from ten digit classes. We use 598388 images for training, another 6000 for validation, and the remaining 26032 for testing. The images are preprocessed with global and local contrast normalization. The model used is:
(2×64C3)−MP2−(2×128C3)−MP2−(2×256C3)−MP2−(2×1024FC)−10SVM. Batch normalization, with a minibatch size of 50, is used. The maximum number of epochs is 50. The learning rate for the weight-binarized (resp. weight-and-activation-binarized) network starts at 0.001 (resp. 0.0005), and decays by a factor of 0.1 at epochs 15 and 25.
Since binarization is a form of regularization (Courbariaux et al., 2015), we do not use other regularization methods (like Dropout). All the weights are initialized as in (Glorot & Bengio, 2010). Adam (Kingma & Ba, 2015) is used as the optimization solver.
Table 1 shows the test classification error rates, and Figure 1 shows the convergence of LAB. As can be seen, the proposed LAB achieves the lowest error on MNIST and SVHN. It even outperforms the full-precision network on MNIST, as weight binarization serves as a regularizer. With the use of curvature information, LAB outperforms BinaryConnect and BWN. On CIFAR-10, LAB is slightly outperformed by BinaryConnect, but is still better than the full-precision network. Among the schemes that binarize both weights and activations, LAB2 also outperforms BNN and the XNOR-Network.
4.2 VARYING THE NUMBER OF FILTERS IN CNN
As in Zhou et al. (2016), we study sensitivity to network width by varying the number of filters K on the SVHN data set. As in Section 4.1, we use the model
(2× KC3)−MP2− (2× 2KC3)−MP2− (2× 4KC3)−MP2− (2× 1024FC)− 10SVM. Results are shown in Table 2. Again, the proposed LAB has the best performance. Moreover, as the number of filters increases, degradation due to binarization becomes less severe. This suggests
that more powerful models (e.g., CNN with more filters, standard feedforward networks with more hidden units) are less susceptible to performance degradation due to binarization. We speculate that this is because large networks often have larger-than-needed capacities, and so are less affected by the limited expressiveness of binary weights. Another related reason is that binarization acts as regularization, and so contributes positively to the performance.
4.3 RECURRENT NEURAL NETWORKS
In this section, we perform experiments on the popular long short-term memory (LSTM) (Hochreiter & Schmidhuber, 1997). Performance is evaluated in the context of character-level language modeling. The LSTM takes as input a sequence of characters, and predicts the next character at each time step. The training objective is the cross-entropy loss over all target sequences. Following Karpathy et al. (2016), we use two data sets (with the same training/validation/test set splitting): (i) Leo Tolstoy’s War and Peace, which consists of 3258246 characters of almost entirely English text with minimal markup and has a vocabulary size of 87; and (ii) the source code of the Linux Kernel, which consists of 6206996 characters and has a vocabulary size of 101.
We use a one-layer LSTM with 512 cells. The maximum number of epochs is 200, and the number of time steps is 100. The initial learning rate is 0.002. After 10 epochs, it is decayed by a factor of 0.98 after each epoch. The weights are initialized uniformly in [0.08, 0.08]. After each iteration, the gradients are clipped to the range [−5, 5], and all the updated weights are clipped to [−1, 1]. For the weight-and-activation-binarized networks, we do not binarize the inputs, as they are one-hot vectors in this language modeling task.
Table 3 shows the testing cross-entropy values. As in Section 4.1, the proposed LAB outperforms other weight binarization schemes, and is even better than the full-precision network on the Linux Kernel data set. BinaryConnect does not work well here because of the problem of exploding gradients (see Section 3.2 and more results in Section 4.4). On the other hand, BWN and the proposed LAB scale the binary weight matrix and perform better. LAB also performs better than BWN as curvature information is considered. Similarly, among schemes that binarize both weights and activations, the proposed LAB2 also outperforms BNN and XNOR-Network.
4.4 VARYING THE NUMBER OF TIME STEPS IN LSTM
In this experiment, we study the sensitivity of the binarization schemes with varying numbers of unrolled time steps (TS) in LSTM. Results are shown in Table 4. Again, the proposed LAB has the best performance. When TS = 10, the LSTM is relatively shallow, and all binarization schemes have similar performance as the full-precision network. When TS ≥ 50, BinaryConnect fails,
while BWN and the proposed LAB perform better (as discussed in Section 3.2). Figure 2 shows the distributions of the hidden-to-hidden weight gradients for TS = 10 and 100. As can be seen, while all models have similar gradient distributions at TS = 10, the gradient values in BinaryConnect are much higher than those of the other algorithms for the deeper network (TS = 100).
Note from Table 4 that as the time step increases, all except BinaryConnect show better performance. However, degradation due to binarization also becomes more severe. This is because the weights are shared across time steps. Hence, error due to binarization also propagates across time.
5 CONCLUSION
In this paper, we propose a binarization algorithm that directly considers its effect on the loss during binarization. The binarized weights are obtained using proximal Newton algorithm with diagonal Hessian approximation. The proximal step has an efficient closed-form solution, and the secondorder information in the Hessian can be readily obtained from the Adam optimizer. Experiments show that the proposed algorithm outperforms existing binarization schemes, has comparable performance as the original full-precision network, and is also robust for wide and deep networks.
ACKNOWLEDGMENTS
This research was supported in part by the Research Grants Council of the Hong Kong Special Administrative Region (Grant 614513). We thank Yongqi Zhang for helping with the experiments, and developers of Theano (Theano Development Team, 2016), Pylearn2 (Goodfellow et al., 2013) and Lasagne. We also thank NVIDIA for the support of Titan X GPU.
A PROOF OF PROPOSITION 3.1
∇`(ŵt−1)>(ŵt − ŵt−1) + 1 2 (ŵt − ŵt−1)>Dt−1(ŵt − ŵt−1)
= 1
2 L∑ l=1 (√ (dt−1l ) > ( ŵtl − (ŵt−1l −∇l`(ŵ t−1) dt−1l ) ))2 + c1
= 1
2 L∑ l=1 (√ (dt−1l ) >(ŵtl −wtl) )2 + c1
= 1
2 L∑ l=1 (√ (dt−1l ) >(αtlb t l −wtl) )2 + c1,
where c1 = − 12 (√ (dt−1l ) >(∇l`(ŵt−1) dt−1l ) )2 . Since αtl > 0,d t l 0,∀l = 1, 2, . . . , L, we
have btl = sign(w t l). Moreover,
1
2 L∑ l=1 ( √ (dt−1l ) >(αtlb t l −wtl))2 + c1 = 1 2 L∑ l=1 (√ (dt−1l ) >( ∣∣αtl1− |wtl |∣∣))2 + c1
= L∑ l=1 1 2 ‖dt−1l ‖1(α t l) 2 − ‖dt−1l w t l‖1αtl + c2,
where c2 = c1 − 12 ‖dt−1l w t l‖ 2 1
‖dt−1l ‖1 . Thus, the optimal αtl is
‖dt−1l w t l‖1
‖dt−1l ‖1 .
B PROOF OF THEOREM 3.1
Let α = [αt1 . . . , α t L] >, and denote the objective in (3) by F (ŵ,α). As ŵt is the minimizer in (6), we have
`(ŵt−1) +∇`(ŵt−1)>(ŵt − ŵt−1) + 1 2 (ŵt − ŵt−1)>Dt−1(ŵt − ŵt−1) ≤ `(ŵt−1). (9)
From Assumption A1, we have
`(ŵt) ≤ `(ŵt−1) +∇`(ŵt−1)>(ŵt − ŵt−1) + β 2 ∥∥ŵt − ŵt−1∥∥2 2 . (10)
Using (9) and (10), we obtain
`(ŵt) ≤ `(ŵt−1)− 1 2 (ŵt − ŵt−1)>(Dt−1 − βI)(ŵt − ŵt−1)
≤ `(ŵt−1)− mink,l([d t−1 l ]k − β) 2 ∥∥ŵt − ŵt−1∥∥2 2 .
Let c3 = mink,l,t([dt−1l ]k − β) > 0. Then,
`(ŵt) ≤ `(ŵt−1)− c3 2 ∥∥ŵt − ŵt−1∥∥2 2 . (11)
From Assumption A2, ` is bounded from below. Together with the fact that {`(ŵt)} is monotonically decreasing from (11), the sequence {`(ŵt)} converges, thus the sequence {F (ŵt,αt)} also converges.
C PROOF OF PROPOSITION 3.2
Let the singulars values of W be λ1(W) ≥ λ2(W) ≥ · · · ≥ λm(W).
λ21(W) ≥ 1
m m∑ i=1 λ2i (W) = 1 m ‖W‖2F = 1 m mn = n.
Thus, λ1(W) ≥ √ n.

The paper presents a second-order method for training a neural networks while ensuring at the same time that weights (and activations) are binary. Through binarization, the method aims to achieve model compression for subsequent deployment on low-memory systems. The method is abbreviated BPN for "binarization using proximal Newton algorithm".
The method incorporates the supervised loss function directly in the binarization procedure, which is an important and desirable property. (Authors mention that existing weight binarization methods ignore the effect of binarization to the loss.) The method is clearly described and related analytically to the previously proposed weight binarization methods.
The experiments are extensive with multiple datasets and architectures, and demonstrate the generally higher performance of the proposed approach.
A minor issue with the feed-forward network experiments is that only test errors are reported. Such information does not really give evidence for the higher optimization performance. (see also comment "RE: AnonReviewer3's questions" stating that all baselines achieve near perfect training accuracy.) Making the optimization problem harder (e.g. by including an explicit regularizer into the training objective, or by using a data extension scheme), and monitoring the training objective instead of the test error could be a more direct way of demonstrating superior optimization performance.
The superiority of BPN is however becoming more clearly apparent in the subsequent LSTM experiments.