Log Hyperbolic Cosine Loss Improves Variational Auto-Encoder

1 INTRODUCTION . Unsupervised generative modeling aims to represent probability distributions over input data . The explicit approach usually employs certain parametric model to estimate and approximate the unknown data distribution pdata , but the effect is sensitive to the model selection , for which expert knowledge is often needed . In a different vein , implicit generative models aim to generate samples from the data distribution without estimating the distribution itself . This greatly weaken the dependency on the model selection , and finds numerous applications such as simulating possible futures in reinforcement learning Finn & Levine ( 2017 ) , predicting the next frame of a video sequence Lotter et al . ( 2015 ) and image super-resolution Ledig et al . ( 2017 ) . Prominent examples of implicit generative models include Boltzmann machines Aarts & Korst ( 1988 ) , Generative Adversarial Network ( GAN ) Goodfellow et al . ( 2014 ) and Variational Auto-Encoder ( VAE ) Kingma & Welling ( 2013 ) . Among these methods , VAE enjoys an efficient and stable training process , and the encoder component of VAE gives dimension reduction and feature learning as a by-product . Meanwhile , it is also observed that images generated by VAE are often blurry . One potential reason is that the model fails to represent high-dimensional probability distributions accurately , which is the target issue we aim to address in this paper . 1.1 VARIATION AUTO-ENCODERS . We now give a very brief account of VAE . In VAE , there is an encoder that converts the real input data x to a latent code z ∼ qφ ( z|x ) , and a decoder that reconstructs an image x̂ ∼ pθ ( x|z ) from the latent code z , where q and p are encoder and decoder , with parameter φ and θ respectively . VAE has two objectives : one is to match the decoded samples x̂ to the input x , and the other is to maintain the posterior qφ ( z|x ) of the hidden code vector z to a given prior distribution p ( z ) . Combining these two , VAE uses the following loss function for datapoint x. L ( θ , φ ; x ) = −Eqφ ( z|x ) log pθ ( x|z ) +DKL ( qφ ( z|x ) ‖ p ( z ) ) . ( 1 ) Here the first term is the reconstruction loss , which reduces to the squared L2 loss ‖x− x̂‖22 if we assume that the decoder predicts a Gaussian distribution at each pixel i.e . pθ ( x|z ) = N ( x̂ , σ2I ) . The second term of the objective matches the distribution of latent code to a prior p ( z ) , which by default is a standard normal distribution . 1.2 RECONSTRUCTION LOSS . To improve accuracy of representing probability distributions , many variants of VAE have been proposed . Much attention has been paid to improving the design of the second term , for example , changing the prior on the latent code ( e.g . S-VAE Davidson et al . ( 2018 ) , VAE with householder flow Tomczak & Welling ( 2016 ) ) , and changing the loss between the latent distribution and the prior ( e.g . AAE Makhzani et al . ( 2016 ) , WAE Tolstikhin et al . ( 2018 ) ) . An alternative approach for improvement is through the reconstruction loss , i.e . the first term in Eq . ( 1 ) , which measures the error between the decoded sample and the input data points . Although the squared L2 loss is widely used as reconstruction loss in VAE setting , it has long been observed that this loss function results in problems such as blurry images Mathieu et al . ( 2015 ) . Several new metrics are proposed . Ridgeway et al . ( 2015 ) applied the structural-similarity ( SSIM ) index Wang et al . ( 2004 ) as a reconstruction metric of an auto-encoder . Larsen et al . ( 2016 ) employs a GAN discriminator to compute the reconstruction loss in VAE . Dosovitskiy & Brox ( 2016 ) combined a GAN discriminator with the perceptual loss which computes distances between image features extracted by deep neural networks . These experiments demonstrate the potential of improving VAE by re-designing the reconstruction loss . However , disadvantages are also evident : The first one Ridgeway et al . ( 2015 ) only demonstrated results on gray-scale images , which are very blurry . The last two Larsen et al . ( 2016 ) ; Dosovitskiy & Brox ( 2016 ) reported that their produce images suffer from severe high-frequency artifacts , possibly due to feature loss as discussed in Mahendran & Vedaldi ( 2015 ) . In addition , as other GAN-based methods , introducing a discriminator makes the training harder and instable . In this paper , we propose an element-wise reconstruction loss based on a simple log-cosh function f ( t ; a ) = 1 a log ( cosh ( at ) ) = 1 a log eat + e−at 2 , ( 2 ) where a ∈ R+ is a parameter and log is the natural logarithm . The plots of f ( t ; a ) in interval [ −1 , 1 ] for different a are shown in Fig . 1 . The function is close to L1 for large |t| and close to L2 for small |t| , thereby combining the smoothness advantage of L2 , and robustness and image sharpness advantage of L1 . In addition , the derivative of the function is simply the sigmoid function , rendering a very efficient training and very simple implementation . In Sec . 4 we show empirically that even if the model trained with the log-cosh reconstruction loss achieved a larger error measured by L2 loss , it significantly outperforms the model trained with squared L2 loss . The result also indicates that squared L2 loss can be ineffective in evaluating high-dimensional data like images , which will be discussed in more details later . 2 LIMITATIONS OF TRADITIONAL LOSS . Given two vectors x and x̂ , the squared L2 loss and the L1 loss are defined as LSquared L2 ( x , x̂ ) = ‖x− x̂‖ 2 2 = ∑ i |xi − x̂i|2 , ( 3 ) LL1 ( x , x̂ ) = ‖x− x̂‖1 = ∑ i |xi − x̂i| , ( 4 ) respectively , where xi stands for the i-th entry of x . The loss functions also extend to matrices by treating them as long vectors obtained from concatenating the rows . When applied in the VAE setting , x and x̂ are the input of encoder and output of decoder , respectively . As mentioned in the last section , the squared L2 loss widely used in VAE comes naturally from multivariate Gaussian distribution pθ ( x|z ) ∝ exp ( −‖x− x̂‖22 ) . Other than the aforementioned issue of causing blurry images , the squared L2 loss is also well known to be sensitive to large noise , and it has long been argued that in image space , it is ineffective in evaluating visual quality Wang & Bovik ( 2009 ) and makes the training to easily get stuck in local minima Zhao et al . ( 2017 ) . If we choose a decoder with zero-mean Laplace distribution pθ ( x|z ) ∝ exp ( −‖x− x̂‖1 ) , we will obtain a loss function in terms of the L1 norm . The L1 loss is used as image to image loss in some variants of VAE and GAN , such as AEGAN Rosca et al . ( 2017 ) and cycle GAN Zhu et al . ( 2017 ) . The gradient of L1 loss is ±1 at differential points . The L1 loss is generally robust to noise and , when used as a regularizer , encourages sparsity in solution . It is found to be quite useful in computer vision tasks such as deblurring Cai et al . ( 2009 ) ; Freeman et al . ( 2009 ) ; Xu & Jia ( 2010 ) , since it appears to enable the solver to escape from local minima , although the exact reason for this is still elusive . In addition , Zhao et al . ( 2017 ) shows that L1 can outperform L2 when directly used as a loss function for neural networks in image restoration applications such as denoising Jain & Seung ( 2009 ) , deblurring Xu et al . ( 2014 ) , demosaicking Wang ( 2014 ) , and super-resolution Dong et al . ( 2014 ) . However , an obvious disadvantage of the L1 loss is that it is not differentiable if any entry xi − x̂i = 0 , causing oscillation between ±1 during the training process . In this paper , we will show that our proposed log-cosh loss takes the advantage of L1 while overcoming the problem of non-differentiability . Experiments show that the log-cosh loss outperforms both L1 and L2 loss . 3 LOG-COSH LOSS . 3.1 PROPERTIES OF LOG-COSH LOSS . Based on the log-cosh function shown in Eq . ( 2 ) , we propose to use the log-cosh loss as the reconstruction error , taking advantages of both squared L2 loss and L1 loss while overcome their limitations . The log-cosh loss is defined as Llog - cosh ( x , x̂ ) = ∑ i f ( xi − x̂i , a ) = 1 a ∑ i log ( cosh ( a ( xi − x̂i ) ) ) , ( 5 ) where a ∈ R+ is a hyper-parameter . We will first show that the log-cosh function behaves like L2 around origin and like L1 at other points . Proposition 1 . The log-cosh function f ( t ; a ) for a > 0 approximates |t| − 1a log 2 when t → ∞ , and approximates 0.5at2 when t→ 0 . More precisely , f ( t ; a ) = 1 a log ( cosh ( at ) ) = 1 a log eat + e−at 2 → |t| − 1 a log 2 , when |t| → ∞ 0.5at2 , when |t| → 0 ( 6 ) Proof . When |t| → ∞ , we have f ( t ; a ) = 1a log ( ( e a|t|+e−a|t| ) /2 ) → 1a log ( e a|t|/2 ) = |t|− 1a log 2 . When |t| → 0 , the Taylor expansion gives f ( t ; a ) = 0.5at2 +O ( t4 ) . For general t , one has bounds |t| − 1a log 2 ≤ f ( t ; a ) ≤ |t| − 1 a log 2 + 1 a . The log-cosh loss can be viewed as a smoothed version of L1 loss differentiable everywhere . The gradient of f ( t ; a ) is df ( t ; a ) dt = tanh ( at ) = eat − e−at eat + e−at = 2σ ( 2at ) − 1 , ( 7 ) where tanh ( t ) = ( et − e−t ) / ( et + e−t ) is the hyperbolic tangent function and σ ( t ) = 1/ ( 1 + e−t ) is the sigmoid function , a commonly used activation function in deep networks . 3.2 MOTIVATION OF APPLYING LOG-COSH LOSS TO VAE . The log-cosh loss applied to VAE can be understood in the following way , which was also how we obtained the method . In general we hope to balance the reconstruction accuracy and latent space optimization in VAE . If the squared L2 loss is used for the reconstruction , then the reconstruction accuracy deteriorates in the region of small reconstruction error , as L2 penalizes these small errors too lightly , making the objection function Eq . ( 1 ) to be dominated by the second term ( the KL divergence term ) . To solve this problem , simply increasing the relative weight of the reconstruction loss does not work well , since it harms the optimization in the latent space when the reconstruction error is large . Putting both large-error and small-error regions into consideration , a feasible solution is to increase the weight of reconstruction loss when it is small , but at the same time to clip the gradient of the loss with an upper bound so that the reconstruction penalty does not always increase linearly with the error . A smooth approximation of this is the tanh function , whose integral turns out to be the log-cosh function . The resulting log-cosh loss becomes more sensitive to small errors because it behaves as a 0.5a-scaled squared L2 loss around the origin . At the same time , it avoids penalizing too much when error is large when it behaves likes the L1 loss . Therefore , we are motivated to implement the log-cosh loss to improve the reconstruction accuracy of VAE without damaging its latent space optimization . The implementation of the log-cosh loss in VAE is quite simple , as the gradient is explicitly available as shown in Eq . ( 7 ) . In tensorflow the implementation is as simple as replacing the reconstruction loss in the back-propagation with x̂ ∗ tf .stop gradient ( 2 ∗ tf .sigmoid ( 2a ∗ ( x̂− x ) ) − 1 ) , where x and x̂ are the input of encoder and output of decoder , respectively . This form also avoids overflow or underflow during the computation , as it directly combines the gradient of the log-cosh loss ( w.r.t decoder samples ) and the gradient of the samples ( w.r.t parameters ) .

Summary: The paper proposes a new loss function which can be used in the reconstruction term of various auto-encoder architectures. The pixel-wise cost function \ell(X, X') = f(X - X'; a) is defined for pairs of two input images X and X' and has one positive real-valued hyperparameter a. For small values of t the function f(t; a) behaves like a quadratic function, while for large t it behaves like |t|. As a consequence, it is smooth, everywhere differentiable (like L2) while not penalizing outliers too hard (like L1). The authors present several experiments conducted on MNIST and Celeba datasets, demonstrating that a simple change of a conventional pixel-wise squared L2 distance with the proposed log-cosh cost function improves the FID scores of generated samples as well as the visual quality of reconstructions (including "the sharpness").