Information Bottleneck Analysis of Deep Neural Networks via Lossy
  Compression
The Information Bottleneck (IB) principle offers an information-theoretic
framework for analyzing the training process of deep neural networks (DNNs).
Its essence lies in tracking the dynamics of two mutual information (MI)
values: between the hidden layer output and the DNN input/target. According to
the hypothesis put forth by Shwartz-Ziv & Tishby (2017), the training process
consists of two distinct phases: fitting and compression. The latter phase is
believed to account for the good generalization performance exhibited by DNNs.
Due to the challenging nature of estimating MI between high-dimensional random
vectors, this hypothesis was only partially verified for NNs of tiny sizes or
specific types, such as quantized NNs. In this paper, we introduce a framework
for conducting IB analysis of general NNs. Our approach leverages the
stochastic NN method proposed by Goldfeld et al. (2019) and incorporates a
compression step to overcome the obstacles associated with high dimensionality.
In other words, we estimate the MI between the compressed representations of
high-dimensional random vectors. The proposed method is supported by both
theoretical and practical justifications. Notably, we demonstrate the accuracy
of our estimator through synthetic experiments featuring predefined MI values
and comparison with MINE (Belghazi et al., 2018). Finally, we perform IB
analysis on a close-to-real-scale convolutional DNN, which reveals new features
of the MI dynamics.
Ivan Butakov, Alexander Tolmachev, Sofia Malanchuk, Anna Neopryatnaya, Alexey Frolov, Kirill Andreev
# Information Bottleneck Analysis of Deep Neural Networks via Lossy Compression

###### Abstract

The Information Bottleneck (IB) principle offers an information-theoretic framework for analyzing the training process of deep neural networks (DNNs). Its essence lies in tracking the dynamics of two mutual information (MI) values: one between the hidden layer and the class label, and the other between the hidden layer and the DNN input. According to the hypothesis put forth by Shwartz-Ziv and Tishby (2017), the training process consists of two distinct phases: fitting and compression. The latter phase is believed to account for the good generalization performance exhibited by DNNs. Due to the challenging nature of estimating MI between high-dimensional random vectors, this hypothesis has only been verified for toy NNs or specific types of NNs, such as quantized NNs and dropout NNs. In this paper, we introduce a comprehensive framework for conducting IB analysis of general NNs. Our approach leverages the stochastic NN method proposed by Goldfeld et al. (2019) and incorporates a compression step to overcome the obstacles associated with high dimensionality. In other words, we estimate the MI between the compressed representations of high-dimensional random vectors. The proposed method is supported by both theoretical and practical justifications. Notably, we demonstrate the accuracy of our estimator through synthetic experiments featuring predefined MI values. Finally, we perform IB analysis on a close-to-real-scale convolutional DNN, which reveals new features of the MI dynamics.

## 1 Introduction

The information-theoretic analysis of deep neural networks (DNNs) is a developing branch of the deep learning theory, which may provide a robust and interpretable way to measure the performance of deep models during training and inference. This type of analysis might complement current non-transparent meta-optimization algorithms for architecture search, like ENAS, DARTS, evolutionary algorithms, and others [1; 2; 3; 4; 5; 6]. For the same reasons, this method may also be preferable (see [7]) to current approaches to explainable AI that rely either on a local analysis of a model [8], or some
complex perturbation of data [9]. We also note that information-theoretic quantities can be considered as training objectives [10].

The information-theoretic analysis of DNNs relies on the _Information Bottleneck_ (IB) principle, which was proposed in [11]. This concept was later developed in [10] and applied to DNNs in [12]. The major idea of the IB approach is to track the dynamics of two _mutual information_ (MI) values: \(I(X;L)\) between the hidden layer (\(L\)) and the DNN input (\(X\)) and \(I(Y;L)\) between the hidden layer and the target of the model (\(Y\)). As a result of the IB analysis, the authors of [12] put forth the so-called _fitting-compression_ hypothesis, which states that the training process consists of two phases: a feature-extraction "fitting" phase (both MI values grow) and a representation compression phase (\(I(Y;L)\) grows while \(I(X;L)\) decreases). The authors of [12] conjectured the compression phase to account for the good generalization performance exhibited by DNNs. However, it is still debated whether empirical confirmations of the compression phase are related to improper mutual information estimators, activation function choice, or other implementation details. For an overview, we refer the reader to [13].

The seminal paper [12] proposed to utilize a quantization (or binning) approach to calculate the MI. The first problem of such an approach lies in the fact that the MI estimate is very sensitive to the bin size. The second problem is even more fundamental, indeed let us consider \(I(X;L)\) at some fixed training epoch. When the training weights are fixed \(L\) is a deterministic function of \(X\) and thus the MI is not a function of the NN parameters (it is infinite for practically all regimes of interest if we speak about continuous case and reasonable activation functions, see e.g. [14]). The subsequent papers addressed the aforementioned problems. To tackle the infinite MI problem it was proposed to consider (a) stochastic NNs [15, 16], (b) quantized NNs [17] or (c) dropout NNs [18]. Simple and inconsistent binning entropy estimators have been replaced with estimators more appropriate for continuous random variables [15, 18, 19, 20].

However, due to the challenging nature of estimating MI between high-dimensional random vectors, the fitting-compression hypothesis has only been verified for toy NNs or special classes of models with trackable information-theoretic quantities [17, 18, 19]. We also mention papers that suggest using lower bounds or other surrogate objectives [21, 22, 23, 24, 25, 26, 27] or even using other definitions of entropy [28, 29].

Unlike all the mentioned approaches, we suggest overcoming the curse of dimensionality problem by compressing the data. As most of the datasets are internally structured (see the _manifold hypothesis_[30]), it may be assumed that it is usually sufficient to estimate information-theoretic quantities of compressed/latent representations of data. This allows conventional and well-developed information-theoretic approaches to be applied to real-scale machine-learning problems. This paper extends the ideas of [31], where we used the compression step to estimate the entropy.

Our contribution is as follows. We introduce a comprehensive framework for conducting IB analysis of general NNs. Our approach leverages the stochastic NN method proposed [15] and incorporates a compression step to overcome the obstacles associated with high dimensionality. In other words, we estimate the MI between the compressed representations of high-dimensional random vectors. We provide a theoretical justification of MI estimation under lossless and lossy compression. The accuracy of our estimator is demonstrated through synthetic experiments featuring predefined MI values. Finally, the experiment with convolutional DNN classifier of the MNIST handwritten digits is performed. The experiment shows that there may be several compression/fitting phases during the training process. It may be concluded that phases revealed by information plane plots are connected to different regimes of learning (i.e. accelerated, stationary, or decelerated drop of loss function).

We note that stochastic DNNs are just a proxy to analyze real DNNs as (a) small noise practically does not affect the accuracy and generalization ability, (b) introduced randomness makes information-theoretic quantities reasonable (MI values do depend on DNN parameters).

The paper is organized as follows. In Section 2, we provide basic concepts and objects, used in this paper. In Section 3, we describe the proposed approach and provide the required theoretical justification, including bounds for mutual information estimate under compression. In Section 4, a general framework to test mutual information estimators on synthetic datasets is developed. This framework is then used in Section 5 to test four selected conventional mutual information estimators, complemented with the proposed compression step. The method performed the best is then used in
Section 6 to conduct information plane analysis of a convolutional classifier of MNIST handwritten digits dataset [32]. The results are discussed in Section 7.

## 2 Preliminaries

Consider an absolutely continuous random vectors \(X\colon\Omega\to\mathbb{R}^{n}\) and \(Y\colon\Omega\to\mathbb{R}^{m}\), here and in what follows by \(\Omega\) we denote the sample space. Let us assume that these random vectors have absolutely continuous probability density functions (PDF) \(\rho(x)\), \(\rho(y)\), and \(\rho(x,y)\), where the last term means a joint PDF. The differential entropy of \(X\) is defined as follows

\[h(X)=-\operatorname{\mathbb{E}}\log\rho(x)=-\int\limits_{\operatorname{supp}X }\rho(x)\log\rho(x)\,dx,\]

where \(\operatorname{supp}X\subseteq\mathbb{R}^{n}\) is the _support_ of \(X\), and \(\log(\cdot)\) means a natural logarithm. Similarly, let us define the joint differential entropy \(h(X,Y)=-\operatorname{\mathbb{E}}\log\rho(x,y)\) and conditional differential entropy \(h(X\mid Y)=-\operatorname{\mathbb{E}}\log\rho\left(X\mid Y\right)=- \operatorname{\mathbb{E}}_{Y}\left(\operatorname{\mathbb{E}}_{X\mid Y=y}\log \rho(X\mid Y=y)\right)\). Finally, mutual information (MI) is given by \(I(X;Y)=h(X)-h(X\mid Y)\), and the following equivalences hold

\[I(X;Y)=h(X)-h(X\mid Y)=h(Y)-h(Y\mid X), \tag{1}\]

\[I(X;Y)=h(X)+h(Y)-h(X,Y). \tag{2}\]

In some cases, \(\operatorname{supp}X\) or \(\operatorname{supp}Y\) can be of measure zero (the distribution is singular). Then, if the supports are manifolds, density functions can be treated as induced probability densities, and \(dx\), \(dy\) -- as area elements of corresponding manifolds; thus all the previous definitions remain valid.

There is an alternative definition of the MI given by the Kullback-Leibler divergence \(D_{KL}(\,\cdot\,)\).

\[I(X;Y)=\int\rho(x,y)\log\frac{\rho(x,y)}{\rho(x)\rho(y)}dxdy=D_{KL}(\rho(x,y)|| \rho(x)\rho(y)).\]

In what follows, we use the quite important property of MI -- invariance under nonsingular mappings between smooth manifolds. We aim to show that one can measure the MI between compressed representations of random vectors. The following statement holds.

**Statement 1**.: _Let \(\xi\colon\Omega\to\mathbb{R}^{n^{\prime}}\) be absolutely continuous random vector, \(f\colon\mathbb{R}^{n^{\prime}}\to\mathbb{R}^{n}\) be injective piecewise-smooth mapping with Jacobian \(J_{f}\) such that \(n\geq n^{\prime}\) and \(\det\left(J_{f}^{T}J_{f}\right)\neq 0\) almost everywhere. Let either \(\eta\) be a discrete random variable, or \((\xi,\eta)\) be an absolutely continuous random vector. Then_

\[I(\xi;\eta)=I\left(f(\xi);\eta\right) \tag{3}\]

**Remark 1**.: _In what follows by \(\xi\colon\Omega\to\mathbb{R}^{n^{\prime}}\) we denote the compressed representation of \(X\), \(n^{\prime}\leq n\)._

Recall that we utilize the stochastic NN approach to deal with the problem of infinite MI \(I(X;f(X))\) for a deterministic mapping \(f\). As it has been shown in [15], adding stochasticity enables proper MI estimation between layers of the network. The stochastic modification of a network is used as a proxy to determine the information-theoretic properties of the original model.

The conventional feedforward NN can be defined as an acyclic computational graph, which can be topologically sorted:

\[L_{0}\triangleq X,\quad L_{1}:=f_{1}(L_{0}),\quad L_{2}:=f_{2}(L_{0},L_{1}), \quad\ldots,\quad\hat{Y}\triangleq L_{n}:=f_{n}(L_{0},\ldots,L_{n-1}),\]

where \(L_{0},\ldots,L_{n}\) are the outputs of the network's layers. The stochastic modification is defined similarly, but using the Markov chain stochastic model:

**Definition 1**.: _The sequence of random vectors \(L_{0},\ldots,L_{n}\) is said to form a stochastic neural network with input \(X\) and output \(\hat{Y}\), if \(L_{0}\triangleq X\), \(\hat{Y}\triangleq L_{n}\), and_

\[L_{0}\ \longrightarrow\ (L_{0},L_{1})\ \longrightarrow\ \ldots\ \longrightarrow\ (L_{0},\ldots,L_{n})\]

_is a Markov chain; \(L_{k}\) represents outputs of the \(k\)-th layer of the network._
Our main goal is to track \(I(L_{i};L_{j})\) during the training process. Due to \(L_{k}\) being of high dimension, it is difficult to acquire the desired estimates [15]. In the following sections, we propose the MI estimation method for high-dimensional random vectors, which utilizes the compression step. Our method works for the datasets for which the manifold hypothesis holds. In what follows we perform all our experiments with raster pictures. As any \(L_{k}\) is an image of \(X\) under some continuous mapping, the manifold hypothesis is likely to be assumed for \(L_{k}\), thus justifying the usage of the proposed method.

## 3 Mutual information estimation via lossy compression

In this section, we discuss the application of lossy compression to MI estimation for high-dimensional random vectors. The task is to estimate \(I(X;Y)\) using \(\{(x_{k},y_{k})\}_{k=1}^{N}\), where \((x_{k},y_{k})\sim(X,Y)\). Conventional estimation approaches rely to some extent on PDF reconstruction in order to acquire the MI or entropy estimate. However, these methods fail to yield a correct estimation for high-dimensional cases. In the following subsections, we analyze the conventional MI estimators, propose and theoretically justify a complementary lossy compression step to tackle the problems of high dimensionality and derive theoretical bounds on mutual information estimate under lossy compression.

### Mutual information estimation

The most straightforward way to estimate mutual information is to estimate all the components in (2) or (1). There are several methods of entropy estimation that require only a set of samples from multivariate distribution: plug-in estimators [33; 34], kernel density-based approaches [33; 35], weighted and non-weighted Kozachenko-Leonenko estimators [36; 37] and other. However, all the mentioned approaches fail to correctly estimate the entropy of high-dimensional random vectors.

We assume the manifold hypothesis (i.e., that data lie along (or close to) some manifold in multidimensional space). This is believed to be true for a large variety of structured data; some datasets are known to satisfy this assumption exactly. In our work we use the following simplified definition of the manifold hypothesis:

**Definition 2**.: _A random vector \(X\colon\Omega\to\mathbb{R}^{n}\) is said to strictly satisfy the manifold hypothesis iff there exist \(\xi\colon\Omega\to\mathbb{R}^{n^{\prime}}\) and \(f\colon\mathbb{R}^{n^{\prime}}\to\mathbb{R}^{n}\) satisfying the conditions of Statement 1, such that \(X=f(\xi)\). A random vector \(X^{\prime}\colon\Omega\to\mathbb{R}^{n}\) is said to loosely satisfy the manifold hypothesis iff \(X^{\prime}=X+Z\), where \(X\) strictly satisfies the manifold hypothesis, and \(Z\) is insignificant in terms of some metric._

To avoid support of measure zero, the probability measure induced on the manifold is considered. As we will show, considering induced PDFs can also solve the problem of high dimensionality, as the small intrinsic dimension of the manifold allows conventional entropy estimators to be applied.

Let \(\{(x_{k},y_{k})\}_{k=1}^{N}\) be a sequence of samples from the joint distribution of random vectors \(X\) and \(Y\). Our goal is to estimate the mutual information between \(X\) and \(Y\), i.e. \(I(X;Y)\), from these samples. To overcome the curse of dimensionality, we suggest learning the manifold with autoencoders [38; 39] and apply conventional estimators to the compressed representations.

First, let us suppose that random vectors \(X\) and \(Y\) are absolutely continuous, and corresponding autoencoders are provided: \(A_{X}=D_{X}\circ E_{X}\), \(A_{Y}=D_{Y}\circ E_{Y}\).

**Corollary 1**.: _Let \(E_{X}^{-1}\colon\mathbb{R}^{n^{\prime}}\supseteq E_{X}(\operatorname{supp}X) \to\mathbb{R}^{n}\), \(E_{Y}^{-1}\colon\mathbb{R}^{m^{\prime}}\supseteq E_{Y}(\operatorname{supp}Y) \to\mathbb{R}^{m}\), \(E_{X}(X)\colon\Omega\to\mathbb{R}^{n^{\prime}}\) and \(E_{Y}(Y)\colon\Omega\to\mathbb{R}^{m^{\prime}}\) exist and satisfy conditions of the Statement 1. Then_

\[I(X;Y)=I(E_{X}(X);E_{Y}(Y)),\]

_so the mutual information estimate can be computed as follows:_

\[\hat{I}(X;Y)=\hat{I}(E_{X}(X);E_{Y}(Y))=\hat{h}(E_{X}(X))+\hat{h}(E_{Y}(Y))- \hat{h}(E_{X}(X),E_{Y}(Y)) \tag{4}\]

Now, the case of absolutely continuous \(X\) and discrete \(Y\) should be considered. In this case, it is impractical to use (2), as the (induced) joined probability distribution is neither absolutely continuous nor discrete. However, (1) is still valid:

\[h(X\mid Y)=\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\!\! \!\!
Probabilities \(p_{Y}\) can be estimated using empirical frequencies: \(\hat{p}_{Y}(y)=\frac{1}{N}\cdot|\{k\mid y_{k}=y\}|\). Conditional entropy \(h(X\mid Y=y)\) can be estimated using corresponding subsets of \(\{x_{k}\}_{k=1}^{N}\): \(\hat{h}(X\mid Y=y)=\hat{h}(\{x_{k}\mid y_{k}=y\})\).

**Corollary 2**.: _Let \(E_{X}^{-1}\colon\mathbb{R}^{n^{\prime}}\supseteq E_{X}(\text{supp}\,X)\to \mathbb{R}^{n}\), \(E_{X}(X)\colon\Omega\to\mathbb{R}^{n^{\prime}}\) and \(Y\colon\Omega\to\text{supp}\,Y\) satisfy conditions of the Statement 1. Then_

\[I(X;Y)=I(E_{X}(X);Y),\]

_so the mutual information estimate can be computed as follows:_

\[\hat{I}(X;Y)=\hat{I}(E_{X}(X);Y)=\hat{h}(E_{X}(X))-\sum_{y\in\text{supp}\,Y} \hat{p}_{Y}(y)\cdot\hat{h}(E_{X}(X)\mid Y=y) \tag{5}\]

According to the strong law of large numbers, \(\hat{p}\xrightarrow{\text{a.s.}}p\). That is why the convergence of the proposed MI estimation methods solely relies on the convergence of the entropy estimator used in (4) and (5).

### Bounds for mutual information estimate

It can be shown that it is not possible to derive non-trivial bounds for \(I(E_{X}(X);E_{Y}(Y))\) and \(I(E_{X}(X);Y)\) in general case if Corollaries 1, 2 do not hold. Let us consider a simple linear autoencoder, optimal in terms of mean squared error (i.e. principal component analysis-based). The following statement shows that there are cases when the proposed method of mutual information estimation via lossy compression fails to estimate the MI.

**Statement 2**.: _For any given \(\varkappa\geq 0\) there exist random variables \(X\colon\Omega\to\mathbb{R}^{n}\), \(Y\colon\Omega\to\mathbb{R}^{m}\) and non-trivial linear autoencoder \(A=D\circ E\) with latent space dimension \(n^{\prime}<n\), optimal in terms of minimizing mean squared error \(\mathbb{E}\,\|X-A(X)\|^{2}\), such that \(I(X;Y)=\varkappa\) and \(I(E(X);Y)=0\)._

Proof.: Let us consider the following three-dimensional Gaussian vector \((X_{1},X_{2},Y)\triangleq(X,Y)\):

\[X\sim\mathcal{N}\left(0,\begin{bmatrix}1&0\\ 0&\sigma\end{bmatrix}\right),\qquad Y\sim\mathcal{N}(0,1)\qquad(X_{1},X_{2},Y) \sim\mathcal{N}\left(0,\begin{bmatrix}1&0&0\\ 0&\sigma&a\\ 0&a&1\end{bmatrix}\right),\]

where \(\text{cov}(X_{2},Y)=a\triangleq\sqrt{1-e^{-2\varkappa}}\), \(\text{cov}(X_{1},Y)=0\) (so \(X_{1}\) and \(Y\) are independent). Let the intrinsic dimension be \(n^{\prime}=1\), and \(\sigma<1\). According to the principal component analysis, the optimal linear encoder is defined up to a scalar factor by the equality \(E(X)=X_{1}\). However, \(I(X;Y)=-\frac{1}{2}\ln\left(1-a^{2}\right)=\varkappa\) (see the Statement 4), but \(I(E(X);Y)=0\), as \(A(X)\) and \(Y\) are independent. 

This statement shows that an arbitrary amount of information can be lost due to the compression. It follows from the fact that "less significant" in terms of metric spaces does not comply with "less significant" in terms of information theory. However, with some additional assumptions, a more useful theoretical result can be achieved.

**Statement 3**.: _Let \(X\), \(Y\) and \(Z\) be random variables such that \(I(X;Y)\) and \(I\left((X,Z);Y\right)\) are defined. Let \(f\) be a function of two arguments such that \(I(f(X,Z);Y)\) is defined. If there exists a function \(g\) such that \(X=g(f(X,Z))\), then the following inequalities hold:_

\[I(X;Y)\leq I\left(f(X,Z);Y\right)\leq I((X,Z);Y)\]

Here \(f(X,Z)\) can be interpreted as (compressed) noisy data, \(X\) -- as denoised data, and \(g\) -- as a perfect denoising autoencoder. This statement justifies the proposed lossy compression method in some cases (namely, when data lost by compression can be considered as an independent random noise).

**Corollary 3**.: _Let \(X\), \(Y\), \(Z\), \(f\) and \(g\) satisfy the conditions of the Statement 3. Let also random variables \((X,Y)\) and \(Z\) be independent. Then \(I(X;Y)=I\left(f(X,Z);Y\right)\)._
These bounds are optimal due to the optimality of the data processing inequality, but additional knowledge about the connection between \(X\), \(Y\), and \(Z\) is required to properly utilize them. Other bounds can also be derived [21; 23; 40], but they do not utilize the fact of compression.

## 4 Synthetic dataset generation

In order to test the proposed mutual information estimator we develop an universal method for synthetic dataset generation with defined information-theoretic properties. This method yields two random vectors \(X\) and \(Y\) which have predefined value of mutual information \(I(X;Y)\). The method requires \(X\) and \(Y\) to be images of normally distributed vectors under some known nonsingular smooth mappings.

The generation consists of two steps: first, a normal vector \((\xi,\eta)\sim\mathcal{N}(0,M)\) is considered, where \(\xi\sim\mathcal{N}(0,I_{n^{\prime}})\), \(\eta\sim\mathcal{N}(0,I_{m^{\prime}})\) and \(n^{\prime}\), \(m^{\prime}\) are dimensions of \(\xi\) and \(\eta\) respectively. The covariance matrix \(M\) is chosen to satisfy \(I(\xi;\eta)=\varkappa\), where \(\varkappa\) is an arbitrary non-negative constant.

**Statement 4**.: _Let \((\xi,\eta)\sim\mathcal{N}(0,M)\) be a Gaussian pair of (scalar) random variables with unit variance such that \(I(\xi;\eta)=\varkappa\). Then_

\[M=\begin{bmatrix}1&a\\ a&1\end{bmatrix},\qquad a=\sqrt{1-e^{-2\varkappa}} \tag{6}\]

**Corollary 4**.: _For every \(\varkappa\geq 0\) and every \(n^{\prime},m^{\prime}\in\mathbb{N}\) exists a matrix \(M\in\mathbb{R}^{(n^{\prime}+m^{\prime})\times(n^{\prime}+m^{\prime})}\) such that \((\xi,\eta)\sim\mathcal{N}(0,M)\), \(\xi\sim\mathcal{N}(0,I_{n^{\prime}})\), \(\eta\sim\mathcal{N}(0,I_{m^{\prime}})\) and \(I(\xi;\eta)=\varkappa\)._

The second step of the method is the application of smooth nonsingular mappings to vectors \(\xi\) and \(\eta\): \(X=f(\xi)\), \(Y=g(\eta)\). Due to the Statement 1, \(I(\xi;\eta)=I(X;Y)\).

## 5 Comparison of the entropy estimators

The MI estimate is acquired according to Subsection 3.1. To estimate the entropy terms in (2) or (1) we leverage conventional entropy estimators, such as kernel density-based and Kozachenko-Leonenko estimators (original and weighted versions), see Table 1. To test the accuracy of these approaches, we use datasets sampled from synthetic random vectors with known MI. We generate these datasets in accordance to Section 4.

To examine the impact of the compression step proposed in Subsection 3.1 we utilize a special type of synthetic datasets. Synthetic data lies on a manifold of small dimension. This is achieved by generating a low-dimensional dataset and then embedding it into a high-dimensional space by a smooth mapping (so the Statement 1 can be applied). Then the acquired datasets are compressed via autoencoders. Finally, the obtained results are fed into a mutual information estimator.

The Algorithm 1 and Figure 2 describe the proposed mutual information estimation quality measurement. We run several experiments with \(f\) and \(g\) mapping normal distributions to rasterised images of geometric shapes (e.g. rectangles) or 2D-plots of smooth functions (e.g. Gaussian functions). The results are presented in the Figures 3 and 4. Blue and green curves correspond to the estimates of MIs marked by the corresponding colors in Figure 2. Thus, we see that the compression step do not lead to bad estimation accuracy especially for KL and WKL estimators, which demonstrate the best performance. In what follows we use WKL estimator.
## 6 Information flow in deep neural networks

This section is dedicated to the information flow estimation in DNNs via the proposed method. We estimate information flow in a convolutional classifier of the MNIST handwritten digits dataset [32]. This neural network is simple enough to be quickly trained and tested, but, at the same time, is complex enough to suffer from the curse of dimensionality. The MNIST dataset consists of images of size \(28\times 28=784\) pixels. It was shown in [44] that these images have relatively low latent space

Figure 3: Maximum-likelihood and Least Squares Error KDE, Non-weighted and Weighted Kozachenko-Leonenko for \(16\times 16\) (first row) and \(32\times 32\) (second row) images of gaussians (\(n^{\prime}=m^{\prime}=3\)), \(5\cdot 10^{3}\) samples. Along \(x\) axes is \(I(X;Y)\), along \(y\) axes is \(\hat{I}(X;Y)\).

Figure 2: Conceptual scheme of Algorithm 1. In order to observe and quantify the loss of information caused by the compression step, we split \(f\colon\mathbb{R}^{n^{\prime}}\to\mathbb{R}^{n}\) into two functions: \(f_{1}\colon\mathbb{R}^{n^{\prime}}\to\mathbb{R}^{n^{\prime}}\) maps \(\xi\) to a structured latent representation of \(X\) (e.g. to parameters of geometric shapes), and \(f_{2}\colon\mathbb{R}^{n^{\prime}}\to\mathbb{R}^{n}\) maps latent representations to corresponding high-dimensional vectors (e.g. rasterised images of geometric shapes). The same goes for \(g=g_{2}\circ g_{1}\). Colors correspond to the Figures 3 and 4.
dimension -- about \(12\)-\(13\). If the preservation of only main features is desired, the latent space can even be narrowed down to \(3\)-\(10\).

It can be concluded from the previous section that the weighted Kozachenko-Leonenko estimator is superior to the other methods tested in this paper. That is why this estimator is used in experiments with the MNIST classifier described in the current section. The analyzed network is designed to return the output of every layer. To avoid the problem of a deterministic relationship between input and output, a small additive Gaussian noise is injected on each layer (we use constant signal-to-noise ratio). Lossy compression of input images \(X\) is performed via convolutional autoencoder with latent dimension \(d_{X}^{\text{latent}}\). Lossy compression of layer outputs \(L_{i}\) is performed via principal component analysis with \(d_{L_{i}}^{\text{latent}}\) as the number of principal components. The general algorithm is described in 2.

```
1:Compress the input dataset \(\{x_{k}\}_{k=1}^{N}\) via the encoder \(E_{X}\): \(c_{k}^{X}=E_{X}(x_{k})\).
2:for epoch : \(1,\dots,\text{number of epochs}\)do
3:for\(L_{i}\) : layers do
4: Collect outputs of the layer \(L_{i}\): \(y_{k}^{L_{i}}=L_{i}(x_{k})\). Each layer must be noisy/stochastic.
5: Compress the outputs via encoder \(E_{L_{i}}\): \(c_{k}^{L_{i}}=E_{L_{i}}(y_{k}^{L_{i}})\).
6: Estimate \(I(E_{X}(X);E_{L_{i}}(L_{i}))\) and \(I(E_{L_{i}}(L_{i});Y(X))\), where \(Y\) maps inputs to their true targets.
7:endfor
8: Perform one-epoch training step of the network.
9:endfor
```

**Algorithm 2** Estimate information flow in the neural network during training

We use the architecture of the classification network provided in Table 2. We train our network with a learning rate equaling \(10^{-5}\) using Nvidia Titan RTX. For other hyperparameters we refer to the source code [43]. The acquired information plane plots are provided by Figure 5. As the direction of the plots with respect to the epoch can be deduced implicitly (the lower left corner of the plot corresponds to the first epochs), we color the lines in accordance with loss function dynamics per epoch. We do this to stress out one of the key observations: the first transition from fitting to compression phase coincides with the acceleration of loss function decrease. It is also evident that there is no clear large-scale compression phase. Moreover, it seems that the number of fitting and compression phases can vary from layer to layer.

\begin{table}
\begin{tabular}{l l} \(L_{1}\): & Conv2d(1, 8, ks=3), LeakyReLU(0.01) \\ \(L_{2}\): & Conv2d(8, 16, ks=3), LeakyReLU(0.01) \\ \(L_{3}\): & Conv2d(16, 32, ks=3), LeakyReLU(0.01) \\ \(L_{4}\): & Dense(32, 32), LeakyReLU(0.01) \\ \(L_{5}\): & Dense(32, 10), LogSoftMax \\ \end{tabular}
\end{table}
Table 2: The architecture of the MNIST convolution-DNN classifier used in this paper.

Figure 4: Maximum-likelihood and Least Squares Error KDE, Non-weighted and Weighted Kozachenko-Leonenko for \(16\times 16\) (first row) and \(32\times 32\) (second row) images of rectangles (\(n^{\prime}=m^{\prime}=4\)), \(5\cdot 10^{3}\) samples. Along \(x\) axes is \(I(X;Y)\), along \(y\) axes is \(\hat{I}(X;Y)\).
## 7 Discussion

An information-theoretic approach to explainable artificial intelligence and deep neural networks analysis seems promising, as it is interpretable, robust, and relies on the well-developed information theory. However, the direct application of information-theoretic analysis still constitutes some problems. The main contribution of our work is overcoming challenges caused by the curse of dimensionality in the information theory of machine learning. We have shown that it is possible to apply information analysis to compressed representations of datasets or models' outputs.

In order to justify our approach, we have acquired several theoretic results regarding mutual information estimation under lossless and lossy compression. These results suggest that this approach is applicable to real data cases. Although it has been shown that an arbitrary amount of information can be lost due to compression, the information required for optimal decompression is still preserved. We have also developed a framework to test conventional mutual information estimators complemented with the proposed lossy compression step. This framework allows generating of pairs of high-dimensional datasets with small internal (latent) dimensions and a predefined quantity of mutual information. Conducted numerical experiments have shown that the proposed method performs well, especially if entropy estimation is done via a weighted Kozachenko-Leonenko estimator; other methods tend to underestimate or overestimate mutual information.

Finally, an information plane experiment with the MNIST dataset classifier has been carried out. This experiment has shown that dynamics of information-theoretic quantities during the training of deep neural networks are, indeed, non-trivial. However, it is not clear whether the original fitting-compression hypothesis holds, as there is no clear large-scale compression phase after the fitting. We suggest that there may be several compression/fitting phases during the training of real-scale neural networks. An interesting observation has also been obtained: the first compression phase coincides with the rapid decrease of loss function (i.e. with the highest absolute value of loss function delta per epoch). It may be concluded that phases revealed by information plane plots are connected

Figure 5: Information plane plots for the MNIST classifier. The lower left parts of the plots correspond to the first epochs of training.
to different regimes of learning (i.e. accelerated, stationary or decelerated drop of loss function). However, further investigation of this phenomenon has to be carried out in order to supply this seeming connection with more evidence and theoretical basis.

As a further research we consider using normalizing flows to improve our approach [45]. Normalizing flows are invertible smooth mappings, providing means of lossless and information-preserving compression. Normalizing flows can also be used to transform the joint distribution to a Gaussian, thus facilitating mutual information estimation.