Good Semi-supervised VAE Requires Tighter Evidence Lower Bound

1 INTRODUCTION . The rise of deep neural networks has led to breakthroughs in computer vision , natural language processing , and many other domains . Most of these models are trained on large labeled datasets via supervised learning . However , in many scenarios , although it is easy to acquire a large amount of the original data , obtaining corresponding labels is often very costly or even infeasible . Semisupervised learning ( Thomas , 2009 ) is proposed to address this problem by training classifiers with sufficient unlabeled data and a small fraction of labeled data . Recent works on semi-supervised learning can be grouped into three categories : ( 1 ) disagreement based learning via data perturbation ( Miyato et al. , 2019 ) and consistency enforcing ( Verma et al. , 2019 ) , ( 2 ) metric learning ( Wu et al. , 2018 ) , ( 3 ) generative approaches via generative adversarial network ( GAN ) ( Springenberg , 2016 ) and variational autoencoder ( VAE ) ( Kingma et al. , 2014 ) . Compared with the first two categories , generative approaches have great advantages in interpretability . Based on the latent variable assumption ( Doersch , 2016 ) , the generative model has an explicit variational inference form , so it can learn the marginal probability distribution of the raw data as well as the conditional distribution of the latent variables given the input data , which makes predictions more reasonable . Besides , generative approaches not only learn the required classification representations , but also capture the semantics-disentangled factors that generate the data , making it easier to generalize to different tasks ( Narayanaswamy et al. , 2017 ) . However , in practice , semi-supervised generative approaches often encounter three major challenges : ( 1 ) The two-stage training process is not robust . Semi-supervised VAE ( Kingma et al. , 2014 ) needs to be trained carefully with a two-stage hierarchical strategy , while the training process of GAN is a two-stage adversarial game ( Chrysos et al. , 2019 ) . ( 2 ) Good semi-supervised learning results and good generative performance can not be obtained at the same time . In GAN , good semi-supervised learning performance will lead to a mismatch between the generated results and the real data distribution ( Dai et al. , 2017 ) . While in VAE , the evidence lower bound ( ELBO ) objective is irrelevant to the classification loss , making it difficult to learn from the labels directly ( Narayanaswamy et al. , 2017 ) . ( 3 ) Even at the expense of sacrificing generative performance , the semi-supervised classification results are still not satisfactory . In practice , disagreement-based meth- ods ( Xie et al. , 2019 ; Berthelot et al. , 2019 ) have dramatically improved the state-of-the-art results on several standard datasets , surpassing generative approaches by a large margin . These challenges naturally raise a question : What limits the performance of generative approaches in semi-supervised learning ? In this work , we propose One-stage Semi-suPervised Optimal Transport VAE ( OSPOT-VAE ) to address these challenges , which consists of two improvements : ( 1 ) a one-stage semi-supervised VAE model that unifies the generation and classification loss in one ELBO framework . ( 2 ) an estimation of the margin between true log-likelihood and the ELBO that exports a tighter evidence lower bound by applying optimal transport ( Ambrosio & Gigli , 2013 ) scheme to the distribution of latent variables . Our model has the following contributions : • We show that OSPOT-VAE can be well trained with a direct one-stage strategy . • We show that OSPOT-VAE can achieve both good generative performance and semi- supervised learning results simultaneously on a series of benchmark datasets . • We point out that it is the large margin between the ELBO and the log-likelihood of the input data that limits the performance of semi-supervised VAE . Besides , we evaluate this assumption across many standard datasets and show that with the proposed tighter ELBO , OSPOT-VAE surpasses the best semi-supervised generative models by a large margin and achieves state-of-the-art performance on Cifar-100 with 10k labels . 2 SEMI-SUPERVISED LEARNING METHODS . In supervised learning ( SL ) , we are facing with training data that appears as input-target pairs ( X , y ) ∈ DL sampled from an unknown distribution p ( X , y ) . Our goal is to learn a function f ( X ; φ ) parameterized by φ that makes the correct inference y for unseen samples from p ( X ) . While in semi-supervised learning ( SSL ) , we can obtain an extra collection of unlabeled data X ∈ DU sampled from the same distribution p ( X ) . We hope to leverage the data from both DL and DU to achieve a more accurate model than what would have been obtained by only using DL . In this section , we review some existing methods for SSL . We mainly focus on those who have reached state-of-the-art results , as well as generative approaches which are strongly connected with our model ; the more comprehensive overview is beyond the scope of this paper , we refer readers to ( Oliver et al. , 2018 ) . 2.1 DISAGREEMENT BASED LEARNING . Disagreement-based learning refers to the general approaches of imposing disagreement among multiple learners on the same task or multiple predictions from a single learner . By eliminating the disagreement , we can enforce the generalization of the model on unseen data . A common technique for creating disagreement is data augmentation , which applies transformations or perturbations on the input data and leaves class semantics unchanged . For X ∈ DU , loss term can be derived as ‖f ( Augment ( X ) ; φ ) − f ( X ; φ ) ‖22 ( 1 ) where the Augment ( X ) is a stochastic function which can be obtained by image transformation ( Xie et al. , 2019 ) , virtual adversarial training ( Miyato et al. , 2019 ) , or mixup method ( Verma et al . 2018 ; Berthelot et al . 2019 ) . Another disagreement construction technique is to train multiple learners on the same dataset and utilize the loss ‖f ( X ; φ1 ) − f ( X ; φ2 ) ‖22 ( 2 ) to enforce the predictive consistency of different models , for example , “ Mean Teacher ” ( Tarvainen & Valpola , 2017 ) and “ Teacher Graph ” ( Luo et al. , 2018 ) . The generalization of the models gets enhanced . 2.2 GENERATIVE APPROACHES . In generative approaches , input X is supposed to have corresponding continuous and discrete latent variables , which we denote by z and c respectively . Feature matching ( FM ) GANs ( Salimans et al. , 2016 ; Dai et al. , 2017 ) apply GANs to semisupervised learning on K-classification tasks by specifying a ( K+1 ) -class objective for the discriminator . Instead of binary classification , true samples are classified into the first K classes respectively and fake samples are classified into the ( K+1 ) -th class . This target function achieves strong empirical results by matching the generator distribution with true data distribution and improves semi-supervised classification performance . Semi-supervised VAEs ( Kingma et al . 2014 ; Narayanaswamy et al . 2017 ) construct a probabilistic model parameterized by θ and φ that respectively describe the generation and inference process between X and latent variables z , c. The generation process of X by z and c is : p ( z ) = N ( z ; 0 , I ) ; p ( c ) = Mult ( c ; K , π ) ; pθ ( X|z , c ) = f ( X ; z , c , θ ) ( 3 ) where Mult ( K , π ) is the multinomial distribution with class K and parameter π. f ( X ; z , c , θ ) is a suitable likelihood function , e.g . a Bernoulli or Gaussian distribution , parameterized by a non-linear transformation of the latent variables z and c. The class label y is treated as c if given . For the inference process , with the following hypothesis qφ ( z , c|X ) = qφ ( z|X ) qφ ( c|X ) ; p ( z , c|X ) = p ( z|X ) p ( c|X ) ; p ( z , c ) = p ( z ) p ( c ) ( 4 ) evidence lower bound ( ELBO ) is used as objective to predict the posterior distribution of latent variables as follows ( see Appendix A.1 for proof ) : log p ( X ) ≥ Eqφ ( z , c|X ) [ log pθ ( X|z , c ) ] −DKL ( qφ ( z|X ) ‖p ( z ) ) −DKL ( qφ ( c|X ) ‖p ( c ) ) = ELBO ( 5 ) For the likelihood log p ( X ) is infeasible , VAE maximizes its evidence lower bound instead , which derives the negative ELBO loss function L ( X ; φ , θ ) = −ELBO . The predictions for the classification label y can be obtained from the inferred posterior distribution qφ ( c|X ) . When class label y is not given , missing label sampling technique is used to sample from qφ ( c|X ) as Eqφ ( c|X ) f ( X ; c , θ ) = K∑ k=1 qφ ( yk|X ) f ( X ; yk , θ ) ( 6 ) Here yk represents a one-hot vector with 1 in k-th dimension . Utilizing this sampling method , the algorithmic complexity of VAE is proportional to K so it is computationally inefficient . Note that in objective ( 4 ) , the label predictive distribution qφ ( c|X ) only contributes to the generative performance . To remedy this , existing models simply add a cross-entropy loss to the negative ELBO loss such that the distribution qφ ( c|X ) can also learn classification rules from the labeled data . The extended objective loss is min φ , θ EX∼DUL ( X ; φ , θ ) + E ( X , y ) ∼DL [ L ( X , c = y ; φ , θ ) − log qφ ( y|X ) ] ( 7 ) Two-stage Training Strategy : In practice , ( Kingma et al. , 2014 ) finds that directly training the one-stage objective ( 7 ) will lead to a bad semi-supervised learning result , so a two-stage training strategy is proposed to improve the model . The two-stage training strategy consists of two parts , M1 and M2 . M1 means to learn a new continuous latent representation z1 first , and M2 means to train a semi-supervised model ( 7 ) with the embedding z1 from M1 instead of the raw data X . This M1+M2 strategy builds a deep VAE with two layers of random variables : pθ ( X , z1 , z2 , c ) = pθ ( X|z1 ) pθ ( z1|z2 , c ) p ( z2 ) p ( c ) , which can dramatically improve the performance of the inference qφ ( c|X ) but is not robust in training . Moreover , GAN ’ s training process can also be considered as two-stage with generator and discriminator competing with each other , and the two-stage adversarial game adds the difficulty in training . 3 ONE-STAGE SEMI-SUPERVISED OPTIMAL TRANSPORT VAE . In this section , we introduce our semi-supervised VAE framework , OSPOT-VAE . Firstly , we derive a one-stage loss function that unifies the generation and classification loss under one ELBO without introducing any additional auxiliary loss items like ( 7 ) . Then , we analyze a phenomenon that good ELBO values do not guarantee good semi-supervised performance and propose the optimal transport estimation to deal with it . At last , combining the two parts , we give the detailed algorithm of OSPOTVAE and discuss some problems in model optimization .

Summary: The paper proposes to combine a VAE model with the Optimal Transport to approximate some components of the model. The authors evaluate their approach on semi-supervised problems and claim to obtain very competitive results compared to literature. Unfortunately, the paper is very unclear and hard to follow. The authors make some claims that are not true, for instance, learning so called M1+M2 architecture end-to-end is hard, however, there are papers that successfully train such model. Moreover, the authors reported their results as SOTA, however, they missed many other papers with much better scores.