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: This paper proposes to change the L2 norm of loss function of VAE into hyperbolic cosh function. The idea  and presentation are clear and straightforward. However, the used cosh function does not convince me since when t=a, f(t,a) will still be very large! Also, they will grow fast with exp|at|. The authors are encouraged to provide more detailed proofs for the advantages of cosh function.