W YNER VAE: A V ARIATIONAL A UTOENCODER WITH S UCCINCT C OMMON R EPRESENTATION L EARNING
Anonymous authors
Paper under double-blind review
A BSTRACT
A new variational autoencoder (VAE) model is proposed that learns a succinct common representation of two correlated data variables for conditional and joint generation tasks.
The proposed
Wyner VAE model is based on two information the- oretic problems—distributed simulation and channel synthesis—in which Wyner’s common information arises as the fundamental limit of the succinctness of the common representation
The Wyner VAE decomposes a pair of correlated data variables into their common representation (e.g., a shared concept) and local representations that capture the remaining randomness (e.g., texture and style) in respective data variables by imposing the mutual information between the data variables and the common representation as a regularization term
The utility of the proposed approach is demonstrated through experiments for joint and conditional generation with and without style control using synthetic data and real images
Experimental results show that learning a succinct common representation achieves better generative performance and that the proposed model outperforms existing
VAE variants and the variational information bottleneck method.
I NTRODUCTION
This paper aims to develop a new probabilistic framework for generation tasks (i.e., learning the distribution of given data and sampling from the learned distributions) for two high-dimensional random vectors.
To motivate the main idea, consider the following cooperative game between Alice and Bob.
Suppose that given an image of a child’s photo, Alice sends its description Z to Bob who draws a portrait of how the child will grow up based on it.
The objective of this game is to draw a nice portrait, and thus Alice needs to help Bob in the process by providing a good description of the child’s photo — any redundant information in the description may confuse Bob in his guessing process.
What description does Alice need to generate and send from the child’s photo?
P. Cuff (2013) formulated this game of conditional generation as the channel synthesis problem in network information theory depicted in Fig.
Given a joint distribution q ( x , y )
=
q ( x ) q ( y | x ) ,
Alice and Bob want to generate Y according to q ( y | x ) based on a sample from q ( x ) .
In this problem,
Alice wishes to ﬁnd the most succinct description Z of X (a child’s photo) such that Y (her adulthood portrait) can be simulated by Bob according to the desired distribution using this description and local randomness V (new features to draw a portrait of adults that are not contained in photos of children).
The minimum description rate for such conditional generation is characterized by Wyner’s common information (Wyner, 1975; El Gamal and Kim, 2011) denoted by J ( X ; Y ) and deﬁned as the optimal value of the optimization problem minimize
I q ( X , Y ; Z ) subject to
X → Z → Y variables q
( z | x , y ) , (1) where X → Z → Y denotes a Markov chain from X to Z to Y and I q
( X , Y ; Z ) denotes the mutual information between ( X , Y ) and Z , where ( X , Y , Z ) ∼ q ( x , y ) q
( z | x , y ) .
The same quantity J ( X ; Y ) arises as the fundamental limit of the distributed simulation of correlated sources studied originally by A. Wyner (1975) in which two distributed agents wish to simulate a target distribution q ( x , y )
(i.e., joint generation of ( X , Y ) )
based on the least possible amount of
[Figure]
shared common randomness.
(See Fig.
1 (c,d).)
In this sense, the joint distribution q ( x , y ) and the conditional distributions q ( y | x ) , q ( x | y ) have the same common information structure characterized by the optimization problem (1), which involves learning the joint distribution in its nature.
We call the joint encoder q ( z | x , y ) (or equivalently, the corresponding random variable Z ) as the common representation of ( X , Y ) , and the mutual information I q ( X , Y ; Z ) then can be viewed as a measure of the complexity of Z .
Based on these information theoretic observations, we apply the idea of learning succinct common representation to design a new generative model for a pair of correlated variables: seeking a succinct representation Z in learning the underlying distribution based on its sample may also help reduce the burden on the decoder’s side and thereby achieve a better generative performance.
The rest of the paper gradually develops our framework as follows.
We ﬁrst deﬁne a probabilistic model based on the motivating problems, which we aim to train and use for generation tasks (Section 2.1), and then establish a general principle for learning the model based on the optimization problem (1) (Section 2.2).
We propose one instantiation of the principle with a standard variational technique by introducing additional encoder distributions (Section 2.3).
The proposed model with its training method can be viewed as a variant of variational autoencoders (VAEs) (Kingma and Welling, 2014;
Rezende et al., 2014), and is thus called Wyner VAE .
(See Appendix A for a brief introduction on VAEs.)
The new encoder components introduced in Wyner VAE allow us to decompose a data vector into the common representation and the local representation, which can be used for sampling with style manipulation (Section 2.4).
We carefully compare our model and show its advantages over the existing VAE variants (Vedantam et al., 2018; Suzuki et al., 2016; Sohn et al., 2015; Wang et al., 2016) and the information bottleneck (IB) principle (Tishby et al., 1999) (Section 3), which is a well-known information theoretic principle in representation learning.
In the experiments, we empirically show the utility of our model in various sampling tasks and its superiority over existing models and that learning a succinct common representation achieves better generative performance in generation tasks (Section 4).
W YNER VARIATIONAL AUTOENCODER 2. P ROBABILISTIC MODEL
We ﬁrst deﬁne a probabilistic model for joint and conditional sampling tasks based on the single-letter characterizations of the motivating problems (Fig. 1 (b,d)).
We assume that all distributions to be introduced below belong to some standard parametric families such as Gaussians, and use ( q, φ ) (or ( p, θ ) ) to denote the parameters and the distribution of encoders (or decoders and priors).
In both channel synthesis and distributed simulation, Z ∼ p θ ( z ) signiﬁes the common randomness fed into the deterministic decoders x θ ( z , u ) and y θ
( z , v ) , while U ∼ p θ ( u ) and V ∼ p θ ( v ) signify the local randomnesses for each decoder.
We take the sources of randomness (or priors) p θ ( z ) , p θ ( u ) , p θ ( v ) as standard parametric distributions such as Gaussian from which a sample can be easily drawn.
U (or V ) is interpreted as a local description of X (or Y ) given the common description Z .
We can also view U (or V ) as an intrinsic randomness of the stochastic decoder ˜
p
θ ( x | z ) (or ˜ p
θ ( y | z ) ), which is the induced distribution by the decoder x θ ( z , u ) (or y θ
( z , v ) ) and the prior p θ ( u ) (or p θ ( v ) ).
To perform joint sampling, we need the priors p θ ( z ) , p θ ( u ) , p θ ( v ) and the decoders x θ ( z , u ) and y θ
( z , v ) (Fig.
(d)).
For conditional sampling of Y given X = x , we need the marginal encoder q φ ( z | x ) to be deﬁned in addition to the prior p θ ( v ) and the decoder
y θ ( z , v ) (Fig. 1 (b)).
We focus on how to learn q φ ( z | x ) in what follows, since q φ ( z | y ) can be dealt with in the same way by the symmetry of the Markov chain X → Z → Y .
After Wyner, we name the entirety of all the components (the marginal encoders q φ ( z | x ) , q φ ( z | y ) , the priors p θ ( z ) , p θ ( u ) , p θ ( v ) , and the decoders p θ ( x | z , u ) , p θ ( y | z , v ) ) that is essential for joint and conditional sampling tasks as the Wyner common representation model or the Wyner model in short.
2. O BJECTIVE FUNCTION
We now propose an objective function for learning each component in the Wyner model based on the optimization problem (1).
In practice, the data distribution q ( x , y ) is replaced by the empirical distribution q emp ( x , y ) deﬁned by the given samples.
To train the components for joint sampling in the Wyner model, we solve the optimization problem (1) by incorporating a new consistency constraint
[Equation]
with the priors p θ ( z ) , p θ ( u ) , p θ ( v ) and the decoders x θ ( z , u ) , y θ ( z , v ) as new optimization variables.
The Markov condition X → Z → Y can be removed under the new constraint (2).
We then relax the constraint (2) with an inequality constraint
[Equation]
for some (cid:15) > 0 , and convert the problem (1) into an unconstrained form as in Zhao et al.
(2018): minimize
[Equation]
Here we use a shorthand notation I q :
= I q ( X , Y ; Z ) , and use λ > 0 to denote the reciprocal of the
Lagrange multiplier.
We can choose D ( p, q ) as any proper distance or divergence measure between distributions such as f -divergences, Jensen–Shannon divergence, Wasserstein distance, or maximum mean discrepancy (Zhao et al., 2018).
To ﬁnd the marginal encoder q φ ( z | x ) that is consistent to q ( y | x ) q φ ( z | x , y ) for conditional generation, we aim to match a joint distribution induced by the Markov chain X
→ Z → Y , i.e., q ( x ) q φ ( z | x )˜
p
θ ( y | z ) , with the decoder distribution p θ ( z )˜
p
θ ( x | z )˜
p
θ ( y | z ) and the encoder distributions q ( x , y )
q φ ( z | x , y ) of the joint model.
That is, we wish to ﬁnd q φ ( z | x ) that minimizes
[Equation]
) and
[Equation]
The objective functions L y → y and L xy → x for q φ ( z | y ) can be deﬁned in a symmetric manner.
The ﬁnal objective function for training the Wyner model with succinct common representation learning then becomes
[Equation]
where the weights α ’s are nonnegative hyperparameters.
Yet, to minimize the objective function in practice, we need to choose which divergence/distance metric D ( p, q ) to use, and also need to address computationally intractable terms—the induced distributions ˜ p
θ ( x | z ) , ˜
p θ ( y | z ) and the mutual information term I q —in (3), (4), and (5).
2. V ARIATIONAL RELAXATION
We propose an instantiation of the objective function with a speciﬁc choice of a divergence function with a standard variational technique.
We choose the metric D ( p, q ) as the reverse KL divergence
D KL ( q (cid:107) p ) , which is a common choice in the variational inference literature (see, e.g., Blei et al.
To remove the intractable induced distributions, we relax the objective function (3) by introducing variational encoders q φ ( u | z , x ) and q φ ( v | z , y ) :
[Equation]
where (8) follows from the chain rule and nonnegativity of KL divergence (see, e.g.,
Cover and Thomas (2006)), ˜
p θ ( u , v | x , y , z ) denotes the induced conditional distribution by p θ ( u ) , p θ ( v ) , p θ ( x | z , u ) , p θ ( y | z , v ) , and p θ ( x | z , u ) (or p θ ( y | z , v ) ) denotes the distribution induced by the decoder x θ ( z , u ) (or y θ
( z , v ) ).
Note that the intractable distributions ˜ p
θ ( x | z ) , ˜
p θ ( y | z ) no longer appear in the upper bound (9).
For the mutual information term I q =
I q ( X , Y ; Z ) with ( X , Y , Z ) ∼ q ( x , y )
q φ ( z | x , y ) , we use the following standard upper bound (see, e.g., Zhao et al.
[Equation]
where ˜ q φ ( z ) denotes the induced marginal distribution by q ( x , y ) and q φ ( z | x , y ) .
Note here that the relaxation gap is D KL (˜ q φ ( z ) (cid:107) p θ ( z )) , which is again upper bounded by the joint KL divergence (7).
Therefore, all relaxation steps in (8) and (11) become tight when the joint distributions induced by data distribution, encoders, priors, and decoders are perfectly consistent with each other, i.e.,
[Equation]
We note that these relaxation techniques are standard in the literature, the tightness of which deserves a separate future study; for recent related work, refer to Cremer et al.
; Poole et al.
To sum up, the objective function for the joint model (3) is relaxed as
[Equation]
Note that the additional information regularization with λ > 0 is only on the common representation
Z , but not on U , V , which we call local representation.
By increasing λ > 0 to a proper degree that does not impedes ﬁtting, we can “reroute” the information ﬂow from ( X , Y ) through the common representation q φ ( z | x , y ) to the local representations q φ ( u | z , x ) , q φ ( v | z , y ) .
Following similar steps in (7), (8), (9), the objective function in (4) can also be upper bounded as
[Equation]
For (5), we choose L xy → y : = D KL ( q ( x , y ) q φ ( z | x , y ) (cid:107) q
x ) q φ ( z | x )˜
p
θ ( y | z )) , which can be viewed as the expected conditional ELBO (Sohn et al., 2015).
It can be also subsequently relaxed as
L xy → y ≤ L xy → y := D KL (
[Equation]
( x ,
[Equation]
q φ ( z | x , y
[Equation]
φ
[Equation]
( y | z , v )) .
After all, the ﬁnal relaxed objective function is given as follows:
[Equation]
See Figure 2 for an overview of each term in the objective function.
We call the overall framework which consists of all the components in the Wyner model and the additional encoders q φ ( z | x , y ) , q φ ( u | z , x ) , q φ ( v | z , y ) together with its training objective (15) as
Wyner common representation VAE or Wyner VAE in short.
After parameterizing each distribution component in Wyner VAE as standard parametric distributions such as Gaussians, whose parameters are again parameterized by deep neural networks, Wyner VAE can be trained efﬁciently by the standard reparameterization trick (Kingma and Welling, 2014) as in the standard VAE.
(See Appendix B for the Gaussian parameterization and the corresponding objective functions.)
In practice, we can either train the objective function (15) jointly with some positive α ’s ( joint training ), or train the joint model ﬁrst
letting α ’s be 0 and then train the marginal encoder in a retrospective manner by freezing the joint model ( two-stage training ) as in Vedantam et al.
The hyperparameter α ’s may be chosen with cross-validation with validation sets.
We elaborate the training schemes used for each experiment in Appendix F.
We remark the role of the encoders q φ ( z | x , y ) , q φ ( u | z , x ) , q φ ( v | z , y ) .
The joint encoder q φ ( z | x , y ) appears in the optimization problem (1) and plays as a reference distribution in learning the components of the Wyner model.
The variational encoders q φ ( u | z , x ) and q φ ( v | z , y ) are introduced to remove the intractable induced distributions ˜
p
θ ( z | x ) and ˜ p
θ ( z | y ) , satisfying the correct conditional independence structure implied by the decoder model p θ ( z ) p θ ( u ) p θ ( v ) p θ ( x | z , u ) p θ ( y | z , v ) , that is, q φ ( z , u , v | x , y )
=
q φ ( z | x , y )
q φ ( u | z , x ) q φ ( v | z , y ) , or equivalently U |
=
( V , Y ) | Z
and
V | =
( U , X )
| Z .
If we learn a succinct common representation q φ ( z | x , y ) (e.g., a shared concept) from ( X , Y ) , then q φ ( u | z , x ) would capture the remaining randomness U of X (e.g., texture and style).
We call this decomposition of the pair ( X , Y ) into the common representation Z and the local representation U , V as common-local information decomposition of ( X , Y ) .
We refer to ( Z , U , V ) as the joint representation of ( X , Y ) , to distinguish it from the common representation Z .
Provided that Wyner VAE achieves a good information decomposition, the variational encoders then can be used to explicitly in ﬁnding the local representations U and V from the data variables X and Y .
2. S AMPLING WITH STYLE CONTROL
As alluded to above, the variational encoders q φ ( u | z , x ) , q φ ( v | z , y ) can be used in sampling tasks with style control as a local representation (i.e., style) extractor.
We illustrate how to perform conditional sampling with style control (Fig.
3 (e)).
Suppose that ( X , Y ) is a pair of correlated images generated from the common concept but from different domains.
Given an image y 0 , we can extract the style information V 0 from y 0 by sampling ( Z 0 , V 0 ) from q φ ( z | y )
q φ ( v | z , y ) (Fig.
3 (d)).
We then generate Y j from an image x j similar to conditional sampling (Fig.
3 (c)), while replacing the randomly drawn local representation V ∼ p θ ( v ) with the previously extracted style V 0 , thereby the generated images Y 0 ,j ’s are of the same style as the reference image y 0 .
In a similar manner, we can also perform joint sampling with a ﬁxed style given a style reference data pair ( x 0 , y 0 ) , by mixing a randomly drawn common representation Z from the prior p θ ( z ) with the extracted style variables ( u 0 , v 0 ) . 2. D EGENERATE CASES IN W YNER VAE
Two degenerate cases may arise in Wyner VAE.
The ﬁrst case is where the common variable Z captures all information of ( X , Y ) , while U and V capture none.
Joint VAE (JVAE) (Vedantam et al., 2018) and joint multimodal VAE (JMVAE) (Suzuki et al., 2016) inherently assume this degenerate case, as it is discssued in the next section.
Wyner VAE is able to avoid such degeneracy by explicitly
[Figure] [Table] having the local variables and controlling the common regularization parameter λ > 0 . On the other extreme, Z may capture no information, while U and V capture all the information of X and Y , respectively.
It may happen in Wyner VAE if the regularization parameter λ is too large.
To avoid the degeneracy, we need to choose a proper λ by cross-validation.
Furthermore, if Wyner VAE cannot ﬁt the necessary common information in Z with the joint objective (10), the conditional objectives
L xy → y and L xy → x in (14) may help by penalizing the ignorance of the common information in Z .
In this section, we compare the proposed Wyner VAE to the existing models, deferring a detailed description of the encoder, decoder, prior components and objective functions of each model to
Appendix C. We provide a summary for capabilities of each model in Table 1.
VAE models.
Wyner VAE can be viewed as a generalization of the probabilistic model (i.e., encoder and decoder) assumed in two existing joint VAEs— JVAE (Vedantam et al., 2018) and JMVAE (Suzuki et al., 2016)—as alluded in the previous paragraph. These models implement a similar idea of performing joint and conditional generation tasks via a symmetric Markov chain X → W → Y , where W is the joint representation of ( X , Y ) .
In other words, these models can be derived by removing the local variables U and V in Wyner VAE. In Section 4, we demonstrated that the local variables in Wyner VAE help generate a variety of samples compared to JVAE.
The same decoder structure of Wyner VAE with the “shared” ( Z ) and the “private” ( U , V ) latent variables has been also studied in the context of multi-view learning (Shon et al., 2006; Ek et al., 2008;
Salzmann et al., 2010; Damianou et al., 2012) mostly based on a linear analysis such as canonical correlation analysis (CCA).
More recently, variational CCA-private (VCCA-private)
(Wang et al., 2016) was proposed to learn the decoder model with variational encoders q φ ( z | x ) , q φ ( u | x ) , and q φ ( v | y ) , with the encoder model q φ ( z , u , v | x , y ) =
q φ ( z | x ) q φ ( u | x ) q φ ( v | y ) to directly capture the conditional model from X to Z to Y .
On the other hand, Wyner VAE relies on the conditional independence structure q φ ( z , u , v | x , y ) =
q φ ( z | x , y ) q φ ( u | z , x ) q φ ( v | z , y ) , which is naturally induced by the decoder model.
We argue that this choice of encoder model in Wyner VAE may capture better semantic meaning of the local (private) random variables U and V , thereby leading a better generative performance; see, e.g., Fig. 6.
Conditional VAE (CVAE)
(Sohn et al., 2015) directly models the conditional distribution q ( y | x ) , obtained by simply conditioning every component in the vanilla VAE for q ( y ) with the conditioning variable X .
If Y is an image and X is an attribute of the image, a latent representation V in
CVAE needs to capture the redundant information of Y , which is not contained in X , i.e., style information of Y given X .
Wyner VAE can be viewed as a combination of two CVAEs with Z as [Table] a common conditioning variable, being capable of bidirectional sampling in its nature.
Yet, if X is high-dimensional, the conditional models like CVAE in general tend to overﬁt the input data of
X (Dutordoir et al., 2018).
To address this problem, a subsequent related work, bottleneck conditional density estimation (BCDE)
(Shu et al., 2017), proposed to learn joint and conditional VAE models simultaneously by softly tying the parameters of the two models for regularization.
We note that Wyner VAE naturally addresses such problem by using a uniﬁed single probabilistic model for both joint and conditional distribution learning, ﬁnding a succinct common representation Z for regularization.
Information bottleneck principle.
The information bottleneck (IB) principle (or method) (Tishby et al., 1999) is a widely known information theoretic approach in representation learning especially for discriminative tasks , i.e., when the target variable Y is a function of X and/or even discrete.
Motivated by lossy compression, the IB principle proposes to ﬁnd a compressed representation Z from the input variable X (i.e., q φ ( z | x ) ) while maximizing the relevance of Z in predicting the target variable Y as the minimizer of the optimization problem minimize q φ ( z | x ) βI q ( X ; Z ) −
I q ( Y ; Z ) , where ( X , Y , Z ) ∼ q ( x , y ) q φ ( z | x ) .
The foremost difference between the IB principle and our approach is in the underlying Markov chains: our symmetric Markov assumption X → Z → Y is more natural than Z → X → Y of
IB, when guessing Y based on Z as a representation of X .
Further, our framework aims to ﬁnd a certain common information structure characterized by Wyner’s common information with proper analogies to generation tasks of our interest, whereas the compressed-from- X yet relevant-to-
Y representation Z in the IB principle lacks its operational meaning, relying on a rather weak analogy to lossy compression.
We summarize other differences in various aspects in Table 2.
In particular, we compare Wyner VAE with variational IB (VIB) (Alemi et al., 2017) in the experiments, which is an instantiation of the IB principle based on a variational technique that can be implemented with neural networks.
As empirically shown below, VIB is not suitable for conditional generative tasks if the target variable Y is high-dimensional.
E XPERIMENTS
We empirically demonstrate that Wyner VAE outperforms JVAE, CVAE, VCCA-private, and VIB, for joint/conditional generation tasks and style manipulation on various datasets.
We defer the implementation details and training schemes used for each experiment to Appendix F.
Synthetic data.
We ﬁrst performed an experiment with a mixture of Gaussians (MoG) dataset as a toy example.
We considered a pair of 10-dim. MoG random vectors ( X , Y ) only correlated through a label Z ∼ Unif([1 , 2 , 3 , 4 , 5]) (common information) and 5-dim.
Gaussian random vectors
U , V ∼ N (0 , I 5 ) (local randomness in each variable).
We used the Gaussian latent variables ( Z , U , V ) of dimensions (10 , 10 , 10) , trained each model for 500 epochs (separate 50 epochs for each marginal encoder for JVAE and Wyner VAE) with a training data of size 50k, and summarized the numerical results in Fig.
4 and Table 3, which were evaluated with a test data of size 10k. See
Fig.
9 in Appendix E.1 for some visualizations.
In Fig.
4 (a,b), Wyner VAE with λ = 0 performed best for ﬁtting joint distributions, but did not excel in conditional log-likelihoods.
We observe that the performance of Wyner VAE on the test data gets improved throughout training by increasing λ : λ = 0 .
05 achieved a good conditional performance without too much sacriﬁce in the joint performance, while a larger value of λ (= 0 . 1) interfered ﬁtting [Table] to the distribution, failing to capture the essential common information structure.
Overall, Wyner VAE with λ control outperformed the other models.
CVAE tends to overﬁt quickly as noted earlier.
JVAE and VIB performed extremely worse in test conditional log-likelihoods compared to others as in Table 3.
JVAE failed to capture the common information structure as the training epochs increased, while VIB was only able to capture the average behaviors — we demonstrate how these models failed in this toy dataset in Appendix E.1.
Although VCCA-private achieved the best conditional likelihood with a comparable joint likelihood performance
[Table]
this model fails to learn more complex distributions as illustrated in Fig.
6 and Appendix E.2, E.3.
Fig.
4 (c) shows that λ can control the common mutual information I q ( X , Y ; Z ) in Wyner VAE — in particular, λ ∈ { 0 . 05 , 0 . 1 } kept the (estimated) mutual information at a constant level. We remark, however, that the MoG dataset has J ( X ; Y ) = log 5 , while the mutual information estimates are not around the true value (Fig. 4).
We attribute this mismatch to modeling the underlying discrete latent variable with a continuous, Gaussian vector.
Note that we deliberately used the standard Gaussian model assuming that we do not have a prior knowledge on the dataset generating process.
MNIST and SVHN datasets. We performed experiments with image datasets MNIST (LeCun, 1998) and SVHN (Netzer et al., 2011), by randomly pairing digit images only through their labels.
In particular, we constructed two dataset, MNIST–MNIST add-1 dataset, where label( Y i ) =
label( X i ) + 1 ( mod 10) , and MNIST–SVHN dataset, where label( Y i ) =
label( X i ) .
For MNIST–MNIST add-1 dataset, we trained Wyner VAE with different choices of λ ∈ { 0 . 0 , 0 .
. , 0 . 5 } to show the effect of λ on the generative performance of Wyner VAE.
Fig.
5 (a) corroborates our main claim that λ > 0 helps learning succinct representation in terms of small
I q ( X , Y ; Z ) and there exists a sweet spot ( λ = 0 . 1 ) that strikes the balance between the ﬁtting and the succinctness.
We also evaluated the label accuracy using the pre-trained Le-Net5 (LeCun et al., 1998) classiﬁcation network of accuracy 99.1%, and the per-pixel variance of samples from conditional generation.
Note that both high accuracy and high variance are desired for good generative models.
Fig.
5 (b) shows that by rerouting the information ﬂow through Z to U , V , λ > 0 helps U , V capture the style of images with a small sacriﬁce in label accuracy.
We present image samples to visualize the effect of λ in Wyner VAE and the superiority of Wyner VAE over the existing models.
We performed four sampling tasks—conditional generation, conditional [Figure] [Figure] n d i t i o n a l n e g a t
i v e l o g li k e li [Figure] t u a l i n [Figure]
r
[Figure]
t
i
[Figure]
n I (
[Figure]
[Figure]
generation with style control, joint stochastic generation, and joint sampling with style control; see
Fig. 3—for both MNIST–MNIST add-1 and MNIST–SVHN datasets, but here we present only a few samples for illustration.
We refer the interested reader to Appendix E.2, E.3 for a full comparison.
We defer all the samples from VIB to the appendix, as it only generated same “average” images.
Fig.
6 presents samples of MNIST–MNIST pairs by conditional generation with and without style control.
Fig.
6 (a1-c1) and Fig.
6 (a2-c2) demonstrated how the variations in the generated images and the style information captured in V are affected by varying λ , respectively.
(We illustrate how we performed conditional generation with style control for CVAE and VCCA-private in Appendix C.) λ
= 0 .
1 showed the best conditional generation results, consistent with Fig.
5 (a).
JVAE generated images without much variation.
VCCA-private erred frequently in guessing the labels, which implies that the shared representation Z in VCCA-private does not capture the “common information”.
Fig.
7 presents samples of MNIST–SVHN pair from conditional sampling with style control (Fig.
7 (a,b)) and two variations of joint sampling tasks (Fig.
7 (c-f)).
Fig.
7 (c,d) show joint stochastic reconstruction: Z is inferred from the label reference data, and samples are generated jointly by drawing local randomness ( U , V ) ∼ p θ ( u ) p θ ( v ) ; see Fig.
3 (b).
Fig.
7 (e,f) show joint sampling with style control: similarly to conditional sampling with style control, we generated joint samples by drawing common randomness Z ∼ p θ ( z ) with a speciﬁed local information from the style reference.
In all cases, we observe that λ = 0 .
1 achieves better style manipulation over λ = 0 , indicating that λ > 0 helps separating style information from common information. [Figure] [Table]
MNIST quadrant prediction dataset.
We performed a quadrant prediction task (Sohn et al., 2015) with a static, binary MNIST dataset (Larochelle and Murray, 2011), using the Bernoulli observation model for decoders.
Speciﬁcally, we split each digit image into two parts into left ( X ; conditioning and right ( Y ; target).
In this case, the most succinct common representation Z is a nontrivial object in contrast to a label (or its bijection) in the previous experiments.
Table 4 summarizes negative log-likelihoods for Wyner VAE, JVAE [Table] and CVAE on the test dataset. Note that Wyner VAE with with λ = 0 performs best, while an additional regularization parameter λ = 0 .
15 only gives a marginal improvement in conditional performance. For this speciﬁc case, having the local variables U and V is enough to capture the local information, as the underlying data seems to have much complex common information structure compared to a discrete label in the previous cases.
CelebA dataset.
We performed an experiment with CelebA dataset (Liu et al., 2015), which is a degenerate case in the sense that X is a function of Y .
While any bijection of X is common representation
Z attaining J ( X ; Y ) in this case, we demonstrated some merits of using Wyner VAE over CVAE in terms of better common representation learning in Appendix E.4.
C ONCLUDING R EMARKS
Cuff’s channel synthesis and Wyner’s distributed simulation are another manifestation of Occam’s razor by ﬁnding the simplest probabilistic structure that connects one random object to another.
The proposed Wyner VAE ﬁnds this succinct structure in a disciplined yet efﬁcient manner, and provides a theoretically sound alternative to the information bottleneck principle.
The experimental results demonstrated the potential of our approach as a new way of learning joint and conditional generation tasks with optimal representation learning that can be further developed and reﬁned for more complex dataset such as auditory, text, or a pair of those.
A few remaining questions from experiments are in order.
First, this paper does not address how close the estimated common information from a learned Wyner VAE model is to the true Wyner’s common information J ( X ; Y ) , and it may be interesting to devise a better estimation technique of
Wyner’s common information from data.
We emphasize that, however, the main goal of this paper is to demonstrate the advantages of the probabilistic structure of Wyner VAE and the regularization parameter λ in various joint and conditional generation tasks. Second, it would be interesting to investigate the relation between the variational approximation gaps and the quality of learned representations in Wyner VAE (Cremer et al., 2018; Poole et al., 2019).
A A Q UICK O VERVIEW ON V ARIATIONAL A UTOENCODERS
Variational autoencoder (VAE) (Kingma and Welling, 2014; Rezende et al., 2014) is a class of deep generative models that aim to simulate the unknown distribution q ( x ) underlying the data x 1 , . .
, x N to generate new samples from this distribution efﬁciently.
Let q emp ( x ) be the empirical distribution deﬁned by the sample.
Assume a generative latent variable model p θ ( z ) p θ ( x | z ) to model the underlying distribution q ( x ) .
One of the standard approach to learn each component in the model, the prior p θ ( z ) and the decoder p θ ( x | z ) , is the maximum likelihood approach that aims to solve
[Equation]
or equivalently, minimize
[Equation]
where ˜ p θ ( x ) is the induced distribution characterized by p θ ( z ) p θ ( x | z ) , i.e., ˜
p θ ( x ) :
= (cid:82) p θ ( z ) p θ ( x | z ) d z .
However, it is often computationally hard to solve the optimization problem directly due to the induced distribution ˜
p
θ ( x ) that involves an integration over a high-dimensional space.
In variational Bayesian learning approach (see, e.g., (Blei et al., 2017)), an approximate posterior q φ ( z | x ) (also called as an encoder) is introduced to relax the objective (16).
Here we present a short derivation of the well-known VAE objective function.
Note that the objective in (16) can be upper bounded as
[Equation]
where ˜ p θ ( z | x ) is the induced posterior characterized by p θ ( z ) p θ ( x | z ) .
(18) follows from the nonnegativity of the KL divergence, and (19) follows from the chain rule of the KL divergence (see, e.g,
Cover and Thomas (2006)).
Note that the ﬁnal relaxed form (20) does not contain the intractable term ˜
p
θ ( x ) .
The variational relaxation is tight if and only if D KL ( q φ ( z | x ) (cid:107) p θ ( z | x ))
=
for all x .
The upper bound (20) is the objective function for the standard VAE model (Kingma and Welling, 2014; Rezende et al., 2014).
That is, the standard VAE model aims to solve the following optimization problem:
[Equation]
To express the objective function in a more standard form in the literature, we add a constant h ( q emp ( x )) , the differential entropy of q emp ( x ) , and then derive
[Equation]
[Equation]
The loss function L θ,φ ( x ) in (23) with negation is called the evidence lower bound (ELBO) in the literature, since −L θ,φ ( x ) lower bounds the evidence log p θ ( x ) for each x .
The KL divergence term and the expected log loss term are called as the regularization term and the reconstruction term, respectively.
Assume that the observed variable X is continuous for simplicity.
The most standard parameterization of the components in the VAE model is the diagonal Gaussian parameterization
[Equation]
where each function may be parameterized by a neural network.
Here, σ 2 denotes a vector of a proper dimension and diag( σ 2 ) denotes a diagonal matrix with diagonal entries σ 2 .
Often, the covariance of the prior is taken to be isotropic and constant such as σ 20 ,θ = 1 , but σ 20 ,θ can also be set as an independent trainable parameter as in this work.
Note that the diagonal Gaussian parameterization p θ ( x | z ) is a formal modeling assumption to have a tractable density, which is required to evaluate the log loss log 1 /p θ ( x | z ) in the reconstruction loss.
This formal Gaussian noise also plays a role in estimating likelihoods.
However, after training, the decoder variance σ 2 θ ( z ) is dropped and the resulting decoder is used as deterministic:
z (cid:55)→ x θ ( z ) .
With this parameterization, the loss function L θ,φ ( x ) in (23) can be estimate efﬁciently for each x via a Monte Carlo approximation by sampling z ∼ q φ ( z | x ) .
The overall objective function then can be minimized with a gradient based optimization algorithm like Adam (Kingma and Ba, 2014) based on the reparameterization trick (Kingma and Welling, 2014;
Rezende and Mohamed, 2015) and backpropagation.
B S TANDARD IMPLEMENTATION OF W YNER VAE B. G AUSSIAN PARAMETERIZATION
As elaborated in Appendix A on VAEs, we use a standard Gaussian parameterization for Wyner VAE in our experiments.
Concretely, we let
[Equation] [Equation]
[Equation] [Equation]
where each function may be parameterized by a neural network
[Equation]
B. O BJECTIVE FUNCTIONS
We can rewrite the objective function L xy → xy in (12) in terms of the expectation of ELBO as in (23):
[Equation]
KL
[Equation]
where
[Equation]
We remark that a β -VAE (Higgins et al., 2017) type regularization in our joint model corresponds to imposing an additional weight β > 1 on L reg θ,φ ( x , y ) such that the objective function becomes
E q
[Equation]
x , y )
[ L rec θ,φ
[Equation]
X , Y ) + β L reg θ,φ ( X , Y )] .
We note that
[Equation]
Therefore, applying β -VAE type regularization imposes an additional regularization on
I q ( X , Y ; Z , U , V ) , which corresponds to the entire latent bottleneck ( U , V , Z ) for the information ﬂow from ( X , Y ) .
B. E STIMATION OF MUTUAL INFORMATION
With the typical Gaussian parameterization as presented above, we have an easy estimate for
I q ( X , Y ; Z ) .
After training, given a test dataset { x i , y i } Ni =1 , the mutual information I q ( X , Y ; Z ) can be estimated as
[Equation]
We can also estimate the variational upper bound of I q ( X , Y ; Z ) in (11) as follows:
[Equation]
In the last expression, each KL divergence can be explicitly obtained from the Gaussian parameters of q φ ( z | x i , y i ) and p θ ( z ) .
C A D EEPER L OOK ON R ELATED W ORK
For the standard VAE (Kingma and Welling, 2014; Rezende et al., 2014), we refer the interested reader to Appendix A. In what follows, we revisit and decompose each VAE-type model into its encoder/prior/decoder components, and express the objective function in the form of the reverse
KL divergence D KL ( q φ (cid:107) p θ ) , where q φ is the joint distribution over the data variables and the latent variables deﬁned by the data distribution and the encoders, and p θ is deﬁned by the priors and the decoders.
JVAE/JMVAE JVAE (Vedantam et al., 2018) and JMVAE (Suzuki et al., 2016) consist of
• the joint encoder q φ ( w | x , y ) , the marginal encoders q φ ( w | x ) , q φ ( w | y ) , • the prior p θ ( w ) , • the decoders p θ ( x | w ) , p θ ( y | w ) .
They share the same joint model objective
[Equation]
[Equation]
but differ in training the marginal encoders.
JMVAE trains the marginal encoder q φ ( z | x ) via minimizing
[Equation]
[Equation] together with the joint model objective (42) (by adding two objective functions with an additional weight as a hyperparameter), while JVAE trains the marginal encoders separately from the joint model by training the marginal VAE as we proposed, i.e., [Equation]
[Equation]
It is worthwhile to compare the regularization term
I q ( X , Y ; Z ) in Wyner VAE with the idea of βVAE (Higgins et al., 2017), which empirically showed that an additional weight on the regularization term ﬁnds a disentangled representation in a VAE model.
If we apply a β -VAE type regularization in our joint model, then it corresponds to an additional weight β > 1 on the mutual information on the joint representation ( I q ( X , Y ; Z , U , V ) ), not only on the common representation ( I q ( X , Y ; Z ) ) as in Wyner VAE.
(See Appendix B.2.)
In words, Wyner VAE provides a ﬁner control on the information ﬂow from ( X , Y ) , by manipulating the common path q φ ( z | x , y ) , while β -VAE blocks the entire latent bottleneck q φ ( z | x , y ) q φ ( u | z , x ) q φ ( v | z , y ) .
CVAE CVAE (Sohn et al., 2015) for modeling q ( y | x ) consists of
• the encoder q φ ( v | y , x ) , • the prior p θ ( v | x ) , • the decoder p θ ( y | v , x ) .
The objective function is then given as
[Equation]
[Equation] Note that without the conditioning variable x , this model boils down to the vanilla VAE for modeling q ( y ) .
If we assume a Markov chain X − V − Y , the decoder p θ ( y | v , x ) can be replaced with p θ ( y | v ) .
We performed conditional generation with style control with CVAE as follows:
Given a style reference y 0 and its corresponding x 0 , we sample and keep v 0 via the encoder q φ ( v | y , x ) .
Given a new x 1 , we take y 0 , 1 = y θ ( v 0 , x 1 ) as a new sample.
Note that this scheme assumes that v 0 ∼ q φ ( v | y 0 , x 0 ) captures only the remaining information in y 0 excluding the information on x 0 .
SSVAE/ADGM Semi-supervised VAE (SSVAE) (Kingma et al., 2014) proposed a similar conditional VAE model for modeling q ( y | x) especially when the conditioning variable x is discrete.
SSVAE consists of
• the encoders q φ ( v | y , x) , q φ (x | y ) (classiﬁer), • the priors p θ ( v ) , p θ (x) , • the decoder p θ ( y | v , x) .
Note that it has an additional encoder q φ (x | y ) (classiﬁer) on top of q φ ( v | y , x ) from CVAE, and assume the priors p θ ( v ) and p θ (x) , which replace the conditional prior p θ ( v | x) and the data distribution q (x) in CVAE, respectively.
With a paired (i.e., labelled) data, SSVAE minimizes
[Equation]
With an unlabeled data, SSVAE minimizes
[Equation]
Note that the label information q (x) in (46) is replaced with q φ (x | y ) in (47) as the label information is missing in the unlabeled data.
Formally, in this degenerate case, identifying the conditioning variable X as the common latent variable Z in Wyner VAE recovers SSVAE.
Yet, SSVAE was proposed in the context of (semisupervised) classiﬁcation problem which aims to learn a classiﬁer q φ (x | y ) , and for a high-dimensional
X , it is not feasible to directly model the conditional distribution q φ ( x | y ) .
Auxiliary deep generative model (ADGM) (Maaløe et al., 2016) adds an auxiliary latent variable Z to
SSVAE to improve the performance.
ADGM consists of
• the encoders q φ ( z | y ) , q φ ( v | z , y , x) , q φ (x | z , y ) (classiﬁer), • the priors p θ ( v ) , p θ (x) , • the decoders p θ ( y | v , x) , p θ ( z | v , x , y ) .
Note the new components q φ ( z | y ) and p θ ( z | v , x , y ) on top of the SSVAE, and the original encoder/decoder components have additional condition on the auxiliary variable Z .
However, as ADGM does not impose any additional conditional independence with Z , it is not directly comparable to
Wyner VAE.
VCCA-private As noted earlier, VCCA-private (Wang et al., 2016) consists of
• the encoders q φ ( z | x ) , (and/or q φ ( z | y ) ), q φ ( u | x ) , q φ ( v | y ) , • the priors p θ ( z ) , p θ ( u ) , p θ ( v ) , • the decoders p θ ( x | z , u ) , p θ ( y | z , v ) .
Note that the prior and the decoder components are same with Wyner VAE.
Hence, the objective for
VCCA-private for the marginal encoder q φ ( z | x ) can be expressed as
[Equation]
To model the other direction of the marginal encoder q φ ( y | z ) , they minimize
[Equation]
[Equation] To learn q φ ( z | x ) , q φ ( z | y ) simultaneously, BiVCCA-private minimizes a convex combination of the two KL divergence terms.
We performed conditional generation with style control with VCCA-private as follows:
Given a style reference y 0 , we sample and keep v 0 via the encoder q φ ( v | y ) .
Given a new attribute x 1 , we take y 0 , 1 = y θ ( v 0 , x 1 ) as a new sample.
VIB VIB (Alemi et al., 2017) proposed a variational relaxation of the following minimization problem posed by the information bottleneck principle (Tishby et al., 1999):
[Equation]
where ( X , Y , Z ) ∼ q ( x , y ) q φ ( z | x ) .
VIB introduces two variational distributions p θ ( z ) and p θ ( y | z ) that approximate ˜ q φ ( z ) and ˜ q
φ ( y | z ) .
Then, based on standard variational bounds on the mutual information terms (see, e.g., Zhao et al.
, we obtain a variational upper bound on βI q ( X ; Z ) −
I q ( Y ; Z ) as
[Equation]
which is the objective function for VIB.
Note that the relaxation gap is given by
[Equation]
D L IKELIHOOD E STIMATION
For each likelihood of our interest, we derive a naive Monte Carlo (MC) estimator and an MC estimator with importance sampling (see, e.g., Rubinstein and Kroese (2016)).
Here we only present the estimators for Wyner VAE, but the estimators for other models can be derived in the same manner.
D. J OINT LIKELIHOOD
We wish to estimate the joint log-likelihood of the model with respect to given test data { ( x ( i ) , y ( i ) )
} Ni =1 , i.e., (cid:80) Ni =1 log p θ ( x ( i ) , y ( i ) ) , where
[Equation]
v
[Equation]
[Equation] [Equation]
[Equation] [Equation]
[Equation] [Equation]
D. C ONDITIONAL LIKELIHOOD
We wish to estimate the conditional log-likelihood of the conditional path of Wyner VAE from x to y , i.e., q φ ( z | x ) p θ ( v ) p θ ( y | z , v ) , with respect to given test data { ( x ( i ) , y ( i ) )
} Ni =1 , i.e.,
[Equation]
56) where r
[Equation]
) (1) Monte Carlo estimation: Given x , let ( z ( s ) , v ( s ) ) ∼ q φ ( z | x ) p θ ( v ) for s = 1 , . . .
, S .
[Equation]
58) (2) Importance sampling: For each ( x , y ) , let ( z ( s ) , v ( s ) ) ∼ q φ ( z | x , y ) q φ ( v | z [Equation] y
[Equation]
for s = 1 , . . .
, S .
[Equation]
=
[Equation]
E A DDITIONAL E XPERIMENTAL R ESULTS E. M IXTURE OF G AUSSIANS
Here we present the test log-likelihood evaluation of JVAE and VIB in conditional generation tasks of the mixture of Gaussians (MoG) dataset; recall Fig. 4.
Both JVAE and VIB performed worse as the training epochs increased, that is, they overﬁt to the training data (Fig.
We visually illustrate how they failed in Fig. 9.
Fig.
9(a) shows the outlook of the MoG dataset in our experiment.
Fig.
9 (e1,e2) show that JVAE captured all the components at the beginning, but then collapsed to a few components afterwards.
On the other hand, Fig. 9 (f1,f2) show that VIB only captured the average behaviors, although gradually adapting to the underlying data.
[Figure]
E. MNIST–MNIST ADD -
[Table]
I npu t
W y n e r V
A E ( 𝝀 = 𝟎 . 𝟎 )
W y n e r V
A E ( 𝝀 = 𝟎 . 𝟏 )
W y n e r V
A E ( 𝝀 = 𝟎 . 𝟐 )
C V A E
V CC A - p r i v a t e
J V A E
V I B
S t y l e r e f .
N / A
N / A
I npu t ( l a b e l r e f . )
W y n e r V
A E ( 𝝀 = 𝟎 . 𝟎 )
W y n e r V
A E ( 𝝀 = 𝟎 . 𝟏 )
W y n e r V
A E ( 𝝀 = 𝟎 . 𝟐 )
C V A E
V CC A - p r i v a t e
J V A E
V I B
L a b e l r e f .
J o i n t s t
o c h a s t i c
r e c o n s t r u c t i o n
S t y l e r e f .
J o i n t g e n e r a t i o n w i
t h s t y l e c o n t r o l
W y n e r V
A E ( 𝝀 = 𝟎 . 𝟎 )
W y n e r V
A E ( 𝝀 = 𝟎 . 𝟏 )
W y n e r V
A E ( 𝝀 = 𝟎 . 𝟐 )
V CC A - p r i v a t e
J V A E
N / A
E.
MNIST–SVHN
I npu t
W y n e r V
A E ( 𝝀 = 𝟎 . 𝟎 )
W y n e r V
A E ( 𝝀 = 𝟎 . 𝟎𝟓 )
W y n e r V
A E ( 𝝀 = 𝟎 . 𝟏 )
C V A E
V CC A - p r i v a t e
J V A E
V I B
S t y l e r e f .
N / A
N / A
I npu t ( l a b e l r e f . )
W y n e r V
A E ( 𝝀 = 𝟎 . 𝟎 )
W y n e r V
A E ( 𝝀 = 𝟎 .
W y n e r V
A E ( 𝝀 = 𝟎 . 𝟏 )
C V A E
V CC A - p r i v a t e
J V A E
V I B
L a b e l r e f .
J o i n t s t
o c h a s t i c
r e c o n s t r u c t i o n
S t y l e r e f .
J o i n t g e n e r a t i o n w i
t h s t y l e c o n t r o l
W y n e r V
A E ( 𝝀 = 𝟎 . 𝟎 )
W y n e r V
A E ( 𝝀 = 𝟎 . 𝟎𝟓 )
W y n e r V
A E ( 𝝀 = 𝟎 . 𝟏 )
V CC A - p r i v a t e
J V A E
N / A
E. (F ACE , ATTRIBUTE ) PAIRS FROM C ELEB A DATASET
CelebA dataset (Liu et al., 2015) consists of pairs of a face image and a 40-dim. binary vector that contains attributes information of the face.
We performed conditional generation of face images ( Y ) given an attribute vector ( X ) .
Since an attribute X is a function of a given face image Y , we let the dimension of the local variable U be 0 in this case.
Fig.
16 presents samples of CelebA faces from
Wyner VAE, JVAE, and CVAE — Wyner VAE with λ = 0 .
1 generated a variety of faces with the correct attributes, while JVAE generated images with little variations as previously observed and
CVAE generated diverse images but often with wrong attributes.
See also Appendix E.4.2 for the results with style control and a numerical evaluation of the performance of Wyner VAE on CelebA dataset.
[Figure]
Note that this is a special case where the target variable Y is a function of a conditioning variable
X , which is a degenerate case in the sense that any random variable Z that is bijective with X can serve as an optimal common representation that achieves J ( X ; Y ) , and Wyner VAE and CVAE may have comparable conditional generation performance.
In practice, however, Wyner VAE outperforms
CVAE since Wyner VAE learns a good representation Z of X that is helpful in generating Y conditionally, while CVAE directly uses the raw X for conditioning.
For a qualitative evidence, we present a few Attribute → Face samples in Fig.
17 from models trained with more data as in the paper, from a truly unseen attribute (likely-female-features + bald)—Wyner
VAE can produce plausible images as it ﬁnds a good representation Z of an unseen attribute, while
CVAE fails.
E.4. N UMERICAL EVALUATION
We present an additional numerical evaluation of Wyner VAE to corroborate the effect of λ > 0 .
With CelebA models, we evaluated both conditional paths, i.e., face to attribute classiﬁcation (face2attribute) and attribute to face generation (attribute2face).
For the attribute classiﬁcation, we counted the number of corrected classiﬁed binary attributes out of 40.
The image variance is evaluated with per pixel, while all pixel values were normalized between 0 and 1.
Table 6 summarizes the results.
We can observe that both the accuracy and the variance for conditional generation from a given attribute were maximized the around λ ∼ 0 . 2 .
Note that JVAE performs worse in classiﬁcation
[Table]
accuracy with a comparably very small per-pixel variance implying much less variations in the generated samples.
E.4. A DDITIONAL ATTRIBUTE 2 FACE GENERATION RESULTS
Here we present additional conditional generated samples (with and without style control) from
CelebA models.
We used the sample images and attributes shown in Fig.
18 for these experiments.
For conditional generation in Fig. 19, samples were generated only based on the sample attributes.
For conditional generation with style control in Fig. 20, Wyner VAE and CVAE ﬁrst extracted style information from sample images in the leftmost column, and then generated new samples from the original attribute added with a new binary attribute speciﬁed in the topmost row for each column.
Hence, the second column corresponds to the reconstruction of the style reference images from the faces and the corresponding attributes.
We remark that JVAE is not capable of style manipulation, and the results from conditional generation with JVAE are given as a reference.
[Table]
[Figure] [Figure]
[Figure]
F E XPERIMENT D ETAILS
We used the same parameterization, same latent dimensions, and the same network architecture across the different models to be a fair comparison.
For simplicity, we used the standard Gaussian parameterization of each component for all the implemented models as in the standrad VAE.
(See
Appendix A and Appendx B.)
For the prior distributions p θ ( z ) , p θ ( u ) , p θ ( v ) , we let the isotropic variances σ 20 ,θ , σ 21 ,θ , σ 22 ,θ be trainable.
With this degree of freedom, the neural networks select necessary dimensions in the latent spaces over Z , U , V by assigning small variances to unused dimensions.
For the MoG experiment, we used a constant 1/2 for the decoder variance, so that the log-loss corresponds to the l 2 -squared loss.
For the rest of the experiments, we allowed the diagonal variances to be trainable: we allocated one trainable decoder variance per channel, independent from the latent inputs.
We found that this trick results in sharper images across the models.
We set the dimension of the latent variable W in JVAE as the sum of the dimensions of Z , U , V as
W corresponds the joint representation.
Similarly, since the latent variable of CVAE corresponds to the local randomness in Wyner VAE, we let the dimension of V in CVAE be equal to the dimension of V in Wyner VAE.
We used (10,10,10), (32,32,32), (128,128,128), and (128,0,128) as the latent dimensions of ( Z , U , V ) for MoG, MNIST–MNIST add-1, MNIST-SVHN, and CelebA, respectively.
All log-likelihood values in the experiments were estimated by importance sampling; see Appendix D.
We used S = 100 importance samples for each data point.
The mutual information I q ( X , Y ; Z ) was estimated in a straightforward manner with the test dataset under the Gaussian parameterization; see
Appendix B.
Computing infrastructure We used NVIDIA TITAN X (Pascal) for our experiments.
Implementation We implemented all models using Keras 3 with tensorﬂow backend.
Datasets For the mixture of Gaussians dataset, we generate the paired dataset ( X i , Y i )
∈ R 10 ×
Ras follows.
Let µ : { 1 , 2 , 3 , 4 , 5 } → R × R be a function deﬁned as
[Equation]
Then, we let
[Equation]
where Z i ∼ Unif( { 1 , 2 , 3 , 4 , 5 } )
, U i , V i ∼ N (0 , I 5 ) are drawn independently.
Here, 1 5 ∈ Rdenotes the all-1 vector.
For MNIST–MNIST add-1, we constructed 50 k add-1 pairs from the MNIST training dataset.
For
MNIST–SVHN domain adapation, we constructed 50 k MNIST-SVHN pairs from MNIST and SVHN training datasets.
For testing, we similarly constructed 1 k paired images from MNIST and SVHN test datasets in each case.
For CelebA experiments, we set aside 5 k samples for test dataset, and used the rest in training.
Network architectures Let c5s1- k -{activation} denote a 5 × 5 Convolution–
BatchNorm–activation with k ﬁlters and stride 1 × 1 .
Let d3s2- k -{activation} / u3s2- k -{activation} denote 3 × 3 Convolution / Deconvolution–BatchNorm–activation with k ﬁlters and stride 2 × 2 , respectively.
Let res- k be a residual block that contains two 3 ×convolutional layers with k ﬁlters in each (i.e., c3s1- k -LReLU , c3s1- k ) and a skip connection from the input to the output.
Let fc- k -{activation} be a fully-connected layer with k units and a non-linear activation.
For MNIST–MNIST add-1, MNIST-SVHN, and CelebA, each Wyner VAE consisted of an outer encoder/decoder pair and a core joint Wyner model.
The outer encoder/decoder pair was introduced to pre-process raw input data.
We summarized the network architectures for the Wyner models in
Table 7, Table 8, and Table 10.
Note that we padded zeros around the 28 × 28 MNIST images to make them of size 32 × 32 .
[Table] [Table]
[Table] [Table]
Training We used the Adam optimizer (Kingma and Ba, 2014) with learning rate 10 − 4 in training
MoG, MNIST–MNIST, MNIST–SVHN and CelebaA models.
We used learning rate 5 · 10 − 4 in training MNIST quadrant prediction experiment.
For MoG dataset, we trained each model for 500 epochs with batch size 100 and trained each marginal encoder for separate 50 epochs for JVAE and Wyner VAE.
For MNIST–MNIST, MNIST–SVHN, and CelebA experiment, each model was trained for 100 epochs with batch size 128 .
For MNIST quadrant prediction experiment, each model was trained for 1000 epochs with batch size 20 .
For the marginal encoder q φ ( z | x ) in JVAE or Wyner VAE, we trained the joint models by setting α x → x = α xy
→ y =
α xy → x = α y → y = 0 and trained only the marginal encoder for everyepochs by minimizing L x → x with freezing all the components in the joint model for 20 epochs.
For
MNIST–MNIST, MNIST–SVHN, and CelebA experiments, we trained q φ ( z | x ) for 1 epoch after every 1 epoch of the joint model training as suggested in Vedantam et al.
Note that the outer encoder/decoder pairs in Tables 8 and 10 were trained with the core joint Wyner model, and they were ﬁxed in the marginal encoder training.
For MNIST quadrant prediction experiment, we trained the joint model and the marginal encoders q φ ( z | x ) and q φ ( z | y ) by minimizing the ﬁnal objective (15) by setting all α ’s to be 1 for simplicity without ﬁne-tuning.
We empirically observed that for some experiments the two-stage training with α xy → y =
α xy → x =is still effective.
Note that we applied the joint training scheme and used α xy → y , α xy → x > 0 only for
MNIST quadrant prediction experiment to achieve a better conditional log-likelihood performance.