SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ravid Shwartz-Ziv* New York University
Randall Balestriero Meta AI, FAIR
Kenji Kawaguchi  
National University of Singapore
Tim G. J. Rudner  
New York University
Yann LeCun  
New York University & Meta AI, FAIR
Variance-Invariance-Covariance Regularization (VICReg) is a self-supervised learning (SSL) method that has shown promising results on a variety of tasks. However, the fundamental mechanisms underlying VICReg remain unexplored. In this paper, we present an information-theoretic perspective on the VICReg objective. We begin by deriving information-theoretic quantities for deterministic networks as an alternative to unrealistic stochastic network assumptions. We then relate the optimization of the VICReg objective to mutual information optimization, highlighting underlying assumptions and facilitating a constructive comparison with other SSL algorithms and derive a generalization bound for VICReg, revealing its inherent advantages for downstream tasks. Building on these results, we introduce a family of SSL methods derived from information-theoretic principles that outperform existing SSL techniques.
Self-supervised learning (SSL) is a promising approach to extracting meaningful representations by optimizing a surrogate objective between inputs and self-generated signals. For example, Variance-Invariance-Covariance Regularization (VICReg) [7], a widely-used SSL algorithm employing a de-correlation mechanism, circumvents learning trivial solutions by applying variance and covariance regularization.
Once the surrogate objective is optimized, the pre-trained model can be used as a feature extractor for a variety of downstream supervised tasks such as image classification, object detection, instance segmentation, or pose estimation [15, 16, 48, 67]. Despite the promising results demonstrated by SSL methods, the theoretical underpinnings explaining their efficacy continue to be the subject of investigation [5, 42].
Information theory has proved a useful tool for improving our understanding of deep neural networks (DNNs), having a significant impact on both applications in representation learning [3] and theoretical explorations [60, 72]. However, applications of information-theoretic principles to SSL have made unrealistic assumptions, making many existing information-theoretic approaches to SSL of limited use. One such assumption is to assume that the DNN to be optimized is stochastic—an assumption that is violated for the vast majority DNNs used in practice. For a comprehensive review on this topic, refer to the work by Shwartz-Ziv and LeCun [62].
*Correspondence to: ravid.shwartz.ziv@nyu.edu.
In this paper, we examine Variance-Invariance-Covariance Regularization (VICReg), an SSL method developed for deterministic DNNs, from an information-theoretic perspective. We propose an approach that addresses the challenge of mutual information estimation in deterministic networks by transitioning the randomness from the networks to the input data—a more plausible assumption. This shift allows us to apply an information-theoretic analysis to deterministic networks. To establish a connection between the VICReg objective and information maximization, we identify and empirically validate the necessary assumptions. Building on this analysis, we describe differences between different SSL algorithms from an information-theoretic perspective and propose a new family of plug-in methods for SSL. This new family of methods leverages existing information estimators and achieves state-of-the-art predictive performance across several benchmarking tasks. Finally, we derive a generalization bound that links information optimization and the VICReg objective to downstream task performance, underscoring the advantages of VICReg.
Our key contributions are summarized as follows:
1. We introduce a novel approach for studying deterministic deep neural networks from an information-theoretic perspective by shifting the stochasticity from the networks to the inputs using the Data Distribution Hypothesis (Section 3)
2. We establish a connection between the VICReg objective and information-theoretic optimization, using this relationship to elucidate the underlying assumptions of the objective and compare it to other SSL methods (Section 4).
3. We propose a family of information-theoretic SSL methods, grounded in our analysis, that achieve state-of-the-art performance (Section 5).
4. We derive a generalization bound that directly links VICReg to downstream task generalization, further emphasizing its practical advantages over other SSL methods (Section 6).
2 Background & Preliminaries
We first introduce the necessary technical background for our analysis.
2.1 Continuous Piecewise Affine (CPA) Mappings.
A rich class of functions emerges from piecewise polynomials: spline operators. In short, given a partition  $\Omega$  of a domain  $\mathbb{R}^D$ , a spline of order  $k$  is a mapping defined by a polynomial of order  $k$  on each region  $\omega \in \Omega$  with continuity constraints on the entire domain for the derivatives of order  $0,\ldots ,k - 1$ . As we will focus on affine splines ( $k = 1$ ), we only define this case for clarity. A  $K$ -dimensional affine spline  $f$  produces its output via  $f(z) = \sum_{\omega \in \Omega}(A_{\omega}z + b_{\omega})\mathbb{1}_{\{z\in \omega \}}$ , with input  $z\in \mathbb{R}^{D}$  and  $A_{\omega}\in \mathbb{R}^{K\times D}$ ,  $b_{\omega}\in \mathbb{R}^{K}$ ,  $\forall \omega \in \Omega$  the per-region slope and offset parameters respectively, with the key constraint that the entire mapping is continuous over the domain  $f\in \mathcal{C}^0 (\mathbb{R}^D)$ .
Deep Neural Networks as CPA Mappings. A deep neural network (DNN) is a (non-linear) operator  $f_{\Theta}$  with parameters  $\Theta$  that map a input  $\pmb{x} \in \mathbb{R}^{D}$  to a prediction  $\pmb{y} \in \mathbb{R}^{K}$ . The precise definitions of DNN operators can be found in Goodfellow et al. [27]. To avoid cluttering notation, we will omit  $\Theta$  unless needed for clarity. For our analysis, we only assume that the non-linearities in the DNN are CPA mappings—as is the case with (leaky-) ReLU, absolute value, and max-pooling operators. The entire input-output mapping then becomes a CPA spline with an implicit partition  $\Omega$ , the function of the weights and architecture of the network [51, 6]. For smooth nonlinearities, our results hold using a first-order Taylor approximation argument.
2.2 Self-Supervised Learning.
SSL is a set of techniques that learn representation functions from unlabeled data, which can then be adapted to various downstream tasks. While supervised learning relies on labeled data, SSL formulates a proxy objective using self-generated signals. The challenge in SSL is to learn useful representations without labels. It aims to avoid trivial solutions where the model maps all inputs to a constant output [11, 36]. To address this, SSL utilizes several strategies. Contrastive methods like SimCLR and its InfoNCE criterion learn representations by distinguishing positive and negative examples [16, 54]. In contrast, non-contrastive methods apply regularization techniques to prevent the collapse [14, 17, 28].
Variance-Invariance-Covariance Regularization (VICReg). A widely used SSL method for training joint embedding architectures [7]. Its loss objective is composed of three terms: the invariance loss, the variance loss, and the covariance loss:
- Invariance loss: The invariance loss is given by the mean-squared Euclidean distance between pairs of embedding and ensures consistency between the representation of the original and augmented inputs.
- Regularization: The regularization term consists of two loss terms: the variance loss—a hinge loss to maintain the standard deviation (over a batch) of each variable of the embedding—and the covariance loss, which penalizes off-diagonal coefficients of the covariance matrix of the embeddings to foster decorrelation among features.
VICReg generates two batches of embeddings,  $\mathbf{Z} = [f(\mathbf{x}_1),\dots ,f(\mathbf{x}_B)]$  and  $Z^{\prime} = [f(\pmb{x}_{1}^{\prime}),\dots ,f(\pmb{x}_{B}^{\prime})]$ , each of size  $(B\times K)$ . Here,  $\pmb {x}_i$  and  $\pmb{x}_i^\prime$  are two distinct random augmentations of a sample  $I_{i}$ . The covariance matrix  $C\in \mathbb{R}^{K\times K}$  is obtained from  $[\mathbf{Z},\mathbf{Z}^{\prime}]$ . The VICReg loss can thus be expressed as follows:
$$
\mathcal {L} = \underbrace {\frac {1}{K} \sum_ {k = 1} ^ {K} \left(\alpha \max  \left(0 , \gamma - \sqrt {\boldsymbol {C} _ {k , k} + \epsilon}\right) + \beta \sum_ {k ^ {\prime} \neq k} \left(\boldsymbol {C} _ {k , k ^ {\prime}}\right) ^ {2}\right)} _ {\text {R e g u l a r i z a t i o n}} + \underbrace {\eta \| \boldsymbol {Z} - \boldsymbol {Z} ^ {\prime} \| _ {F} ^ {2} / N} _ {\text {I n v a r i a n c e}}. \tag {1}
$$
2.3 Deep Neural Networks and Information Theory
Recently, information-theoretic methods have played an essential role in advancing deep learning by developing and applying information-theoretic estimators and learning principles to DNN training [3, 9, 34, 56, 64, 65, 69, 72]. However, information-theoretic objectives for deterministic DNNs often face a common issue: the mutual information between the input and the DNN representation is infinite. This leads to ill-posed optimization problems.
Several strategies have been suggested to address this challenge. One involves using stochastic DNNs with variational bounds, where the output of the deterministic network is used as the parameters of the conditional distribution [43, 61]. Another approach, as suggested by Dubois et al. [22], assumes that the randomness of data augmentation among the two views is the primary source of stochasticity in the network. However, these methods assume that randomness comes from the DNN, contrary to common practice. Other research has presumed a random input but has made no assumptions about the network's representation distribution properties. Instead, it relies on general lower bounds to analyze the objective [71, 75].
3 Self-Supervised Learning in DNNs: An Information-Theoretic Perspective
To analyze information within deterministic networks, we first need to establish an information-theoretic perspective on SSL (Section 3.1). Subsequently, we utilize the Data Distribution Hypothesis (Section 3.2) to demonstrate its applicability to deterministic SSL networks.
3.1 Self-Supervised Learning from an Information-Theoretic Viewpoint
Our discussion begins with the MultiView InfoMax principle, which aims to maximize the mutual information  $I(Z; \overline{X}')$  between a view  $X'$  and the second representation  $Z$ . As demonstrated in Federici et al. [24], we can optimize this information by employing the following lower bound:
$$
I (Z, X ^ {\prime}) = H (Z) - H (Z \mid X ^ {\prime}) \geq H (Z) + \mathbb {E} _ {x ^ {\prime}} [ \log q (z \mid x ^ {\prime}) ]. \tag {2}
$$
Here,  $H(Z)$  represents the entropy of  $Z$ . In supervised learning, the labels  $Y$  remain fixed, making the entropy term  $H(Y)$  a constant. Consequently, the optimization is solely focused on the log-loss,  $\mathbb{E}_{x'}[\log q(z|x')]$ , which could be either cross-entropy or square loss.
However, for joint embedding networks, a degenerate solution can emerge, where all outputs "collapse" into an undesired value [16]. Upon examining Equation (2), we observe that the entropies are not fixed and can be optimized. As a result, minimizing the log loss alone can lead the representations to collapse into a trivial solution and must be regularized.
3.2 Understanding the Data Distribution Hypothesis
Previously, we mentioned that a naive analysis might suggest that the information in deterministic DNNs is infinite. To address this point, we investigate whether assuming a dataset is a mixture of Gaussians with non-overlapping support can provide a manageable distribution over the neural network outputs. This assumption is less restrictive compared to assuming that the neural network itself is stochastic, as it concerns the generative process of the data, not the model and training process. For a detailed discussion about the limitations of assuming stochastic networks and a comparison between stochastic networks vs stochastic input, see Appendix N. In Section 4.2, we verify that this assumption holds for real-world datasets.
The so-called manifold hypothesis allows us to treat any point as a Gaussian random variable with a low-rank covariance matrix aligned with the data's manifold tangent space [25], which enables us to examine the conditioning of a latent representation with respect to the mean of the observation, i.e.,  $X|\pmb{x}^* \sim \mathcal{N}(\pmb{x};\pmb{x}^*,\pmb{\Sigma}_{\pmb{x}^*})$ . Here, the eigenvectors of  $\Sigma_{\pmb{x}^*}$  align with the tangent space of the data manifold at  $\pmb{x}^*$ , which varies with the position of  $\pmb{x}^*$  in space. In this setting, a dataset is considered a collection of distinct points  $\pmb{x}_n^*$ ,  $n = 1,\dots,N$ , and the full data distribution is expressed as a sum of Gaussian densities with low-rank covariance, defined as:
$$
X \sim \sum_ {n = 1} ^ {N} \mathcal {N} \left(\boldsymbol {x} _ {n} ^ {*}, \Sigma_ {\boldsymbol {x} _ {n} ^ {*}}\right) ^ {\mathbb {I} _ {\{T = n \}}} \quad \text {w i t h} \quad T \sim \operatorname {C a t} (N). \tag {3}
$$
Here,  $T$  is a uniform Categorical random variable. For simplicity, we assume that the effective support of  $\mathcal{N}(\boldsymbol{x}_i^*,\Sigma_{\boldsymbol{x}_i^*})$  do not overlap (for empirical validation of this assumption see Section 4.2). The effective support is defined as  $\{x\in \mathbb{R}^{D}:p(x) > \epsilon \}$ . We can then approximate the density function as follows:
$$
p (\boldsymbol {x}) \approx \mathcal {N} \left(\boldsymbol {x}; \boldsymbol {x} _ {n (\boldsymbol {x})} ^ {*}, \Sigma_ {\boldsymbol {x} _ {n (\boldsymbol {x})} ^ {*}}\right) / N, \tag {4}
$$
where  $\mathcal{N}(\pmb{x};.,.)$  is the Gaussian density at  $\pmb{x}$  and  $n(\pmb{x}) = \arg \min_{n}(\pmb{x} - \pmb{x}_n^*)^T\Sigma_{\pmb{x}_n^*}(\pmb{x} - \pmb{x}_n^*)$ .
3.3 Data Distribution Under the Deep Neural Network Transformation
Let us consider an affine spline operator  $f$ , as illustrated in Section 2.1, which maps a space of dimension  $D$  to a space of dimension  $K$ , where  $K \geq D$ . The image or the span of this mapping is expressed as follows:
$$
\operatorname {I m} (f) \triangleq \left\{f (\boldsymbol {x}): \boldsymbol {x} \in \mathbb {R} ^ {D} \right\} = \bigcup_ {\omega \in \Omega} \operatorname {A f f} (\omega ; \boldsymbol {A} _ {\omega}, \boldsymbol {b} _ {\omega}) \tag {5}
$$
In this equation,  $\mathrm{Aff}(\omega; A_{\omega}, b_{\omega}) = \{A_{\omega}x + b_{\omega}: x \in \omega\}$  denotes the affine transformation of region  $\omega$  by the per-region parameters  $A_{\omega}, b_{\omega}$ .  $\Omega$  denotes the partition of the input space where  $\pmb{x}$  resides. To practically compute the per-region affine mapping, we set  $A_{\omega}$  to the Jacobian matrix of the network at the corresponding input  $x$ , and  $b$  to be defined as  $f(x) - A_{\omega}x$ . Therefore, the DNN mapping composed of affine transformations on each input space partition region  $\omega \in \Omega$  based on the coordinate change induced by  $A_{\omega}$  and the shift induced by  $b_{\omega}$ .
When the input space is associated with a density distribution, this density is transformed by the mapping  $f$ . Calculating the density of  $f(X)$  is generally intractable. However, under the disjoint support assumption from Section 3.2, we can arbitrarily increase the density's representation power by raising the number of prototypes  $N$ . As a result, each Gaussian's support is contained within the region  $\omega$  where its means lie, leading to the following theorem:
Theorem 1. Given the setting of Equation (4), the unconditional DNN output density,  $Z$ , can be approximated as a mixture of the affinely transformed distributions  $\pmb{x} | \pmb{x}_{n(\pmb{x})}^*$ :
$$
Z \sim \sum_ {n = 1} ^ {N} \mathcal {N} \left(\boldsymbol {A} _ {\omega (\boldsymbol {x} _ {n} ^ {*})} \boldsymbol {x} _ {n} ^ {*} + \boldsymbol {b} _ {\omega (\boldsymbol {x} _ {n} ^ {*})}, \boldsymbol {A} _ {\omega (\boldsymbol {x} _ {n} ^ {*})} ^ {T} \Sigma_ {\boldsymbol {x} _ {n} ^ {*}} \boldsymbol {A} _ {\omega (\boldsymbol {x} _ {n} ^ {*})}\right) ^ {1 _ {\{T = n \}}},
$$
where  $\omega (\pmb{x}_n^*) = \omega \in \Omega \iff \pmb{x}_n^*\in \omega$  is the partition region in which the prototype  $\pmb{x}_n^*$  lives in.
In other words, Theorem 1 implies that when the input noise is small, we can simplify the conditional output density to a single Gaussian:  $(Z'|X' = x_n) \sim \mathcal{N}(\mu(x_n), \Sigma(x_n))$ , where  $\mu(x_n) = A_{\omega(\boldsymbol{x}_n)}\boldsymbol{x}_n + b_{\omega(\boldsymbol{x}_n)}$  and  $\Sigma(x_n) = A_{\omega(\boldsymbol{x}_n)}^T\Sigma_{\boldsymbol{x}_n}A_{\omega(\boldsymbol{x}_n)}$ .
Figure 1: Left: The network output for SSL training is more Gaussian for small input noise. The  $p$ -value of the normality test for different SSL models trained on ImageNet for different input noise levels. The dashed line represents the point at which the null hypothesis (Gaussian distribution) can be rejected with  $99\%$  confidence. Right: The Gaussians around each point are not overlapping. The plots show the  $l2$  distances between raw images for different datasets. As can be seen, the distances are largest for more complex real-world datasets.
4 Information Optimization and the VICReg Optimization Objective
Building on our earlier discussion, we used the Data Distribution Hypothesis to model the conditional output in deterministic networks as a Gaussian mixture. This allowed us to frame the SSL training objective as maximizing the mutual information,  $I(Z;X^{\prime})$  and  $I(Z^{\prime};X)$ .
However, in general, this mutual information is intractable. Therefore, we will use our derivation for the network's representation to obtain a tractable variational approximation using the expected loss, which we can optimize.
The computation of expected loss requires us to marginalize the stochasticity in the output. We can employ maximum likelihood estimation with a Gaussian observation model. For computing the expected loss over  $x'$  samples, we must marginalize the stochasticity in  $Z'$ . This procedure implies that the conditional decoder adheres to a Gaussian distribution:  $(Z|X' = x_n) \sim \mathcal{N}(\mu(x_n), I + \Sigma(x_n))$ .
However, calculating the expected log loss over samples of  $Z$  is challenging. We thus focus on a lower bound – the expected log loss over  $Z'$  samples. Utilizing Jensen's inequality, we derive the following lower bound:
$$
\mathbb {E} _ {x ^ {\prime}} \left[ \log q \left(z \mid x ^ {\prime}\right) \right] \geq \mathbb {E} _ {z ^ {\prime} \mid x ^ {\prime}} \left[ \log q \left(z \mid z ^ {\prime}\right) \right] = \frac {1}{2} \left(d \log 2 \pi - \left(z - \mu \left(x ^ {\prime}\right)\right) ^ {2} - \operatorname {T r} \log \Sigma \left(x ^ {\prime}\right)\right). \tag {6}
$$
Taking the expectation over  $Z$ , we get
$$
\mathbb {E} _ {z \mid x} \left[ \mathbb {E} _ {z ^ {\prime} \mid x ^ {\prime}} \left[ \log q \left(z \mid z ^ {\prime}\right) \right] \right] = \frac {1}{2} \left(d \log 2 \pi - \left(\mu (x) - \mu \left(x ^ {\prime}\right)\right) ^ {2} - \log \left(| \Sigma (x) | \cdot | \Sigma \left(x ^ {\prime}\right) |\right)\right). \tag {7}
$$
Combining all of the above, we obtain
$$
I (Z; X ^ {\prime}) \geq H (Z) + \frac {d}{2} \log 2 \pi - \frac {1}{2} \mathbb {E} _ {x, x ^ {\prime}} [ (\mu (x) - \mu (x ^ {\prime})) ^ {2} + \log (| \Sigma (x) | \cdot | \Sigma (x ^ {\prime}) |) ]. \tag {8}
$$
The full derivations are presented in Appendix A. To optimize this objective in practice, we approximate  $p(x,x^{\prime})$  using the empirical data distribution
$$
L \left(x _ {1} \dots x _ {N}, x _ {1} ^ {\prime} \dots x _ {N} ^ {\prime}\right) \approx \frac {1}{N} \sum_ {i = 1} ^ {N} \underbrace {H (Z) - \log \left(\left| \Sigma \left(x _ {i}\right) \right| \cdot \left| \Sigma \left(x _ {i} ^ {\prime}\right) \right|\right)} _ {\text {R e g u l a r i z e r}} - \underbrace {\frac {1}{2} \left(\mu \left(x _ {i}\right) - \mu \left(x _ {i} ^ {\prime}\right)\right) ^ {2}} _ {\text {I n v a r i a n c e}}. \tag {9}
$$
4.1 Variance-Invariance-Covariance Regularization: An Information-Theoretic Perspective
Next, we connect the VICReg to our information-theoretic-based objective. The "invariance term" in Equation (9), which pushes augmentations from the same image closer together, is the same term used in the VICReg objective. However, the computation of the regularization term poses a significant challenge. Entropy estimation is a well-established problem within information theory, with Gaussian mixture densities often used for representation. Yet, the differential entropy of Gaussian mixtures lacks a closed-form solution [55].
A straightforward method for approximating entropy involves capturing the distribution's first two moments, which provides an upper bound on the entropy. However, minimizing an upper bound doesn't necessarily optimize the original objective. Despite reported success from minimizing an upper bound [46, 53], this approach may induce instability during the training process.
Let  $\Sigma_Z$  denote the covariance matrix of  $Z$ . We utilize the first two moments to approximate the entropy we aim to maximize. Because the invariance term appears in the same form as the original VICReg objective, we will look only at the regularizer. Consequently, we get the approximation
$$
\bar {L} \left(x _ {1} \dots x _ {N}, x _ {1} ^ {\prime} \dots x _ {N} ^ {\prime}\right) \approx \sum_ {i = 1} ^ {N} \log \frac {\left| \Sigma_ {Z} \left(x _ {1} \dots x _ {N}\right) \right|}{\left| \Sigma \left(x _ {i}\right) \right| \cdot \left| \Sigma \left(x _ {i} ^ {\prime}\right) \right|}. \tag {10}
$$
Theorem 2. Assuming that the eigenvalues of  $\Sigma (x_{i})$  and  $\Sigma (x_i^{\prime})$ , along with the differences between the Gaussian means  $\mu (x_{i})$  and  $\mu (x_i^{\prime})$ , are bounded, the solution to the maximization problem
$$
\max  _ {\Sigma_ {Z}} \left\{\sum_ {i = 1} ^ {N} \log \frac {\left| \Sigma_ {Z} \left(x _ {1} \dots x _ {N}\right) \right|}{\left| \Sigma \left(x _ {i}\right) \right| \cdot \left| \Sigma \left(x _ {i} ^ {\prime}\right) \right|} \right\} \tag {11}
$$
involves setting  $\Sigma_Z$  to a diagonal matrix.
Proof. See Appendix J.
According to Theorem 2, we can maximize Equation (10) by diagonalizing the covariance matrix and increasing its diagonal elements. This goal can be achieved by minimizing the off-diagonal elements of  $\Sigma_Z$ -the covariance criterion of VICReg--and by maximizing the sum of the log of its diagonal elements. While this approach is straightforward and efficient, it does have a drawback: the diagonal values could tend towards zero, potentially causing instability during logarithm computations. A solution to this issue is to use an upper bound and directly compute the sum of the diagonal elements, resulting in the variance term of VICReg. This establishes the link between our information-theoretic objective and the three key components of VICReg.
4.2 Empirical Validation of Assumptions About Data Distributions
Validating our theory, we tested if the conditional output density  $P(Z|X)$  becomes a Gaussian as input noise lessens. We used a ResNet-50 model trained with SimCLR or VICReg objectives on CIFAR-10, CIFAR-100, and ImageNet datasets. We sampled 512 Gaussian samples for each image from the test dataset, examining whether each sample remained Gaussian in the DNN's penultimate layer. We applied the D'Agostino and Pearson's test to ascertain the validity of this assumption [19].
Figure 1 (left) displays the  $p$ -value as a function of the normalized std. For low noise levels, we reject that the network's conditional output density is non-Gaussian with an  $85\%$  probability when using VICReg. However, the network output deviates from Gaussian as the input noise increases.