Wavelet Loss Function for Auto-Encoder

Abstract

In the field of image generation, especially for auto-encoder models, how to extract better features and obtain better quality reconstruction samples by modifying network structure and training algorithms has always been the focus of attention. For example, Variational Auto-Encoder (VAE), which is a very popular auto-encoder model, is a theoretically rigorously derived image generation model. For commonly used auto-encoders, such as VAE, Regularized Auto-Encoder (RAE), and Wasserstein Auto-Encoder (WAE), the mean square error (MSE) is all used as the loss function of the reconstructed process, which makes the blur problem of the reconstruction samples unavoidable. Especially for larger-sized images, the blur phenomenon is more obvious. To solve this problem, Perceptual Loss Function is used in some cases. Although it can improve the image quality to a certain extent, the amount of calculation is large and the image quality improvement in auto-encoder is also relatively limited. For this reason, we try to propose a new loss function, Wavelet loss function, to better generate and reconstruct images. Wavelet transform is applied to the reconstructed image loss function of the auto-encoder, and the frequency characteristics of the decomposed image are used to constrain it. We conducted comparative experiments on two larger-size image datasets (FaceSrub, COIL20) and a small-size image dataset (Fashion_MNIST), and proved the effectiveness of the wavelet loss function. At the same time, we propose a new image quality index: wavelet high-frequency signal-to-noise ratio (WHF-SNR), which can better measure the quality of the reconstructed image of the auto-encoder.