Visualizing Information Bottleneck through Variational Inference

Given any two random variables, $X$ and $Y$, with a joint distribution $p(x,y)$, their Mutual Information is defined as:

$$I(X;Y)=D_{KL}[p(x,y)||p(x)p(y)]=\sum\limits_{x\in X,y\in Y}p(x,y)\log\left(\frac{p(x,y)}{p(x)p(y)}\right)$$

(1)

$$=\sum\limits_{x\in X,y\in Y}p(x,y)\log\left(\frac{p(x|y)}{p(x)}\right)=H(X)-H(X|Y)$$

(2)

where $D_{KL}[p||q]$ is the Kullback-Liebler divergence of the distributions $p$ and $q$, and $H(X)$ and $H(X|Y)$ are the entropy and conditional entropy of $X$ and $Y$, respectively.

When analyzing representations of $X$ w.r.t $Y$, the classical notion of minimal sufficient statistics provide good candidates for optimal representation. Sufficient statistics, in our context, are maps or partitions of $X$, $S(X)$, that capture all the information that $X$ has on $Y$. Namely, $I(S(X);Y)=I(X;Y)$.

Minimal sufficient statistics, $T(X)$, are the simplest sufficient statistics and induce the coarsest sufficient partition on $X$. A simple of way of formulating this is through the Markov chain: $Y\rightarrow X\rightarrow S(X)\rightarrow T(X)$, which should hold for a minimal sufficient statistics $T(X)$ with any other sufficient statistics $S(X)$. Since exact minimal sufficient statistics only exist for very special distributions, (i.e. exponential families), Tishby et al. [2001] relaxed this optimization problem by first allowing the map to be stochastic, defined as an encoder $P(T|X)$, and then, by allowing the map to capture as much as possible of $I(X;Y)$, not necessarily all of it.

This leads to the Information Bottleneck (IB) tradeoff [Tishby et al. [2001]], which provides a computational framework for finding approximate minimal sufficient statistics, or the optimal tradeoff between compression of $X$ and prediction of $Y$.

Any representation variable, $T$, defined as a (possibly stochastic) map of the input $X$, is characterized by its joint distributions with $X$ and $Y$, or by its encoder and decoder distributions, $P(T|X)$ and $P(Y|T)$, respectively. Given $P(X;Y)$, $T$ is uniquely mapped to a point in the Information Plane with coordinates $I(X;T),I(T;Y)$.

Tishby and Zaslavsky [2015] proposed to analyze DNNs in the Information Plane and suggested that the goal of the neural network is to optimize the Information Bottleneck tradeoff between compression and prediction, successively, for each layer.

Building on top of this work Shwartz-Ziv and Tishby [2017] analyzed the information plane of a network with 7 fully connected hidden layers, and widths 12-10-7-5-4-3-2 neurons with hyperbolic tangent activations. The analysis showed that the Stochastic Gradient Descent (SGD) optimization has two main phases, in the first and shorter phase the layers increase the information on the input (fitting), while in the second much longer phase the layers reduce the information on the input (compression phase). The tasks were chosen as binary decision rules which are invariant under $O(3)$ rotations of the sphere, with 12 binary inputs that represent 12 uniformly distributed points on the sphere. As the network size was small they calculated the mutual information exhaustively by binning the output activations into 30 buckets and computing the joint distribution. Figure 1 shows the information plane for the described setup.

Anonymous [2018] showed that results of Shwartz-Ziv and Tishby [2017] are not applicable for networks with ReLU activations and also claim that the shape of information plane is a result of saturating hyperbolic tangent non-linearities and for ReLU activation the same shape is not observed. Anonymous [2018] follows a binning strategy to estimate the mutual information for ReLU activations. This method of estimating mutual information has been refuted by authors of Shwartz-Ziv and Tishby [2017], ICLR-2018 Openreview¹¹1Review comments and discussion available at https://goo.gl/U24Kfp.

In section 4 we discuss how we estimate mutual information using variational inference. Using these results we draw the information plane for a neural network learned to classify on the MNIST dataset²²2http://yann.lecun.com/exdb/mnist. We use ReLU non-linearities and a deeper neural network for our task. In our results we observe the two distinct phases of SGD, fitting and compression as originally observed by Shwartz-Ziv and Tishby [2017].

The MNIST classification task consists of images of digits and the task is to predict the label of digit. An example of the task is shown in in Figure 2. We consider $X$ to be the input space (images) and $Y$ to be the output label space (digit labels). We modify this task into the teacher-student setting, where the digit images are generated by a pre-trained teacher model.

The IB setup views the data generating process as a Markov chain $Y\rightarrow X\rightarrow Z\rightarrow\hat{Y}$. $Y$ is a signal we wish to extract from observations $X$. $Z$ is some intermediate representation computed from $Z$ by the model on the way to computing the predictions $\hat{Y}$.

The VIB [Alemi et al., 2016] method is a generalization of the Variational Autoencoder [Kingma and Welling, 2013] into the supervised setting. Its objective function is to maximize the rate distortion tradeoff:

$$\mathcal{L}(\theta)=I(Y;Z)-\beta I(X;Z)$$

These two quantities correspond to the position in the information plane, so computing an approximation of the loss will give us an estimate of the network’s position. The MI of $Y$ (target variable) and $Z$ can be computed as the difference between entropy and conditional entropy (the remaining uncertainty after observing $Z$):

	$\displaystyle I(Y;Z)$	$\displaystyle=H(Y)-H(Y|Z)$
		$\displaystyle=H(Y)+\sum_{y,z}p(y,z)\log p(y|z)$

Under our setup $p(y|z)$ is difficult to compute, because we would need to invoke Bayes’ rule on $p(z|y)$, which is what we have (after marginalizing over all $x$). So we introduce an approximate distribution $q(y|z)$:

	$\displaystyle I(Y;Z)$	$\displaystyle=H(Y)+\sum_{y,z}p(y,z)\log p(y|z)\frac{q(y|z)}{q(y|z)}$
		$\displaystyle=H(Y)+\sum_{y,z}p(y,z)\log q(y|z)+\sum_{y,z}p(y,z)\log\frac{p(y|z)}{q(y|z)}$
		$\displaystyle=H(Y)+\mathbb{E}_{y,z}\log q(y|z)+D_{KL}(p(y|z)||q(y|z))$
		$\displaystyle\geq H(Y)+\mathbb{E}_{y,z}\log q(y|z)$

The second term is the average log-likelihood of $y$ under $q$. In practice we will use the learned decoder of the VIB model as the approximate $q$. The second term can be thought of as a conditional cross entropy term, which is an upper bound on the conditional entropy. To calculate $I(X;Z)$ we decompose it the same way:

	$\displaystyle I(X;Z)$	$\displaystyle=H(Z)-H(Z|X)$
		$\displaystyle=H(Z)-\mathbb{E}_{x}H[p(z|x)]$

$\mathbb{E}_{x}H[p(z|x)]$ is the average entropy of the conditional distribution $p(z|x)$ for a fixed $x$. Intuitively, this makes sense. If the entropy of the encoding is high, then there is high uncertainty of choosing $z$ from $x$, so mutual information should be low. Because $H(Z)$ is hard to compute, we substitute it with the cross-entropy of $p(z)$ and some variational approximation $r(z)$. Cross entropy is always greater than entropy, so:

	$\displaystyle I(X;Z)$	$\displaystyle\leq-\sum_{x,z}p(x,z)\log r(z)+\sum_{x,z}p(x,z)\log p(z|x)$
		$\displaystyle=\sum_{x,z}p(x,z)\log\frac{p(z|x)}{r(z)}$
		$\displaystyle=\mathbb{E}_{x}D_{KL}(p(z|x)||r(z))$

Suppose we assume that $X$ can be sampled using the process $Z_{v}\rightarrow X$. In practice, this would mean pre-training an unsupervised latent variable model which we will define as $p(x|z_{v})$. We can think of this as an instance of the teacher-student setup, where the student is trained to mimic the teacher. Given this, we can use another approach to calculate $H(Z)$ by estimating $p(z)$ for a given $z$:

$\displaystyle p(z)$

$\displaystyle=\sum\limits_{x}p(x,z)=\sum\limits_{x}p(z|x)p(x)=\sum\limits_{x}\sum_{z_{v}}p(z|x)p(z_{v},x)$

To perform variational inference, we introduce a distribution $q(z_{v},x)$:

$$p(z)=\sum_{z_{v},x}q(z_{v},x)p(z|x)\frac{p(z_{v},x)}{q(z_{v},x)}=\mathbb{E}_{z_{v},x\sim Q}\,\,\,\,p(z|x)\frac{p(z_{v},x)}{q(z_{v},x)}$$

(3)

We can use this to lower bound $\log p(z)$:

	$\displaystyle\log p(z)$	$\displaystyle=\log\mathbb{E}_{z_{v},x\sim Q}p(z|x)\frac{p(z_{v},x)}{q(z_{v},x)}$		(4)
		$\displaystyle\geq\mathbb{E}_{x\sim Q}\log p(z|x)-\mathbb{E}_{z_{v},x\sim Q}\log\frac{q(z_{v},x)}{p(z_{v},x)}$
		$\displaystyle=\mathbb{E}_{x\sim Q}\log p(z|x)-D_{\text{KL}}(Q(z_{v},x)||P(z_{v},x))$

If we define $q(z_{v},x)$ to use the teacher model:

$$q(z_{v},x)=q(z_{v})p(x|z_{v})$$

we can simplify the KL divergence:

	$\displaystyle D_{\text{KL}}[q(z_{v},x)\mid\mid p(z_{v},x)]$	$\displaystyle=\sum_{z_{v},x}q(z_{v})p(x\mid z_{v})\log\frac{q(z_{v})p(x\mid z_{v})}{p(z_{v})p(x\mid z_{v})}$
		$\displaystyle=\sum_{z_{v}}q(z_{v})\sum_{x}p(x\mid z_{v})\log\frac{q(z_{v})}{p(z_{v})}$
		$\displaystyle=\sum_{z}q(z)\log\frac{q(z_{v})}{p(z_{v})}\left(\sum_{x}p(x\mid z_{v})\right)$
		$\displaystyle=D_{\text{KL}}[q(z_{v})\mid\mid p(z_{v})]$

In other words, if we have access to the "real" data generating distribution $p(x|z_{v})$, we can use it during variational inference. The final equation is as follows:

	$\displaystyle I(X,Z)\leq$	$\displaystyle-\mathbb{E}_{x}H[p(z|x)]$		(5)
		$\displaystyle-\mathbb{E}_{z,x}[\mathbb{E}_{z_{v},x^{\prime}\sim Q(z)}\log p(z|x^{\prime})-D_{KL}(Q(z_{v}|x)||P(z_{v}))]$

The variational distribution $Q$ is computed for a specific $z$. Algorithmically the inference process is as follows:

1. 1.

   Sample $z_{v}$, $x$ from the teacher model $p(x|z_{v})$

2. 2.

   Run the student encoder to get $p(z|x)$

3. 3.

   Sample $z$ from $p(z|x)$

4. 4.

   Run the inference network to get $q(z_{v}|z)$

5. 5.

   Sample $z_{v}^{\prime}$ from $q$, and re-run the teacher model $p(x^{\prime}|z_{v}^{\prime})$

6. 6.

   Sample $x^{\prime}$ from this distribution to compute $p(z|x^{\prime})$

7. 7.

   Use all the samples to compute an MI upper bound via Equation 5

8. 8.

   Use the upper bound as a loss function to train the inference network $q(z_{v}|z)$

The classification model we analyze is trained to maximize the IB tradeoff directly. However, the objective does not specify how this optimization will take place. We want to see whether the two training modes observed by Shwartz-Ziv and Tishby [2017] will happen for a VIB model, or whether it will train to fit and compress in an interleaved manner [Anonymous, 2018].

Our secondary goal is to compare our two approaches to estimate mutual information of $X$ and $Z$. Both methods upper bound the true quantity $I(X,Z)$, we would like to see which methods will produce a lower result. The VIB objective has the advantage of being fast to calculate, whereas our method is optimization-based that runs over multiple iterations. However we might get a better result due to the network having access to the data generating distribution.

To verify our implementation, we ran training on a dataset where the images are generated independently from the labels, thereby having $I(X,Y)=0$. We would expect to see $I(X,Z)$ and $I(Y,Z)$ to drop rapidly once the network has determined that the image contains no useful information for predicting the digit.

Our results show that the direct estimation method from Alemi et al. [2016] is better at bounding $I(X,T)$ during the later stages of training. We suspect that this is because at later stages of training, the marginal $p(z)$ is sufficiently close to the approximation $p(z)$, which produces a better estimate than doing it the roundabout way. In contrast, the inference network was able to reach a lower bound while mutual information is still rising.

In the following sections, we will take the minimum of both estimates.

Our experiments reveal that we were able to see an inflection point in the mutual information $I(X,T)$ as the VIB model is training (see figures 5 and 6). We see that the regularization parameter strongly affects how much the model will fit to the data before starting compression.

VIB models can declare their uncertainty of an encoding given a model $p(z|x)$ by enlarging the variance. Surprisingly, during compression phase we are still able to see the variance given to real samples go down.