SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Next, we verified our assumption of non-overlapping effective support in the model's data distribution. We calculate the distribution of pairwise  $l_{2}$  distances between images across seven datasets: MNIST [41], CIFAR10, CIFAR100 [40], Flowers102 [52], Food101 [12], and FGVAircraft [45]. Figure 1 (right) reveals that the pairwise distances are far from zero, even for raw pixels. This implies that we can use a small Gaussian around each point without overlap, validating our assumption as realistic.
5 Self-Supervised Learning Models through Information Maximization
The practical application of Equation (8) involves several key "design choices". We begin by comparing how existing SSL models have implemented it, investigating the estimators used, and discussing the implications of their assumptions. Subsequently, we introduce new methods for SSL that incorporate sophisticated estimators from the field of information theory, which outperform current approaches.
5.1 VICReg vs. SimCLR
In order to evaluate their underlying assumptions and strategies for information maximization, we compare VICReg to contrastive SSL methods such as SimCLR along with non-contrastive methods like BYOL and SimSiam.
Contrastive Learning with SimCLR. In their work, Lee et al. [43] drew a connection between the SimCLR objective and the variational bound on information regarding representations by employing the von Mises-Fisher distribution. By applying our analysis for information in deterministic networks with their work, we compare the main differences between SimCLR and VICReg:
(i) Conditional distribution: SimCLR assumes a von Mises-Fisher distribution for the encoder, while VICReg assumes a Gaussian distribution. (ii) Entropy estimation: SimCLR approximated it based on the finite sum of the input samples. In contrast, VICReg estimates the entropy of  $Z$  solely based on the second moment. Developing SSL methods that integrate these two distinctions form an intriguing direction for future research.
Empirical comparison. We trained ResNet18 on CIFAR-10 for VICReg, SimCLR, and BYOL and compared their entropies directly using the pairwise distances entropy estimator. (For more details, see Appendix K.) This estimator was not directly optimized by any method and was an independent validation. The results (Figure 2), showed that entropy increased for all methods during training, with SimCLR having the lowest and VICReg the highest entropy.
Figure 2: VICReg has higher Entropy during training. The entropy along the training for different SSL methods. Experiments were conducted with ResNet-18 on CIFAR-10. Error bars represent one standard error over 5 trials.
5.2 Family of alternative Entropy Estimators
Next, we suggest integrating the invariance term of current SSL methods with plug-in methods that optimize entropy.
Entropy estimators. The VICReg objective seeks to approximate the log determinant of the empirical covariance matrix through its diagonal terms. As discussed in Section 4.1, this approach has its drawbacks. An alternative is to employ different entropy estimators. The LogDet Entropy Estimator [74] is one such option, offering a tighter upper bound. This estimator employs the differential entropy of  $\alpha$  order with scaled noise and has been previously shown to be a tight estimator for high-dimensional features, proving robust to random noise. However, since this estimator provides an upper bound on entropy, maximizing this bound doesn't guarantee optimization of the original objective. To counteract this, we also introduce a lower bound estimator based on the pairwise distances of individual mixture components [38]. In this family, a function determining pairwise distances between component densities is designated for each member. These estimators are computationally efficient and typically straightforward to optimize. For additional entropy estimators, see Appendix F. Beyond VICReg, these methods can serve as plug-in estimators for numerous SSL algorithms. Apart from VICReg, we also conducted experiments integrating these estimators with the BYOL algorithm.
Table 1: The proposed entropy estimators outperform previous methods. CIFAR-10, CIFAR-100, and Tiny-ImageNet top-1 accuracy under linear evaluation using ResNet-18, ConvNetX and VIT as backbones. Error bars correspond to one standard error over three trials.
Method	CIFAR-10
ResNet-18	Tiny-ImageNet	CIFAR-100
ConvNetX	VIT	ConvNetX	VIT
SimCLR	89.72 ± 0.05	50.86 ± 0.13	51.16 ± 0.13	67.21 ± 0.24	67.31 ± 0.18
Barlow Twins	88.81 ± 0.10	51.34 ± 0.10	51.40 ± 0.16	68.54 ± 0.15	68.02 ± 0.12
SwAV	89.12 ± 0.13	50.76 ± 0.14	51.54 ± 0.20	68.93 ± 0.14	67.89 ± 0.21
MoCo	89.46 ± 0.08	52.36 ± 0.21	53.06 ± 0.21	70.32 ± 0.15	69.89 ± 0.14
VICReg	89.32 ± 0.09	51.02 ± 0.26	52.12 ± 0.25	70.09 ± 0.20	70.12 ± 0.17
BYOL	89.21 ± 0.11	52.24 ± 0.17	53.44 ± 0.20	70.01 ± 0.27	69.59 ± 0.22
VICReg + PairDist (ours)	90.37 ± 0.09	52.61 ± 0.15	53.70 ± 0.13	71.10 ± 0.16	70.50 ± 0.19
VICReg + LogDet (ours)	90.27 ± 0.08	52.91 ± 0.17	54.89 ± 0.20	71.23 ± 0.18	70.61 ± 0.17
BYOL + PairDist (ours)	90.19 ± 0.14	53.47 ± 0.22	54.33 ± 0.21	71.39 ± 0.25	71.09 ± 0.24
BYOL + LogDet (ours)	90.11 ± 0.16	53.19 ± 0.25	54.67 ± 0.27	71.20 ± 0.21	70.79 ± 0.26
Setup. Experiments were conducted on three image datasets: CIFAR-10, CIFAR-100 [39], and Tiny-ImageNet [20]. For CIFAR-10, ResNet-18 [31] was used. In contrast, both ConvNeXt [44] and Vision Transformer [21] were used for CIFAR-100 and Tiny-ImageNet. For comparison, we examined the following SSL methods: VICReg, SimCLR, BYOL, SwAV [14], Barlow Twins [73], and MoCo [33]. The quality of representation was assessed through linear evaluation. A detailed description of different methods can be found in Appendix H.
Results. As evidenced by Table 1, the proposed entropy estimators surpass the original SSL methods. Using a more precise entropy estimator enhances the performance of both VICReg and BYOL, compared to their initial implementations. Notably, the pairwise distance estimator, being a lower bound, achieves superior results, resonating with the theoretical preference for maximizing a true entropy's lower bound. Our findings suggest that the astute choice of entropy estimators, guided by our framework, paves the way for enhanced performance.
6 A Generalization Bound for Downstream Tasks
In earlier sections, we linked information theory principles with the VICReg objective. Now, we aim to extend this link to downstream generalization via a generalization bound. This connection further aligns VICReg's generalization with information maximization and implicit regularization.
Notation. Consider input points  $x$ , outputs  $y \in \mathbb{R}^r$ , labeled training data  $S = ((x_i, y_i))_{i=1}^n$  of size  $n$  and unlabeled training data  $\bar{S} = ((x_i^+, x_i^{++}))_{i=1}^m$  of size  $m$ , where  $x_i^+$  and  $x_i^{++}$  share the same (unknown) label. With the unlabeled training data, we define the invariance loss
$$
I _ {\bar {S}} (f _ {\theta}) = \frac {1}{m} \sum_ {i = 1} ^ {m} \| f _ {\theta} (x _ {i} ^ {+}) - f _ {\theta} (x _ {i} ^ {+ +}) \|,
$$
where  $f_{\theta}$  is the trained representation on the unlabeled data  $\bar{S}$ . We define a labeled loss  $\ell_{x,y}(w) = \| Wf_{\theta}(x) - y\|$  where  $w = \operatorname{vec}[W] \in \mathbb{R}^{dr}$  is the vectorization of the matrix  $W \in \mathbb{R}^{r \times d}$ . Let  $w_{S} = \operatorname{vec}[W_{S}]$  be the minimum norm solution as  $W_{S} = \mathrm{minimize}_{W^{\prime}} \| W^{\prime}\|_{F}$  such that
$$
W ^ {\prime} \in \underset {W} {\arg \min } \frac {1}{n} \sum_ {i = 1} ^ {n} \| W f _ {\theta} (x _ {i}) - y _ {i} \| ^ {2}.
$$
We also define the representation matrices
$$
Z _ {S} = \left[ f \left(x _ {1}\right), \dots , f \left(x _ {n}\right) \right] \in \mathbb {R} ^ {d \times n} \quad \text {a n d} \quad Z _ {\bar {S}} = \left[ f \left(x _ {1} ^ {+}\right), \dots , f \left(x _ {m} ^ {+}\right) \right] \in \mathbb {R} ^ {d \times m},
$$
and the projection matrices
$$
\mathbf {P} _ {Z _ {S}} = I - Z _ {S} ^ {\top} (Z _ {S} Z _ {S} ^ {\top}) ^ {\dagger} Z _ {S} \quad \text {a n d} \quad \mathbf {P} _ {Z _ {\bar {S}}} = I - Z _ {\bar {S}} ^ {\top} (Z _ {\bar {S}} Z _ {\bar {S}} ^ {\top}) ^ {\dagger} Z _ {\bar {S}}.
$$
We define the label matrix  $Y_{S} = [y_{1},\ldots ,y_{n}]^{\top}\in \mathbb{R}^{n\times r}$  and the unknown label matrix  $Y_{\bar{S}} = [y_1^+, \dots, y_m^+]^\top \in \mathbb{R}^{m\times r}$ , where  $y_{i}^{+}$  is the unknown label of  $x_{i}^{+}$ . Let  $\mathcal{F}$  be a hypothesis space of  $f_{\theta}$ . For a given hypothesis space  $\mathcal{F}$ , we define the normalized Rademacher complexity
$$
\tilde {\mathcal {R}} _ {m} (\mathcal {F}) = \frac {1}{\sqrt {m}} \mathbb {E} _ {\bar {S}, \xi} \left[ \sup _ {f \in \mathcal {F}} \sum_ {i = 1} ^ {m} \xi_ {i} \| f (x _ {i} ^ {+}) - f (x _ {i} ^ {+ +}) \| \right],
$$
where  $\xi_1, \ldots, \xi_m$  are independent uniform random variables in  $\{-1, 1\}$ . It is normalized such that  $\tilde{\mathcal{R}}_m(\mathcal{F}) = O(1)$  as  $m \to \infty$ .
6.1 A Generalization Bound for Variance-Invariance-Covariance Regularization
Now we will show that the VICReg objective improves generalization on supervised downstream tasks. More specifically, minimizing the unlabeled invariance loss while controlling the covariance  $Z_{\bar{S}}Z_{\bar{S}}^{\top}$  and the complexity of representations  $\tilde{\mathcal{R}}_m(\mathcal{F})$  minimizes the expected labeled loss:
Theorem 3. (Informal version). For any  $\delta >0$ , with probability at least  $1 - \delta$
$$
\mathbb {E} _ {x, y} \left[ \ell_ {x, y} \left(w _ {S}\right) \right] \leq I _ {\bar {S}} \left(f _ {\theta}\right) + \frac {2}{\sqrt {m}} \| \mathbf {P} _ {Z _ {\bar {S}}} Y _ {\bar {S}} \| _ {F} + \frac {1}{\sqrt {n}} \| \mathbf {P} _ {Z _ {S}} Y _ {S} \| _ {F} + \frac {2 \tilde {\mathcal {R}} _ {m} (\mathcal {F})}{\sqrt {m}} + \mathcal {Q} _ {m, n}, \tag {12}
$$
where  $\mathcal{Q}_{m,n} = O(G\sqrt{\ln(1 / \delta) / m} +\sqrt{\ln(1 / \delta) / n})\to 0$  as  $m,n\to \infty$ . In  $\mathcal{Q}_{m,n}$ , the value of  $G$  for the term decaying at the rate  $1 / \sqrt{m}$  depends on the hypothesis space of  $f_{\theta}$  and  $w$  whereas the term decaying at the rate  $1 / \sqrt{n}$  is independent of any hypothesis space.
Proof. The complete version of Theorem 3 and its proof are presented in Appendix I.
The term  $\| \mathbf{P}_{Z_{\bar{S}}}Y_{\bar{S}}\| _F$  in Theorem 3 contains the unobservable label matrix  $Y_{\bar{S}}$ . However, we can minimize this term by using  $\| \mathbf{P}_{Z_{\bar{S}}}Y_{\bar{S}}\| _F\leq \| \mathbf{P}_{Z_{\bar{S}}}\| _F\| Y_{\bar{S}}\| _F$  and by minimizing  $\| \mathbf{P}_{Z_{\bar{S}}}\| _F$ . The factor  $\| \mathbf{P}_{Z_{\bar{S}}}\| _F$  is minimized when the rank of the covariance  $Z_{\bar{S}}Z_{\bar{S}}^{\top}$  is maximized. This can be enforced by maximizing the diagonal entries while minimizing the off-diagonal entries, as is done in VICReg.
The term  $\| \mathbf{P}_{Z_S}Y_S\| _F$  contains only observable variables, and we can directly measure the value of this term using training data. In addition, the term  $\| \mathbf{P}_{Z_S}Y_S\| _F$  is also minimized when the rank of the covariance  $Z_{S}Z_{S}^{\top}$  is maximized. Since the covariances  $Z_{S}Z_{S}^{\top}$  and  $Z_{\bar{S}}Z_{\bar{S}}^{\top}$  concentrate to each other via concentration inequalities with the error in the order of  $O(\sqrt{(\ln(1 / \delta)) / n} +$ $\tilde{\mathcal{R}}_m(\mathcal{F})\sqrt{(\ln(1 / \delta)) / m})$  , we can also minimize the upper bound on  $\| \mathbf{P}_{Z_S}Y_S\| _F$  by maximizing the diagonal entries of  $Z_{\bar{S}}Z_{\bar{S}}^{\top}$  while minimizing its off-diagonal entries, as is done in VICReg.
Thus, VICReg can be understood as a method to minimize the generalization bound in Theorem 3 by minimizing the invariance loss while controlling the covariance to minimize the label-agnostic upper bounds on  $\| \mathbf{P}_{Z_S}Y_{\bar{S}}\| _F$  and  $\| \mathbf{P}_{Z_S}Y_S\| _F$ . If we know partial information about the label  $Y_{\bar{S}}$  of the unlabeled data, we can use it to minimize  $\| \mathbf{P}_{Z_S}Y_{\bar{S}}\| _F$  and  $\| \mathbf{P}_{Z_S}Y_S\| _F$  directly.
6.2 Comparison of Generalization Bounds
The SimCLR generalization bound [58] requires the number of labeled classes to go infinite to close the generalization gap, whereas the VICReg bound in Theorem 3 does not require the number of label classes to approach infinity for the generalization gap to go to zero. This reflects that, unlike SimCLR, VICReg does not use negative pairs and thus does not use a loss function based on the implicit expectation that the labels of a negative pair are different. Another difference is that our VICReg bound improves as  $n$  increases, while the previous bound of SimCLR [58] does not depend on  $n$ . This is because Saunshi et al. [58] assumes partial access to the true distribution per class for setting, which removes the importance of labeled data size  $n$  and is not assumed in our study.
Consequently, the generalization bound in Theorem 3 provides a new insight for VICReg regarding the ratio of the effects of  $m$  v.s.  $n$  through  $G\sqrt{\ln(1 / \delta) / m} + \sqrt{\ln(1 / \delta) / n}$ . Finally, Theorem 3 also illuminates the advantages of VICReg over standard supervised training. That is, with standard training, the generalization bound via the Rademacher complexity requires the complexities of hypothesis spaces,  $\tilde{\mathcal{R}}_n(\mathcal{W}) / \sqrt{n}$  and  $\tilde{\mathcal{R}}_n(\mathcal{F}) / \sqrt{n}$ , with respect to the size of labeled data  $n$ , instead of the size of unlabeled data  $m$ . Thus, Theorem 3 shows that using SSL, we can replace the complexities of hypothesis spaces in terms of  $n$  with those in terms of  $m$ . Since the number of unlabeled data points is typically much larger than the number of labeled data points, this illuminates the benefit of SSL.
6.3 Understanding Theorem 2 via Mutual Information Maximization
Theorem 3, together with the result of the previous section, shows that, for generalization in the downstream task, it is helpful to maximize the mutual information  $I(Z;X^{\prime})$  in SSL via minimizing the invariance loss  $I_{\bar{S}}(f_{\theta})$  while controlling the covariance  $Z_{\bar{S}}Z_{\bar{S}}^{\top}$ . The term  $2\tilde{\mathcal{R}}_m(\mathcal{F}) / \sqrt{m}$  captures the importance of controlling the complexity of the representations  $f_{\theta}$ . To understand this term further, let us consider a discretization of the parameter space of  $\mathcal{F}$  to have finite  $|\mathcal{F}| < \infty$ . Then, by Massart's Finite Class Lemma, we have that  $\tilde{\mathcal{R}}_m(\mathcal{F}) \leq C\sqrt{\ln|\mathcal{F}|}$  for some constant  $C > 0$ . Moreover, Shwartz-Ziv [60] shows that we can approximate  $\ln |\mathcal{F}|$  by  $2^{I(Z;X)}$ . Thus, in Theorem 3, the term  $I_{\bar{S}}(f_{\theta}) + \frac{2}{\sqrt{m}}\| \mathbf{P}_{Z_{\bar{S}}}Y_{\bar{S}}\|_{F} + \frac{1}{\sqrt{n}}\| \mathbf{P}_{Z_{S}}Y_{S}\|_{F}$  corresponds to  $I(Z;X^{\prime})$  which we want to maximize while compressing the term of  $2\tilde{\mathcal{R}}_m(\mathcal{F}) / \sqrt{m}$  which corresponds to  $I(Z;X)$  [23, 63, 66].
Although we can explicitly add regularization on the information to control  $2\tilde{\mathcal{R}}_m(\mathcal{F}) / \sqrt{m}$ , it is possible that  $I(Z;X|X^{\prime})$  and  $2\tilde{\mathcal{R}}_m(\mathcal{F}) / \sqrt{m}$  are implicitly regularized via implicit bias through design choices [29, 68, 30]. Thus, Theorem 3 connects the information-theoretic understanding of VICReg with the probabilistic guarantee on downstream generalization.
In our paper, we proposed novel methods for SSL premised on information maximization. Although our methods demonstrated superior performance on some datasets, computational constraints precluded us from testing them on larger datasets. Furthermore, our study hinges on certain assumptions that, despite rigorous validation efforts, may not hold universally. While we strive for meticulous testing and validation, it's crucial to note that some assumptions might not be applicable in all scenarios or conditions. These limitations should be taken into account when interpreting our study's results.
We analyzed the Variance-Invariance-Covariance Regularization for self-supervised learning through an information-theoretic lens. By transferring the stochasticity required for an information-theoretic analysis to the input distribution, we showed how the VICReg objective can be derived from information-theoretic principles, used this perspective to highlight assumptions implicit in the VICReg objective, derived a VICReg generalization bound for downstream tasks, and related it to information maximization.
Building on these findings, we introduced a new VICReg-inspired SSL objective. Our probabilistic guarantee suggests that VICReg can be further improved for the settings of partial label information by aligning the covariance matrix with the partially observable label matrix, which opens up several avenues for future work, including the design of improved estimators for information-theoretic quantities and investigations into the suitability of different SSL methods for specific data characteristics.
[73] Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stephane Deny. Barlow twins: Self-supervised learning via redundancy reduction. In International Conference on Machine Learning, pages 12310-12320. PMLR, 2021.
[74] Zhanghao Zhouyin and Ding Liu. Understanding neural networks with logarithm determinant entropy estimator. arXiv preprint arXiv:2105.03705, 2021.
[75] Roland S Zimmermann, Yash Sharma, Steffen Schneider, Matthias Bethge, and Wieland Brendel. Contrastive learning inverts the data generating process. In International Conference on Machine Learning, pages 12979-12990. PMLR, 2021.
This appendix is organized as follows:
- In Appendix A, we provide a detailed derivation of the lower bound of Equation (8).
- In Appendix B, we provide full proof of our Theorem 1 on the network's representation distribution.
- In Appendix C, we provide additional empirical validations. Specifically, we empirically check if the optimal solution to the information maximization problem in Section 4.1 is the diagonal matrix.
- In Appendix D, we show the collapse phenomenon under Gaussian Mixture Model (GMM) using Expectation Maximization (EM) and demonstrate how it is related to SSL and how we can prevent it.
- Appendix E provides additional details on the SimCLR method.
- Appendix F provides a detailed review of entropy estimators, their implications, assumptions, and limitations.
- In Appendix G, we provide proofs for known lemmas that we are using throughout our paper.
- In Appendix H, we provide detailed information on the hyperparameters, datasets, and architectures used in our experiments in Section 5.2.
- In Appendix I, we provide full proof of our generalization bound for downstream tasks from Section 6.
- In Appendix J, we provide full proof of the theorems for Section 4.1 on the connection between information optimization and the VICReg objective.
- Appendix K provides experimental details on experiments conducted in Section 5.1 for entropy comparison between different SSL methods.
- In Appendix L, we provide detailed information on the reproducibility of our study.
- In Appendix O, we discuss the broader impact of our work. This section explores the implications, significance, and potential applications of our findings beyond the scope of the immediate study.
Appendix A Lower bounds on  $\mathbb{E}_{x'}[\log q(z|x')]$
In this section of the supplementary material, we present the full derivation of the lower bound on  $\mathbb{E}_{x'}[\log q(z|x')]$ . Because  $Z'|X'$  is a Gaussian, we can write it as  $Z' = \mu(x') + L(x')\epsilon$  where  $\epsilon \sim \mathcal{N}(0,1)$  and  $L(x')^T L(x') = \Sigma(x')$ . Now, setting  $\Sigma_r = I$ , will give us:
$$
\begin{array}{l} \mathbb {E} _ {x ^ {\prime}} \left[ \log q \left(z \mid x ^ {\prime}\right) \right] (13) \\ \geq \mathbb {E} _ {z ^ {\prime} \mid x ^ {\prime}} [ \log q (z \mid z ^ {\prime}) ] (13) \\ = \mathbb {E} _ {z ^ {\prime} | x ^ {\prime}} \left[ \frac {d}{2} \log 2 \pi - \frac {1}{2} (z - z ^ {\prime}) ^ {T} (I)) ^ {- 1} (z - z ^ {\prime}) \right] (14) \\ = \frac {d}{2} \log 2 \pi - \frac {1}{2} \mathbb {E} _ {z ^ {\prime} \mid x ^ {\prime}}, \left[ \left(z - z ^ {\prime}\right) ^ {2} \right] (15) \\ = \frac {d}{2} \log 2 \pi - \frac {1}{2} \mathbb {E} _ {\epsilon} \left[ \left(z - \mu \left(x ^ {\prime}\right) - L \left(x ^ {\prime}\right) \epsilon\right) ^ {2} \right] (16) \\ = \frac {d}{2} \log 2 \pi - \frac {1}{2} \mathbb {E} _ {\epsilon} \left[ (z - \mu \left(x ^ {\prime}\right)) ^ {2} - 2 (z - \mu \left(x ^ {\prime}\right) * L \left(x ^ {\prime}\right) \epsilon) + \left(\left(L \left(x ^ {\prime}\right) \epsilon\right) ^ {T} \left(L \left(x ^ {\prime}\right) \epsilon\right)\right) \right] (17) \\ = \frac {d}{2} \log 2 \pi - \frac {1}{2} \mathbb {E} _ {\epsilon} \left[ (z - \mu \left(x ^ {\prime}\right)) ^ {2} \right] + (z - \mu \left(x ^ {\prime}\right) L \left(x ^ {\prime}\right)) \mathbb {E} _ {\epsilon} [ \epsilon ] - \frac {1}{2} \mathbb {E} _ {\epsilon} \left[ \epsilon^ {T} L \left(x ^ {\prime}\right) ^ {T} L \left(x ^ {\prime}\right) \epsilon \right] (18) \\ = \frac {d}{2} \log 2 \pi - \frac {1}{2} (z - \mu \left(x ^ {\prime}\right)) ^ {2} - \frac {1}{2} T r \log \Sigma \left(x ^ {\prime}\right) (19) \\ \end{array}
$$
where  $\mathbb{E}_{x'}[\log q(z|x')] = \mathbb{E}_{x'}[\log \mathbb{E}_{z'|x'}[q(z|z')]]\geq \mathbb{E}_{z'}[\log q(z|z')]]$  by Jensen's inequality,  $\mathbb{E}_{\epsilon}[\epsilon ] = 0$  and  $\mathbb{E}_{\epsilon}\left[\epsilon \left(L(x^{\prime})^{T}L(x^{\prime})\right.\epsilon \right] = Tr\log \Sigma (x^{\prime})$  by the Hutchinson's estimator.
$$
\begin{array}{l} \mathbb {E} _ {z \mid x} \left[ \mathbb {E} _ {z ^ {\prime} \mid x ^ {\prime}} \left[ \log q \left(z \mid z ^ {\prime}\right) \right] \right] = \mathbb {E} _ {z \mid x} \left[ \frac {d}{2} \log 2 \pi - \frac {1}{2} (z - \mu \left(x ^ {\prime}\right)) ^ {2} - \frac {1}{2} T r \log \Sigma \left(x ^ {\prime}\right) \right] (20) \\ = \frac {d}{2} \log 2 \pi - \frac {1}{2} \mathbb {E} _ {z | x} \left[ \left(z - \mu \left(x ^ {\prime}\right)\right) ^ {2} \right] - \frac {1}{2} T r \log \Sigma \left(x ^ {\prime}\right) (21) \\ = \frac {d}{2} \log 2 \pi - \frac {1}{2} \mathbb {E} _ {\epsilon} \left[ \left(\mu (x) + L (x) \epsilon - \mu \left(x ^ {\prime}\right)\right) ^ {2} \right] - \frac {1}{2} T r \log \Sigma \left(x ^ {\prime}\right) (22) \\ = \frac {d}{2} \log 2 \pi - \frac {1}{2} \mathbb {E} _ {\epsilon} \left[ \left(\mu (x) - \mu \left(x ^ {\prime}\right)\right) ^ {2} \right] + \mathbb {E} _ {\epsilon} \left[ \left(\mu (x) - \mu \left(x ^ {\prime}\right)\right) L (x) \epsilon \right] \\ - \frac {1}{2} \mathbb {E} _ {\epsilon} \left[ \epsilon^ {T} L (x) ^ {T} L (x) \epsilon \right] - \frac {1}{2} T r \log \Sigma \left(x ^ {\prime}\right) (23) \\ = \frac {d}{2} \log 2 \pi - \frac {1}{2} (\mu (x) - \mu \left(x ^ {\prime}\right)) ^ {2} - \frac {1}{2} T r \log \Sigma (x) - \frac {1}{2} T r \log \Sigma \left(x ^ {\prime}\right) (24) \\ = \frac {d}{2} \log 2 \pi - \frac {1}{2} (\mu (x) - \mu \left(x ^ {\prime}\right)) ^ {2} - \frac {1}{2} \log \left(| \Sigma (x) | \cdot | \Sigma \left(x ^ {\prime}\right) |\right) (25) \\ \end{array}
$$
Appendix B Data Distribution after Deep Network Transformation
Theorem 4. Given the setting of Equation (4), the unconditional DNN output density denoted as  $Z$  approximates (given the truncation of the Gaussian on its effective support that is included within a single region  $\omega$  of the DN's input space partition) a mixture of the affinely transformed distributions  $\pmb{x} | \pmb{x}_{n(\pmb{x})}^{*}$  e.g. for the Gaussian case
$$
Z \sim \sum_ {n = 1} ^ {N} \mathcal {N} \left(\boldsymbol {A} _ {\omega \left(\boldsymbol {x} _ {n} ^ {*}\right)} \boldsymbol {x} _ {n} ^ {*} + \boldsymbol {b} _ {\omega \left(\boldsymbol {x} _ {n} ^ {*}\right)}, \boldsymbol {A} _ {\omega \left(\boldsymbol {x} _ {n} ^ {*}\right)} ^ {T} \Sigma_ {\boldsymbol {x} _ {n} ^ {*}} \boldsymbol {A} _ {\omega \left(\boldsymbol {x} _ {n} ^ {*}\right)}\right) ^ {T = n},
$$