M$^2$VAE - Derivation of a Multi-Modal Variational Autoencoder Objective from the Marginal Joint Log-Likelihood

TL;DR

This paper provides a detailed derivation of the training objective for Multi-Modal Variational Autoencoders (M$^2$VAE) based on the marginal joint log-likelihood, facilitating better understanding and implementation.

Contribution

It offers a comprehensive derivation of the evidence lower bound for M$^2$VAE, clarifying the training process for multi-modal generative models.

Findings

Derived a clear training objective for M$^2$VAE

Enhanced understanding of multi-modal variational autoencoders

Facilitated implementation of M$^2$VAE models

Abstract

This work gives an in-depth derivation of the trainable evidence lower bound obtained from the marginal joint log-Likelihood with the goal of training a Multi-Modal Variational Autoencoder (M$^2$VAE).

Figures

Equations

$$L$$

$$=\sum _{z}q\left(z|a\right)\log \left(p\left(a\right)\right)$$

$$=\sum _{z}q\left(z|a\right)\log \left(\frac{p\left(z,a\right)}{p\left(z|a\right)}\right)$$

$$=\sum _{z}q\left(z|a\right)\log \left(\frac{p\left(z,a\right)}{p\left(z|a\right)}\frac{q\left(z|a\right)}{q\left(z|a\right)}\right)$$

$$=\sum _{z}q\left(z|a\right)\log \left(\frac{p\left(z,a\right)}{q\left(z|a\right)}\frac{q\left(z|a\right)}{p\left(z|a\right)}\right)$$

$$=\sum _{z}q\left(z|a\right)\log \left(\frac{p\left(z,a\right)}{q\left(z|a\right)}\right)+\sum _{z}q\left(z|a\right)\log \left(\frac{q\left(z|a\right)}{p\left(z|a\right)}\right)$$

$$=\mathcal{L}+{\mathrm{D}}_{\mathrm{KL}}\left(q\left(z|a\right)\Vert p\left(z|a\right)\right)$$

$$\ge \mathcal{L}$$

$$\mathcal{L}$$

$$=\sum _{z}q\left(z|a\right)\log \left(\frac{p\left(a|z\right)p\left(z\right)}{q\left(z|a\right)}\right)$$

$$=\sum _{z}q\left(z|a\right)\log \left(\frac{p\left(z\right)}{q\left(z|a\right)}\right)+\sum _{z}q\left(z|a\right)\log \left(p\left(a|z\right)\right)$$

$$=-{\mathrm{D}}_{\mathrm{KL}}\left(q\left(z|a\right)\Vert p\left(z\right)\right)+E_{q\left(z|a\right)}\log \left(p\left(a|z\right)\right)$$

$$\mathcal{L}$$

$$L_{\text{J}}$$

$$=\sum _{z}q\left(z|a,b\right)\log \left(p\left(a,b\right)\right)$$

$$=\sum _{z}q\left(z|a,b\right)\log \left(\frac{p\left(z,a,b\right)}{p\left(z|a,b\right)}\right)$$

$$=\sum _{z}q\left(z|a,b\right)\log \left(\frac{p\left(z,a,b\right)}{p\left(z|a,b\right)}\frac{q\left(z|a,b\right)}{q\left(z|a,b\right)}\right)$$

$$=\sum _{z}q\left(z|a,b\right)\log \left(\frac{p\left(z,a,b\right)}{q\left(z|a,b\right)}\frac{q\left(z|a,b\right)}{p\left(z|a,b\right)}\right)$$

$$=\sum _{z}q\left(z|a,b\right)\log \left(\frac{p\left(z,a,b\right)}{q\left(z|a,b\right)}\right)+\sum _{z}q\left(z|a,b\right)\log \left(\frac{q\left(z|a,b\right)}{p\left(z|a,b\right)}\right)$$

$$={\mathcal{L}}_{\text{J}}+{\mathrm{D}}_{\mathrm{KL}}\left(q\left(z|a,b\right)\Vert p\left(z|a,b\right)\right)$$

$$\ge {\mathcal{L}}_{\text{J}}$$

$${\mathcal{L}}_{\text{J}}$$

$$=\sum _{z}q\left(z|a,b\right)\log \left(\frac{p\left(a,b|z\right)p\left(z\right)}{q\left(z|a,b\right)}\right)$$

$$=\sum _{z}q\left(z|a,b\right)\log \left(\frac{p\left(z\right)}{q\left(z|a,b\right)}\right)+\sum _{z}q\left(z|a,b\right)\log \left(p\left(a,b|z\right)\right)$$

$$=\sum _{z}q\left(z|a,b\right)\log \left(\frac{p\left(z\right)}{q\left(z|a,b\right)}\right)+\sum _{z}q\left(z|a,b\right)\log \left(p\left(a|z\right)\right)+\sum _{z}q\left(z|a,b\right)\log \left(p\left(b|z\right)\right)$$

$$=-{\mathrm{D}}_{\mathrm{KL}}\left(q\left(z|a,b\right)\Vert p\left(z\right)\right)+E_{q\left(z|a,b\right)}\log \left(p\left(a|z\right)\right)+E_{q\left(z|a,b\right)}\log \left(p\left(b|z\right)\right)$$

$${\mathcal{L}}_{\text{J}}=\underset{\text{Regularization}}{\underbrace{-{\mathrm{D}}_{\mathrm{KL}}\left(q_{{\varphi }_{\text{a}\text{b}}}\left(z|a,b\right)\Vert p\left(z\right)\right)}}+\underset{\text{Reconstruction wrt. a}}{\underbrace{E_{q_{{\varphi }_{\text{a}\text{b}}}\left(z|a,b\right)}\log \left(p_{{\theta }_{\text{a}}}\left(a|z\right)\right)}}+\underset{\text{Reconstruction wrt. b}}{\underbrace{E_{q_{{\varphi }_{\text{a}\text{b}}}\left(z|a,b\right)}\log \left(p_{{\theta }_{\text{b}}}\left(b|z\right)\right)}}$$

$$p\left(a|b\right)$$

$$=\frac{1}{p\left(z|a,b\right)}p\left(z,a|b\right)$$

$$=\frac{1}{p\left(z|a,b\right)}\frac{p\left(z,a,b\right)}{p\left(b\right)}$$

$$=\frac{1}{p\left(z|a,b\right)}\frac{p\left(a,b|z\right)p\left(z\right)}{p\left(b\right)}$$

$$=\frac{1}{p\left(z|a,b\right)}\frac{p\left(a|z\right)p\left(b|z\right)p\left(z\right)}{p\left(b\right)}$$

$$=\frac{1}{p\left(z|a,b\right)}\frac{p\left(a|z\right)p\left(z|b\right)\frac{p\left(b\right)}{p\left(z\right)}p\left(z\right)}{p\left(b\right)}$$

$$=\frac{p\left(a|z\right)p\left(z|b\right)}{p\left(z|a,b\right)}$$

$$L_{{\text{M}}_{\text{a}}}$$

$$=\sum _{z}q\left(z|a,b\right)\log \left(p\left(a|b\right)\right)$$

$$=\sum _{z}q\left(z|a,b\right)\log \left(\frac{p\left(z,a|b\right)}{p\left(z|a,b\right)}\right)$$

$$=\sum _{z}q\left(z|a,b\right)\log \left(\frac{p\left(z,a|b\right)}{p\left(z|a,b\right)}\frac{q\left(z|a,b\right)}{q\left(z|a,b\right)}\right)$$

$$=\sum _{z}q\left(z|a,b\right)\log \left(\frac{p\left(z,a|b\right)}{q\left(z|a,b\right)}\frac{q\left(z|a,b\right)}{p\left(z|a,b\right)}\right)$$

$$=\sum _{z}q\left(z|a,b\right)\log \left(\frac{p\left(z,a|b\right)}{q\left(z|a,b\right)}\right)+\sum _{z}q\left(z|a,b\right)\log \left(\frac{q\left(z|a,b\right)}{p\left(z|a,b\right)}\right)$$

$$={\mathcal{L}}_{{\text{M}}_{\text{a}}}+{\mathrm{D}}_{\mathrm{KL}}\left(q\left(z|a,b\right)\Vert p\left(z|a,b\right)\right)$$

$$\ge {\mathcal{L}}_{{\text{M}}_{\text{a}}}$$

$$L_{\text{M}}$$

$$=\log \left(p\left(a|b\right)\right)+\log \left(p\left(b|a\right)\right)$$

$$={\mathcal{L}}_{{\text{M}}_{\text{a}}}+{\mathcal{L}}_{{\text{M}}_{\text{b}}}+2{\mathrm{D}}_{\mathrm{KL}}\left(q\left(z|a,b\right)\Vert p\left(z|a,b\right)\right)$$

$$\ge {\mathcal{L}}_{{\text{M}}_{\text{a}}}+{\mathcal{L}}_{{\text{M}}_{\text{b}}}$$

$${\mathcal{L}}_{{\text{M}}_{\text{a}}}+{\mathcal{L}}_{{\text{M}}_{\text{b}}}$$

$$=\sum _{z}q\left(z|a,b\right)\log \left(\frac{p\left(a|z\right)p\left(z|b\right)}{q\left(z|a,b\right)}\right)+\sum _{z}q\left(z|a,b\right)\log \left(\frac{p\left(b|z\right)p\left(z|a\right)}{q\left(z|a,b\right)}\right)$$

$$=\sum _{z}q\left(z|a,b\right)\log \left(\frac{p\left(a|z\right)}{q\left(z|a,b\right)}\right)+\sum _{z}q\left(z|a,b\right)\log \left(\frac{p\left(z|b\right)}{q\left(z|a,b\right)}\right)$$

$$+\sum _{z}q\left(z|a,b\right)\log \left(\frac{p\left(b|z\right)}{q\left(z|a,b\right)}\right)+\sum _{z}q\left(z|a,b\right)\log \left(\frac{p\left(z|a\right)}{q\left(z|a,b\right)}\right)$$

Full text

I Introduction

Variational auto encoder (VAE) combine neural networks with variational inference to allow unsupervised learning of complicated distributions according to the graphical model shown in Fig. 1 (left). A $D_{a}$-dimensional observation $a$ is modeled in terms of a $D_{z}$-dimensional latent vector $z$ using a probabilistic decoder $p_{\theta_{\text{a}}}\mathopen{}\left(z\right)\mathclose{}$ with parameters $\theta$. To generate the corresponding embedding $z$ from observation $a$, a probabilistic encoder network with $q_{\phi_{\text{a}}}\mathopen{}\left(z\right)\mathclose{}$ is being provided which parametrizes the posterior distribution from which $z$ is sampled. The encoder and decoder, given by neural networks, are trained jointly to bring $a$ close to an $a^{\prime}$ under the constraint that an approximate distribution needs to be close to a prior $p\mathopen{}\left(z\right)\mathclose{}$ and hence inference is basically learned during training.

The specific objective of VAEs is the maximization of the marginal distribution $p\mathopen{}\left(a\right)\mathclose{}=\int p_{\theta}\mathopen{}\left(a|z\right)\mathclose{}p\mathopen{}\left(z\right)\mathclose{}\operatorname{d}\!a$. Because this distribution is intractable, the model is instead trained via stochastic gradient variational Bayes (SGVB) by maximizing the evidence lower bound (ELBO) $\mathcal{L}$ of the marginal log-likelihood as described in Sec. II This approach proposed by [1] is used in settings where only a single modality $a$ is present in order to find a latent encoding $z$ (c.f. Fig. 1 (left)).

This work gives an in-depth derivation of the trainable evidence lower bound (ELBO) obtained from the marginal joint log-Likelihood, that satisfies all plate models as depicted in Fig. 1, we are with the goal of training a multi-modal variational Autoencoder (M²VAE).

II Variational Autoencoder

First, the derivation of the vanilla Variational Autoencoder by [1] is recaped.

II-A The Variational Bound

[TABLE]

$\operatorname{D_{KL}}$ is the Kulbeck-Leibler divergence, with $\operatorname{D_{KL}}\geq 0$, thats depends on how good $q\mathopen{}\left(z|a\right)\mathclose{}$ can approximate $p\mathopen{}\left(z|a\right)\mathclose{}$. $\mathcal{L}$ is the lower variational bound of the marginal log-likelihood, also called the evidence lower bound (ELBO). If and only if the two distributions $q$ and $p$ are identical, $\operatorname{D_{KL}}$ becomes 0 ($q=p\Leftrightarrow\operatorname{D_{KL}}=0$). $\mathcal{L}=L$ means on the other hand therefore implicitly, that $q$ perfectly approximates $p$. It is beacuse $\mathcal{L}$ and $\operatorname{D_{KL}}$ are in equilibrium so that minimizing $\operatorname{D_{KL}}$ is identical to the maximization of $\mathcal{L}$ ($\min\operatorname{D_{KL}}\Leftrightarrow\max\mathcal{L}$). Minimizing $\operatorname{D_{KL}}$ is not feasible, because we don’t know the true posterior $p\mathopen{}\left(z|a\right)\mathclose{}$. Therefore, $\mathcal{L}$ is further investigated.

II-B Approximate Inference (i.e. rewriting $\mathcal{L}$)

[TABLE]

If the variable $a$ is replaced by some real valued sample $a^{\mathopen{}\left(i\right)\mathclose{}}$ (e.g. image or LiDAR scan), two terms can be identified:

[TABLE]

The first term is just a regularize that punishes the variational distribution $q$, that is the approximatior of the posterior distribution, if it deviates from some prior $p\mathopen{}\left(z\right)\mathclose{}$. The reconstruction term on the other hand compares the difference between the data $a^{\mathopen{}\left(i\right)\mathclose{}}$ of $q_{\phi}\mathopen{}\left(z|a^{\mathopen{}\left(i\right)\mathclose{}}\right)\mathclose{}$ wrt. the sampled data $a^{\mathopen{}\left(i\right)\mathclose{}}$ from the likelihood function $p_{\theta}\mathopen{}\left(a^{\mathopen{}\left(i\right)\mathclose{}}|z\right)\mathclose{}$.

That means if $\mathcal{L}$ is going to be maximized, the posterior function has to be equal to some prior and the data that is used to sample the latent feature $z$ from the posterior should be equal to the data from the likelihood function $p_{\theta}$. The objective is now, to find a function $q_{\phi}$ and $p_{\theta}$ which own these properties. Luckily, if some parametrized functions $q_{\phi}$ and $p_{\theta}$ (with parameters $\phi$ and $\theta$) are applied, every increase of Equation 13 means that the variational approximator $q_{\phi}\mathopen{}\left(z|a\right)\mathclose{}$ comes closer to the real posterior functions $p\mathopen{}\left(z|a\right)\mathclose{}$. Therefore, numerical optimization techniques like gradient descent can be applied to this issue. Commonly, two neuronal network, where each tries to find the optimal parameters $\phi$ and $\theta$, are applied to approximate the functions $q_{\phi}$ (i.e. the encoder) and $p_{\theta}$ (i.e. the decoder). However, since the true value of $L$ remains unknown, maximization of $\mathcal{L}$ can only be done until convergence. Thus, the overall procedure has the property of finding local optima.

III Joint Variational Autoencoder

Second, we expand the VAE from Sec. II to the marginal joint log-likelihood and derive the variational bound as follows:

[TABLE]

Approximate Inference (i.e. rewriting $\mathcal{L}_{\text{J}}$):

[TABLE]

Three different terms can be identified:

[TABLE]

A regularization for the joint encoder $q_{\phi_{\text{a}\text{b}}}$ and two reconstruction terms, one for each decoder $p_{\theta_{\text{a}}}$ and $p_{\theta_{\text{b}}}$.

IV Joint Multi-Modal Variational Autoencoder via Variation of Information

The issue with the joint VAE is the lacking possibility of encoding just one modality $a$ or $b$. Thus, Suzuki et al. [2] exploit the Variation of Information (VI) and derive the evidence lower bound wrt. the VI.

First, the conditional probability is investigated

[TABLE]

Further, the marginal log-likelihood of a conditional distribution can be written as:

[TABLE]

Further, the log-likelihood of the VI can be written as

[TABLE]

[TABLE]

With respect to Equation 27, the following regularization terms can be identified:

[TABLE]

IV-A Conclusion

The introduced KL regularization by Suzuki et al. [2] tries to find a mean representative of all parameters between clusters in the latent space. This is absolutely correct for the mean values but is insufficient for all other statistics of the distribution that are estimated by $q_{\phi_{\text{*}}}$. The approach would be sufficient, iff the true latent joint probability would be known, because only then the KL divergence is able to adapt to it in a correct manner. Unfortunately, this is not the case and thus all the estimated statistic values of the uni-modal encoders $q_{\phi_{\text{*}}}$, except the mean, are questionable.

V Proposed Joint Multi-Modal Autoencoder

V-A The Variational Bound

[TABLE]

V-B Approximating Inference (i.e. rewriting $\mathcal{L}_{\text{PJ}}$)

[TABLE]

Applying the corresponding function approximators, the formula can be written as:

[TABLE]

Investigating every line of the formula, the following properties can be identified: Equation 76 is the common multi-modal VAE loss derived from the joint probability, while Equation 77 adds the feature introduced by Suzuki et al. [2]. It introduces the KL regularization that brings the posterior distribution of an uni-modal encoder close to the distribution of the multi-modal case. The drawback of this approach is discussed in Section IV-A. The new lines, i.e. Equation 75 and 78, introduce the regularization of the uni-modal encoders wrt. the common prior and the reconstruction loss. The regularizer cares about the fact, that the uni-modal distribution does not deviate to much from the common prior while the reconstruction term shapes the remaining statistics including the mean. However, the last fact is very important, while the mean value in latent space might not be the best representative of the likelihood (i.e. the decoded data). This property cannot be respected by the KL divergence, but by the introduced reconstruction term.

V-B1 of General Expression for Arbitrary Number of Modalities

Comprising the applied steps to derive $\mathcal{L}_{\text{PJ}}$ from the former section, we can identify that by successively applying logarithm and Bayes rules, we derive the ELBO for the proposed multi-modal VAE as follows: First, given the independent set of observable modalities $\mathcal{M}=\{\text{a},\text{b},\text{c},\ldots\}$, its marginal log-likelihood $\log p\mathopen{}\left(\mathcal{M}\right)\mathclose{}=:L_{\text{M\textsuperscript{2}}}$ is multiplied by the cardinality of the set as the neutral element $1=\frac{|\mathcal{M}|}{|\mathcal{M}|}$. Second, applying logarithm multiplication rule, the nominator is written as the argument’s exponent. Third, Bayes rule is applied to each term wrt. the remaining observable modalities to derive their conditionals. Therefore, we can write

[TABLE]

The expression $\sum_{m\in\mathcal{M}}\log p\mathopen{}\left(m|\mathcal{M}\setminus m\right)\mathclose{}$ is the general form of the marginal log-likelihood for the variation of information (VI), as introduced by [2] for the JMVAE, for any set $\mathcal{M}$. Thus, it can be directly substituted with $L_{\text{M}_{\mathcal{M}}}$. The expression $\sum_{m\in\mathcal{M}}\log p\mathopen{}\left(\mathcal{M}\setminus m\right)\mathclose{}$ is the combination of all joint log-likelihoods of the subsets of $\mathcal{M}$ which have one less element. Therefore, this term can be rewritten as

[TABLE]

with $\widetilde{\mathcal{M}}=\left\{m|m\in\mathcal{P}\mathopen{}\left(\mathcal{M}\right)\mathclose{},|m|=|\mathcal{M}|-1\right\}$ Finally, $\log p\mathopen{}\left(\widetilde{m}\right)\mathclose{}$ can be substituted by $L_{\text{M\textsuperscript{2}}_{\widetilde{m}}}$ without loss of generality. However, it is worth noticing that substitution stops at the end of recursion and therefore, all final expressions $\log p\mathopen{}\left(\widetilde{m}\right)\mathclose{}\ \forall\ |\widetilde{m}|\equiv 1$ remain. $\square$

This results in the final recursive log-likelihood expression from which the ELBO can be directly derived as follows:

[TABLE]

VI Appendix

Variants of Bayes equation:

[TABLE]

[TABLE]

Logarithm rules:

[TABLE]

[TABLE]

Evidence lower bound:

[TABLE]

Marginal likelihood:

[TABLE]

Independent and identically distributed random variables (i.i.d. or iid or IID):

[TABLE]

VI-A Kulbeck-Leibler Divergence

[TABLE]

[TABLE]

more tbd.

VI-B Variation of Information

VI-B1 Operator Names

• •

$\operatorname{VI}\mathopen{}\left(A,B\right)\mathclose{}$: Variation of Information between some properties $A$ and $B$

• •

$\operatorname{I}\mathopen{}\left(A\right)\mathclose{}$: Information of $A$ (or mutual information)

• •

$\operatorname{I}\mathopen{}\left(A,B\right)\mathclose{}$: Mutual Information (MI) of $A$ and $B$

• •

$\operatorname{I}\mathopen{}\left(A,B|C\right)\mathclose{}$: Mutual Conditional Information (MCI) of $A$ and $B$ given $C$

• •

$\operatorname{H}\mathopen{}\left(A\right)\mathclose{}$: Entropy of $A$

• •

$\operatorname{H}\mathopen{}\left(A,B\right)\mathclose{}$: Joint Entropy (JE) of $A$ and $B$

• •

$\operatorname{H}\mathopen{}\left(A|B\right)\mathclose{}$: Conditional Entropy (CE) of $A$ given $B$

The Variation of Information (VI) between some random variables can be written as

[TABLE]

and

[TABLE]

and so on ….

VI-C Extension to three Modalities

It should be clear that both approaches, proposed and Suzuki’s [2], can be extended to multiple modalities. In the following, an example for three modalities is given.

First, the conditional probability is investigated:

[TABLE]

The log-likelihood of a single joint distribution can be written as:

[TABLE]

VI-C1 JMVAE for three Modalities

The log-likelihood of the VI between three distributions can be written as:

[TABLE]

The combined ELBO can then be rewritten as

[TABLE]

$\mathcal{L}_{\widetilde{\text{J}}}$ is the joint ELBO of a joint probability distribution having three arguments (i.e. $a$, $b$, $c$). The derivation is analog to Section III. The next steps are the application of encoders and decoders for this network which is straight forward and should be clear to the reader.

However, if we investigate the last equations, the following properties can be identified: There are the common reconstruction terms ($\operatorname{\mathbb{E}}$) for each decoder $p\mathopen{}\left(*|z\right)\mathclose{}$ wrt. the full multi-modal decoder $q\mathopen{}\left(z|a,b,c\right)\mathclose{}$. The KL terms show the drawback of the VI approach. As before, these regularizer tend to bring the encoders’ distribution to match each other. But only pairwise encoders (e.g. $p\mathopen{}\left(z|a,b\right)\mathclose{}$) remain and thus, uni-modal encoders are neglected.

This means from a practical point of view, that when we have $N$ modalities in a setup, we only can build derived setups having $N-1$ modalities.

VI-C2 Proposed JMVAE for three Modalities

The derivation from the joint log likelihood can be written analogously:

[TABLE]

It is now straight forward, by applying all former mentioned equations, to derive the ELBO for the above marginal log-likelihood. As one can imagine, the above equation results in a pretty heavy loss term but with the big advantage of respecting all permutations of modalities.

This means again from a practical point of view, in comparison to the approach by Suzuki et al. [2], that we can build arbitrary sensor setups having $1$ to $N$ modalities.