Deep Variational Multivariate Information Bottleneck -
A Framework for Variational Losses

Large dimensional multi-modal datasets are abundant in multimedia systems utilized for language modeling (Xu et al., 2016; Zhou et al., 2018; Rohrbach et al., 2017; Sanabria et al., 2018; Goyal et al., 2017; Wang et al., 2018; Hendrycks et al., 2020), neural control of behavior studies (Steinmetz et al., 2021; Urai et al., 2022; Krakauer et al., 2017; Pang et al., 2016), multi-omics approaches in systems biology (Clark et al., 2013; Zheng et al., 2017; Svensson et al., 2018; Huntley et al., 2015; Lorenzi et al., 2018), and many other domains. Such data come with the curse of dimensionality, making it hard to learn the relevant statistical correlations from samples. The problem is made even harder by the data often containing information that is irrelevant to the specific questions one asks. To tackle these challenges, a myriad of dimensionality reduction (DR) methods have emerged. By preserving certain aspects of the data while discarding the remainder, DR can decrease the complexity of the problem, yield clearer insights, and provide a foundation for more refined modeling approaches.

DR techniques span linear methods like Principal Component Analysis (PCA) (Hotelling, 1933), Partial Least Squares (PLS) (Wold et al., 2001), Canonical Correlations Analysis (CCA) (Hotelling, 1936), and regularized CCA (Vinod, 1976; Årup Nielsen et al., 1998), as well as nonlinear approaches, including Autoencoders (AE) (Hinton and Salakhutdinov, 2006), Deep CCA (Andrew et al., 2013), Deep Canonical Correlated AE (Wang et al., 2015), Correlational Neural Networks (Chandar et al., 2016), Deep Generalized CCA (Benton et al., 2017), and Deep Tensor CCA (Wong et al., 2021). Of particular interest to us are variational methods, such as Variational Autoencoders (VAE) (Kingma and Welling, 2014), beta-VAE (Higgins et al., 2016), Joint Multimodal VAE (JMVAE) (Suzuki et al., 2016), Deep Variational CCA (DVCCA) (Wang et al., 2016), Deep Variational Information Bottleneck (DVIB) (Alemi et al., 2017), Variational Mixture-of-experts AE (Shi et al., 2019), and Multiview Information Bottleneck (Federici et al., 2020b). These DR methods use deep neural networks and variational approximations to learn robust and accurate representations of the data, while, at the same time, often serving as generative models for creating samples from the learned distributions.

There are many theoretical derivations and justifications for variational DR methods (Kingma and Welling, 2014; Higgins et al., 2016; Suzuki et al., 2016; Wang et al., 2016; Karami and Schuurmans, 2021; Qiu et al., 2022; Alemi et al., 2017; Bao, 2021; Lee and Van der Schaar, 2021; Wang et al., 2019; Wan et al., 2021; Federici et al., 2020a; Huang et al., 2022; Hu et al., 2020). This diversity of derivations, while enabling adaptability, often leaves researchers with no principled ways for choosing a method for a particular application, for designing new methods with distinct assumptions, or for comparing methods to each other.

Here, we introduce the Deep Variational Multivariate Information Bottleneck (DVMIB) framework, offering a unified mathematical foundation for many variational DR methods. Our framework is grounded in the multivariate information bottleneck loss function (Tishby et al., 2000; Friedman et al., 2013). This loss, amenable to approximation through upper and lower variational bounds, provides a system for implementing diverse DR variants using deep neural networks. We demonstrate the framework’s efficacy by deriving the loss functions of many existing variational DR methods starting from the same principles. Furthermore, our framework naturally allows the adjustment of trade-off parameters, leading to generalizations of these existing methods. For instance, we generalize DVCCA to $\beta$-DVCCA. The framework further allows us to introduce and implement in software novel DR methods. We view the DVMIB framework, with its uniform information bottleneck language, conceptual clarity of translating statistical dependencies in data via graphical models of encoder and decoder structures into variational losses, the ability to unify existing approaches, and easy adaptability to new scenarios as one of the main contributions of our work.

Beyond its unifying role, our framework offers a principled approach for deriving problem-specific loss functions using domain-specific knowledge. Thus, we anticipate its application for multi-view representation learning across diverse fields. To illustrate this, we use the framework to derive a novel dimensionality reduction method, the Deep Variational Symmetric Information Bottleneck (DVSIB), which compresses two random variables into two distinct latent variables that are maximally informative about one another. This new method produces better representations of classic datasets than previous approaches. The introduction of DVSIB is another major contribution of our paper.

In summary, our paper makes the following contributions to the field:

1. 1.

   Introduction of the Variational Multivariate Information Bottleneck Framework: We provide both intuitive and mathematical insights into this framework, establishing a robust foundation for further exploration.

2. 2.

   Rederivation and Generalization of Existing Methods within a Common Framework: We demonstrate the versatility of our framework by systematically rederiving and generalizing various existing methods from the literature, showcasing the framework’s ability to unify diverse approaches.

3. 3.

   Design of a Novel Method — Deep Variational Symmetric Information Bottleneck (DVSIB): Employing our framework, we introduce DVSIB as a new method, contributing to the growing repertoire of techniques in variational dimensionality reduction. The method constructs high-accuracy latent spaces from substantially fewer samples than comparable approaches.

The paper is structured as follows. First, we introduce the underlying mathematics and the implementation of the DVMIB framework. We then explain how to use the framework to generate new DR methods. In Tbl. 1, we present several known and newly derived variational methods, illustrating how easily they can be derived within the framework. As a proof of concept, we then benchmark simple computational implementations of methods in Tbl. 1 against the Noisy MNIST dataset. Appendices present detailed treatment of all terms in variational losses in our framework, discussion of multi-view generalizations, and more details —including visualizations— of the performance of many methods on the Noisy MNIST.

We represent DR problems similar to the Multivariate Information Bottleneck (MIB) of Friedman et al. (2013), which is a generalization of the more traditional Information Bottleneck algorithm (Tishby et al., 2000) to multiple variables. The reduced representation is achieved as a trade-off between two Bayesian networks. Bayesian networks are directed acyclic graphs that provide a factorization of the joint probability distribution, $P(X_{1},X_{2},X_{3},..,X_{N})=\prod_{i=1}^{N}P(X_{i}|Pa_{X_{i}}^{G})$, where $Pa_{X_{i}}^{G}$ is the set of parents of $X_{i}$ in graph $G$. The multiinformation (Studenỳ and Vejnarová, 1998) of a Bayesian network is defined as the Kullback-Leibler divergence between the joint probability distribution and the product of the marginals, and it serves as a measure of the total correlations among the variables, $I(X_{1},X_{2},X_{3},...,X_{N})=D_{KL}(P(X_{1},X_{2},X_{3},...,X_{N})\|P(X_{1})P(X_{2})P(X_{3})...P(X_{N}))$. For a Bayesian network, the multiinformation reduces to the sum of all the local informations $I(X_{1},X_{2},..X_{N})=\sum_{i=1}^{N}I(X_{i};Pa_{X_{i}}^{G})$ (Friedman et al., 2013).

The first of the Bayesian networks is an encoder (compression) graph, which models how compressed (reduced, latent) variables are obtained from the observations. The second network is a decoder graph, which specifies a generative model for the data from the compressed variables, i.e., it is an alternate factorization of the distribution. In MIB, the information of the encoder graph is minimized, ensuring strong compression (corresponding to the approximate posterior). The information of the decoder graph is maximized, promoting the most accurate model of the data (corresponding to maximizing the log-likelihood). As in IB (Tishby et al., 2000), the trade-off between the compression and reconstruction is controlled by a trade-off parameter $\beta$:

$$L=I_{\text{encoder}}-\beta I_{\text{decoder}}.$$

(1)

In this work, our key contribution is in writing an explicit variational loss for typical information terms found in both the encoder and the decoder graphs. All terms in the decoder graph use samples of the compressed variables as determined from the encoder graph. If there are two terms that correspond to the same information in Eq. (1), one from each of the graphs, they do not cancel each other since they correspond to two different variational expressions. For pedagogical clarity, we do this by first analyzing the Symmetric Information Bottleneck (SIB), a special case of MIB. We derive the bounds for three types of information terms in SIB, which we then use as building blocks for all other variational MIB methods in subsequent Sections.

The variational bounds used in DVSIB can be used to implement loss functions that correspond to other encoder-decoder graph pairs and hence to other DR algorithms. The simplest is the beta variational auto-encoder. Here $G_{\rm encoder}$ consists of one term: $X$ compressed into $Z_{X}$. Similarly $G_{\rm decoder}$ consists of one term: $X$ decoded from $Z_{X}$ (see Table 1). Using this simple set of Bayesian networks, we find the variational loss:

$$L_{\text{beta-VAE}}=\tilde{I}^{E}(X;Z_{X})-\beta\tilde{I}^{D}(X;Z_{X}).$$

(14)

Both terms in Eq. (14) are the same as Eqs. (10, 11) and can be approximated and implemented by neural networks.

Similarly, we can re-derive the DVCCA family of losses (Wang et al., 2016). Here $G_{\rm encoder}$ is $X$ compressed into $Z_{X}$. $G_{\rm decoder}$ reconstructs both $X$ and $Y$ from the same compressed latent space $Z_{X}$. In fact, our loss function is more general than the DVCCA loss and has an additional compression-reconstruction trade-off parameter $\beta$. We call this more general loss $\beta$-DVCCA, and the original DVCCA emerges when $\beta=1$:

$$L_{\rm DVCCA}=\tilde{I}^{E}(X;Z_{X})-\beta(\tilde{I}^{D}(Y;Z_{X})+\tilde{I}^{D}(X;Z_{X})).$$

(15)

Using the same library of terms as we found in DVSIB, Eqs. (10, 11), we find:

$$L_{\text{DVCCA}}\approx\frac{1}{N}\sum_{i=1}^{N}D_{\rm KL}(p(z_{x}|x_{i})\|r(z_{x}))\\ -\beta\left(\frac{1}{N}\sum_{i=1}^{N}\int dz_{x}p(z_{x}|x_{i})\ln(q(y_{i}|z_{x}))+\frac{1}{N}\sum_{i=1}^{N}\int dz_{x}p(z_{x}|x_{i})\ln(q(x_{i}|z_{x}))\right).$$

(16)

This is similar to the loss function of the deep variational CCA (Wang et al., 2016), but now it has a trade-off parameter $\beta$. It trades off the compression into $Z_{X}$ against the reconstruction of $X$ and $Y$ from the compressed variable $Z_{X}$.

Table 1 shows how our framework reproduces and generalizes other DR losses (see Appendix A). Our framework naturally extends beyond two variables as well (see Appendix B).

To test our methods, we created a dataset inspired by the noisy MNIST dataset (LeCun et al., 1998; Wang et al., 2015, 2016), consisting of two distinct views of data, both with dimensions of $28\times 28$ pixels, cf. Fig. 2. The first view comprises the original image randomly rotated by an angle uniformly sampled between $0$ and $\frac{\pi}{2}$ and scaled by a factor uniformly distributed between $0.5$ and $1.5$. The second view consists of the original image with an added background Perlin noise (Perlin, 1985) with the noise factor uniformly distributed between $0$ and $1$. Both image intensities are scaled to the range of $[0,1)$. The dataset was shuffled within labels, retaining only the shared label identity between two images, while disregarding the view-specific details, i.e., the random rotation and scaling for $X$, and the correlated background noise for $Y$. The dataset, totaling $70,000$ images, was partitioned into training ($80\%$), testing ($10\%$), and validation ($10\%$) subsets. Visualization via t-SNE (Hinton and Roweis, 2002) plots of the original dataset suggest poor separation by digit, and the two digit views have diverse correlations, making this a sufficiently hard problem.

The DR methods we evaluated include all methods from Tbl. 1. PCA and CCA (Hotelling, 1933, 1936) served as a baseline for linear dimensionality reduction. Multi-view Information Bottleneck Federici et al. (2020b) was included for a specific comparison with DVSIB (see Appendix C). We emphasize that none of the algorithms were given labeled data. They had to infer compressed latent representations that presumably should cluster into ten different digits based simply on the fact that images come in pairs, and the (unknown) digit label is the only information that relates the two images.

Each method was trained for 100 epochs using fully connected neural networks with layer sizes $(\text{input\_dim},1024,1024,(k_{Z},k_{Z}))$, where $k_{Z}$ is the latent dimension size, employing ReLU activations for the hidden layers. The input dimension $(\text{input\_dim})$ was either the size of $X$ (784) or the size of the concatenated $[X,Y]$ (1568). The last two layers of size $k_{Z}$ represented the means and $\log(\text{variance})$ learned. For the decoders, we employed regular decoders, fully connected neural networks with layer sizes $(k_{Z},1024,1024,\text{output\_dim})$, using ReLU activations for the hidden layers and sigmoid activation for the output layer. Again, the output dimension $(\text{output\_dim})$ could either be the size of $X$ (784) or the size of the concatenated $[X,Y]$ (1568). The latent dimension $(k_{Z})$ could be $k_{Z_{X}}$ or $k_{Z_{Y}}$ for regular decoders, or $k_{Z_{X}}+k_{W_{X}}$ or $k_{Z_{Y}}+k_{W_{Y}}$ for decoders with private information. Additionally, another decoder denoted as decoder_MINE, based on the MINE estimator for estimating $I(Z_{X},Z_{Y})$, was used in DVSIB and DVSIB with private information. The decoder_MINE is a fully connected neural network with layer sizes $(k_{Z_{X}}+k_{Z_{Y}},1024,1024,1)$ and ReLU activations for the hidden layers. Optimization was conducted using the ADAM optimizer with default parameters.

To evaluate the methods, we trained them on the training portions of $X$ and $Y$ without exposure to the true labels. Subsequently, we utilized the trained encoders to compute $Z_{\text{train}}$, $Z_{\text{test}}$, and $Z_{\text{validation}}$ on the respective datasets. To assess the quality of the learned representations, we revealed the labels of $Z_{\text{train}}$ and trained a linear SVM classifier with $Z_{\text{train}}$ and $\text{labels}_{\text{train}}$. Fine-tuning of the classifier was performed to identify the optimal SVM slack parameter ($C$ value), maximizing accuracy on $Z_{\text{test}}$. This best classifier was then used to predict $Z_{\text{validation}}$, yielding the reported accuracy. We also conducted classification experiments using fully connected neural networks, with detailed results available in the Appendix D. For both SVM and the fully connected network, we find the baseline accuracy on the original training data and labels $(X_{\text{train}},\text{labels}_{\text{train}})$ and $(Y_{\text{train}},\text{labels}_{\text{train}})$, fine-tuning with the test datasets, and reporting the results of the validation datasets. Using Linear SVM enables us to assess the linear separability of the clusters of $Z_{X}$ and $Z_{Y}$ obtained through the DR methods. While neural networks excel at uncovering nonlinear relationships that could result in higher classification accuracy, the comparison with a linear SVM establishes a level playing field. It ensures a fair comparison among different methods and is independent of the success of the classifier used for comparison in detecting nonlinear features in the data, which might have been missed by the DR methods. Here, we focus on the results of the $Y$ datasets (MNIST with correlated noise background); results for $X$ are in the Appendix D. A parameter sweep was performed to identify optimal $k_{Z}$ values, ranging from $2^{1}$ to $2^{8}$ dimensions on $\log_{2}$ scale, as well as optimal $\beta$ values, ranging from $2^{-5}$ to $2^{10}$. For methods with private information, $k_{W_{X}}$ and $k_{W_{Y}}$ were varied from $2^{1}$ to $2^{6}$. The highest accuracy is reported in Tbl. 2, along with the optimal parameters used to obtain this accuracy. Additionally, for every method we find the range of $\beta$ and the dimensionality $k_{Z}$ of the latent variable $Z_{Y}$ that gives 95% of the method’s maximum accuracy. If the range includes the limits of the parameter, this is indicated by an asterisk.

Figure 3 shows a t-SNE plot of DVSIB’s latent space, $Z_{Y}$, colored by the identity of digits. The resulting latent space has 10 clusters, each corresponding to one digit. The clusters are well separated and interpretable. Further, DVSIB’s $Z_{Y}$ latent space provides the best classification of digits using a linear method such as an SVM showing the latent space is linearly separable. DVSIB maximum classification accuracy obtained for the linear SVM is 97.8%. Crucially, DVSIB maintains accuracy of at least 92.9% (95% of 97.8%) for $\beta\in[2,1024^{*}]$ and $k_{Z}\in[8,256^{*}]$. This accuracy is high compared to other methods and has a large range of hyperparameters that maintain its ability to correctly capture information about the identity of the shared digit. DVSIB is a generative method, we have provided sample generated digits from the decoders that were trained from the model graph.

In Fig. 4, we show the highest SVM classification accuracy curves for each method. DVSIB and DVSIB-private tie for the best classification accuracy for $Y$. Together with $\beta$-DVCCA-private they have the highest accuracy for all dimensions of the latent space, $k_{Z}$. In theory, only one dimension should be needed to capture the identity of a digit, but our data sets also contain information about the rotation and scale for $X$ and the strength of the background noise for $Y$. $Y$ should then need at least two latent dimensions to be reconstructed and $X$ should need at least three. Since DVSIB, DVSIB-private, and $\beta$-DVCCA-private performed with the best accuracy starting with the smallest $k_{Z}$, we conclude that methods with the encoder-decoder graphs that more closely match the structure of the data produce higher accuracy with lower dimensional latent spaces.

Next, in Fig. 5, we compare the sample training efficiency of DVSIB and $\beta$-VAE by training new instances of these methods on a geometrically increasing number of samples $n=[256,339,451,\dots,\sim 42k,\sim 56k]$, consisting of 20 subsamples of the full training data ($X,Y$) to get ($X_{\text{train}_{n}},Y_{\text{train}_{n}}$), where each larger subsample includes the previous one. Each method was trained for 60 epochs, and we used $\beta=1024$ (as defined by the DVMIB framework). Further, all reported results are with the latent space size $k_{Z}=64$. We explored other numbers of training epochs and latent space dimensions (see Appendix D.6), but did not observe qualitative differences. We follow the same procedure as outlined earlier, using the 20 trained encoders for each method to compute $Z_{\text{train}_{n}}$, $Z_{\text{test}}$, and $Z_{\text{validation}}$ for the training, test, and validation datasets. As before, we then train and evaluate the classification accuracy of SVMs for the $Z_{Y}$ representation learned by each method. Fig. 5, inset, shows the classification accuracy of each method as a function of the number of samples used in training. Again, CCA and PCA serve as linear methods baselines. PCA is able to capture the linear correlations in the dataset consistently, even at low sample sizes. However, it is unable to capture the nonlinearities of the data, and its accuracy does not improve with the sample size. Because of the iterative nature of the implementation of the PCA algorithm (Pedregosa et al., 2011), it is able to capture some linear correlations in a relatively low number of dimensions, which are sufficiently sampled even with small-sized datasets. Thus the accuracy of PCA barely depends on the training set size. CCA, on the other hand, does not work in the under-sampled regime (see Abdelaleem et al. (2023) for discussion of this). DVSIB performs uniformly better, at all training set sizes, than the $\beta$-VAE. Furthermore, DVSIB improves its quality faster, with a different sample size scaling. Specifically, DVSIB and $\beta$-VAE accuracy ($A$, measured in percent) appears to follow the scaling form $A=100-c/n^{m}$, where $c$ is a constant, and the scaling exponent $m=0.345\pm 0.007$ for DVSIB, and $0.196\pm 0.013$ for $\beta$-VAE. We illustrate this scaling in Fig. 5 by plotting a log-log plot of $100-A$ vs $1/n$ and observing a linear relationship.

We developed an MIB-based framework for deriving variational loss functions for DR applications. We demonstrated the use of this framework by developing a novel variational method, DVSIB. DVSIB compresses the variables $X$ and $Y$ into latent variables $Z_{X}$ and $Z_{Y}$ respectively, while maximizing the information between $Z_{X}$ and $Z_{Y}$. The method generates two distinct latent spaces—a feature highly sought after in various applications—but it accomplishes this with superior data efficiency, compared to other methods. The example of DVSIB demonstrates the process of deriving variational bounds for terms present in all examined DR methods. A comprehensive library of typical terms is included in Appendix A for reference, which can be used to derive additional DR methods. Further, we (re)-derive several DR methods, as outlined in Table 1. These include well-known techniques such as $\beta$-VAE, DVIB, DVCCA, and DVCCA-private. MIB naturally introduces a trade-off parameter into the DVCCA family of methods, resulting in what we term the $\beta$-DVCCA DR methods, of which DVCCA is a special case. We implement this new family of methods and show that it produces better latent spaces than DVCCA at $\beta=1$, cf. Tbl. 2.

We observe that methods that more closely match the structure of dependencies in the data can give better latent spaces as measured by the dimensionality of the latent space and the accuracy of reconstruction (see Figure 4). This makes DVSIB, DVSIB-private, and $\beta$-DVCCA-private perform the best. DVSIB and DVSIB-private both have separate latent spaces for $X$ and $Y$. The private methods allow us to learn additional aspects about $X$ and $Y$ that are not important for the shared digit label, but allow reconstruction of the rotation and scale for $X$ and the background noise of $Y$. We also found that DVSIB can make more efficient use of data when producing latent spaces as compared to $\beta$-VAEs and linear methods.

Our framework may be extended beyond variational approaches. For instance, in the deterministic limit of VAE, autoencoders can be retrieved by defining the encoder/decoder graphs as nonlinear neural networks $z=f(x)$ and $x=g(z)$. Additionally, linear methods like CCA can be viewed as special cases of the information bottleneck (Chechik et al., 2003) and hence must follow from our approach. Similarly, by using specialized encoder and decoder neural networks, e.g., convolutional ones, our framework can implement symmetries and other constraints into the DR process. Overall, the framework serves as a versatile and customizable toolkit, capable of encompassing a wide spectrum of dimensionality reduction methods. With the provided tools and code, we aim to facilitate the adaptation of the approach to diverse problems.