GOOD SEMI-SUPERVISED VAE REQUIRES TIGHTER EVIDENCE LOWER BOUND
Year: 2019

Abstract:
Semi-supervised learning approaches based on generative models have now encountered 3 challenges: (1) The two-stage training strategy is not robust. (2) Good semi-supervised learning results and good generative performance can not be obtained at the same time. (3) Even at the expense of sacrificing generative performance, the semi-supervised classification results are still not satisfactory. To address these problems, we propose One-stage Semi-suPervised Optimal Transport VAE (OSPOT-VAE), a one-stage deep generative model that theoretically unifies the generation and classification loss in one ELBO framework and achieves a tighter ELBO by applying the optimal transport scheme to the distribution of latent variables. We show that with tighter ELBO, our OSPOT-VAE surpasses the best semi-supervised generative models by a large margin across many benchmark datasets. For example, we reduce the error rate from 14.41% to 6.11% on Cifar-10 with 4k labels and achieve state-of-the-art performance with 25.30% on Cifar-100 with 10k labels. We also demonstrate that good generative models and semi-supervised results can be achieved simultaneously by OSPOT-VAE.

heading: 1 INTRODUCTION
text: 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.

heading: 2 SEMI-SUPERVISED LEARNING METHODS
text: 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).

heading: 2.1 DISAGREEMENT BASED LEARNING
text: 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.

heading: 2.2 GENERATIVE APPROACHES
text: 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.

heading: 3 ONE-STAGE SEMI-SUPERVISED OPTIMAL TRANSPORT VAE
text: 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.

heading: 3.1 ONE-STAGE SEMI-SUPERVISED VAE
text: Following the notations and assumptions (3, 4) in Section 2.2, we derive our one-stage semisupervised VAE. With the empirical distribution pemp(X;D) = 1 |D| ∑ X′∈D 1X=X′ , we utilize the decomposition in (Zhao et al., 2017) and rewrite the second part of (5) into (proof in Appendix A.2) Epemp(X)DKL(qφ(z|X)‖p(z)) = Iqφ(X; z) +DKL(qφ(z)‖p(z)) ≥ Iqφ(X; z) (8)
where qφ(z) = 1 |D| ∑ X∈D qφ(z|x) and Iqφ(X; z) is the mutual information between X and z. The left part of (8) equals to 0 when X and z are independent. This is undesirable, so Iqφ(X; z) can be regarded as the lower bound of controlled mutual information. The continuous variables in (8) can be easily extend to discrete variables c. We can use Iz and Ic to denote the controlled information capacity and derive the objective for the unlabeled dataset DU
LDU (X;θ,φ) = Eqφ(z,c|X)[− log pθ(X|z, c)] + βz|DKL(qφ(z|X)‖p(z)− Iz| + βc|DKL(qφ(c|X)‖p(c))− Ic|
(9)
where β, Iz, Ic are all hyper-parameters forcing the KL divergence term to match the mutual information capacities of z and c.
For the labeled subset DL, instead of directly employing class label y as sampled c, we view it as the parameter of the true posterior distribution, i.e. p(c|X) = Mult(c;K,y) and derive the following one-stage ELBO form:
log p(X) = logEqφ(z|X),p(c|X) p(X, z, c)
qφ(z|X)p(c|X) ≥ Eqφ(z|X),p(c|X) log
p(X, z, c)
qφ(z|X)p(c|X)
= Eqφ(z|X),p(c|X)[log p(X|z, c) + log p(z)p(c)
qφ(z|X)p(c|X) ] = Eqφ(z|X),p(c|X) log p(X|z, c)
−DKL(qφ(z|X)‖p(z))−DKL(p(c|X)‖qφ(c|X)) + Ep(c|X) log p(c)
qφ(c|X)
(10)
Notice that DKL(p(c|X)‖qφ(c|X)) is equal to the common cross-entropy loss for y is a one-hot vector. In this respect, the margin between p(c|X) and qφ(c|X) can be significantly small when the suitable optimization method is chosen. This allows us to utilize the approximation p(c|X) ≈
qφ(c|X) to modify Ep(c|X) log p(c)
qφ(c|X) in (10), resulting in a consist ELBO with LDU (X;θ,φ):
Ep(c|X) log p(c)
qφ(c|X) ≈(when p(c|X)≈qφ(c|X)) Eqφ(c|X) log
p(c)
qφ(c|X) = DKL(qφ(c|X)‖p(c)) (11)
Combining the ELBO form (10) of DL with the mutual information decomposition (8) and the approximation (11), the new objective for semi-supervised VAE is: LDL(X,y;θ,φ) = Eqφ(z|X),p(c|X)[− log pθ(X|z, c)] + βz|DKL(qφ(z|X)‖p(z))− Iz|
+ βc|DKL(qφ(c|X)‖p(c))− Ic|+DKL(p(c|X)‖qφ(c|X)) (12)
With (9) and (12), the objective for the entire dataset is now min φ,θ EX∼pemp(X;DU )LDU (X;θ,φ) + E(X,y)∼pemp((X,y);DL)LDL(X,y;θ,φ) (13)
This one-stage objective with a simple approximate transformation (9) unifies the generation loss as well as the target of SSL and results in improved performance of semi-supervised learning, which we demonstrate in Section 4.1.
Algorithm 1 Optimal transport estimation ingests a batch of observation X as well as the representation qφ(z|X), qφ(c|X) inferred from the original VAE and returns the estimation of the margin DKL(qφ(z|X)‖p(z|X)) and DKL(qφ(c|X)‖p(c|X)). Input:
Batch of observation X sampled from pemp(X); Inferred parameter (µ, diag(σ2)) of qφ(z|X) = N (z;µ, diag(σ2)); Inferred parameter π of qφ(c|X) = Mult(c;K,π); Hyperparameter α for mixup vicinal distribution pmixup(X)
Output: X̃ sampled from pmixup(X); Estimation LMz of the margin DKL(qφ(z|X̃)‖p(z|X̃)); Estimation LMc of the margin DKL(qφ(c|X̃)‖p(c|X̃))
1: X′,π′,µ′,σ′2 = RandomPermutation(X,π,µ,σ2) 2: X̃ = λ ∗X + (1− λ) ∗X′, λ ∈ β(α, α) 3: π̃ = OptimalTransportC(π,π′, λ) 4: (µ̃, σ̃2) = OptimalTransportZ((µ,σ2), (µ′,σ′2), λ) 5: qφ(z|X̃), qφ(c|X̃) = VAE(X̃) 6: p̃(z|X̃) = N (z; µ̃, diag(σ̃2)) 7: p̃(c|X̃) = Mult(c;K, π̃) 8: LMz = DKL(qφ(z|X̃)‖p̃(z|X̃)) 9: LMc = DKL(qφ(c|X̃)‖p̃(c|X̃))
10: return X̃, LMz , LMc

heading: 3.2 OPTIMAL TRANSPORT ESTIMATION
text: To summarize the above, VAE aims to learn the useful representation qφ(z|X) and qφ(c|X) by reducing the KL divergence between the empirical distribution pemp(X) and the model marginal p(X) = ∫ z ∫ c p(X)p(z|X)p(c|X)dzdc. Instead of minimizing DKL(pemp(X)‖p(X)) directly, VAE models use the expected ELBO mentioned in (5) as target via the following inequality
DKL(pemp(X)‖p(X)) ≤ H(pemp(X))− Epemp(X)ELBO (14)
However, one phenomenon is that good ELBO values do not imply accurate inference. A typical example has been discussed in (Zhao et al., 2017). Here we mainly focus on the cause of this phenomenon and propose optimal transport estimation to alleviate this problem in semi-supervised learning. Following the work in (Rezende et al., 2014), we write down the closed form of the expected margin between true log-likelihood and ELBO as (proof in Appendix A.3):
Epemp(X)[log p(X)− ELBO] = Epemp(X)[DKL(qφ(z|X)‖p(z|X)) +DKL(qφ(c|X)‖p(c|X))] (15) Combined with the decomposition (8), training the expected ELBO target can only reduce the difference between marginal distributions qφ(c),qφ(z) and p(c),p(z). It means that even with a good ELBO, the margin Epemp(X)DKL(qφ(z|X)‖p(z|X)) and Epemp(X)DKL(qφ(c|X)‖p(c|X)) in (15) can still be large. In this scenario, the consistent optimization of ELBO will contribute no more to the semi-supervised classification performance. However, optimizing the margin in (15) directly is impossible, for p(c|X) and p(z|X) are unknown. To remedy this, we extend the empirically effective approximation in (Zhang et al., 2018) to our VAE framework with the form
Epmixup(X)DKL(qφ(z|X)‖p(z|X)) ≈ Epemp(X)DKL(qφ(z|X)‖p(z|X))(α→ 0) Epmixup(X)DKL(qφ(c|X)‖p(c|X)) ≈ Epemp(X)DKL(qφ(c|X)‖p(c|X))(α→ 0)
(16)
where pmixup(X) is the mixup vicinal distribution (Zhang et al., 2018) and α is the related parameter. Then we propose optimal transport estimation to construct the estimations of Epmixup(X)DKL(qφ(z|X)‖p(z|X)) as well as Epmixup(X)DKL(qφ(z|X)‖p(z|X)) by applying optimal transport scheme to latent variables z and c. The computation steps of optimal transport estimation are provided in Algorithm 1, and we present the details of the optimal transport scheme in the rest of this section.
Algorithm 2 OSPOT-VAE training process with epoch t. Input:
Batch of labeled pairs (XL,yL) ∈ DL,Batch of unlabeled examples XU ∈ DU ; ELBO hyperparameters: βz, βc, Iz, Ic = ELBOScheduler(t); Optimal transport estimation weights: wMz , wMc = WeightScheduler(t); Model parameters: θ(t−1),φ(t−1); Model optimizer: SGD
Output: Updated parameters: θ(t),φ(t)
1: LL = LDL(XL,yL;θ(t−1),φ(t−1);βz, βc, Iz, Ic) 2: LU = LDU (XU ;θ(t−1),φ(t−1);βz, βc, Iz, Ic) 3: LMz , LMc = OptimalTransportEstimation(XU , qφ(z|XU ), qφ(c|XU )) 4: L = LL + LU + wMzLMz + wMcLMc 5: θ(t),φ(t) = SGD(θ(t−1),φ(t−1), ∂L∂θ , ∂L ∂φ ) 6: return θ(t),φ(t)
Optimal Transport Scheme: The mixup vicinal distribution can be understood as applying linear transport between the points X,X′ ∈ D, extending the original dataset with new points falling on one straight line X̃ = λ ∗X + (1− λ) ∗X′, λ ∈ [0, 1]. For X̃, it is a natural thought that this linear transformation could associate with the shortest-path transport in the latent space. Based on this, we calculate the distributions p̃(z|X̃), p̃(c|X̃) of z, c and consider them as the estimation of the true posterior distributions. Following the work of (Ambrosio & Gigli, 2013), the norm-2 based optimal transport scheme γ(x, y) between two distributions p(x) and p(y) satisfy:
inf γ(x,y) ∫ x ∫ y ‖x− y‖22γ(x, y)dxdy
s.t. ∫
y γ(x, y)dy = p(x); ∫ x γ(x, y)dx = p(y)
(17)
For the continuous variable z ∼ N (µ, diag(σ2)) and discrete variable c ∼ Mult(K,π), the following 2 propositions are proposed to calculate the shortest-path based on optimal transport scheme (see Appendix A.4 for proof). Proposition 3.1. The shortest-path derived from optimal transport scheme (17) between z1 ∼ N (µ1, diag(σ21)) and z2 ∼ N (µ2, diag(σ22)) with λ ∈ [0, 1] is
µ̃ = λµ1 + (1− λ)µ2 σ̃ = λσ1 + (1− λ)σ2
(18)
Proposition 3.2. The shortest-path derived from KL divergence based optimal transport scheme between c1 ∼ Mult(K,π1) and c2 ∼ Mult(K,π2) with λ ∈ [0, 1] is
π̃ = λπ1 + (1− λ)π2 (19)
Algorithm 1 yields the optimal transport estimation of the margin in (15), which leads to a tighter ELBO. In Section 4.2, we demonstrate that with this tighter ELBO, the inference performance of semi-supervised VAE is significantly improved on many benchmark datasets.

heading: 3.3 OPTIMIZATION OF OSPOT-VAE
text: Combining one-stage semi-supervised VAE and optimal transport estimation, we can get the complete OSPOT-VAE model. The full OSPOT-VAE algorithm is provided in Algorithm 2, and a schematic is shown in Figure 1. Note that the conditions for the approximations used in Algorithm 1,2 satisfy (1) qφ(c|X) ≈ p(c|X) and (2) the VAE model has already achieved a good ELBO. Therefore, the warm-up schedule (Higgins et al., 2017) is used to set parameters Iz, Ic, βz, βc and wMz , wMc . We list the details of “ELBOScheduler(t)” and “WeightScheduler(t)” in Appendix A.5.
In Algorithm 2, we apply stochastic gradient descent (SGD) as optimizer, which needs to calculate the gradient ∇θ,φL. The target loss L consists of KL divergence and the expected loglikelihood . The derivation of KL divergence part has a closed form, while calculating the gradient
of Eqφ(z|X),qφ(c|X) log pθ(X|z, c) is difficult. To this end, we follow the work of (Rezende et al., 2014) and (Jang et al., 2017), using the reparameterization trick as
∇θ,φEqφ(z|X) log pθ(X|z) = EN ( ;0,I)∇θ,φ log pθ(X|µ+ σ · )
EGumbel( ;0,1))∇θ,φ log pθ(X|Softmax( logπ +
τ ))→ ∇θ,φEqφ(c|X) log pθ(X|c)(τ → 0)
(20)
Note that with (20), the algorithmic complexity of one-stage semi-supervised VAE is independent with the class number K, making it easier to extend to large-scale classification tasks.

heading: 4 EXPERIMENTS
text: In this section, we demonstrate the 3 contributions of our OSPOT-VAE model with sufficient experiments on 4 standard SSL benchmark datasets, that is, MNIST, SVHN, Cifar10, and Cifar100. In Section 4.1, we show the validity of our one-stage semi-supervised VAE objective (13) by comparing with other one-stage and two-stage VAE models. Then, we evaluate the performance of OSPOT-VAE under “WideResNet”(Zagoruyko & Komodakis, 2016) backbone and compare with other state-of-the-art SSL models mentioned in Section 2. Besides, We provide an ablation study to verify the contribution of the optimal transport estimation. As an additional application, we show that good generative models and semi-supervised results can be obtained at the same time by OSPOT-VAE (Section 4.3). The source code is available at https: //github.com/PaperCodeSubmission/OSPOT-VAE; more details are available in Appendix A.6.

heading: 4.1 ONE-STAGE SEMI-SUPERVISED VAE
text: We evaluate the effectiveness of the one-stage semi-supervised VAE objective on 2 standard benchmarks, MNIST and SVHN. As for baseline models, we consider two VAE-based SSL models, which are one-stage disentangled VAE (Narayanaswamy et al., 2017) and two-stage VAE(M1+M2) (Kingma et al., 2014). For fairness, except the target loss functions, all models use the same structure as is used in M1+M2 (Kingma et al., 2014). The results are presented in Table 1, and our model achieves the best performance.

heading: 4.2 OSPOT-VAE
text: We compare the results of OSPOT-VAE with two categories of state-of-the-art models mentioned in Section2. In all experiments, we use the “WideResNet-28” model or other deep models with a comparable amount of parameters as the backbone. The results in Table 2,3 demonstrate that our model outperforms most of the existing methods and surpasses state-of-the-art semi-supervised generative models (Dai et al., 2017) by a large margin. Notice that recently, data-augmentation based method, MixMatch, (Berthelot et al., 2019) achieves the absolute state-of-the-art results in all benchmarks. It uses pre-designed sophisticated data augmentation strategies for different datasets and outperforms our model. We list its results fairly as a comparison, while OSPOT-VAE surpasses it in Cifar100 dataset.
Ablation Study: The OSPOT-VAE model consists of two parts: (1) a one-stage VAE objective and (2) an optimal transport estimation. In ablation study, we analyze the effect of each component in our model with the backbone “WideResNet-28-2”. To study the independent effects of transport estimation, we combine it with the encoder part of OSPOT-VAE to build a classifier with loss function LMc . The improved classification error rates in Table 4 show that, with optimal transport estimation, the posterior inference qφ(c|X) gets closer to the true distribution p(c|X). It indicates that our optimal transport estimation does reduce the gap between ELBO and the log-likelihood of the input data and yield a tighter ELBO, which leads to a better semi-supervised performance.

heading: 4.3 GENERATIVE PERFORMANCE
text: Epemp(X)ELBO measures the margin between the true data distribution and the distribution learned by generation models (Doersch, 2016). By comparing the Epemp(X) value of pure unsupervised VAE and our semi-supervised VAE model under the same “WideResNet-28-2” backbone, we demonstrate that good generative models and semi-supervised results can be obtained at the same time in OSPOTVAE. The results in Table 5 show that the data generative distribution learned by our OSPOT-VAE model is as good as the pure VAE model. Further generated results are available in Appendix A.7.

heading: 5 CONCLUSION
text: In this work, we pointed out that it was the large margin between ELBO and the true log-likelihood of the raw data that limits the performance of semi-supervised VAE. To this end, we introduced OSPOT-VAE, a one-stage generative model that unified the classification and generation objective and achieved a tighter ELBO by optimal transport estimation. We demonstrated our assertion through extensive experiments, and our semi-supervised results significantly outperform former state-of-the-art generative SSL methods by a large margin on Cifar10 and Cifar100.

heading: A APPENDIX

heading: A.1 BASIC INEQUALITY OF ELBO
text: Proposition A.1. The Basic inequality of ELBO is log p(X) ≥ Eqφ(z,c|X)[log pθ(X|z, c)]−DKL(qφ(z|X)‖p(z))−DKL(qφ(c|X)‖p(c))
proof
log p(X) = log ∫ z,c qφ(z, c|X) p(X, z, c) qφ(z, c|X) = logEqφ(z,c|X) p(X, z, c) qφ(z, c|X)
≥ Eqφ(z,c|X) log p(X, z, c)
qφ(z, c|X)
= ∫ z,c qφ(z, c|X) log p(z, c) qφ(z, c|X) + ∫ z,c qφ(z, c|X) log pθ(X|z, c)
= −DKL(qφ(z, c|X)‖p(z, c)) + Eqφ(z,c|X)[log pθ(X|z, c)] =with assumption (3,4) Eqφ(z,c|X)[log pθ(X|z, c)]−DKL(qφ(z|X)‖p(z))−DKL(qφ(c|X)‖p(c))

heading: A.2 DECOMPOSITION OF ELBO IN INFO-VAE
text: Proposition A.2. The expected ELBO with empirical distribution satisfies Epemp(X)DKL(qφ(z|X)‖p(z)) = Iqφ(X; z) +DKL(qφ(z)‖p(z))
proof
Epemp(X)DKL(qφ(z|X)‖p(z)) = ∫ X pemp(X) ∫ z qφ(z|X) qφ(z|X) p(z) dzdX
= ∫ X pemp(X) ∫ z qφ(z|X) qφ(z|X) qφ(z) dzdX + ∫ X pemp(X) ∫ z qφ(z|X) qφ(z) p(z) dzdX
= ∫ X ∫ z qφ(z,X) log qφ(z,X) qφ(z)pemp(X) dzdX +DKL(qφ(z)‖p(z))
= Iqφ(X; z) +DKL(qφ(z)‖p(z))

heading: A.3 THE EQUATION FORM OF ELBO
text: Proposition A.3. The expected margin between the true log-likelihood Epemp(X) log p(X) and Epemp(X)ELBO is
Epemp(X)[log p(X)− ELBO] = Epemp(X)[DKL(qφ(z|X)‖p(z|X)) +DKL(qφ(c|X)‖p(c|X))]
proof
We just need to prove the following equation
log p(X)− ELBO = DKL(qφ(z|X)‖p(z|X)) +DKL(qφ(c|X)‖p(c|X)) and the proof under assumption (3, 4) is
log p(X) = ∫ z,c qφ(z|X)qφ(c|X) log p(X)dzdc = ∫ z,c qφ(z|X)qφ(c|X) log p(X, z, c) p(z|X)p(c|X) dzdc
= Eqφ(z|X)qφ(c|X) log p(X, z, c)
qφ(z|X)qφ(c|X) + ∫ z,c qφ(z|X)qφ(c|X) log qφ(z|X)qφ(c|X) p(z|X)p(c|X)
= ELBO +DKL(qφ(z|X)‖p(z|X)) +DKL(qφ(c|X)‖p(c|X))

heading: A.4 OPTIMAL TRANSPORT SCHEME

heading: Proposition A.4.
text: 1. The shortest-path derived from optimal transport scheme (17) between z1 ∼ N (µ1, diag(σ21)) and z2 ∼ N (µ2, diag(σ22)) with λ ∈ [0, 1] is
µ̃ = λµ1 + (1− λ)µ2 σ̃ = λσ1 + (1− λ)σ2
(A.1)
proof
Utilizing the conclusions in Kuang & Tabak (2017), the closed form of optimal transport from one multi-normal distribution N (z1;µ1,Σ1) to another normal distribution N (z2;µ2,Σ2) is
z→ T (z) = µ2 + T(z− µ1); T = Σ − 12 1 (Σ 1 2 1 Σ2Σ 1 2 1 ) 1 2 Σ − 12 1
Utilize the diag matrix assumption, the optimal transport scheme with λ is
zλ = (1− λ)z1 + λT (z1) T = diag(σ1/σ2)
(A.2)
Utilize (A.2), we can get (A.1) as
zλ ∼ N (µ̃, diag(σ̃2)); µ̃ = λµ1 + (1− λ)µ2; σ̃ = λσ1 + (1− λ)σ2 2. The shortest-path derived from KL divergence based optimal transport scheme between c1 ∼ Mult(K,π1) and c2 ∼ Mult(K,π2) with λ ∈ [0, 1] is
π̃ = λπ1 + (1− λ)π2 (A.3) proof
As the definition (17) has no closed-form solution for multinomial distribution, we use KL divergence instead. The KL divergence based optimal transport target cλ ∼ Mult(K, π̃) between two multinomial distribution c1 ∼ Mult(K,π1) and c2 ∼ Mult(K,π2) with λ ∈ [0, 1] satisfy
min π̃ λDKL(π1‖π̃) + (1− λ)DKL(π2‖π̃) s.t. K∑ i=1 π̃i = 1 (A.4)
The Lagrange multiplier form of (A.4) is
L(π̃, t) = λDKL(π1‖π̃) + (1− λ)DKL(π2‖π̃) + t ∗ ( K∑ i=1 π̃i − 1)
and the related KKT conditions are
∂L(π̃, t) ∂π̃ = t− λπ1 + (1− λ)π2 π̃ = 0
t ∗ ( K∑ i=1 π̃i − 1) = 0 (A.5)
Solve the equation (A.5), we can get the closed form of π̃ as
π̃ = λπ1 + (1− λ)π2

heading: A.5 SCHEDULE ANALYSIS
text: The warm-up scheduler aims to slowly increase the parameters until they reach their maximum. For a certain hyperparameter h, there are 2 parameters control its warm-up process, the target value hmax and the total epoch tmax to reach the target value. We use the exponential function to get the middle value ht as
ht = hmax × exp [−5 ∗ (1−min(1, t/tmax))2] (A.6)
The curve of (A.6) is shown in Figure 2, and we list the scheduler parameters of 4 benchmark datasets, i.e. MNIST, SVHN, Cifar10, Cifar100, in Table 6.

heading: A.6 DETAILS OF TRAINING PROCESS
text: Here we list some important items need to set in the training process for different process. We classify these items into 2 categories: (1) items related to the loss function and (2) items related to the optimization strategy. The details are as follows and we list the exact value in Table 7:
1. Items related to the loss function
• Latent dim The latent dim of discrete variable c is the same as the number of classifications, that is, K. For continuous variable z, the latent dim is determined by experiments.
• Mutual information We find the value of continuous mutual information will affect the generative performance, but have little impact on semi-supervised learning results, so we choose a suitable value to get the best generative performance. For discrete information, we choose the value in the ideal scene, that is, Ic = log(K).
• Calculation of Eqφ(z|X),qφ(c|X) − log pθ(X|c, z) pθ(X|c, z) has two forms: (1) normal distribution with N (fθ(c, z), I) and (2) multinomial distribution with Mult(dim(X), fθ(c, z)). For the two forms, the loss function of − log pθ(X|c, z) is mean square error(MSE) and binary cross entropy(BCE) respectively. We use reparameterization trick in (20) to approximate expectation, and the sampling frequency is 1.
• α of pmixup(X) The β(α, α) for mixup vicinal distribution will strongly affect SSL performance. We set it to 2 in all experiments.
2. Items related to the optimization strategy
• Optimizer In one-stage VAE, we use Adam. In OSPOT-VAE, we use SGD with momentum 0.9.
• Learning rate We set the initial learning rate to 0.1 in OSPOT-VAE, which obtains best SSL performance.
• The scheduler of adjusting learning rate We decay the learning ratio in some milestones with the specified decay rate.
• Weight decay Weight decay controls the strength of L2 regularization.
We also list some standard training curves on benchmark datasets in Figure 3 and 4.
Left: classification performance. Right: Eqφ(z|X),qφ(c|X) − log pθ(X|c, z)
Left: classification performance. Right: Eqφ(z|X),qφ(c|X) − log pθ(X|c, z)

heading: A.7 GENERATIVE PERFORMANCE
text: Following figures5-7 show the generation performance of our OSPOT-VAE.

Review summaries:
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.