Learning Consistent Deep Generative Models from Sparse Data via Prediction Constraints

We develop a new framework for learning variational autoencoders and other deep generative models that balances generative and discriminative goals . Our framework optimizes model parameters to maximize a variational lower bound on the likelihood of observed data , subject to a task-specific prediction constraint that prevents model misspecification from leading to inaccurate predictions . We further enforce a consistency constraint , derived naturally from the generative model , that requires predictions on reconstructed data to match those on the original data . We show that these two contributions – prediction constraints and consistency constraints – lead to promising image classification performance , especially in the semi-supervised scenario where category labels are sparse but unlabeled data is plentiful . Our approach enables advances in generative modeling to directly boost semi-supervised classification performance , an ability we demonstrate by augmenting deep generative models with latent variables capturing spatial transformations . 1 INTRODUCTION . We develop broadly applicable methods for learning flexible models of high-dimensional data , like images , that are paired with ( discrete or continuous ) labels . We are particularly interested in semisupervised learning ( Zhu , 2005 ; Oliver et al. , 2018 ) from data that is sparsely labeled , a common situation in practice due to the cost or privacy concerns associated with data annotation . Given a large and sparsely labeled dataset , we seek a single probabilistic model that simultaneously makes good predictions of labels and provides a high-quality generative model of the high-dimensional input data . Strong generative models are valuable because they can allow incorporation of domain knowledge , can address partially missing or corrupted data , and can be visualized to improve interpretability . Prior approaches for the semi-supervised learning of deep generative models include methods based on variational autoencoders ( VAEs ) ( Kingma et al. , 2014 ; Siddharth et al. , 2017 ) , generative adversarial networks ( GANs ) ( Dumoulin et al. , 2017 ; Kumar et al. , 2017 ) , and hybrids of the two ( Larsen et al. , 2016 ; de Bem et al. , 2018 ; Zhang et al. , 2019 ) . While these all allow sampling of data , a major shortcoming of these approaches is that they do not adequately use labels to inform the generative model . Furthermore , GAN-based approaches lack the ability to evaluate the learned probability density function , which can be important for tasks such as model selection and anomaly detection . This paper develops a framework for training prediction constrained variational autoencoders ( PC-VAEs ) that minimize application-motivated loss functions in the prediction of labels , while simultaneously learning high-quality generative models of the raw data . Our approach is inspired by the prediction-constrained framework recently proposed for learning supervised topic models of “ bag of words ” count data ( Hughes et al. , 2018 ) , but differs in four major ways . First , we develop scalable algorithms for learning a much larger and richer family of deep generative models . Second , we capture uncertainty in latent variables rather than simply using point estimates . Third , we allow more flexible specification of loss functions . Finally , we show that the generative model structure leads to a natural consistency constraint vital for semi-supervised learning from very sparse labels . Our experiments demonstrate that consistent prediction-constrained ( CPC ) VAE training leads to prediction performance competitive with state-of-the-art discriminative methods on fully-labeled datasets , and excels over these baselines when given semi-supervised datasets where labels are rare . 2 BACKGROUND : DEEP GENERATIVE MODELS AND SEMI-SUPERVISION . We now describe VAEs as deep generative models and review previous methods for semi-supervised learning ( SSL ) of VAEs , highlighting weaknesses that we later improve upon . We assume all SSL tasks provide two training datasets : an unsupervised ( or unlabeled ) dataset DU of N feature vectors x , and a supervised ( or labeled ) dataset DS containing M pairs ( x , y ) of features x and label y ∈ Y . Labels are often sparse ( N M ) and can be discrete or continuous . 2.1 UNSUPERVISED GENERATIVE MODELING WITH THE VAE . The variational autoencoder ( Kingma & Welling , 2014 ) is an unsupervised model with two components : a generative model and an inference model . The generative model defines for each example a joint distribution pθ ( x , z ) over “ features ” ( observed vector x ∈ RD ) and “ encodings ” ( hidden vector z ∈ RC ) . The “ inference model ” of the VAE defines an approximate posterior qφ ( z | x ) , which is trained to be close to the true posterior ( qφ ( z | x ) ≈ pθ ( z | x ) ) but much easier to evaluate . As in Kingma & Welling ( 2014 ) , we assume the following conditional independence structure : pθ ( x , z ) = N ( z | 0 , IC ) · F ( x | µθ ( z ) , σθ ( z ) ) , qφ ( z | x ) = N ( z | µφ ( x ) , σφ ( x ) ) . ( 1 ) The likelihood F is often multivariate normal , but other distributions may give robustness to outliers . The ( deterministic ) functions µθ and σθ , with trainable parameters θ , define the mean and covariance of the likelihood . Given any observation x , the posterior of z is approximated as normal with mean µφ and ( diagonal ) covariance σφ parameterized by φ . These functions can be represented as multi-layer perceptrons ( MLPs ) , convolutional neural networks ( CNNs ) , or other ( deep ) neural networks . We would ideally learn generative parameters θ by maximizing the marginal likelihood of features x , integrating latent variable z . Since this is intractable , we instead maximize a variational lower bound : max θ , φ ∑ x∈D LVAE ( x ; θ , φ ) , LVAE ( x ; θ , φ ) = Eqφ ( z|x ) [ log pθ ( x , z ) qφ ( z|x ) ] ≤ log pθ ( x ) . ( 2 ) This expectation can be evaluated via Monte Carlo samples from the inference model qφ ( z|x ) . Gradients with respect to θ , φ can be similarly estimated by the reparameterization “ trick ” of representing qφ ( z | x ) as a linear transformation of standard normal variables ( Kingma & Welling , 2014 ) . Throughout this paper , we denote variational parameters by φ . Because the factorization of q changes for more complex models , we will write φz|x to denote the parameters specific to factor q ( z|x ) . 2.2 TWO-STAGE SSL : MAXIMIZE FEATURE LIKELIHOOD THEN TRAIN PREDICTOR . One way to employ the VAE for a semi-supervised task is a two-stage “ VAE-then-MLP ” . First , train a VAE to maximize the unsupervised likelihood ( 2 ) of all observed features x ( both labeled DS and unlabeled DU ) . Second , we define a label-from-code predictor ŷw ( z ) that maps each learned code representation z to a predicted label y ∈ Y . We use an MLP with weights w , though any predictor could do . Let ` S ( y , ŷ ) be a loss function , such as cross-entropy , appropriate for the prediction task . We train the predictor to minimize the loss : minw ∑ x , y∈DS Eqφ ( z|x ) [ ` S ( y , ŷw ( z ) ) ] . Importantly , this second stage uses only the small labeled dataset and relies on fixed parameters φ from stage one . While “ VAE-then-MLP ” is a simple common baseline ( Kingma et al. , 2014 ) , it has a key disadvantage : Labels are only used in the second stage , and thus a misspecified generative model in stage one will likely produce inferior predictions . Fig . 1 illustrates this weakness . 2.3 SEMI-SUPERVISED VAES : MAXIMIZE JOINT LIKELIHOOD OF LABELS AND FEATURES . To overcome the weakness of the two-stage approach , previous work by Kingma et al . ( 2014 ) presented a VAE-inspired model called “ M2 ” focused on the joint generative modeling of labels y and data x. M2 has two components : a generative model pθ ( x , y , z ) and an inference model qφ ( y , z | x ) . Their generative model is factorized to sample labels ( with frequencies π ) first , and then features x : pθ ( x , y , z ) = N ( z | 0 , IC ) · Cat ( y | π ) · F ( x | µθ ( y , z ) , σθ ( y , z ) ) . ( 3 ) The M2 inference model sets qφ ( y , z | x ) = qφy|x ( y | x ) qφz|x , y ( z | x , y ) , where φ = ( φy|x , φz|x , y ) . To train M2 , Kingma et al . ( 2014 ) maximize the likelihood of all observations ( labels and features ) : max θ , φy|x , φz|x , y ∑ x , y∈DS LS ( x , y ; θ , φz|x , y ) + ∑ x∈DU LU ( x ; θ , φy|x , φz|x , y ) . ( 4 ) The first , “ supervised ” term in Eq . ( 4 ) is a variational bound for the feature-and-label joint likelihood : LS ( x , y ; θ , φz|x , y ) = Eq φz|x , y ( z|x , y ) [ log pθ ( x , y , z ) q φz|x , y ( z|x , y ) ] ≤ log pθ ( x , y ) . ( 5 ) The second , “ unsupervised ” term is a variational lower bound for the features-only likelihood log pθ ( x ) ≥ LU , where LU = Eqφ ( y , z|x ) [ log pθ ( x , y , z ) qφ ( y , z|x ) ] can be simply expressed in terms of LS : LU ( x ; θ , φy|x , φz|x , y ) = ∑ y∈Y qφy|x ( y | x ) ( LS ( x , y ; θ , φz|x , y ) − log qφy|x ( y | x ) ) . ( 6 ) As with the unsupervised VAE , both terms in the objective can be computed via Monte Carlo sampling from the variational posterior , and gradients can be estimated via the reparameterization trick . M2 ’ s prediction dilemma and heuristic fix . After training parameters θ , φ , we need to predict labels y given test data x. M2 ’ s structure assumes we make predictions via the inference model ’ s discriminator density qφy|x ( y | x ) . However , the discriminator ’ s parameter φy|x is only informed by the unlabeled data when using the objective above ( it is not used to compute LS ) . We can not expect accurate predictions from a parameter that does not touch any labeled examples in the training set . To partially overcome this issue , Kingma et al . ( 2014 ) and later work use a weighted objective : max θ , φ ∑ x , y∈DS ( α log qφy|x ( y | x ) + λLS ( x , y ; θ , φz|x , y ) ) + ∑ x∈DU LU ( x ; θ , φy|x , φz|x , y ) . ( 7 ) This objective biases the inference model ’ s discriminator to do well on the labeled set via an extra loss term ( weighted by hyperparameter α > 0 ) . We can further include λ > 0 to balance the supervised and unsupervised terms . Originally , Kingma et al . ( 2014 ) fix λ = 1 and tune α to achieve good performance . Later , Siddharth et al . ( 2017 ) tuned λ to improve performance . Maaløe et al . ( 2016 ) used this same α log q ( y | x ) term for labeled data to train VAEs with auxiliary variables . Disadvantage : What Justification ? While the LS and LU terms in Eq . ( 7 ) have a rigorous justification as maximizing the data likelihood under the assumed generative model , the first term ( α log q ( y | x ) ) is not justified by the generative or inference model . In particular , suppose the training data were fully labeled : we would ignore the LU terms altogether , and the remaining terms would decouple the parameters θ , φz|x , y from the discriminator parameters φy|x . This is deeply unsatisfying : We want a single model guided by both generative and discriminative goals , not two separate models . Even in partially-labeled scenarios , including this α term does not adequately balance generative and discriminative goals , as we demonstrate in later examples . An overly flexible yet misspecified generative model may go astray and compromise predictions . Disadvantage : Runtime Cost . Another disadvantage is that the computation of LU in Eq . ( 6 ) is expensive . If labels are discrete , computing this sum exactly is possible but requires a sum over all L = |Y| possible class labels , computing a Monte Carlo estimate of LS for each one . While Monte Carlo approximations can avoid the explicit sum in Eq . ( 6 ) , they may make gradients too noisy . Extensions . Siddharth et al . ( 2017 ) showed how LS and LU could be extended to any desired conditional independence structure for qφ ( y , z | x ) , generalizing the label-then-code factorization qφ ( y | x ) qφ ( z | x , y ) of Kingma et al . ( 2014 ) . While importance sampling leads to likelihood bounds , the overall objective still has two undesirable traits . First , it is expensive , requiring either marginalization of y to compute LU in Eq . ( 6 ) or marginalization of z to compute q ( y|x ) = ∫ qφ ( y , z|x ) dz . Second , the approach requires the heuristic inclusion of the discriminator loss α log q ( y | x ) . While recent parallel work by Gordon & Hernández-Lobato ( 2020 ) also tries to improve SSL for VAEs , their approach couples discriminative and generative terms only distantly through a joint prior over parameters and still requires expensive sums over labels when computing generative likelihoods .

Summary: The paper proposes a framework for semi-supervised settings to leverage both unlabeled data and (limited) labeled data where VAEs are trained subject to regularization terms from label information. More specifically, the proposed method trains a VAE and a NN classifier simultaneously by optimizing an objective that consists of the usual (unsupervised) variational lower bound, classification error for the labeled data based on the latent space, and consistency term for all data encouraging the same prediction for latent representations corresponding to the original and reconstructed version of a data point. The proposed method is compared with a few other deep generative semi-supervised learning methods on three image datasets.