\section{Introduction}
A central goal for representation learning models is that the resulting latent representation should be \emph{compact} yet \emph{disentangled}. \emph{Compact} requires the representation $\mathbf{z}$ does not contain any nuance factors in the input signal $\mathbf{x}$ that are not relevant for the desired response $\mathbf{y}$~\cite{achille2018information}, whereas \emph{disentangled} means that $\mathbf{z}$ is factorizable and has consistent semantics associated to different generating factors of the underlying data generation process. To achieve this goal in an unsupervised fashion, different variational autoencoder (VAE)~\cite{kingma2014auto} based models have been developed in recent years by formulating the objective as finding an informative latent representation that can minimize reconstruction error under proper regularizations. Well-known regularizations include the channel capacity~\cite{burgess2018understanding} and the total correlation~\cite{watanabe1960information} associated with each dimension of $\mathbf{z}$~\cite{kim2018disentangling,chen2018isolating}.


Despite these recent efforts, the fundamental properties of the dynamics of learning for most existing VAE models still remain unknown and under-investigated. Information theory provides a solid methodology to analyze the dynamics of learning and different trade-offs in deep neural networks (DNNs)~\cite{shwartz2017opening}. This is because the encoding, the decoding, and the compression are precisely among the core problems information theory was made to solve. Therefore, a natural idea is to analyze existing VAE models by inspecting the evolution of some key information-theoretic quantities associated with the latent representations across training epochs. Albeit easy-to-understand, estimating information-theoretic quantities (such as mutual information) exactly in high-dimensional space is problematic~\cite{paninski2003estimation}. As a result, current attempts (e.g.,~\cite{alemi2018information,achille2018information}) on interpreting VAE models still lie on using rate-distortion (RD) theory~\cite[Chapter~10]{cover2012elements} to approximately compare the RD trade-off of different VAEs or illuminating its connection to the famed Information Bottleneck (IB)~\cite{tishby99information}. 




In this work, we address aforementioned problems by deriving a novel VAE objective from a \emph{first principle}. We also suggest a novel estimator to evaluate mutual information and total correlation in VAEs; and illuminate its usage on analyzing and interpreting VAEs. Our contributions are fourfold:


\begin{itemize}
    \item We propose a novel VAE objective based on the principle of relevant information~\cite[Chapter~8]{principe2010information}. We term it PRI-VAE and establish its connections to other VAEs.
    \item We introduce the recently developed matrix-based R{\'e}nyi's $\alpha$-entropy functional estimator~\cite{giraldo2014measures,yu2019multivariate} to measure the mutual information $\mathbf{I}(\mathbf{x};\mathbf{z})$ and the total correlation $\mathbf{T}(\mathbf{z})$ of different VAEs across training epochs.
    \item We observed that, with the common stochastic gradient descent (SGD) optimization, both $\mathbf{I}(\mathbf{x};\mathbf{z})$ and $\mathbf{T}(\mathbf{z})$ increase rapidly at the first few epochs and then continuously decrease in the remaining epochs. We also observed a positive correlation between $\mathbf{I}(\mathbf{x};\mathbf{z})$ and $\mathbf{T}(\mathbf{z})$.
   
    \item Experiments on the four benchmark data sets suggest that PRI-VAE encourages more disentangled representations and reasonable reconstruction quality with fewer training epochs. We also demonstrate the flexibility of PRI-VAE and generalize its idea to a higher level.
\end{itemize}


\section{Preliminary Knowledge}
\subsection{VAE and the Evidence Lower Bound}

A VAE operates with two probabilistic mappings, an encoder $X\mapsto Z$ (represented by a neural network with parameter $\phi$), and a decoder ($Z\mapsto X$ represented by another neural network with parameter $\theta$). Normally, we assume a fixed prior distribution $p(\mathbf{z}
)$ over $\mathbf{z}$. Since the distribution of $\mathbf{x}$ is also fixed (i.e., the data distribution $q(\mathbf{x})$), the encoder and decoder induce joint distributions $q(\mathbf{x},\mathbf{z})=q_\phi(\mathbf{z}|\mathbf{x})q(\mathbf{x})$ and $p(\mathbf{x},\mathbf{z})=p_\theta(\mathbf{x}|\mathbf{z})p(\mathbf{z})$, respectively. An ideal VAE objective is to maximize the marginalized log-likelihood:
\begin{equation}\label{eq_ideal_VAE}
    \mathbb{E}_{p(\mathbf{x})}[\log p_\theta(\mathbf{x})].\\[3pt]
\end{equation}

Eq.~(\ref{eq_ideal_VAE}) is, however, not tractable and is approximated by the evidence lower bound (ELBO)~\cite{kingma2014auto}:
\begin{equation}\label{eq_VAE_obj}
    \mathbb{E}_{q(\mathbf{z}|\mathbf{x})}[\log p(\mathbf{x}|\mathbf{z})]-\mathbb{E}_{p(\mathbf{x})}[D_{KL}(q_\phi(\mathbf{z}|\mathbf{x})\|p(\mathbf{z}))], \\[3pt]
\end{equation}
where the first term measures the reconstruction loss, and the second one is the regularization term, which corresponds to the Kullback-Leibler (KL) divergence between the latent distribution $q_\phi(\mathbf{z}|\mathbf{x})$ and the prior distribution $p(\mathbf{z})$.
Usually, we assume $p(\mathbf{z})$ follows a standard isotropic multivariate normal distribution $\mathcal{N}(0,I)$ where $I$ is the identify matrix.


\subsection{PRI: the General Idea and its Objective}\label{sec:PRI_introduction}
PRI is an unsupervised information-theoretic principle that aims to perform mode decomposition of a random variable $X$ with a known (and fixed) probability distribution $g$. Suppose we obtain a reduced statistical representation characterized by a random variable $Y$ with probability distribution $f$. 

The PRI casts this problem as a trade-off between the entropy $\mathbf{H}(f)$ of $Y$ and its descriptive power about $X$ in terms of their divergence $D(f\|g)$:
\begin{equation}\label{eq_PRI_obj}
    J(f) = \argmin\limits_{f}\mathbf{H}(f) + \gamma D(f\|g),
\end{equation}
where $\gamma$ is a hyper-parameter controlling the amount of relevant information that $Y$ can extract from $X$. The minimization of entropy can be viewed as a means of reducing uncertainty (or redundancy) and finding the statistical regularities in the outcomes of a process, whereas the minimization of information divergence ensures that such regularities are closely related to $X$. The PRI is similar in spirit to the IB approach, but the formulation is different because PRI does not require a relevant auxiliary variable $Z$ and the optimization is done directly on the random variable $X$.


Note that the choice of entropy and divergence is application-specific and depends mostly on the simplicity of optimization. Although this work still uses the basic Shannon's differential entropy and the KL divergence for a fair comparison, the $2$-order R{\'e}nyi's entropy functional~\cite{renyi1961measures} and the Parzen window density estimator~\cite{parzen1962estimation} evidence a long track record of usefulness in various machine learning applications~\cite{principe2010information}.
\textcolor{black}{Fig.~\ref{fig:pri_demo} illustrates a set of solutions revealed by PRI that are related to the principal curves or surfaces.}





\begin{comment}
can be simplified as:
\begin{equation}\label{eq_PRI_renyi}
\begin{aligned}
J(\mathbf{Y})&=&\argmin\limits_{\mathbf{Y}}[-(1-\gamma)\log\big(\frac{1}{N^2}\sum_{i,j=1}^NG_{\sigma}(y_i-y_j)\big)\\
&&-2\gamma\log\big(\frac{1}{N^2}\sum_{i,j=1}^NG_{\delta}(y_i-x_j)\big)].
\end{aligned}
\end{equation}
where $x_i\in X$, $y_i\in Y$, and $N$ is the cardinality of $X$. Eq.~(\ref{eq_PRI_renyi}) is differentiable with respect to $X$, and provides a set of solutions that are related to the principal curves~\cite{hastie1989principal} of $g$, as demonstrated in Fig.~\ref{fig:pri_demo}.
\end{comment}

\begin{figure}
    \centering
    \subfloat[3d Gaussian.]{\includegraphics[width=0.28\textwidth]{Figs/original_update.png}}
    \subfloat[$\gamma=0$]{\includegraphics[width=0.28\textwidth]{Figs/PRI_0_update.png}}
    \subfloat[$\gamma=1$]{\includegraphics[width=0.28\textwidth]{Figs/PRI_1_update.png}}\\
    \subfloat[$\gamma=2$]{\includegraphics[width=0.28\textwidth]{Figs/PRI_2_update.png}}
    \subfloat[$\gamma=5$]{\includegraphics[width=0.28\textwidth]{Figs/PRI_5_update.png}}
    \subfloat[$\gamma=100$]{\includegraphics[width=0.28\textwidth]{Figs/PRI_100_update.png}}
    \caption{Illustration of the structures revealed by the PRI for (a) a 3d isotropic Gaussian. As the values of $\gamma$ increase, the solution passes through (b) a single point, (c) mode, (d) principal curves, (e) principal surfaces, and in the extreme case of (f) $\gamma\rightarrow\infty$ we get back the data themselves as the solution.}
    \label{fig:pri_demo}
\end{figure}










\section{PRI-VAE}
As mentioned earlier, a good latent representation $\mathbf{z}$ should be \emph{compact} yet \emph{disentangled}. We thus propose the following VAE objective:
\begin{equation}
\begin{split}
    \mathcal{L}_{\mathrm{PRI-VAE}} &= \mathbb{E}_{p(\mathbf{x})}\left[\mathbb{E}_{q(\mathbf{z} | \mathbf{x})}\left[\log p(\mathbf{x} | \mathbf{z})\right]\right. \\
    &- \alpha \mathbf{H}(\mathbf{z}) - \beta D_{KL}[q_{\phi}(\mathbf{z}) \| p(\mathbf{z})]. \\[3pt]
\end{split}
\label{eq:PRI_vae}
\end{equation}

Eq.~(\ref{eq:PRI_vae}) has three terms, but each term has its own physical meaning. 
$\mathbb{E}_{p(\mathbf{x})}\left[\mathbb{E}_{q(\mathbf{z} | \mathbf{x})}\left[\log p(\mathbf{x} | \mathbf{z})\right]\right.$ guarantees a reliable reconstruction of input $\mathbf{x}$ from latent representation $\mathbf{z}$.
$\alpha\mathbf{H}(\mathbf{z})$ regularizes the uncertainty or the degree of \emph{compactness} in $\mathbf{z}$: a large $\alpha$ encourages our model to learn a more compact latent representation, whereas a small $\alpha$ may lead to the latent representation contains nuisance factors that are irrelevant for downstream tasks. On the other hand, $\beta$ controls the closeness between $q_\phi(\mathbf{z})$ and $p(\mathbf{z})$, which can be interpreted as the extent of \emph{disentanglement} of $\mathbf{z}$: a large $\beta$ enforces the independence in each dimension of $\mathbf{z}$, and should, therefore, be preferred. One should note that, the combination of $\alpha \mathbf{H}(\mathbf{z})+\beta D_{KL}[q_\phi(\mathbf{z})\| p(\mathbf{z})]$ is exactly the objective of PRI, in which the ratio $\beta/\alpha$ plays the same role as $\gamma$ in Eq.~(\ref{eq_PRI_obj}). We therefore term Eq.~(\ref{eq:PRI_vae}) the principle-of-relevant-information variational autoencoder (PRI-VAE).




\subsection{Optimizing the PRI-VAE}
To optimize the PRI-VAE objective, we need to compute the inference marginal $q(\mathbf{z}) = \mathbb{E}_{p(\mathbf{x})}[q(\mathbf{z} | \mathbf{x})]$.
However, in practice, exactly computing its value is intractable in the training phase, since it depends on the entire data set. 
In PRI-VAE, we use the weighted Monte Carlo approximation suggested in~\cite{chen2018isolating} to estimate $q(\mathbf{z})$, as it does not require additional hyper-parameters or inner optimization loops.



\subsection{Relation to Prior Art}\label{sec:PRI_VAE_relation}
In this section, we establish the connections between PRI-VAE and other recently developed VAEs.

$\beta$-VAE~\cite{higgins2017beta} learns disentangled representation by introducing a hyper-parameter $\beta$ (usually $\beta>1$) to heavily penalize the term
$D_{KL}(q_\phi(\mathbf{z}|\mathbf{x})\|p(\mathbf{z}))$ in Eq.~(\ref{eq_VAE_obj}). The emergence of disentanglement can be simply explained if we decompose $\mathbb{E}_{p(\mathbf{x})}[D_{KL}(q_\phi(\mathbf{z}|\mathbf{x})\|p(\mathbf{z}))]$ as (see proof in~\cite{kim2018disentangling}):
\begin{equation}\label{eq_decomp_KL}
   
    \underbrace{\mathbf{I}(\mathbf{x};\mathbf{z})}_{\circled{\tiny 1}}
    +\underbrace{D_{KL}(q_\phi(\mathbf{z})\|p(\mathbf{z})}_{\circled{\tiny 2}}). \\[3pt]
\end{equation}

To address this limitation, DIP-VAE~\cite{kumar2018variational} and InfoVAE~\cite{zhao2019infovae} suggest assigning different weights to the two terms in Eq.~(\ref{eq_decomp_KL}).
Moreover, suppose $p(\mathbf{z})$ is factorizable, \circled{\footnotesize 2} can be further decomposed as~\cite{kim2018disentangling,chen2018isolating}:
{\small
\begin{equation}\label{eq_tc_decompose}
\begin{split}
&\circled{\footnotesize 2} = 
\underbrace{D_{KL}(q(\mathbf{z})\|\prod\limits_{j=1}q(\mathbf{z}_j)}_{\circled{\tiny A}} + \sum_{d} \underbrace{\operatorname{D_{KL}}\left(q_{\phi}\left(\mathbf{z}_{j}\right) \| p\left(\mathbf{z}_{j}\right)\right)}_{\circled{\tiny B}}), \\[3pt]
\end{split}
\end{equation}
}
where term \circled{\footnotesize A} is the total correlation regularization associated with each dimension of $\mathbf{z}$ and term \circled{\footnotesize B} refers to the dimension-wise KL divergence ($d$ is the dimension of $\mathbf{z}$), which controls the penalty on the deviation between each latent dimension to the prior.
FactorVAE~\cite{kim2018disentangling} and $\beta$-TCVAE~\cite{chen2018isolating} add an external penalty on \circled{\footnotesize A} to encourage the independence in each dimension of $\mathbf{z}$.



Before our work, IB approach (e.g., \cite{achille2018information}) suggests a restriction on $\mathbf{I}(\mathbf{x};\mathbf{z})$ in Eq.~(\ref{eq_decomp_KL}) to keep $\mathbf{z}$ from containing nuance factors, whereas~\cite{alemi2018information,burgess2018understanding} introduce a hard constraint on the channel capacity to make sure $\mathbf{I}(\mathbf{x};\mathbf{z})$ is consistently larger than a predefined value.
PRI-VAE absorbs both proposals by flexibly adjusting the value of $\beta/\alpha$: a large $\alpha$ penalizes more on $\mathbf{H}(\mathbf{z})$ (and hence $\mathbf{I}(\mathbf{x};\mathbf{z})$, because $\mathbf{I}(\mathbf{x};\mathbf{z})$ is always saturated~\cite{esmaeili2019structured}) in practice, whereas a small $\alpha$ avoids $\mathbf{H}(\mathbf{z})$ reducing to zero, which also prevents a small channel capacity $\mathbf{I}(\mathbf{x};\mathbf{z})$. 
In an extreme case, when $\alpha=0$, we get back to the InfoVAE objective. 
We provide, in Table~\ref{tb:regularization}, an overview of objectives in different VAE models.


\begin{comment}
On the other hand, several works (e.g.,~\cite{achille2018information,alemi2018information}) have interpreted ELBO objective as an unsupervised rate-distortion problem. In this sense, reconstruction of the input is the distortion, whereas KL divergence between the encoder and prior distributions gives an upper bound on the information rate. \cite{brekelmans2019exact} introduces echo noise to exactly optimize the information rate (i.e., $\mathbf{I}(\mathbf{x};\mathbf{z})$) for arbitrary input distributions. When $\alpha=\beta=1$, the PRI term can be viewed as a tight upper bound to the information rate. This is because $\mathbf{I}(\mathbf{x};\mathbf{z})$ is always observed to be equal to $\mathbf{H}(\mathbf{z})$~\cite{saxe2018information}.
\end{comment}


\begin{table}[h]
\caption{Objectives of different VAE models. $\protect\circled{\footnotesize 1}^{*}$ represents the $\mathbf{H}(\mathbf{z})$ term, which is a tight upper bound of the term \protect\circled{\footnotesize 1}.}
\label{tb:regularization}
\renewcommand{\arraystretch}{1.2}
\begin{center}
\scalebox{1.0}{
\begin{tabular}{lcl}
\hline
Method & \multicolumn{1}{l}{Reconstruction term} & Regularization term \\ \hline
$\beta$-VAE~\cite{higgins2017beta} & \multirow{7}{*}{$\mathbb{E}_{q(\mathbf{z}|\mathbf{x})}[\log p(\mathbf{x}|\mathbf{z})]$} & $\beta(\circled{\footnotesize 1} + \circled{\footnotesize 2})$ \\
AnnealedVAE~\cite{burgess2018understanding} &  & $\gamma |(\circled{\footnotesize 1} + \circled{\footnotesize 2}) - C|$ \\
DIP-VAE~\cite{kumar2018variational} &  & $\circled{\footnotesize 1} + \lambda \circled{\footnotesize 2}$ \\
InfoVAE~\cite{zhao2019infovae} &  & $\lambda\circled{\footnotesize 2}$ \\
FactorVAE~\cite{kim2018disentangling} &  & $\circled{\footnotesize 1} + \gamma\circled{\footnotesize A} + \circled{\footnotesize B}$ \\
$\beta$-TCVAE~\cite{chen2018isolating} &  & $\circled{\footnotesize 1} + \beta\circled{\footnotesize A} + \circled{\footnotesize B}$ \\ \cline{1-1} \cline{3-3} 
PRI-VAE &  & $\alpha \circled{\footnotesize 1}^{*} + \beta\circled{\footnotesize 2}$ \\ \hline
\end{tabular}}
\end{center}
\end{table}


\begin{comment}
\textcolor{red}{We illustrate the rationality of such setting theoretically and empirically in Appendix 3.1.}.
\end{comment}





\section{Empirical behavior of $\mathbf{I}(\mathbf{x};\mathbf{z})$ and $\mathbf{T}(\mathbf{z})$}\label{sec:Empirical_behavior}




Motivated by~\cite{shwartz2017opening}, we study disentanglement behavior and the learning dynamics of VAEs by measuring the evolution of $\mathbf{I}(\mathbf{x};\mathbf{z})$ and $\mathbf{T}(\mathbf{z})$ across training epochs. 

Given input $\mathbf{x}\in\mathbb{R}^p$ and latent representation $\mathbf{z}\in\mathbb{R}^d$, the standard Shannon (differential) entropy functional defines $\mathbf{I}(\mathbf{x};\mathbf{z})$ and $\mathbf{T}(\mathbf{z})$ with Eq.~(\ref{eq_mutual_information}) and Eq.~(\ref{eq_total_corr}), respectively:
\begin{equation}\label{eq_mutual_information}
\begin{split}
\mathbf{I}(\mathbf{x}; \mathbf{z}) & = D_{\text{KL}}\left(p(\mathbf{x},\mathbf{z})\|p(\mathbf{x})p(\mathbf{z})\right) \\
& = \int\int p(\mathbf{x},\mathbf{z})\log\left(\frac{p(\mathbf{x},\mathbf{z})}{p(\mathbf{x})p(\mathbf{z})}\right)d\mathbf{x} d\mathbf{z}
\end{split}
\end{equation}

\begin{equation}\label{eq_total_corr}
\begin{split}
\mathbf{T}(\mathbf{z}) & = D_{\text{KL}}\left(p(\mathbf{z})\|p(z_1)\cdots p(z_d)\right) \\
& = \int\cdots\int p(\mathbf{z})\log\left(\frac{p(\mathbf{z})}{p(z_1)\cdots p(z_d)}\right)dz_1\cdots dz_d
\end{split}
\end{equation}

As can be seen, a precise estimation to the joint distribution 
$p(\mathbf{x},\mathbf{z})$ or $p(z_1,z_2,\cdots,z_d)$ is a prerequisite to obtain a reliable evaluation to both $\mathbf{I}(\mathbf{x};\mathbf{z})$ and $\mathbf{T}(\mathbf{z})$. Unfortunately, the density estimation in high or ultra-high dimensional space is always problematic and impractical. 
Taking the simplest $dSprites$ data set as an example, we have $\mathbf{x}\in\mathbb{R}^{4096}$ and $\mathbf{z}\in\mathbb{R}^{10}$ in the general setup.

To circumvent this issue, we use the recently proposed matrix-based R{\'e}nyi's $\alpha$-order entropy functional~\cite{giraldo2014measures,yu2019multivariate} to estimate $\mathbf{I}(\mathbf{x};\mathbf{z})$ and $\mathbf{T}(\mathbf{z})$. Unlike the standard Shannon entropy functional, the novel estimator is defined over the normalized eigenspectrum of a Hermitian matrix of the projected data in a reproducing kernel Hilbert space (RKHS). In this way, it avoids the explicit estimation of the underlying distributions of data. In the remaining of this section, we only demonstrate our observations related to the fundamental properties of VAEs. We leave the implementation details and the robustness of our estimator with respect to different hyper-parameter setting to the supplementary material.




\begin{comment}
{\small
\begin{equation}
\mathbf{S}_\alpha(A)=\frac{1}{1-\alpha}\log_2\left(\mathrm{tr}(A^\alpha)\right)=
\frac{1}{1-\alpha}\log_2\big[\sum_{i=1}^N\lambda_i(A)^\alpha\big] \label{eq_entropy}
\end{equation}}
where $A_{ij}=\frac{K_{ij}}{\mathrm{tr}(K)}$ and $\lambda_i(A)$ denotes its $i$-th eigenvalue.
\end{comment}



\begin{comment}
\begin{definition} Let $\kappa:\mathcal{X}\times\mathcal{X}\mapsto\mathbb{R}$ be a real valued positive definite kernel that is also infinitely divisible~\cite{bhatia2006infinitely}. Given $X=\{x_1,x_2,...,x_n\}$ and the Gram matrix $K$ obtained from evaluating a positive definite kernel $\kappa$ on all pairs of exemplars, that is $(K)_{ij}=\kappa(x_i,x_j)$, a matrix-based analogue to R{\'e}nyi's $\alpha$-entropy for a normalized positive definite (NPD) matrix $A$ of size $n\times n$,  such that $\mathrm{tr}(A)=1$, can be given by the following functional:
\begin{equation}
\mathbf{S}_\alpha(A)=\frac{1}{1-\alpha}\log_2\left(\mathrm{tr}(A^\alpha)\right)=
\frac{1}{1-\alpha}\log_2\big[\sum_{i=1}^n\lambda_i(A)^\alpha\big] \label{eq_entropy}
\end{equation}
where $A_{ij}=\frac{K_{ij}}{\mathrm{tr}(K)}$ and $\lambda_i(A)$ denotes its $i$-th eigenvalue.
\end{definition}

\begin{definition}
Given a collection of $n$ samples $\{s_i=(x_1^i,x_2^i,\cdots, x_k^i)\}_{i=1}^n$, where the superscript $i$ denotes the sample index, each sample contains $k$ ($k\geq2$) measurements $x_1\in \mathcal{X}_1$, $x_2\in \mathcal{X}_2$, $\cdots$, $x_k\in \mathcal{X}_k$ obtained from the same realization, and the positive definite kernels $\kappa_1:\mathcal{X}_1\times \mathcal{X}_1\mapsto\mathbb{R}$, $\kappa_2:\mathcal{X}_2\times \mathcal{X}_2\mapsto\mathbb{R}$, $\cdots$, $\kappa_k:\mathcal{X}_k\times \mathcal{X}_k\mapsto\mathbb{R}$, a matrix-based analogue to R{\'e}nyi's $\alpha$-order joint entropy among $k$ variables can be defined as:
\begin{equation}
\mathbf{S}_\alpha(A_1,A_2,\cdots,A_k)=\mathbf{S}_\alpha\left(\frac{A_1\circ A_2\circ\cdots\circ A_k}{\mathrm{tr}(A_1\circ A_2\circ\cdots\circ A_k)}\right) \label{eq_joint}
\end{equation}
where $(A_1)_{ij}=\kappa_1(x_1^i,x_1^j)$, $(A_2)_{ij}=\kappa_2(x_2^i,x_2^j)$, $\cdots$, $(A_k)_{ij}=\kappa_k(x_k^i,x_k^j)$, and $\circ$ denotes the Hadamard product.
\end{definition}

Given \textbf{Definition~1} and \textbf{Definition~2}, the matrix-based R{\'e}nyi's $\alpha$-order mutual information and total correlation can be measured, in analogy to the Shannon's information, with Eq.~(\ref{eq_mutual_information}) and Eq.~(\ref{eq_total_correlation}), respectively. See~\cite{giraldo2014measures,yu2019multivariate} on the theoretical foundations.
\begin{equation}
\mathbf{I}_\alpha(A;B)=\mathbf{S}_\alpha(A)+\mathbf{S}_\alpha(B)-\mathbf{S}_\alpha(A,B). \label{eq_mutual_information}
\end{equation}

\begin{equation}
\begin{split}
\mathbf{T}_\alpha(A_1,A_2,\cdots,A_k)&=\mathbf{S}_\alpha(A_1)+\mathbf{S}_\alpha(A_2)+\cdots+
\mathbf{S}_\alpha(A_k)\\
&-\mathbf{S}_\alpha\big(\frac{A_1\circ A_2\circ\cdots\circ A_k}{\mathrm{tr}(A_1\circ A_2\circ\cdots\circ A_k)}\big). \label{eq_total_correlation}
\end{split}
\end{equation}
\end{comment}




\begin{comment}
\textcolor{red}{The new estimator requires an infinite divisible kernel. The RBF kernel meets this requirement and was recommended by its inventors.
One should note that Eq.~(\ref{eq_entropy}) reduces the von Neumann entropy (same expression) when $\alpha\to1$ (see proof in Appendix 1.2), a measure of entanglement for pure states in quantum information theory~\cite{nielsen2002quantum}.}
\end{comment}




\begin{comment}
\begin{figure}
  \begin{subfigure}[t]{0.23\textwidth}
    \includegraphics[width=\textwidth]{MI_x_z.png}
   
   
  \end{subfigure}
  %
  \begin{subfigure}[t]{0.23\textwidth}
    \includegraphics[width=\textwidth]{Figs/TC_z.png}
   
   
  \end{subfigure}
  \caption{The learning curve of $I(\mathbf{x}; \mathbf{z})$ (left) and $TC(\mathbf{z})$ (right) of a $\beta$-VAE model trained on dSprites}
  \label{fig:learning_curve_of_MI_TC}
\end{figure}
\end{comment}

  

\begin{comment}
\begin{figure}
  \begin{subfigure}[t]{0.23\textwidth}
    \includegraphics[width=\textwidth]{Figs/correlation_mi_tc_1.png}
   
   
  \end{subfigure}
  %
  \begin{subfigure}[t]{0.23\textwidth}
    \includegraphics[width=\textwidth]{Figs/correlation_mi_tc_2.png}
   
   
  \end{subfigure}
  \caption{The correlation between $I(\mathbf{x}; \mathbf{z})$ and $TC(\mathbf{z})$ of a $\beta$-TCVAE model (left) trained on the dSprites and $beta$-VAE model (right) trained on the MNIST.}
  \label{fig:correlation_mi_tc}
\end{figure}
\end{comment}

\subsection{Compression and Positive Correlation}

Fig.~\ref{fig:learning_curve_of_MI_TC_alpha-1_1} demonstrates the evolution of $\mathbf{I}(\mathbf{x};\mathbf{z})$ and $\mathbf{T}(\mathbf{z})$ for two representative VAEs across training iterations. 
Obviously, during the common SGD optimization, both $\mathbf{I}(\mathbf{x};\mathbf{z})$ and $\mathbf{T}(\mathbf{z})$ show two separate phases: an early ``fitting" phase, in which both $\mathbf{I}(\mathbf{x};\mathbf{t})$ and $\mathbf{I}(\mathbf{t};\mathbf{y})$ increase rapidly, and a later ``compression" phase, in which there is a reversal such that $\mathbf{I}(\mathbf{x};\mathbf{t})$ and $\mathbf{I}(\mathbf{t};\mathbf{y})$ continually decrease.
Our observation fits well with the arguments of the mainstream information theory community.




More interestingly, we also observe a significant positive correlation between $\mathbf{I}(\mathbf{x};\mathbf{z})$ and $\mathbf{T}(\mathbf{z})$, as manifested by their Pearson correlation coefficient.
Note that, a small value of $\mathbf{T}(\mathbf{z})$ is always a good indicator to model disentanglement~\cite{chen2018isolating}. 
According to our observation, if $\mathbf{T}(\mathbf{z})$ is optimized to a small value, $\mathbf{I}(\mathbf{x};\mathbf{z})$ should be small accordingly (even zero if fully disentangled). However, it is counter-intuitive that an extremely small value of $\mathbf{I}(\mathbf{x};\mathbf{z})$ is able to faithfully reconstruct the input~\cite{higgins2017beta,chen2018isolating}. \textcolor{black}{In this sense, our observation provides empirical evidence to support the argument in~\cite{zhao2019infovae}, which states that existing VAE models tend to learn uninformative latent representation. It also suggests there is always a trade-off between reconstruction and disentanglement~\cite{lezama2018overcoming,esmaeili2019structured,kim2018disentangling}}. 

\begin{figure}
    \centering
    \includegraphics[width=0.9\textwidth]{Figs/learning_curve_v4.png}
    \caption
    The evolution of the value of $\mathbf{I}(\mathbf{x};\mathbf{z})$ (left column) and the value of $\mathbf{T}(\mathbf{z})$ (middle column) across training iterations for FactorVAE, $\beta$-TCVAE and PRI-VAE on \textit{Cars3d} data set. The right column shows the positive correlation between these two values.}
    \label{fig:learning_curve_of_MI_TC_alpha-1_1}
\end{figure}

\subsection{Observations with Other Estimators}
Our aim here is to verify if the general observation made by our suggested matrix-based R{\'e}nyi's $\alpha$-entropy functional are consistent with other existing Shannon estimators. To this end, we estimate $\mathbf{T}(\mathbf{z})$ with the widely used $k$NN estimator~\cite{kraskov2004estimating} and kernel density estimator (KDE)~\cite{moon1995estimation}. We also apply the recently proposed  ensemble dependency graph estimator (EDGE)~\cite{noshad2019scalable} to measure $\mathbf{I}(\mathbf{x};\mathbf{z})$. Fig.~\ref{fig:edge} suggests that our general observation is in accordance with other popular estimators. However, note that both $k$NN and KDE suffer from the curse of dimensionality, and it still remains a problem on estimating total correlation with EDGE. 


\begin{figure}
    \centering
    \subfloat{\includegraphics[width=0.5\textwidth]{edge/cars3d/beta_tc_vae-edge-4_2.png}} \hspace{0.05\textwidth}
    \subfloat{\includegraphics[width=0.4\textwidth]{consistency_estimators/total_correlation_new.png}}
    \caption{The evolution of value of $\mathbf{I}(\mathbf{x};\mathbf{z})$ (left) measured by EDGE estimator and $\mathbf{T}(\mathbf{z})$ (right) measured by $k$NN and KDE.}
    \label{fig:edge}
\end{figure}







\begin{comment}
\subsection{The Saturation of $\mathbf{I}(\mathbf{x};\mathbf{z})$}

Another interesting phenomenon we observed is the saturation of $\mathbf{I}(\mathbf{x};\mathbf{z})$, i.e., $\mathbf{I}(\mathbf{x};\mathbf{z})$ is always equal to $\mathbf{H}(\mathbf{z})$ (see Fig.~\ref{fig:saturation}). Theoretically, $\mathbf{I}(\mathbf{x};\mathbf{z})=\mathbf{H}(\mathbf{z})-\mathbf{H}(\mathbf{z}|\mathbf{x})\leq\mathbf{H}(\mathbf{z})$, where $\mathbf{H}(\mathbf{z}|\mathbf{x})$ denotes the conditional entropy of $\mathbf{z}$ given $\mathbf{x}$. Therefore, the approximate equivalence also indicates a near-zero value of $\mathbf{H}(\mathbf{z}|\mathbf{x})$.


A zero $\mathbf{H}(\mathbf{z}|\mathbf{x})$ is certain for a deterministic autoencoder with discrete input signal. Specifically, since the neural network is deterministic, the distribution of the layered representation $\mathbf{t}$ is discrete as well and is entirely determined by $\mathbf{x}$, from which follows that $\mathbf{I}(\mathbf{x};\mathbf{t})=\mathbf{H}(\mathbf{t})$~\cite{amjad2019learning,goldfeld2019estimating}. However, the existence of saturation of $\mathbf{I}(\mathbf{x};\mathbf{z})$ in VAEs with stochastic mappings may suggest the feasibility of learning meaningful latent space by substituting stochasticity with explicit regularization schemes, although some initial works have been done (e.g.,~\cite{ghosh2019variational,wu2019quantization}).



\begin{figure}[ht]
  \begin{subfigure}[t]{0.23\textwidth}
    \includegraphics[width=\textwidth]{Figs/dsprites_single_MIxz-Hz-1_1_rmtd7.png}
    \caption{\textit{dSprites}}
   
  \end{subfigure}
  %
  \begin{subfigure}[t]{0.23\textwidth}
    \includegraphics[width=\textwidth]{Figs/cars3d_single_MIxz-Hz-1_2_rmtd7.png}
    \caption{\textit{Cars3D}}
   
  \end{subfigure}
  \caption{The box plots of $\mathbf{I}(\mathbf{x};\mathbf{z})$ and $\mathbf{H}(\mathbf{z})$ across training iterations for different VAE models in (a) \textit{dSprites} and (b) \textit{Cars3D} data sets.}
  \label{fig:saturation}
\end{figure}
\end{comment}




\begin{comment}
In addition, as the measurement of the dependency between the input data $\mathbf{x}$ and latent variable $\mathbf{z}$, the mutual information $I(\mathbf{x}; \mathbf{z})$ can be expressed as: 
\begin{equation}
    I(\mathbf{x};\mathbf{z}) = H(\mathbf{x}) + H(\mathbf{z}) - H(\mathbf{x}, \mathbf{z})
    \label{eq:mi}
\end{equation}
where $H(\mathbf{x})$ and $H(\mathbf{z})$ are the marginal entropies of $\mathbf{x}$ and $\mathbf{z}$ respectively, and the $H(\mathbf{x}, \mathbf{z})$ is their joint entropy.
\end{comment}



\section{Experiments}\label{sec:experiment}
We evaluate the performances of our proposed PRI-VAE on four benchmark data sets to demonstrate its capability to learn disentangled, interpretable latent representations, and reconstruct input signal with high fidelity. We compare PRI-VAE with six SOTA VAE models, namely $\beta$-VAE~\cite{higgins2017beta}, AnnealedVAE~\cite{burgess2018understanding}, DIP-VAE~\cite{kumar2018variational}, InfoVAE~\cite{zhao2019infovae}, FactorVAE~\cite{kim2018disentangling}, and $\beta$-TCVAE~\cite{chen2018isolating}.
We evaluate the disentanglement and reconstruction performances of all competing models quantitatively (in Section~\ref{sec:quantitative}) and qualitatively (in Section~\ref{sec:qualitative}). We also illustrate the potential implications of our methodology for future work and generalize the idea of PRI-VAE to a higher level in Section~\ref{sec:implication}.
For a fair comparison, all competing methods are trained with the same network architecture, optimizer, and mini-batch and evaluated by unified disentangling metrics~\cite{locatello2019challenging}. 





\subsection{Quantitative Evaluation} \label{sec:quantitative}
We perform quantitative evaluation on \textit{dSprites}~\cite{dSprites17} and \textit{Cars3D}~\cite{reed2015deep}. Samples in both data sets are generated with ground-truth independent latent factors. 
Since most of the existing metrics on disentanglement evaluation are positively correlated \cite{locatello2019challenging}, we consider the most widely used Mutual Information Gap (MIG)~\cite{chen2018isolating} and DCI Disentanglement metric~\cite{ridgeway2018learning,locatello2019challenging} in this work.




\subsubsection{Comparison of PRI-VAE and InfoVAE Family}\label{sec:PRI_synthetic_comparison1}

The most similar objective to PRI-VAE is the InfoVAE family~\cite{zhao2019infovae} with the regularization term $\alpha\circled{\footnotesize 1} + \lambda\circled{\footnotesize 2}$.
Despite reaching similar objectives, the routes to this result are rather different. 
Our PRI-VAE is motivated by an information-theoretic perspective in the sense that we want to learn a compact (regularized by $\alpha\mathbf{H}(\mathbf{z})$) latent representation that is able to reconstruct input signal (regularized by $\mathbb{E}_{p(\mathbf{x})}\left[\mathbb{E}_{q(\mathbf{z}|\mathbf{x})}\left[\log p(\mathbf{x}|\mathbf{z})\right]\right]$) and is also disentangled (regularized by $\beta D_{KL}[q_{\phi}(\mathbf{z})\| p(\mathbf{z})]$). By contrast, InfoVAE family originates from a generative model perspective that attempts to address two fundamental issues in the original VAE (i.e., the inaccurate amortized inference distributions and the vanishing relevance of latent variable and input signal). 
\textcolor{black}{Although authors of InfoVAE suggest setting $\alpha=0$ to completely drop \circled{\footnotesize 1}, we demonstrate here the superiority gained from the difference, i.e.,  $\alpha(\mathbf{H}(\mathbf{z}) - \mathbf{I}(\mathbf{x};\mathbf{z}))$.} 


\begin{figure}[t]
    \centering
    \subfloat{\includegraphics[width=0.45\textwidth]{compare_infovae/mig-5-151-merge_diff_runs_v2_compress.png}} \hspace{0.03\textwidth}
    \subfloat{\includegraphics[width=0.47\textwidth]{compare_infovae/run_rec-5-151_sm0-8_compress.png}} 
    \caption{The evolutions of MIG score (left) and reconstruction loss (right, after exponential moving average with weight $0.8$) for PRI-VAE and InfoVAE family with respect to different values of $\alpha$ when $\beta$ is fixed to $4$.}
    \label{fig:comparison_infovae_dSprites}
    \vspace{-1em}
\end{figure}

\begin{figure}
    \centering
    \subfloat{\includegraphics[width=0.45\textwidth]{compare_infovae/mig-5-161-merge_diff_runs_v2_compress.png}} \hspace{0.03\textwidth}
    \subfloat{\includegraphics[width=0.47\textwidth]{compare_infovae/run_rec-5-161_sm0-8_compress.png}} 
    \caption{The evolutions of MIG score (left) and reconstruction loss (right, after exponential moving average with weight $0.8$) for PRI-VAE and InfoVAE family with respect to different values of $\beta$ when $\alpha$ is fixed to $1$.}
    \label{fig:comparison_infovae_2_dSprites}
\end{figure}



We first evaluate the performances of PRI-VAE and InfoVAE family with respect to different values of $\alpha$ by fixing $\beta=4$ (a common setting in previous works~\cite{higgins2017beta,esmaeili2019structured}).
Fig.~\ref{fig:comparison_infovae_dSprites} demonstrates the evolution of the MIG score and the reconstruction loss across training iterations on \textit{dSprites} data set. Under the same hyper-parameter setting (i.e., same value of $\alpha$), PRI-VAE can quickly obtain more disentangled latent representation with smaller reconstruction error than its InfoVAE counterpart. This is not surprising. Note that $\mathbf{H}(\mathbf{z}) - \mathbf{I}(\mathbf{x};\mathbf{z})\geq0$, i.e., $\mathbf{H}(\mathbf{z})$ is always a heavy penalty than $\mathbf{I}(\mathbf{x};\mathbf{z})$. Therefore, a faster convergence is expected. 





Fig.~\ref{fig:comparison_infovae_2_dSprites} reaffirms the advantage of PRI-VAE over InfoVAE family concerning different values of $\beta$. Another interesting observation is that, with the increase of $\beta$, the MIG score keeps increasing, whereas the reconstruction quality is reduced. This is probably because there is a trade-off between reconstruction and disentanglement~\cite{lezama2018overcoming,kim2018disentangling,esmaeili2019structured}. 



\subsubsection{Comparison of PRI-VAE and other SOTA VAE Models} \label{sec:PRI_synthetic_comparison2}
Next, we compare PRI-VAE with VAE, $\beta$-VAE, AnnealedVAE, DIP-VAE-\RNum{1}, DIP-VAE-\RNum{2}, FactorVAE and $\beta$-TCVAE.
As recommended in~\cite{locatello2019challenging}, we finalize the hyperparameters for each methods by sweeping over a wide enough regularization weight range. We then fix the hyperparameters of all methods on different data sets.
Specifically, we select $\beta=4$ for $\beta$-VAE and $\beta$-TCVAE, $\gamma=30$ for FactorVAE, $\lambda_{od}=5$, $\lambda_{d}=10\lambda_{od}$ for DIP-VAE-\RNum{1} and $\lambda_{od}=5$, $\lambda_{d}=\lambda_{od}$ for DIP-VAE-\RNum{2}.
For AnnealedVAE, we fix $\gamma=1000$ and linearly increase the channel capacity $C$ from $0$ to $25$ over the course of $10,000$ training steps.
For our model, we select $\alpha=0.6$ and $\beta=6.0$\footnote{\textcolor{black}{We recommend $\beta\geq 2\alpha$. This is because $\mathbf{z}$ converges to a single point if $\beta\leq \alpha$ (see Fig.~\ref{fig:pri_demo}), from which it is hard to reconstruct $\mathbf{x}$.}}. Ten independent runs are performed.

\begin{figure}[t]
    \centering
    \subfloat{\includegraphics[width=0.44\textwidth]{compare_baseline/dsprites_mean_mig_4-1.png}} \hspace{0.05\textwidth}
    \subfloat{\includegraphics[width=0.44\textwidth]{compare_baseline/cars3d_mean_mig_4-1.png}} 
    \caption{The MIG score comparison on \textit{dSprites} (left) and \textit{Cars3D} (right) data sets.}
    \label{fig:compare_baseline_mig}
\end{figure}


\begin{comment}
\begin{figure}
  \begin{subfigure}[t]{0.23\textwidth}
  \includegraphics[width=\textwidth]{compare_baseline/dsprites_mean_dci_4-1.png}
   
   
   
   
  \end{subfigure}
  %
  \begin{subfigure}[t]{0.23\textwidth}
  \includegraphics[width=\textwidth]{compare_baseline/dsprites_mean_mig_4-1.png}
   
   
   
   
  \end{subfigure}
  \caption{The DCI disentanglement metric (left) and the MIG score (right) on \textit{dSprites} data set.}
  \label{fig:dsprites_compare}
    \vspace{-1em}
\end{figure}

\begin{figure}
  \begin{subfigure}[]{0.23\textwidth}
  \includegraphics[width=\textwidth]{compare_baseline/cars3d_mean_dci_4-1.png}
   
   
   
   
  \end{subfigure}
  %
  \begin{subfigure}[]{0.23\textwidth}
    \includegraphics[width=\textwidth]{compare_baseline/cars3d_mean_mig_4-1.png}
   
   
   
   
  \end{subfigure}
  \caption{The DCI disentanglement metric (left) and the MIG score (right) on \textit{Cars3D} data set.}
    \vspace{-0.3em}
  \label{fig:cars3d_compare}
\end{figure}
\end{comment}


Fig.~\ref{fig:compare_baseline_mig} demonstrates the MIG score for all competing models in \textit{dSprites} and \textit{Cars3D}, respectively\footnote{Same to \cite{locatello2019challenging,esmaeili2019structured}, we also observe that the hyper-parameter tuning and the random seed have substantial impacts to the performances of different models.}. 
The comparison result on DCI disentanglement metric is similar, which is shown in Appendix Fig.~\ref{fig:compare_baseline_dci}.
As can be seen, PRI-VAE achieves consistent better disentanglement performance than VAE, $\beta$-VAE, AnnealedVAE, and DPI-VAE, and preserves a satisfactory reconstruction quality suggested by our visualization results (see Appendix Fig.~\ref{fig:privae_reconstruction}).
However, $\beta$-TCVAE and FactorVAE  obtain higher disentanglement scores than ours. 
This is not surprising. Both FactorVAE and $\beta$-TCVAE add a new regularization term $\mathbf{T}(\mathbf{z})$ with extra weight to the original VAE objective. 
One should note that $\mathbf{T}(\mathbf{z})$ is always a more explicit and heavy penalty than $D_{KL}(q_\phi(\mathbf{z})\|p(\mathbf{z}))$ to encourage independence of each dimension of $\mathbf{z}$~\cite{chen2018isolating}.
We will show in Section~\ref{sec:implication_2} that the merits of PRI-VAE and $\beta$-TCVAE can be merged to significantly improve the performances of both methods.  
 
\subsection{Qualitative Evaluation} \label{sec:qualitative}

\begin{comment}
\begin{figure*}
    \centering
    \subcaptionbox{Chair size\label{fig:chair_size}}{\includegraphics[width=0.25\textwidth]{Figs/chairs_size.png}}\hspace{3em}%
    \subcaptionbox{Azimuth\label{fig:Azimuth}}{\includegraphics[width=0.25\textwidth]{Figs/chars_Azimuth.png}}\hspace{3em}%
    \subcaptionbox{Backrest\label{fig:Backrest}}{\includegraphics[width=0.25\textwidth]{Figs/chairs_Backrest.png}
    
    \centering
    \subcaptionbox{Leg style\label{fig:leg_style}}{\includegraphics[width=0.25\textwidth]{Figs/chairs_leg_style.png}}\hspace{3em}%
    \subcaptionbox{Leg length\label{fig:leg_length}}{\includegraphics[width=0.25\textwidth]{Figs/chairs_leg_length.png}}\hspace{3em}%
    \subcaptionbox{Material\label{fig:material}}{\includegraphics[width=0.25\textwidth]{Figs/chairs_material.png}}%

   
   
   
   
   
   
   
   
   
   
   
   
   
   
   
   
   
   
   
   
   
   
   
   
   
   
   
   
   
   
    \caption{Latent traversals of the model trained on \textit{3D chairs}. Each column corresponds to varying a single latent unit. The traversal is over the $[-2, 2]$ range.}
    \label{fig:charis_traverals}
\end{figure*}
\end{comment}

We then quantitatively evaluate PRI-VAE on \textit{3D Chairs}~\cite{aubry2014seeing}, and \textit{Fashion-MNIST}~\cite{xiao2017fashion}.
Fig.~\ref{fig:chairs_traversal_partial} and Fig.~\ref{fig:FashionMNIST} show the traversal results of the latent variables on both data sets, respectively. Similar to $\beta$-TCVAE, PRI-VAE is able to discover six factors of variation in \textit{3D Chairs}, such as chair size, azimuth, backrest, leg style, leg length, and chair material. By contrast, $\beta$-VAE can only learn four factors~\cite{chen2018isolating}. Moreover, PRI-VAE discovers a unique factor (i.e., leg separation) in~\textit{Fashion-MNIST}, which has never been mentioned in previous works. 
\begin{figure}[h]
\vspace{-0em}
    \centering
    \includegraphics[width=0.8\textwidth]{Figs/charis_partial_2.png}
    \vspace{-1em}
    \caption{Latent traversals for the PRI-VAE ($\alpha=0.6, \beta=4$) trained on \textit{3D chairs}. Each column corresponds to varying a single latent unit. The traversal is over the $[-2, 2]$ range.}
    \label{fig:chairs_traversal_partial}
\end{figure}

\begin{figure}[h]
    \centering
    \includegraphics[width=0.78\textwidth]{Figs/FashionMNIST_a0_6_b8_t3_v4_m}
    \caption{Latent traversals for the PRI-VAE ($\alpha=0.6, \beta=8$) trained on \textit{Fashion-MNIST}. Each column corresponds to varying a single latent unit. The traversal is over the $[-3, 3]$ range.}
    \label{fig:FashionMNIST}
\end{figure}





\begin{comment}
\subsection{Hyperparameter Analysis} \label{sec:hyperparameter}
In this section, we investigate the effects of hyper-parameters $\alpha$ and $\beta$ to the performance of PRI-VAE to corroborate our analysis in Sections~\ref{sec:PRI_introduction} and \ref{sec:PRI_VAE_relation}. 
\begin{figure}
    \centering
    \includegraphics[width=0.45\textwidth]{Figs/PriVAE_hyperparameters.png}
    \caption{Caption}
    \label{fig:PriVAE_hyperparameters}
\end{figure}
\end{comment}

\vspace{-1em}
\subsection{Implications for Future Work} \label{sec:implication}

We finally present two positive implications of our methodology for future work with convincing empirical validation.

\subsubsection{Merging the Merits of PRI-VAE and $\beta$-TCVAE}\label{sec:implication_2}

\begin{comment}
\begin{figure}[ht]
  \begin{subfigure}[]{0.23\textwidth}
    \includegraphics[width=\textwidth]{Figs/dsprites_PirVAE_star_mig_2_3.png}
   
   
  \end{subfigure}
  %
  \begin{subfigure}[]{0.23\textwidth}
    \includegraphics[width=\textwidth]{Figs/cars3d_PirVAE_star_mig_2_1.png}
   
   
  \end{subfigure}
  \caption{The MIG score comparison of $\beta$-TCVAE, PRI-VAE and PRI-VAE$^\star$ on \textit{dSprites} (left) and \textit{Cars3D} (right) data set.}
  \label{fig:pri_vae_star}
\end{figure}
\end{comment}

We continue our discussion in Section~\ref{sec:PRI_synthetic_comparison2} to demonstrate the benefits gained from merging the merits of PRI-VAE and $\beta$-TCVAE.
As discussed earlier, FactorVAE and $\beta$-TCVAE have better disentanglement performance than others by explicitly regularizing an extra term $\mathbf{T}(\mathbf{z})$. 
On the other hand, the PRI is geometrically interpretable and enables faster convergence. 
Therefore, a natural idea to combine the merits of both methodologies is via the following objective:
\textcolor{black}{
\begin{equation}
\begin{split}
&\mathcal{L}_{\mathrm{PRI}-\mathrm{VAE}^\star} \\
    &=\mathbb{E}_{p(\mathbf{x})}\left[\mathbb{E}_{q(\mathbf{z} | \mathbf{x})}\left[\log p(\mathbf{x} | \mathbf{z})\right]\right] - \alpha H(\mathbf{z}) \\
    &-\beta D_{KL}[q(\mathbf{z}) \| p(\mathbf{z})] - \gamma D_{KL}(q(\mathbf{z}) \| \prod_{j} q\left(\mathbf{z}_{j}\right)).
\end{split}
\label{eq:pri_vae_star}
\end{equation}}

We term Eq.~(\ref{eq:pri_vae_star}) PRI-VAE$^\star$ since it just adds a weighted term $\gamma\mathbf{T}(\mathbf{z})$ to the original PRI-VAE objective.
\textcolor{black}{In Eq.~(\ref{eq:pri_vae_star}), the term $\alpha H(\mathbf{z}) + \beta D_{KL}q(\mathbf{z})||p(\mathbf{z}))$ can be interpreted as an upper bound of the exact information rate $\mathbf{I}(\mathbf{x};\mathbf{z})$, whereas the new regularization term $\gamma D_{KL}(q(\mathbf{z}) \| \prod_{j} q\left(\mathbf{z}_{j}\right))$ further enforces the disentanglement. This way, we are actually optimizing a regularized rate-distortion problem. Given that minimizing $\mathbf{T}(\mathbf{z})$ is likely to incur the decrease of $\mathbf{I}(\mathbf{x};\mathbf{z})$, the distortion may increase accordingly. }

We compare PRI-VAE$^\star$ with PRI-VAE and $\beta$-TCVAE on \textit{dSprites} and \textit{Cars3D}.
We keep the same hyper-parameter setting for PRI-VAE and $\beta$-TCVAE as shown in Section~\ref{sec:PRI_synthetic_comparison2}. 
For PRI-VAE$^\star$, we set $\alpha=0.5, \beta=1.0$ and $\gamma=4.0$. 
Table~\ref{tb:pri_vae_star} summarizes the quantitative evaluation results and DCI comparison is shown in Appendix Table~\ref{tb:pri_vae_star_dci}.
As can be seen, PRI-VAE$^\star$ has improved disentanglement than PRI-VAE and it is slightly better than $\beta$-TCVAE. 
Note that the hyperparameters of all the methods are not additionally tuned.

\begin{comment}
\begin{table}[ht]
\caption{The MIG and DCI disentanglement scores of $\beta$-TCVAE, PRI-VAE and PRI-VAE$^\star$ on \textit{dSprites} and \textit{Cars3D}. We report the mean value over $10$ runs and the standard deviation in parentheses.} \label{tb:pri_vae_star}
\renewcommand{\arraystretch}{1.5}
\begin{center}
\scalebox{0.9}{
\begin{tabular}{ccccc}
\hline
\multicolumn{1}{l}{} & \multicolumn{2}{c}{\textit{dSprites}} & \multicolumn{2}{c}{\textit{Cars3D}} \\[-2pt]
 & MIG & DCI & MIG & DCI \\ \hline
$\beta$-TCVAE & 0.194 (0.030) & 0.318 (0.040) & 0.113 (0.038) & 0.343 (0.037) \\ \hline
PRI-VAE & 0.155 (0.042) & 0.234 (0.058) & 0.112 (0.021) & 0.281 (0.050) \\ \hline
PRI-VAE$^\star$ & \mathbf{0.198} (0.063) & \mathbf{0.345} (0.088) & \mathbf{0.132} (0.033) & \mathbf{0.321} (0.048) \\ \hline
\end{tabular}}
\end{center}
\end{table}
\end{comment}

\begin{table}[h]
\caption{The MIG values of $\beta$-TCVAE, PRI-VAE and PRI-VAE$^\star$ on \textit{dSprites} and \textit{Cars3D}. We report the mean value over $10$ runs and the standard deviation in parentheses.} \label{tb:pri_vae_star}
\renewcommand{\arraystretch}{1.5}
\begin{center}
\begin{tabular}{llll}
\hline
 & $\beta$-TCVAE & PRI-VAE & PRI-VAE$^\star$ \\ \hline
\textit{dSprites} & 0.194 (0.030) & 0.155 (0.042) & \textbf{0.198 (0.063)} \\ \hline
\textit{Cars3D} & 0.113 (0.038) & 0.112 (0.021) & \textbf{0.132 (0.033)} \\ \hline
\end{tabular}
\end{center}
\end{table}

\subsubsection{Replacing $\mathbf{I}(\mathbf{x};\mathbf{z})$ with $\mathbf{H}(\mathbf{z})$ in SOTA VAE models}

As demonstrated in Section~\ref{sec:PRI_synthetic_comparison1}, PRI-VAE achieves an obvious performance gain over InfoVAE family by replacing $\mathbf{I}(\mathbf{x};\mathbf{z})$ with $\mathbf{H}(\mathbf{z})$. 
We now investigate if this modification can be generalized to other SOTA VAE models. 
Again, we use the $\beta$-TCVAE as an example, and suggest the following objective (term it $\beta$-TCVAE$^\star$):
{\small
\begin{equation}
\begin{split}
&\mathcal{L}_{\beta-\mathrm{TCVAE}^\star} \\
    &=\mathbb{E}_{p(\mathbf{x})}\left[\mathbb{E}_{q(\mathbf{z} | \mathbf{x})}\left[\log p(\mathbf{x} | \mathbf{z})\right]\right] - H(\mathbf{z}) \\
    &-\beta D_{KL}(q(\mathbf{z}) \| \prod_{j} q\left(\mathbf{z}_{j}\right)) - \sum_{j} D_{KL}(q\left(\mathbf{z}_{j}\right) \| p\left(\mathbf{z}_{j}\right)).
\end{split}
\end{equation}
}
We compare $\beta$-TCVAE$^\star$ with the original $\beta$-TCVAE with the same value of $\beta$. The comparison result on \textit{dSprites} is shown in Fig.~\ref{fig:pri_beta_tcvae}.
It is interesting to find that $\beta$-TCVAE$^\star$ is consistently better than classical $\beta$-TCVAE in terms of both disentanglement and reconstruction. 




\begin{figure}[hb]
    \centering
    \subfloat{\includegraphics[width=0.44\textwidth]{Figs/PRI_beta_tcvae_mig-10_2-merged.png}} \hspace{0.03\textwidth}
    \subfloat{\includegraphics[width=0.44\textwidth]{Figs/PRI_beta_tcvae_run_rec-10_2_sm0_6.png}} 
    \caption{The comparison of disentanglement (left) and reconstruction (right) between $\beta$-TCVAE$^\star$ and $\beta$-TCVAE on \textit{dSprites}.}
    \label{fig:pri_beta_tcvae}
\end{figure}


\begin{comment}
\begin{figure}[ht]
  \begin{subfigure}[]{0.23\textwidth}
    \includegraphics[width=\textwidth]{Figs/dsprites_full_compare_mig_2_2.png}
   
   
  \end{subfigure}
  %
  \begin{subfigure}[]{0.23\textwidth}
    \includegraphics[width=\textwidth]{Figs/dsprites_full_compare_factor_score_2_2.png}
   
   
  \end{subfigure}
  \caption{The comparison of MIG score (left) and FactorVAE score (right) for all the studied models in this paper on \textit{dSprites} data set.}
  \label{fig:dsprites_all_compare}
\end{figure}
\end{comment}

\begin{comment}
\subsubsection{Learning compressed latent representation via PRI}

We finally show that the potential of PRI also applies to deterministic AEs. As an example, we train a simpler, deterministic AE regularized with PRI (term it PRI-AE), and demonstrate its capability to learn compressed latent representations in the presence of channel noise (i.e., noise in $\mathbf{z}$). This is
an old research topic~\cite{doi2007robust} that has re-attracted increasing attention in both machine learning and communication communities (e.g.,~\cite{grover2019uncertainty,choi2019neural}). Fig.~\ref{fig:PRI_AE_channel_noise} demonstrates the $2$-dimensional latent representations of AE and PRI-AE on MNIST~\cite{lecun1998gradient} data set. Table~\ref{Tab:channel_noise} reports their reconstruction errors (measured with the mean square error or MSE) with respect to injected channel noise under different signal-to-noise ratios (SNRs). Obviously, the latent codes in PRI-AE has less uncertainty and is more robust to channel noise.


\begin{figure}[ht]
  \begin{subfigure}[t]{0.23\textwidth}
    \includegraphics[width=\textwidth]{Figs/mse_latent.png}
   
   
  \end{subfigure}
  %
  \begin{subfigure}[t]{0.23\textwidth}
    \includegraphics[width=\textwidth]{Figs/pri_latent_2.png}
   
   
  \end{subfigure}
  \caption{Two-dimensional learned representations of a basic deterministic autoencoder (left) and a deterministic autoencoder regularized with PRI (right) on MNIST. Each color indicates one digit class.}
  \label{fig:PRI_AE_channel_noise}
\end{figure}

\begin{table}[!hbpt]
\begin{center}
\caption{MSE comparison with respect to channel noise with different SNRs (dB).}\label{Tab:channel_noise}
\begin{tabular}{lcccccccc}\hline
& 4 & 6 & 8 & 10 & 12 \\ \hline
AE & 0.0713 & 0.0592 & 0.0619 &  0.0571 & 0.0521  \\
PRI-AE & 0.0663 & 0.0549 & 0.0516 & 0.0490 & 0.0469 \\\hline
\end{tabular}
\end{center}
\end{table}
\end{comment}


\section{Related Work}
\subsection{Disentangled Representation Learning with VAEs}
Before our work, many efforts have been made for disentangled representation learning with the VAE framework. 
In Section~\ref{sec:PRI_VAE_relation} we illustrate the connections between our PRI-VAE and some prevalent VAEs.
More recently, HFVAE~\cite{esmaeili2019structured} decomposes the VAE objective into four terms, which provides a unified view to most VAE variants.
Interested readers can refer to~\cite{tschannen2018recent} for a comprehensive survey.






\vspace{-1em}
\subsection{Interpreting DNNs with Information Theory}
There has been a growing interest in understanding DNNs using information theory. According to~\cite{shwartz2017opening}, a DNN can be analyzed by measuring the amount of information that each hidden layer's representation $\mathbf{t}$ preserves about the input signal $\mathbf{x}$ with respect to the desired response $\mathbf{y}$ (i.e., $\mathbf{I}(\mathbf{x};\mathbf{t})$ with respect to $\mathbf{I}(\mathbf{t};\mathbf{y})$). 
This technique has been applied to various DNN architectures, such as the the multilayer perceptrons (e.g.,~\cite{chelombiev2019adaptive,shwartz2017opening}), the deterministic autoencoders (e.g.,~\cite{yu2019understanding}) and the convolutional neural networks (e.g.,~\cite{noshad2019scalable,yu2018understanding}).
In general, most of these recent studies suggest that both $\mathbf{I}(\mathbf{x};\mathbf{t})$ and $\mathbf{I}(\mathbf{t};\mathbf{y})$ undergo separate ``fitting'' and ``compression'' phases, which is consistent with our observations as described in Section~\ref{sec:Empirical_behavior}.





Before our work,~\cite{burgess2018understanding} interpreted the emergence of disentangled representation in $\beta$-VAE with the rate-distortion (RD) theorem~\cite[Chapter~10]{cover2012elements}. Similarly,~\cite{alemi2018information} uses the RD function to explore the Pareto front of existing VAEs in terms of both rate (i.e., the average number of additional nats necessary to encode samples from the encoder) and distortion (i.e., the reconstruction error). On the other hand,~\cite{achille2018information} illuminates the connection between VAE and IB approach.


\section{Conclusion}

In this paper, we introduced the recently proposed matrix-based R{\'e}nyi's $\alpha$-entropy functional to interpret the dynamics of learning of VAE models. We also developed a novel VAE learning objective and establish its connections to prior art. Extensive experimental results on four benchmark data sets suggest that our objective can obtain more interpretable factors and comparable disentangled representation while preserving a high reconstruction fidelity. Moreover, our objective can be easily tailored to expand its capacity; and the merit of our objective can be generalized to other SOTA models.





\bibliographystyle{ieeetrans}
