Metropolis-CVAE: Bootstrapping Labels for Bayesian Inference via Semi-Supervised Conditional Variational Autoencoders

In Bayesian parameter estimation , models make simplifying assumptions to make1 parameter inference feasible . If learned inference methods are trained using data2 simulated by models , however , distributional differences between simulated and3 observed data may lead to biased inference results on the observed data . In this4 work , we introduce a semi-supervised learned Bayesian inference method which5 makes use of both simulated data – for which the underlying parameters are known6 by construction – and unlabeled data , which may depend on nuissance parameters7 not captured by the simulation procedure . A conditional variational autoencoder8 ( CVAE ) is trained to perform approximate inference simultaneously on the sets of9 labeled simulated data and unlabeled data , where the unlabeled data is initialized10 with arbitrary pseudo labels . At each training iteration , new candidate pseudo11 labels are drawn from the CVAE posterior and the pseudo labels are updated using12 the Metropolis-Hastings algorithm . This process results in a Markov chain of13 bootstrapped pseudo labels for each unlabeled datum , effectively performing online14 Markov chain Monte Carlo ( MCMC ) inference wherein the proposal distribution15 is a CVAE informed by labeled simulated data , producing proposals which are16 increasingly likely to be accepted as training proceeds . The resulting CVAE is17 able to efficiently produce samples from the posterior distributions of both the18 simulated and unlabeled data , implicitly marginalizing over nuissance parameters19 in the unlabeled data . We demonstrate the effectiveness of this method in magnetic20 resonance imaging ( MRI ) where MCMC is computationally impractical to due the21 ( 3+1 ) D nature of the images , showing improvement against traditional MCMC22 inference in both speed and posterior quality.23 1 Introduction24 Bayesian parameter estimation methods are robust techniques for quantifying properties of a system25 that can not be observed directly [ 1 ] . In order to estimate such parameters , one first needs to develop26 a model of the phenomena to be studied . This process requires deep domain-specific knowledge . For27 all but the most basic of systems , a series of simplifying assumptions on the system are required to28 make parameter inference tractable . Typically , tractable is synonymous with not being unreasonably29 expensive to compute . Examples of such methodologies include perturbation theory , in which models30 containing Taylor series expansions drop higher order terms ; mean-field theory , in which interactions31 involving many degrees of freedom are replaced by averaged approximations ; and ( stochastic ) 32 differential equations , when used as continuous limits of discrete stochastic processes . The price one33 pays for utilizing a given approximation is highly problem dependent . When approximate models34 are used for parameter inference , the mismatch between model predictions and data may propagate35 Submitted to 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Do not distribute . through the inference process and lead to bias and other misestimations of the inferred parameters . In36 this work , we are interested in the following research question : can a deep learning model be trained37 to perform inference , while incorporating both labeled data generated from a well-understood model,38 and unlabeled data which contains additional complex structure which can not be modeled ? 39 Here , we present Metropolis conditional variational autoencoders ( Metropolis-CVAEs ) . Metropolis-40 CVAEs combine traditional CVAEs with the Metropolis-Hastings Markov chain monte carlo ( MCMC ) 41 inference method . The resulting networks combine the ability of CVAEs to learn to perform rapid42 approximate Bayesian inference from labeled data with traditional MCMC methodology which43 requires only a likelihood model and a prior distribution over the parameters of interest . The44 Metropolis-CVAE initializes the unlabeled data with pseudo labels drawn from the prior distribution45 and , informed by the labeled data , iteratively improves the pseudo labels throughout training . We46 demonstrate the effectiveness of the Metropolis-CVAE network compared to traditional MCMC47 methods on an inference problem from magnetic resonance imaging ( MRI ) which is inherently48 computationally challenging due to the ( 3+1 ) D nature of the data.49 2 Methods50 2.1 Related work51 In the pioneering work by Sohn et al . [ 2 ] , conditional variational autoencoders were introduced52 by considering a variational lower bound to the conditional log-likelihood log p ( y|x ) of labels y53 given corresponding data x . Ideally , to perform data-driven approximate inference , one would like54 to train a network to learn to maximize the conditional log-likelihood directly . However , this is55 known to be an intractable problem . To mitigate this issue , the stochastic gradient variational Bayes56 ( SGVB ) framework is employed . In SGVB , a latent space which factorizes the marginal likelihood is57 introduced via p ( y|x ) = ∫ z dz p ( z|x ) p ( y|z , x ) , as well as a recognition distribution q ( z|x , y ) . The58 conditional log-likelihood is then maximized indirectly via maximizing the variational lower bound59 log p ( y|x ) ≥−KL ( q ( z|x , y ) || p ( z|x ) ) + Eq ( z|x , y ) [ log p ( y|x , z ) ] : =−L ( x , y ) ( 1 ) where KL is the Kullback-Leibler divergence . The interpretation as a conditional autoencoder is60 as follows : let p ( z|x ) , q ( z|x , y ) , and p ( y|x , z ) be parameterized by deep neural networks E1 ( z|x ) ,61 E2 ( z|x , y ) , and D ( y|x , z ) , respectively . Then , E1 and E2 can be viewed as encoders which map62 their inputs into distributions over the latent variables z . D can be viewed as a decoder which maps63 stochastic latent representations z and data x into posterior distributions over the labels y . Hence , y64 is conditionally autoencoded via the encoder-decoder pipeline z ∼ E2 ( z|x , y ) → y ∼ D ( y|x , z ) .65 At inference time , posterior samples are similarly drawn via z ∼ E1 ( z|x ) → y ∼ D ( y|x , z ) . The66 CVAE – the triplet of networks ( E1 , E2 , D ) – is trained by minimizing L over the CVAE parameters,67 thereby maximizing the variational lower bound on log p ( y|x ) .68 The CVAE approach by Sohn et al . [ 2 ] is designed for supervised learning problems . Earlier work of69 a similar vein by Kingma et al . [ 3 ] introduces a semi-supervised framework in which labeled data70 ( x ` , y ) is used to infer labels for unlabeled data xu via the minimization of a two-term loss function.71 The first term of the loss function is a supervised loss over ( x ` , y ) samples , similar to Equation 1,72 derived from a variational lower bound on the joint log-likelihood log p ( x ` , y ) . The second term is an73 unsupervised loss over xu samples which treats label inference as a data imputation task . Specifically,74 the unknown label is treated as a parameter over which posterior inference is performed ; the resulting75 loss is a variational lower bound on the log-likelihood log p ( xu ) .76 While the semi-supervised method of Kingma et al . is an elegant approach to discovering labels,77 it is not quite suitable for inferring labels for out of distribution data . This is due to the implicit78 assumption that relationships between x ` and y learned from the joint distribution p ( x ` , y ) generalize79 to data from p ( xu ) . In fact , this assumption is made explicit via an extended objective function which80 adds a regularization term Ep̃ ` ( x , y ) [ −log q ( y|x ` ) ] over the empirical distribution p̃ ` ( x , y ) of labeled81 data . This penalty encourages the learned posterior distribution q ( y|x ) to generate labels for xu by82 extrapolating from the relationships it discovers between ( x ` , y ) pairs . This will naturally lead to83 biased labels for xu when the distribution underlying xu differs from that of x ` .84 Gabbard et al. , who made use of CVAEs to accelerate inference for an application in gravitational85 wave astronomy [ 4 ] , presented an alternate view of Equation 1 . Gabbard et al . begin by aiming to86 minimize the expected cross-entropy87 H : = − Ep ( x ) [ ∫ y dy p ( y|x ) log p̂ ( y|x ) ] ( 2 ) over the data distribution p ( x ) between the true posterior p ( y|x ) and approximate posterior p̂ ( y|x ) .88 Employing the SGVB framework and letting p̂ ( y|x ) = ∫ z dz p̂ ( z|x ) p̂ ( y|z , x ) , it follows from 1 that89 H ≤ Ep ( x ) [ ∫ y dy p ( y|x ) L ( x , y ) ] . ( 3 ) Applying Bayes ’ theorem , we equivalently have that90 H ≤ Ep ( x ) Ep ( y|x ) [ L ( x , y ) ] ( 4 ) = Ep ( y ) Ep ( x|y ) [ L ( x , y ) ] ( 5 ) = Ep ( x , y ) [ L ( x , y ) ] . ( 6 ) Therefore , maximizing the variational lower bound to log p ( y|x ) over a dataset of ( x , y ) pairs , as in91 Equation 1 , is equivalent to minimizing the expected cross-entropy via Equation 6 . The interpretation92 of Equations 4 and 5 , however , will prove useful for remedying the issue of label inference for out of93 distribution data.94 2.2 Theoretical contributions95 Equation 5 is a natural framework for using CVAEs to perform inference on simulated data with96 known labels [ 4 ] . Suppose labels y ∼ p ( y ) are sampled from a prior distribution and x ∼ p ( x|y ) 97 is subsequently given by a ( possibly stochastic ) model function x = f ( y ) . Then , Equation 598 corresponds to minimizing the average CVAE loss L ( x , y ) over pairs of simulated data ( x=f ( y ) , y ) .99 The novel contribution of this work stems from of Equation 4 . Using this formulation directly would100 require sampling from the posterior p ( y|x ) , which is our stated objective . However , this can be101 circumvented by making the observation that if we could construct a Markov chain of pseudo labels102 ỹ during training , such that the stationary distribution of the sequence ( ỹk ) k∈N was p ( y|x ) , then103 Equation 4 could be approximated as104 Ep ( x ) Ep ( y|x ) [ L ( x , y ) ] ≈ Ep ( x ) [ 1 Lc Lc−1∑ i=0 L ( x , ỹn−i ) ] ( 7 ) where ( ỹn−Lc+1 , . . . , ỹn ) are the Lc most recent samples in the Markov chain . This approach to105 distribution sampling – Markov chain Monte Carlo ( MCMC ) sampling – is considered the gold106 standard in parameter inference.107 In this work , we consider the Metropolis-Hastings ( MH ) algorithm [ 5 , 6 ] . In the context of Bayesian108 inference for recovering labels y from data x , MH sampling begins with a prior distribution p ( y ) , a109 likelihood function p ( x|y ) , and a proposal distribution Q ( y′|y ) which quantifies the probability of110 transitioning from y to y′ in the space of possible labels . Given a sample yn of a Markov chain , the111 MH update rule is given by112 y′ ∼ Q ( y′|yn ) α = min ( 1 , p ( y′ ) p ( yn ) · p ( x|y ′ ) p ( x|yn ) ) u ∼ Uniform ( 0 , 1 ) yn+1 = { y′ u ≤ α yn otherwise . ( 8 ) The proposal distribution Q ( y′|y ) is a free parameter of the MH algorithm . In general , choosing Q113 to resemble the true posterior as closely as possible improves the efficiency of the MH algorithm.114 Therefore , we propose to use the approximate posterior p̂ ( y|x ) of the CVAE itself as the proposal115 distribution . Recalling that p̂ ( y|x ) = ∫ z dz p̂ ( z|x ) p̂ ( y|z , x ) : = ∫ z dzE1 ( z|x ) D ( y|z , x ) , we let116 Q ( y′|yn ) = Q ( y′ ) = ∫ z dzE1 ( z|x ) D ( y′|z , x ) ( 9 ) ≈ 1 Lz Lz∑ i=1 D ( y′|zi , x ) where zi ∼ E1 ( z|x ) . ( 10 ) By choosing this proposal function , we are able to bootstrap pseudo labels ỹ onto unlabeled data xu.117 In particular , we minimize the semi-supervised hybrid loss118 Lhybrid = Lsuper + Lself ( 11 ) Lsuper = E ( x ` , y ) ∼p̃ ` ( x , y ) [ L ( x ` , y ) ] ( 12 ) Lself = E ( xu , ỹ ) ∼p̃u ( x , ỹ ) [ L ( xu , ỹ ) ] ( 13 ) where p̃ ` ( x , y ) and p̃u ( x , ỹ ) are the empirical distributions over the pairs of labeled data ( x ` , y ) 119 and unlabeled data with bootstrapped pseudo labels ( xu , ỹ ) , respectively . The pseudolabels ỹ are120 initialized uniformly from the prior space p ( y ) and are updated according to the MH algorithm 8 at121 each training iteration.122 Note that in the MH update step , the MH acceptance ratio α can be interepreted as computing a123 Bayesian goodness of fit check relative to the current label ỹn . Therefore , p̃ ` ( x ) and p̃u ( x ) need124 not be identical distributions , merely close enough such that decreasing the supervised loss Lsuper125 improves the proposal quality for the self-supervised loss Lself early in training.126 2.3 MRI physics application127 We consider an application in magnetic resonance imaging . In MRI , data is typically acquired in128 the form of ( 3+1 ) D spatio-temporal grids , with 1D magnetic resonance time signals measured in129 each voxel of the three spatial dimensions . Advanced imaging methods typically involve voxelwise130 parameter inference for each time signal for the computation of quantitative maps . Modeling the131 individual time signals for inference , however , is challenging due to imperfections in the magnetic132 field generated by the scanner and other sources of signal corruption . Effectively , in MRI there is a133 distributional mismatch problem : simulated data p̃ ` ( x ) , with labels corresponding to well understood134 physics parameters , does not contain the full distribution of measured data p̃u ( x ) which , while135 generated by the same physics in principle , depends on additional nuissance parameters which cannot136 be modeled . Machine learning models which intended to generalize to p̃u ( x ) should therefore not be137 trained only on data from p̃ ` ( x ) .138 In this work , we acquire multi spin-echo ( MSE ) MRI images . The MSE time signals are modeled139 using a two-component extended phase graph ( EPG ) model using the algorithm detailed in Prasloski140 et al . [ 7 ] . Using this model , the j-th time point for each signal is given by141 x̂j = f ( j · TE ) f ( t ) = 2∑ ` =1 A ` EPG ( t , α , β , T2 , ` , T1 ) ( 14 ) where α is the spin flip angle , β the refocusing control angle , A1 , A2 the component amplitudes,142 T2,1 ≤ T2,2 the short and long transverse relaxation times , and T1 the longitudinal relaxation time.143 The EPG ( t , . . . ) terms are approximately exponentially decaying in t with time constants T2 , ` , with144 additional modifications due to MRI physics determined by α , β , and T1 . The echo time TE is the145 uniform spacing between time points . We reparameterize A1 , A2 , T2,1 , and T2,2 in terms of the146 unconstrained parameters η , δ1 , and δ2 as follows : A1 = η , A2 = 1 − η , log T2,1 = log T2 , min +147 ( log T2 , max−log T2 , min ) ·δ1 , and log T2,2 = log T2 , min+ ( log T2 , max−log T2 , min ) · ( δ1+δ2 · ( 1−δ1 ) ) ,148 where T2 , min = 10 ms and T2 , max = 1 s. The longitudinal relaxation time is fixed T1 = 1.0 s.149 MRI signal noise can be modeled as Rician [ 8 , 9 ] . Given data x normalized to have maximum value150 1 , the likelihood p ( x|y ) under Rician noise is given by151 log p ( x|y ) = Nx∑ j=1 log pRice ( xj ∣∣∣∣ s · x̂ ( θ ) jmaxk x̂ ( θ ) k , s · ) , ( 15 ) where pRice ( ξ|ν , σ ) = ξ σ2 exp ( −ξ 2 + ν2 2σ2 ) I0 ( ξν σ2 ) ( 16 ) is the Rician probability density function with location parameter ν and scale parameter σ ; I0 is152 the modified Bessel function of the first kind with order zero . We have introduced two additional153 parameters : a scale parameter s to account for signal normalization , and a noise level relative to this154 scale . Additionally , we denote θ = ( α , β , η , δ1 , δ2 ) the parameters of the EPG model 14.155 In total , there are 7 labels to be inferred : y = ( α , β , η , δ1 , δ2 , log , log s ) . We place the following156 priors on the parameters:157 α ∼ TN ( 180◦ , 45◦ , 90◦ , 180◦ ) β ∼ TN ( 180◦ , 45◦ , 90◦ , 180◦ ) η ∼ TN ( 0.0 , 0.5 , 0.0 , 1.0 ) δ1 ∼ TN ( 0.0 , 0.5 , 0.0 , 1.0 ) δ2 ∼ TN ( 1.0 , 0.5 , 0.0 , 1.0 ) log ∼ U ( log 10−5 , log 10−1 ) log s ∼ TN ( 0.0 , 0.5 , −2.5 , 2.5 ) ( 17 ) where TN ( µ , σ , a , b ) is the normal distribution with parameters ( µ , σ ) truncated to the interval [ a , b ] ,158 and U ( a , b ) is the uniform distribution on [ a , b ] . The priors were chosen to align with the expectations159 that : α , β are typically near 180◦ ; the short component amplitude η is typically less than the long160 amplitude 1− η ; δ1 and δ2 should prefer to represent the shortest and longest components ; the noise161 level is chosen uniformly from signal-to-noise ratios between 20 and 100 ; the scale parameter s162 should prefer to be 1.163 3 Experiments164 3.1 Data sets165 MRI data The MRI data used for this study consists of two anonymized brain scans acquired using166 a Carr-Purcell-Meiboom-Gill ( CPMG ) [ 10 , 11 ] multi spin-echo sequence [ 12 ] . The first data set,167 denoted p̃u,1 ( x ) , contains signals with Nx,1 = 48 samples at times ti = i · TE , with echo spacing168 TE = 8 ms , repetition time TR = 1073 ms , matrix size 240× 240× 48 , and spatial resolution169 0.96× 0.96× 2.5 mm3 . The second data set , denoted p̃u,2 ( x ) , contains signals with Nx,2 = 56170 samples at times ti = i · TE , with echo spacing TE = 7 ms , repetition time TR = 1066 ms,171 matrix size 240× 240× 113 , and spatial resolution 1.0× 1.0× 3.0 mm3 . Following the extraction172 of image volumes containing the brain , p̃u,1 ( x ) and p̃u,2 ( x ) contain 821 145 and 1 265 306 signals,173 respectively . MRI data was acquired on a 3 T MR system ( Ingenia Elition , Philips Medical Systems,174 Best , The Netherlands ) from healthy volunteers giving written and informed consent , and approved175 by our university ethics board.176 Simulated data Using the EPG physics model 14 with Rician noise , we consider two simulated177 data sets . First , the labeled data set p̃ ` ( x ) which is generated on demand during training using labels178 y ∼ p ( y ) drawn from the prior distributions 17 . Second , a precomputed simulated data set p̃u,3 ( x ) 179 used for validation of the method , where the labels are held out during training . All simulated signals180 are generated with Nx,3 = 64 samples and TE = 10 ms.181 MCMC data MCMC is performed using the No-U-Turn Sampler [ 13 ] algorithm to generate182 posterior samples ŷ ∼ p ( y|xu ) which can be compared with the ( Metropolis- ) CVAE posterior183 samples . MCMC was performed twice : once to draw 100 posterior samples for every time signal in184 p̃u,1 ( x ) , p̃u,2 ( x ) , and p̃u,3 ( x ) , and once to draw 3000 posterior samples for a subset of 5000 signals185 from the training , validation , and testing partitions of all three data sets . In total , performing the186 above MCMC analysis took approximately 72 h using an AMD Ryzen 9 3950X 16-Core CPU.187 3.2 Model architecture188 CVAE components Let x ∈ RNx , y ∈ RNy , and z ∈ RNz be input data , corresponding labels,189 and latent space samples , respectively . The encoders E1 and E2 are chosen to be multivariate normal190 distributions : E1 ( z|x ) = N ( µz1 , σz1 ) and E2 ( z|x , y ) = N ( µz2 , σz2 ) , where µz1 , µz2 ∈ RNz and191 σz1 , σz2 ∈ R Nz + . Similarly , the decoder is given byD ( y|x , z ) = TN ( µy , σy , 0 , 1 ) , where µy ∈ RNy192 and σy ∈ R Ny + parameterize independent multivariate normal distributions truncated to [ 0 , 1 ] . The193 labels y are scaled linearly from the prior domains 17 to [ 0 , 1 ] Ny in order to better condition the194 network during training . Similarly , as we are interested only in the relative x values and not their195 absolute scale , inputs x are normalized to [ 0 , 1 ] Nx .196 Each of the E1 , E2 , and D networks are composed of fully connected layers with ReLU activation197 functions and H = 2 hidden layers , with hidden dimension NH = 512 , for a total of H + 2 = 4198 layers . The dimensions of the data , labels , and latent space are Nx = 64 , Ny = 7 , and Nz = 12,199 respectively . The encoder networks output µz1 , log σz1 , µz2 , and log σz2 vectors . In order to avoid200 latent space collapse during early training stages – that is , one or both of |µzi | , | log σzi | → ∞ –201 each µzi was bounded to ( −3 , 3 ) using the activation function x→ 3 tanh ( x ) , and each log σzi was202 bounded to ( −6 , 0 ) using the activation function x→ 3 tanh ( x ) − 3 . These bounds were determined203 by observing empirical µzi and log σzi values and choosing intervals which clipped only the tails of204 the distributions . The decoder network outputs µy and log σy vectors without further nonlinearities.205 Loss functions The terms Lsuper and Lself in Equation 11 are identical apart from their inputs,206 with each term consisting of the KL-divergence and evidence lower bound ( ELBO ) terms from the207 variational lower bound of Equation 1 . The KL-divergence can be computed in closed closed-form,208 and is given by209 KL ( E2 ( z|x , y ) ||E1 ( z|x ) ) = Nz∑ j=1 σ2z2 , j + ( µz2 , j − µz1 , j ) 2 2σ2z1 , j + log σz1 , j σz2 , j − 1 2 ( 18 ) Following sampling z2 ∼ E2 ( z|x , y ) and subsequently computing ( µy , log σy ) = D ( y|x , z2 ) , the210 ELBO component is approximated as211 EE2 ( z|x , y ) [ logD ( y|x , z ) ] ≈ Ny∑ j=1 { log φ ( yj − µy , j σy , j ) − log ( Φ ( 1− µy , j σy , j ) − Φ ( 0− µy , j σy , j ) ) − log σy , j } ( 19 ) where φ ( ξ ) = 1√ 2π exp ( −1 2 ξ2 ) Φ ( ζ ) = 1 2 ( 1 + erf ( ζ√ 2 ) ) . φ ( ξ ) and Φ ( ζ ) are the probability density and cumulative distribution functions of the standard212 normal distribution , respectively . Note that each term in the sum of Equation 19 is the log-likelihood213 of a normal distribution with parameters ( µy , i , σy , i ) truncated to the unit interval [ 0 , 1 ] .214 3.3 Training215 All data sets were split into training/validation/testing with proportions 50 % /25 % /25 % . Input data216 was padded with zeros to length Nx = 64 , if necessary . To support data with differing unpadded217 lengths , random masking was performed during training : for each x , a random integer jmask was218 sampled from 32–64 and all elements xj > jmask were set to zero . The ADAM optimizer was used with219 an initial learning rate of 10−4 . The learning rate was decreased every 1000 epochs by a factor of √ 10.220 Training was completed after 5000 epochs , where an epoch is defined as 100 iterations . Each iteration,221 a batch of 1024 data and ( pseudo- ) label pairs are drawn from either a labeled or an unlabeled dataset,222 with equal probability . If labeled data are sampled , the loss component Lsuper ( x ` , y ) is descended223 on . If unlabeled data are sampled , the pseudo labels ỹ corresponding to the sampled xu are updated224 using Equation 8 before descending on the loss component Lself ( xu , ỹ ) . In the Metropolis-Hastings225 update step , we set Lc = 1 in equation 7 and Lz = 1 in equation 9 . Training was performed on a226 single Nvidia GeForce RTX 3080 GPU with 10 GB of VRAM ; approximately 15 hours was required227 to train for 5000 epochs.228 4 Results and discussion229 In the first experiment , a Metropolis-CVAE is trained on online simulated labeled data x ` ∼ p̃ ` ( x ) as230 well as precomputed simulated data xu ∼ p̃u,3 ( x ) with labels held out during training . Figure 1 ( A-G ) 231 shows the distributions of prediction errors by the trained network for each label . For each method,232 the prediction error is defined as the difference between the true label and the mean of the posterior233 samples . The labels have been normalized to [ 0 , 1 ] . The histograms for the Metropolis-CVAE with234 100 posterior samples are more tightly clustered around zero than MCMC with either 100 or 3000235 posterior samples – denoted MCMC-100 and MCMC-3000 – in all cases except for one ( Figure 1F ) .236 Figure 1H shows a p-p plot : the fraction of posterior samples greater than or equal to the true label , p,237 is plotted against the cumulative distribution of p-values across the data set ; the Metropolis-CVAE238 produces similar curves as MCMC-3000.239 In the second experiment , a Metropolis-CVAE is similarly trained using labeled simulated x ` ∼ p̃ ` ( x ) ,240 but now with the unlabeled data xu drawn from the MRI data sets p̃u,1 ( x ) and p̃u,2 ( x ) . For241 comparison , a second traditional CVAE is trained on the simulated data only . Figure 2A shows the242 acceptance rate of the proposed pseudo labels ỹ′ ∼ Q ( y′ ) for each data set during training . As the243 Metropolis-CVAE continually learns from both data sets , we find that the network quickly enters244 a negative feedback loop in which the acceptance rate continually increases throughout training.245 Noteworthy is that the acceptance rate of ỹ converges to a value between 0.6 and 0.7 for both datasets.246 The default target acceptance rate of the No-U-Turn sampler is 0.65 , a value which originates from247 a theoretical result pertaining to Hamiltonian Monte Carlo ( HMC ) [ 14 ] . Theoretical work would248 be required to demonstrate a direct connection between Metropolis-CVAEs and the HMC result.249 However , this illustrates that the network neither accepts nor rejects more proposals than is typically250 desired . Figure 2 ( B-C ) shows the average Wasserstein distance between empirical label distributions251 from the Metropolis-CVAE with MCMC-100 and MCMC-3000 , respectively . During training , the252 Metropolis-CVAE quickly reaches minimum Wasserstein distances with respect to MCMC-100.253 Minima are reached with respect to MCMC-3000 similarly quickly . We can extrapolate from this254 result that the Metropolis-CVAE likely produces higher quality posterior samples than MCMC-3000.255 This may be expected , as the Metropolis-CVAE updates its pseudo labels at every training iteration,256 continuously performing MCMC throughout training using the MH update rule 8.257 Figure 3 ( A-G ) compares the empirical distributions of a traditional CVAE trained only on simulated258 data with a Metropolis-CVAE . Histograms of the Wasserstein distances between the two networks and259 MCMC-100 do not show large differences , due to the low-quality of MCMC-100 . Compared with260 MCMC-3000 , parameters which are traditionally easier to infer , such as α and β , show little difference261 between the networks . The CVAE trained on supervised data alone shows significant deviation from262 MCMC-3000 for several parameters , particularly η , δ1 , and log . This illustrates the ability of the263 Metropolis-CVAE to generalize well to the unlabeled test data , as unlabeled training data from p̃u ( x ) 264 has been explicitly incorporated into its training process . Figure 3 ( H ) shows quantile-quantile plots265 for distributions of all label samples across the data sets . All methods agree.266 Limitations A limitation of this method is the requirement of a likelihood function which is fast to267 compute in order to perform the Metropolis-Hastings update step 8 . For physics models , computing268 the likelihood of a set of parameters often involves costly forward simulations of complex models.269 Therefore , computationally intensive models which require solving non-trivial integral or differential270 equations will not be suitable.271 Another limitation is that the pseudo label sampling procedure may converge to a stable local272 minimum which is far from the globally optimal labels . This is inherent to the MH update step ; while273 MCMC methods often provide convergence guarantees in the limit of large numbers of update steps,274 one can not predict how many update steps this will require in practice . However , by training the275 Metropolis-CVAE on both simulated data generated from the prior space as well as real unlabeled276 data , we have not found this to be a practical concern.277 Potential negative societal impacts This work describes a framework for ascribing labels to unla-278 beled data . Potential malicious and unintended uses could occur if this framework were significantly279 extended beyond the MRI physics inference problem which we considered . For example , this method-280 ology could be used to infer missing data generally , with the inferred data then presented under the281 pretense that it were true data . Further , we have not shown that our method of data inference for282 unlabeled data is inherently fair , nor that the clinical MRI data under which the Metropolis-CVAE283 model is trained on can not be recovered from the network weights.284 References285 [ 1 ] Udo von Toussaint . Bayesian inference in physics . Reviews of Modern Physics , 83 ( 3 ) :943–999,286 September 2011. doi : 10.1103/RevModPhys.83.943 . URL https : //link.aps.org/doi/287 10.1103/RevModPhys.83.943 . Publisher : American Physical Society.288 [ 2 ] Kihyuk Sohn , Honglak Lee , and Xinchen Yan . Learning Structured Output Representation289 using Deep Conditional Generative Models . In C. Cortes , N. Lawrence , D. Lee , M. Sugiyama,290 and R. Garnett , editors , Advances in Neural Information Processing Systems , volume 28 . Curran291 Associates , Inc. , December 2015 . URL https : //proceedings.neurips.cc/paper/2015/292 file/8d55a249e6baa5c06772297520da2051-Paper.pdf.293 [ 3 ] Diederik P. Kingma , Danilo J. Rezende , Shakir Mohamed , and Max Welling . Semi-Supervised294 Learning with Deep Generative Models . arXiv:1406.5298 [ cs , stat ] , October 2014 . URL295 http : //arxiv.org/abs/1406.5298 . arXiv : 1406.5298.296 [ 4 ] Hunter Gabbard , Chris Messenger , Ik Siong Heng , Francesco Tonolini , and Roderick297 Murray-Smith . Bayesian parameter estimation using conditional variational autoencoders298 for gravitational-wave astronomy . arXiv:1909.06296 [ astro-ph , physics : gr-qc ] , September299 2019 . URL http : //arxiv.org/abs/1909.06296 . arXiv : 1909.06296.300 [ 5 ] Nicholas Metropolis , Arianna W. Rosenbluth , Marshall N. Rosenbluth , Augusta H. Teller , and301 Edward Teller . Equation of State Calculations by Fast Computing Machines . The Journal of302 Chemical Physics , 21 ( 6 ) :1087–1092 , June 1953 . ISSN 0021-9606. doi : 10.1063/1.1699114.303 URL https : //aip.scitation.org/doi/10.1063/1.1699114 . Publisher : American In-304 stitute of Physics.305 [ 6 ] W. K. Hastings . Monte Carlo sampling methods using Markov chains and their applications.306 Biometrika , 57 ( 1 ) :97–109 , April 1970 . ISSN 0006-3444. doi : 10.1093/biomet/57.1.97 . URL307 https : //doi.org/10.1093/biomet/57.1.97.308 [ 7 ] Thomas Prasloski , Burkhard Mädler , Qing-San Xiang , Alex MacKay , and Craig Jones.309 Applications of stimulated echo correction to multicomponent T2 analysis . Magnetic310 Resonance in Medicine , 67 ( 6 ) :1803–1814 , 2012 . ISSN 1522-2594. doi : 10.1002/311 mrm.23157 . URL https : //onlinelibrary.wiley.com/doi/abs/10.1002/mrm.23157.312 _eprint : https : //onlinelibrary.wiley.com/doi/pdf/10.1002/mrm.23157.313 [ 8 ] S. O . Rice . Mathematical Analysis of Random Noise . Bell System Technical Journal , 23 ( 3 ) :314 282–332 , 1944 . ISSN 1538-7305. doi : 10.1002/j.1538-7305.1944.tb00874.x . URL https:315 //onlinelibrary.wiley.com/doi/abs/10.1002/j.1538-7305.1944.tb00874.x.316 _eprint : https : //onlinelibrary.wiley.com/doi/pdf/10.1002/j.1538-7305.1944.tb00874.x.317 [ 9 ] S. O . Rice . Mathematical Analysis of Random Noise . Bell System Technical Journal , 24 ( 1 ) :318 46–156 , 1945 . ISSN 1538-7305. doi : 10.1002/j.1538-7305.1945.tb00453.x . URL https:319 //onlinelibrary.wiley.com/doi/abs/10.1002/j.1538-7305.1945.tb00453.x.320 _eprint : https : //onlinelibrary.wiley.com/doi/pdf/10.1002/j.1538-7305.1945.tb00453.x.321 [ 10 ] H. Y. Carr and E. M. Purcell . Effects of Diffusion on Free Precession in Nuclear Magnetic322 Resonance Experiments . Physical Review , 94 ( 3 ) :630–638 , May 1954. doi : 10.1103/PhysRev.94.323 630 . URL https : //link.aps.org/doi/10.1103/PhysRev.94.630 . Publisher : American324 Physical Society.325 [ 11 ] S. Meiboom and D. Gill . Modified Spin-Echo Method for Measuring Nuclear Relaxation326 Times . Review of Scientific Instruments , 29 ( 8 ) :688–691 , August 1958 . ISSN 0034-6748. doi:327 10.1063/1.1716296 . URL https : //aip.scitation.org/doi/abs/10.1063/1.1716296.328 Publisher : American Institute of Physics.329 [ 12 ] Kenneth P. Whittall , Alex L. Mackay , Douglas A. Graeb , Robert A. Nugent , David K. B.330 Li , and Donald W. Paty . In vivo measurement of T2 distributions and water contents in331 normal human brain . Magnetic Resonance in Medicine , 37 ( 1 ) :34–43 , 1997 . ISSN 1522-2594.332 doi : 10.1002/mrm.1910370107 . URL https : //onlinelibrary.wiley.com/doi/abs/10.333 1002/mrm.1910370107.334 [ 13 ] Matthew D Hoffman and Andrew Gelman . The No-U-Turn sampler : adaptively setting path335 lengths in Hamiltonian Monte Carlo . J. Mach . Learn . Res. , 15 ( 1 ) :1593–1623 , 2014.336 [ 14 ] Alexandros Beskos , Natesh Pillai , Gareth Roberts , Jesus-Maria Sanz-Serna , and An-337 drew Stuart . Optimal tuning of the hybrid Monte Carlo algorithm . Bernoulli , 19338 ( 5A ) :1501–1534 , November 2013 . ISSN 1350-7265. doi : 10.3150/12-BEJ414.339 URL https : //projecteuclid.org/journals/bernoulli/volume-19/issue-5A/340 Optimal-tuning-of-the-hybrid-Monte-Carlo-algorithm/10.3150/12-BEJ414.341 full . Publisher : Bernoulli Society for Mathematical Statistics and Probability.342 Checklist343 1 . For all authors ... 344 ( a ) Do the main claims made in the abstract and introduction accurately reflect the paper ’ s345 contributions and scope ? [ Yes ] 346 ( b ) Did you describe the limitations of your work ? [ Yes ] See the Limitations paragraph of347 the Results and discussion section.348 ( c ) Did you discuss any potential negative societal impacts of your work ? [ Yes ] See the349 Potential negative societal impacts paragraph of the Results and discussion section.350 ( d ) Have you read the ethics review guidelines and ensured that your paper conforms to351 them ? [ Yes ] 352 2 . If you are including theoretical results ... 353 ( a ) Did you state the full set of assumptions of all theoretical results ? [ N/A ] 354 ( b ) Did you include complete proofs of all theoretical results ? [ N/A ] 355 3 . If you ran experiments ... 356 ( a ) Did you include the code , data , and instructions needed to reproduce the main ex-357 perimental results ( either in the supplemental material or as a URL ) ? [ Yes ] Code to358 reproduce the main results is included with the supplemental material . The MRI data359 can not be included , as it is medical data . However , the simulated data is provided in360 the supplemental material.361 ( b ) Did you specify all the training details ( e.g. , data splits , hyperparameters , how they362 were chosen ) ? [ Yes ] See the Training subsection of the Experiments section.363 ( c ) Did you report error bars ( e.g. , with respect to the random seed after running experi-364 ments multiple times ) ? [ Yes ] We provide entire histograms of relevant error metrics ; 365 see Figures 1 and 3.366 ( d ) Did you include the total amount of compute and the type of resources used ( e.g. , type367 of GPUs , internal cluster , or cloud provider ) ? [ Yes ] See the MCMC data paragraph of368 the Data sets subsection of the Experiments section , as well as the Training subsection369 of the Experiments section.370 4 . If you are using existing assets ( e.g. , code , data , models ) or curating/releasing new assets ... 371 ( a ) If your work uses existing assets , did you cite the creators ? [ N/A ] 372 ( b ) Did you mention the license of the assets ? [ N/A ] 373 ( c ) Did you include any new assets either in the supplemental material or as a URL ? [ Yes ] 374 The trained models are supplied in the supplemental material.375 ( d ) Did you discuss whether and how consent was obtained from people whose data you ’ re376 using/curating ? [ Yes ] See the MRI data paragraph of the Data sets subsection of the377 Experiments section.378 ( e ) Did you discuss whether the data you are using/curating contains personally identifiable379 information or offensive content ? [ Yes ] The MRI data used in this work is anonymized ; 380 see the MRI data paragraph of the Data sets subsection of the Experiments section.381 5 . If you used crowdsourcing or conducted research with human subjects ... 382 ( a ) Did you include the full text of instructions given to participants and screenshots , if383 applicable ? [ No ] This was a prospective study ; this information was not available.384 ( b ) Did you describe any potential participant risks , with links to Institutional Review385 Board ( IRB ) approvals , if applicable ? [ Yes ] These MRI scans were acquired with386 approval from our university ethics board ; see the MRI data paragraph of the Data sets387 subsection of the Experiments section.388 ( c ) Did you include the estimated hourly wage paid to participants and the total amount389 spent on participant compensation ? [ N/A ] 390

Summary: This paper introduces an online MCMC framework for semi-supervised learning with variational autoencoders. The idea of the paper stems from physics problems, where one typically wants to infer certain parameters of a data generating process from empirical data. An example of this is quantitative MRI.   The authors consider the model of data generating process known, and also assume a distribution mismatch between the simulated data, as well as the empirical data, which may include certain confounding factors that may be impossible to model. During the the training of a VAE, the authors progressively improve the psudolabels via a Metropolis-Hastings scheme. The key novelty of the paper is that the proposal distribution for the sampler is an posterior distribution $\hat p(\mathbf{y} | \mathbf{x})$, modeled by the VAE itself.