What Do We Maximize in Self-Supervised Learning?
Abstract
This paper analyzes self-supervised learning (SSL) methods, VICReg in particular, to provide for the first time an information-theoretical understanding of its construction.
As a first step, we demonstrate how information-theoretic quantities can be obtained for a deterministic network, offering a possible alternative to prior work that relies on stochastic models.
This enables us to demonstrate how VICReg can be (re)discovered from first principles and its assumptions about data distribution.
Furthermore, we empirically demonstrate the validity of our assumptions, confirming our novel understanding of VICReg.
Finally, we believe that the derivation and insights we obtain can be generalized to many other SSL methods, opening new avenues for theoretical and practical understanding of SSL and transfer learning.
1. Introduction
Self-Supervised Learning (SSL) algorithms (Bromley et al., 1993) learn representations using a proxy objective (i e.,
SSL objective) between inputs and self-defined signals.
The results indicate that the learned representations can generalize well to a wide range of downstream tasks (Chen et al., 2020; Misra & Maaten, 2020), even when the SSL objective does not use downstream supervision during training.
In
SimCLR (Chen et al., 2020), for example, a contrastive loss is defined between images with different augmentations (i e., one as input and the other as a self-supervised signal).
Then, we take our pre-learned model as a feature extractor and adopt the features to various applications, including image classification, object detection, instance segmentation, and pose estimation (Caron et al., 2021).
However, despite the success in practice, only a few works (Arora et al., 2019; Lee et al., 2021a) provide theoretical insights into the learning efficacy of SSL.
In recent years, information theory methods have played a key role in several notable deep learning achievements, from practical applications in representation learning as the variational information bottleneck (Alemi et al., 2016), to theoretical investigations (e.g., the generalization bound induced by mutual information (Xu & Raginsky, 2017; Steinke &
Zakynthinou, 2020; Shwartz-Ziv, 2022).
Moreover, different deep learning problems have been successfully approached by developing and applying novel estimators and learning principles derived from information-theoretic quantities, such as mutual information estimation.
Many works have attempted to analyze SSL from an information theory perspective.
An example is the use of the mutual information neural estimator (MINE) (Belghazi et al., 2018) in representation learning (Hjelm et al., 2018) in conjunction with the renowned information maximization (InfoMax) principle (Linsker, 1988).
However, looking at these works may be confusing.
Numerous objective functions are presented, some contradicting each other, as well as many implicit assumptions.
Moreover, these works rely on a crucial assumption: a stochastic (often Gaussian) DN mapping, which is rarely the case in most cases these days.
This paper presents a unified framework for SSL methods from an information theory perspective which can be applied to deterministic DN training.
We summarize our contributions into two points: (i)
First, we develop a method to study deterministic DNs from an information theory perspective by shifting the stochasticity to the DN input, which is a much more faithful assumption currently valid for most methods.
(ii)
Second, based on this formulation, we analyze how current SSL methods that use deterministic networks optimize information-theoretic quantities.
2. Background
Continuous Piecewise Affine (CPA) Mappings.
A rich class of functions emerges from piecewise polynomials: spline operators.
In short, given a partition, Ω of a domain R D , a spline of order k is a mapping defined by a polynomial of order k on each region ω ∈ Ω with continuity constraints on the entire domain for the derivatives of order 0 ,..., k − 1 .
As we will focus on affine splines ( k = 1 ), we define this case only for concreteness.
An K -dimensional affine spline f produces its output via
[Equation]
with input z ∈ R D and A ω ∈ R K × D , b ω ∈ R K , ∀ ω ∈ Ω the per-region slope and offset parameters respectively, with the key constraint that the entire mapping is continuous over the domain f ∈ C
( R D ) .
Spline operators and especially affine spline operators have been widely used in function approximation theory (Cheney & Light, 2009), optimal control (Egerstedt & Martin, 2009), statistics (Fantuzzi et al., 2002), and related fields.
Deep Networks.
A deep network (DN) is a (non-linear) operator f Θ with parameters Θ that map a input x ∈ R D to a prediction y ∈ R K .
The precise definitions of DNs operators can be found in (Goodfellow et al., 2016).
We will omit the Θ notation for clarity unless needed.
The only assumption we require for our study is that the nonlinearities present in the DN are CPA, as is the case with (leaky-)
ReLU, absolute value, and max-pooling.
In that case, the entire input-output mapping becomes a CPA spline with an implicit partition Ω , the function of the weights and architecture of the network (Montufar et al., 2014; Balestriero & Baraniuk, 2018).
For smooth nonlinearities, our results hold from a first-order Taylor approximation argument.
Self-Supervised Learning .
Joint embedding methods learn the DN parameters Θ without supervision and input reconstruction.
Due to this formulation, the difficulty of SSL is to produce a good representation for downstream tasks whose labels are not available during training —while avoiding a trivially simple solution where the model maps all inputs to constant output.
Many methods have been proposed to solve this problem.
Contrastive methods learn representations by contrasting positive and negative examples, e.g., SimCLR (Chen et al., 2020) and its the InfoNCE criterion (Oord et al., 2018).
Other recent work introduced non-contrastive methods that employ different regularization methods to prevent collapsing of the representation.
Several papers used stopgradients and extra predictors to avoid collapse (Chen & He, 2021;
Grill et al., 2020) while (Caron et al., 2020) uses an additional clustering step.
As opposed to contrastive methods, noncontrastive methods do not explicitly rely on negative samples.
Of particular interest to us is the VICReg method (Bardes et al., 2021) that considers two embedding batches
Z =
[ f ( x 1 ) , . .
f ( x N )] and Z ′ =
[ f ( x ′ 1 ) , . . .
, f ( x ′ N )] each of size ( N × K ) .
Denoting by C the ( K × K ) covariance matrix obtained from [ Z , Z ′ ] we obtain the VICReg triplet loss
[Equation]
Our goal will now be to formulate SSL as an informationtheoretic problem from which we can precisely relate VICReg to known methods even with a deterministic network.
Deep Networks and Information-Theory.
Recently, information-theoretic methods have played a key role in several remarkable deep learning achievements (Alemi et al., 2016; Xu & Raginsky, 2017; Steinke & Zakynthinou, 2020;
Shwartz-Ziv & Tishby, 2017).
Moreover, different deep learning problems have been successfully approached by developing and applying information-theoretic estimators and learning principles (Hjelm et al., 2018; Belghazi et al., 2018; Piran et al., 2020; Shwartz-Ziv et al., 2018).
There is, however, a major problem when it comes to analyzing information-theoretic objectives in deterministic deep neural networks: the source of randomness.
The mutual information between the input and the representation in such networks is infinite, resulting in ill-posed optimization problems or piecewise constant, making gradient-based optimization methods ineffective (Amjad & Geiger, 2019).
To solve these problems, researchers have proposed several solutions.
For SSL, stochastic deep networks with variational bounds could be used, where the output of the deterministic network is used as the parameter of the conditional distribution (Lee et al., 2021b; Shwartz-Ziv & Alemi, 2020).
(Dubois et al., 2021) suggested another option, which assumed that the randomness of data augmentation among the two views was the source of stochasticity.
For supervised learning, (Goldfeld et al., 2018) introduced an auxiliary (noisy) DN framework by injecting additive noise into the model and demonstrated that it is a good proxy for the original (deterministic) DN in terms of both performance and representation.
Finally, Achille and Soatto (Achille &
Soatto, 2018) found that minimizing a stochastic network with a regularizer is equivalent to minimizing cross-entropy over deterministic DNs with multiplicative noise.
However, it is important to note that all of these methods assume that the noise comes from the model itself, thus constricting current training methods.
In this work, we explicitly assume that the stochasticity comes from the data itself, which is a less restrictive assumption and does not require changing the current algorithms.
3. Information Maximization of Deep Networks Outputs
This section first sets up notation and assumption on the information-theoretic challenges in self-supervised learning (Section 3.1) and on our assumptions regarding the data distribution (Section 3.2) so that any training sample x can be seen as coming from a single Gaussian distribution as in x ∼ N ( µ x , Σ x ) .
From this we obtain that the output of any deep network f ( x ) corresponds to a mixture of truncated
Gaussian (Section 3.3).
In particular, it can fall back to a single Gaussian under some small noise ( det(Σ) → ϵ ) as- sumptions.
These results will enable information measures to be applied to deterministic DNs.
We then recover the known SSL methods (Bardes et al., 2021) by making different assumptions about the data distribution and estimating their information.
3.1. SSL as an Information-Theoretic Problem
To better grasp the difference between key SSL methods, we first formulate the general SSL goal from an informationtheoretical perspective.
We start with the MultiView InfoMax principle, i.e., maximizing the mutual information between the representations of the two views.
To do so, as shown in (Federici et al., 2020), we need to maximize I ( Z ; X ′ ) and I ( Z ′ ; X ) .
We can do so by a lower bound using
[Equation]
where H ( Z ) is the entropy of Z .
In supervised learning, where we need to maximize I ( Z ; Y ) , the labels ( Y ) are fixed, the entropy term H ( Y ) is constant, and you only need to optimize the log-loss E [log q ( z | x )]
(crossentropy or square loss).
However, in SSL, the entropy
H ( Z ) and H ( Z ′ ) are not constant and can be optimized throughout the learning process.
Therefore, only maximizing E [log q ( z | x ′ )] will cause it to collapse to the trivial solution of making the representations constant (where the entropy goes to zero).
To regularize these entropies, i.e., prevent collapse, different methods utilize different approaches to implicit regularizing information.
To recover them in
Section 4, we must first introduce the notation and results around the data distribution (Section 3.2) and how a DN transforms that distribution (Section 3.3).
3.2. Data Distribution Hypothesis
Our first step is to assess how the output random variables of the network are represented, assuming a distribution on the data itself.
Under a manifold hypothesis, any point can be seen as a Gaussian random variable with a low-rank covariance matrix in the direction of the manifold tangent space of the data.
Therefore, we will consider throughout this study the conditioning of a latent representation with respect to the mean of the observation, i.e., X | x ∗ ∼ N ( x ∗ , Σ x ∗ ) where the eigenvectors of Σ x ∗ are in the same linear subspace than the tangent space of the data manifold at x ∗ , which varies with the position of x ∗ in space.
Hence a dataset is considered to be a collection of { x ∗ n , n = 1 , . .
. , N } and the full data distribution to be a sum of lowrank covariance Gaussian densities as in
[Equation]
with T the uniform Categorical random variable.
To keep things simple and without loss of generality, we consider that the effective support of N ( x ∗ i , Σ x ∗ i ) and N ( x ∗ j , Σ x ∗ j ) do not overlap.
This keeps things general, as it is enough to cover the domain of the data manifold overall, without overlap between different Gaussians.
Hence, in general, we have that.
[Equation]
where N ( x ; .
, . ) is the Gaussian density at x and with n ( x ) =
arg min n ( x − x ∗ n ) T Σ x ∗ n ( x − x ∗ n ) .
This assumption that a dataset is a mixture of Gaussians with nonoverlapping support will simplify our derivations below, which could be extended to the general case if needed.
Note that this is not restrictive since, given a sufficiently large N , the above can represent any manifold with an arbitrarily good approximation.
3.3. Data Distribution After Deep Network Transformation
Consider an affine spline operator f (Eq. 1) that goes from a space of dimension D to a space of dimension K with
K ≥ D .
The span —that we denote as image— of this mapping is given by
[Equation]
with Aff ( ω ; A ω , b ω ) =
{ A ω x + b ω : x ∈ ω } the affine transformation of region ω by the per-region parameters
A ω , b ω , and with Ω the partition of the input space in which x lives in.
We also provide an analytical form of the perregion affine mappings in Section 2.
Hence, the DN mapping consists of affine transformations on each input space partition region ω ∈ Ω based on the coordinate change induced by A ω and the shift induced by b ω .
When the input space is equipped with a density distribution, this density is transformed by the mapping f .
In general, finding the density of f ( X ) is an intractable task.
However, given our disjoint support assumption provided in
Section 3.2, we can arbitrarily increase the representation power of the density by increasing the number of prototypes
N .
In doing so, the support of each Gaussian is included with the region ω in which its means lie in, leading to the following result.
Theorem 1.
Given the setting of Equation (4) the unconditional DN output density denoted as Z is a mixture of the affinely transformed distributions x | x ∗ n ( x ) e.g. for the
Gaussian case
[Equation]
where ω ( x ∗ n ) = ω ∈ Ω ⇐⇒ x ∗ n ∈ ω is the partition region in which the prototype x ∗ n lives in.
The proof of the above involves the fact that if (cid:82) ω p ( x | x ∗ n ( x ) ) d x ≈ 1 then f is linear within the effective support of p .
Therefore, any sample from p will almost surely lie within a single region ω ∈ Ω and therefore, the entire mapping can be considered linear with respect to p .
Thus, the output distribution is a linear transformation of the input distribution based on the per-region affine mapping.
4. Information Optimization and Optimality
As we saw in Section 3.1, our objective is to maximize
I ( Z ; X ′ ) and I ( Z ′ ; X ) .
Based on our analysis, we now show how specific algorithms can be derived for SSL.
As we saw, when the noise in our input is small, the conditional output density p z | x ∗ reduces to a single Gaussian:
[Equation]
where we abbreviated the parameters.
Using that, and the result from Section 3.1, we see that one should optimize both H ( Z | X ′ ) and H ( Z ) .
As in a standard regression task, assume a Gaussian observation model, which means p ( z | z ′ ) ∼ N ( z ′ , Σ r ) .
Using the mean square error as our log-loss function in training regression tasks is a specific case of this assumption, where Σ r =
I
To compute our expected loss, we need to marginalize out the stochasticity in Z ′ , to obtain the result that the conditional decoding map is a Gaussian.
[Equation]
which gives the distribution N ( µ n [Equation] , Σ r + Σ n ′ ))
meaning that we can lower the bound of mutual information with
I ( Z ; X
[Equation]
−
N (cid:88) n =
[Equation]
What happens if we attempt to optimize this objective?
The only intractable component is the entropy of Z .
We will begin by examining Z itself.
It is natural to ask why the entropy of Z will not increase to infinity.
Intuitively, the answer is that H ( Z ) and H ( Z | X ′ ) are tied together and one cannot increase without the other.
Now, recalling that under our distribution assumption, Z is a mixture of Gaussian (recall Thm. 1), we can see how the existing upper and lower bounds could be used for this case; for example, the ones in (Moshksar & Khandani, 2016).
4.1. Deriving VICReg From First Principles
We now propose to recover VICReg from the first principles per the above information-theoretic principle.
Recall that our goal is to estimate the entropy H ( Z ) in 7, where Z is a Gaussian mixture.
This quantity is not known for a mixture of Gaussians due to the logarithm of a sum of exponential functions, except for the special case of a single Gaussian density.
There are, however, several approximations in the literature that include both upper and lower bounds.
Among the methods, some use the logarithmic sum of the probability (Kolchinsky & Tracey, 2017), and some use an entropy-adjusted log probabilities (Huber et al., 2008).
An even simpler solution is to approximate the entire mixture as a single Gaussian by capturing only the first two moments of the mixture distribution.
Since the Gaussian distribution maximizes the entropy for a given covariance matrix, this method provides an upper bound approximation of our entropy of interest H ( Z ) .
In this case, denoting by Σ Z is the covariance matrix of Z , we find that we should maximize the following objective:
[Equation]
where Σ r is constant with respect to our optimization process, and the second term is the prediction performance of one representation from the other.
A key result from Shi et al.
(2009) connects the eigenvectors and eigenvalues of Σ Z and those of each component
Σ
i , ∀ i , and showed that under the assumption that the separation ( µ i − µ j )Σ − 1 i ( µ i − µ j ) T between the different components is large enough —which holds true in our case as per our data distribution model— the eigenfunctions of Σ i , ∀ i are approximately the eigenfunctions of Σ Z .
Hence, in our case, this means that since all those eigenvalues are tied, we only need to find the most efficient way to maximize | Σ Z | .
We know that the determinant of the matrix is the product of its eigenvalues.
For every positive matrix, the maximum eigenvalue is bigger than or equal to each diagonal element.
Therefore, under the constraint of the eigenvalues of the matrix, the most efficient way is to decrease the off-diagonal terms and increase the diagonal terms.
By setting Σ r
= I , we therefore fully recover the VICReg objective .
4.2. Empirical Evaluation
The next step is to verify the validity of our assumptions.
Based on the theory presented in Section 3.3, the conditional output density p z | T
= i reduces to a single Gaussian with decreasing input noise.
We validated it using a ResNetmodel trained with SimCLR or VICReg on the CIFAR10 dataset (Krizhevsky, 2009).
From the test dataset, we sample
Gaussian samples for each image and analyze whether these samples remain Gaussian (for each image) at the penultimate layer of the DN, i.e., before the linear classification layer, independently for each output dimension.
Then, we employ D’Agostino and Pearson’s test (D’Agostino, 1971) to compute the p-value of the normality test under the null hypothesis that the sample represents a normal distribution.
The process is repeated for different noise standard deviations.
Fig.
1 shows the p-value as a function of the normalized standard deviation.
We can observe that the network’s output is indeed Gaussian with a high probability for small input noise.
VICReg is, interestingly, more ”Gaussian” than SimCLR, which may have to do with the fact that it optimizes only the second moments of the density distribution to regularize H ( Z ) .
5. Conclusions
In this study, we examine SSL’s objective function from an information-theoretic perspective.
Our analysis, based on transferring the required stochasticity to the input distribution, shows how to derive SSL objectives.
Therefore, it is possible to obtain an information-theoretic analysis even when using deterministic DNs.
In the second part, we rediscovered VICReg’s loss function from first principles and showed its implicit assumptions.
In short, VICReg performs a crude lower bound estimate of the output density entropy by approximating this distribution with a Gaussian matching the first two moments.
Finally, we empirically validated that our assumptions are valid in practice, thus confirming the validity of our novel understanding of VICReg.
Our work opens many new paths for future research; A better estimation of information-theoretic quantities fits our assumptions.
Another exciting research direction is to identify which SSL method is preferred based on data properties.