RMVAE: one-class classification via divergence regularization and maximization mutual information

One-class classification aims to learn the classifier from only one class of data. Variational auto-encoder (VAE) has been widely used in it. Trained on the normal samples, all the images reconstructed by the VAE in the test stage are similar to the normal samples. Thus, the VAE can produce higher reconstruction errors for abnormal samples than normal ones, which can be used as a classification criterion. However, the VAE can reconstruct abnormal samples well and produce lower reconstruction errors due to the model generalization. It leads to the wrong classification for the normal images. To alleviate this shortcoming of the VAE, we propose to use mutual information module and divergence regularization to enhance the VAE. The new model is called RMVAE. Firstly, we refer to the idea of contrast learning to maximize the mutual information between the input image and the corresponding latent representation so that the encoder can express the unique characteristics of the normal class. Besides, the attention mechanism is used in the encoder to enhance the feature extraction capabilities of the model. Secondly, we introduce divergence regularization to make the latent representation of the normal samples evenly distributed in the latent space. Extensive experiments demonstrate that the proposed method achieves a better effect against other state-of-the-art methods on the three public benchmark datasets.