Self-Supervised Learning with an Information Maximization Criterion
Anonymous Author(s)
AfﬁliationAddress email
Abstract
Self-supervised learning allows AI systems to learn effective representations fromlarge amounts of data using tasks that do not require costly labeling.
Mode collapse, i.e., the model producing identical representations for all inputs, is a centralproblem to many self-supervised learning approaches, making self-supervisedtasks, such as matching distorted variants of the inputs, ineffective.
In this article,we argue that a straightforward application of information maximization amongalternative latent representations of the same input naturally solves the collapseproblem and achieves competitive empirical results.
We propose a self-supervisedlearning method, CorInfoMax, that uses a second-order statistics-based mutualinformation measure that reﬂects the level of correlation among its arguments.
Maximizing this correlative information measure between alternative representations of the same input serves two purposes: (1) it avoids the collapse problemby generating feature vectors with non-degenerate covariances; (2) it establishesrelevance among alternative representations by increasing the linear dependenceamong them.
An approximation of the proposed information maximization objective simpliﬁes to a Euclidean distance-based objective function regularized by thelog-determinant of the feature covariance matrix.
The regularization term acts as anatural barrier against feature space degeneracy.
Consequently, beyond avoidingcomplete output collapse to a single point, the proposed approach also preventsdimensional collapse by encouraging the spread of information across the wholefeature space.
Numerical experiments demonstrate that CorInfoMax achieves betteror competitive performance results relative to the state-of-the-art SSL approaches.
Introduction
Self-supervised learning (SSL) is an important paradigm with signiﬁcant impact in artiﬁcial intelligence [1–6].
In particular, SSL can signiﬁcantly reduce the need for expensively-obtained labeleddata.
Furthermore, the internal representations learned by SSL approaches are more suitable fortask generalization in transfer learning applications [7].
Recently, SSL methods have even beendemonstrated to achieve comparable performance to their supervised counterparts [8, 9].
The primary approach for SSL is to deﬁne a pretext objective based on unlabeled data, whoseoptimization leads to rich and useful representations for downstream applications.
For example, forvisual processing, a common way to construct a pretext task is to obtain similar features for twodifferent augmentations of the same input, which exploits the expected invariance of representationsto particular transformations of the input.
However, the corresponding optimization setting shouldavoid the (total) “collapse problem,” which is described as generating the same feature output forall inputs.
Figure 1.(a) illustrates a toy view of the total collapse problem, where all latent vectoroutputs converge to a single point during training.
We can also deﬁne its less severe version, which
[Table]
/gAn9qz9q59ad+zaE7LZg7BHLSfX/bXrLg=</latexit> r z (1)  z (1)  z (2)  <latexit sha1_base64="jZHb6u2dMVFquqE1HkE2IlyxnXI=">AAACKXicbVDLSsNAFJ3UV62vqEs3wSK0m5KIoispuHFZxT6giWEymbRDJ5MwMxFqyO+48VfcKCjq1h9x0kbQtgcGDuecy9x7vJgSIU3zUystLa+srpXXKxubW9s7+u5eR0QJR7iNIhrxngcFpoThtiSS4l7MMQw9irve6DL3u/eYCxKxWzmOsRPCASMBQVAqydWbNoMehW5qh1AOvSB9yO7SGq1nmU2jge1jWft1brIFqbqrV82GOYExT6yCVEGBlqu/2n6EkhAziSgUom+ZsXRSyCVBFGcVOxE4hmgEB7ivKIMhFk46uTQzjpTiG0HE1WPSmKh/J1IYCjEOPZXMNxWzXi4u8vqJDM6dlLA4kZih6UdBQg0ZGXlthk84RpKOFYGIE7WrgYaQQyRVuRVVgjV78jzpHDes04Z5fVJtXhR1lMEBOAQ1YIEz0ARXoAXaAIFH8AzewLv2pL1oH9rXNFrSipl98A/a9w9TC6h5</latexit> r z ( l ) logdet( R z ( l ) ) is referred to as "dimensional collapse," where all latent vectors lie in a strict subspace of the wholelatent space as illustrated by Figure 1.(b).
This article proposes an SSL approach based on the maximization of a form of mutual informationamong the alternative latent representations obtained from the same input.
The mutual informationmaximization approach would serve two fundamental purposes for SSL:
(i).
( Similarity )
It would enforce the representations obtained for the same input source to be related toor dependent on each other,
(ii).
( No Collapse )
Since the entropy of each alternative representation is lower bounded by the mutual information, maximizing mutual information ensures that latent representations have non-degeneratedistributions avoiding collapse.
Therefore, the information maximization approach provides a natural solution to the problems targetedby the existing SSL approaches, which are brieﬂy surveyed in Sec. 2.2.
The conventional choice for mutual information measure is Shannon Mutual Information (SMI) [10].
For SSL, the fundamental limitations on the precise computation of SMI [11–13] impose seriouschallenges on model training. Especially when mutual information is large, we require drasticallymore samples (exponentially related to SMI) to obtain a reliable estimate of SMI [13].
However,the use of large batch sizes to improve precision may not be a favorable solution, as (i) it mayhave an adverse impact on generalization performance [14, 15], (ii) it can signiﬁcantly increasememory requirements.
Furthermore, estimating and optimizing SMI can signiﬁcantly increase thecomputational burden.
Finally, SMI maximization may induce a nonlinear dependence between therepresentations corresponding to the same input, where the resultant partitioning of the latent spacemay not be suitable as an input to a simple linear classiﬁer.
To address these issues, we can consider a second-order-statistics-based mutual information measure,referred to as log-determinant mutual information (LDMI)
The joint entropy measurecorresponding to LDMI is the log-determinant of the diagonally perturbed covariance matrix.
Thezero perturbation case boils down to the Shannon differential entropy for Gaussian distributedvectors.
Unlike SMI, which reﬂects general dependence between its arguments, LDMI reﬂects lineardependence and is computationally more efﬁcient.
In this article, we propose a novel SSL approach, correlative information maximization (CorInfoMax),derived from the LDMI measure.
The loss function for CorInfoMax is obtained through the ﬁrst-orderapproximation of LDMI and restricting the linear dependence to the identity map.
The resultingoptimization objective is a Euclidean distance-based loss function regularized by the log-determinantof the latent vector covariance matrix.
This regularization term encourages the spreading of thefeature vectors and acts as a natural barrier against the dimensional degeneracy and collapse in thefeature space.
Our numerical experiments conﬁrm that CorInfoMax learns effective representationsthat perform well in downstream tasks.
The main contributions can be summarized as follows:
• Based on information theoretical grounds, we propose a novel framework, CorInfoMax, with aninterpretable loss function that has explicit components for minimizing variance of positive samples(‘the attraction factor’) and for making full use of the embedding space to avoid dimensionalcollapse (‘the big-bang factor’).• We introduce a computationally-efﬁcient approximation to LDMI that simpliﬁes it by using anidentity mapping instead of a general linear mapping.
This approximation directly minimizesrepresentation variance, which is desired for self-supervised learning.• CorInfoMax relies only on second-order statistics, does not require negative samples, and does notforce embeddings to be uncorrelated.•
In addition to these theoretically appealing features, the proposed CorInfoMax framework achieveseither better or similar performance results compared to the state-of-the-art SSL approaches.
The following is the organization of the article: Section 2, provides a discussion of the relevantliterature related to the proposed approach.
We introduce the log-determinant entropy and mutualinformation measures by highlighting LDMI as a measure of linear dependence in Section 3.
Inthe same section, we derive the objective function of CorInfoMax as a variation on the LDMI measure.
Section 4 introduces the proposed correlative information maximization-based self-supervisedlearning method.
Section 5 provides the numerical experiments illustrating the performance of theproposed approach.
The appendix provided in the supplementary document contains the detailsrelated to these experiments.
Finally, Section 6 presents the discussion and conclusions.
Related work
2.1 Determinant maximization for unsupervised learning
The determinant maximization criterion utilized in our framework has been used as an effectivealgorithmic tool in unsupervised matrix factorization methods such as nonnegative matrix factorization(NMF)
[18, 19], simplex-structured matrix factorization (SSMF)
[20], sparse component analysis(SCA) [21], bounded component analysis (BCA) [22] and polytopic matrix factorization (PMF) [23].
In the generative models of these frameworks, the input data is assumed to be linear transformationsof some latent vectors.
Furthermore, these latent vectors are assumed to be sufﬁciently scattered intheir domain.
Maximizing the determinant of the latent covariance matrix spreads the latent vectorestimates to capture this presumed scattering of the generative model samples.
Similarly, the logdeterminant of the latent vector covariance matrix in the CorInfoMax objective causes the spreadingof latent vectors in their ambient space and, therefore, avoids collapse.
According to [17], the determinant maximization criterion used in these matrix factorization frameworks is essentially equivalent tocorrelative information maximization between input and its latent factors based on the LDMI objective,where the latent vectors are constrained to lie in their domain determined by the generative model.
In the CorInfoMax approach, we utilize determinant maximization for correlative information maximization among latent space vectors rather than between inputs and their latent space representations.
2.2 Handling collapse for SSL
Collapse is a central concern for SSL with the same input in computer vision, and we can categorizethe existing approaches according to how they handle this issue.
For example, contrastive methods,such as SimCLR [8] and MoCo [24], deﬁne an objective that pushes features for different inputs(negative samples) away from each other while keeping the representations corresponding to the sameinput (positive samples) close to each other.
The number and the choice of the negative samples appearas critical factors for the performance and scalability of contrastive methods.
Another category of SSLis the distillation methods, such as BYOL [4] and SimSiam [25], which avoid collapse through theasymmetry of alternative encoder branches and algorithmic tuning [26].
Other SSL approaches suchas DeepCluster [27], SeLa [28] and SvAW [9] enforce a cluster structure in the feature space to prevent constant output.
As a different line of algorithms, decorrelated batch normalization (DBN) [29] Barlow-Twins [5], Whitening MSE (W-MSE) [30] and VICReg [6], use feature decorrelation [31] as ameans to avoid information collapse.
The effective whitening mechanisms used in these methods aimfor the isotropic spread of information inside the feature space which also prevents "dimensional collapse" illustrated in Figure 1.(b).
The reference [32] constructs an SSL loss function using the HilbertSchmidt Independence Criterion, a kernel-based independence measure, to maximize the dependence between alternative representations of the same input.
Finally, [33] proposed graph representationsfor the embeddings of positive samples and the corresponding spectral decomposition algorithm.
The proposed CorInfoMax approach is most related to the decorrelation-based methods discussedabove, mainly due to correlation-based measures.
However, unlike the decorrelation methods,
CorInfoMax does not constrain latent vectors to be uncorrelated.
Instead, it avoids covariance matrixdegeneracy by using its log-determinant as a regularizer loss function.
Furthermore, the informationmaximization principle is more direct and explicit for the CorInfoMax algorithm.
2.3 Information maximization for unsupervised learning
The maximization of SMI has been proposed as an unsupervised learning mechanism in differentbut related contexts.
As one of the earliest approaches, Linsker proposed maximum informationtransfer from input data to its latent representation and showed that it is equivalent to maximizingthe determinant of the output covariance under the Gaussian distribution assumption [34].
Aroundthe same time frame, Becker & Hinton [35]put forward a representation learning approach based onthe maximization of (an approximation of) the SMI between the alternative latent vectors obtainedfrom the same image.
The most well-known application is the Independent Component Analysis(ICA)
Infomax algorithm [36] for separating independent sources from their linear combinations.
The
ICA-Infomax algorithm targets to maximize the mutual information between mixtures and sourceestimates while imposing statistical independence among outputs.
The Deep Infomax approach [37]extends this idea to unsupervised feature learning by maximizing the mutual information between theinput and output while matching a prior distribution for the representations.
It is essential to underline that the proposed method clearly distinguishes itself from the Deep Infomaxin [37]: Our objective is not to maximize the mutual information between inputs and outputs of a deepnetwork.
Instead, we maximize the mutual information content of the alternative latent representationsof the same input.
From this point of view, our approach is closer to what is aimed at by [35].
However,we use a different (correlative) information measure, which is computationally more efﬁcient andinduces a special form of linear dependence among alternative latent representations of the sameinput, which may be more desirable considering the goal of generating features for a linear classiﬁer.
CorInfoMax as a criterion based on log-determinant mutual information
This section derives the optimization objective for the CorInfoMax framework as a variation on the
LDMI measure.
As described in Appendix A, the LDMI between random vectors x and y is given by
[Equation]
where R x and R y are auto-covariance matrices for x and y , respectively, and R xy is theircross-covariance matrix.
There are three potential advantages of using LDMI in (1) over SMI in (A.1) for SSL:
i.
It is
[-• I ( " ) LD ( x ; y )  0 , • I ( " ) LD ( x ; y ) = 
 if and-] only based on second-order statistics.
On the other hand, SMI is a statistic based on the joint PDFs of the argument vectors, and its accurate estimation has high sample complexity.
In contrast, itis more practical to obtain (the estimates of) the auto-covariance and the cross-covariance matrices, as discussed in Section
ii. There are beneﬁts of information maximization principle in connection with self-supervised training.
The maximization of SMI of two vectors induces a general, potentially nonlinear dependencebetween them whereas the maximization of LDMI increases correlation, or linear dependence
Therefore, LDMI-based information maximization is expected to organize the feature space beingmore favorable for data-efﬁcient and low-complexity linear or shallow supervised classiﬁers astargeted in
SSL applications. iii. The resulting LDMI-based objective functions are more interpretable – they involve the log-
determinant of the projector-space covariance
The primary motivation for obtaining an LDMI variant is to increase the similarity between thealternative latent representations of the same input by enforcing the identity transformation betweenthem.
Let z (1) and z (2) represent alternative latent representations corresponding to the same input.
One potential SSL approach would be to pursue direct maximization of I ( " ) LD ( z (1) ; z (2) ) in (1).
As discussedearlier, this would maximize the correlation between z (1) and z (2) , therefore inducing a linear dependence between them.
For the SSL application, we would prefer the identity mapping to an arbitrarylinear relationship such that the alternative latent representations corresponding to the same inputconcentrate in the same neighborhood of the latent space.
Therefore, we modify the LDMI expressionin (1) to impose this constraint.
For this purpose, we apply the ﬁrst-order Taylor series approximation,log det [Equation] C + D ) ⇡ log det( C )+
Tr ( D T C  1 ) , on the third term in the right side of (1), which provides
[Equation]
Therefore, the expression in (2) corresponds to the mean square error of the best linear (afﬁne) MMSE estimator of z (1) from z (2) , multiplied by "  1 .
Similarly, if we apply the same approximationto the fourth term on the right side of (1), we obtain
[Equation]
In order to induce the identity mapping between z (1) and z (2) , we constrain A 1 =
A 2 =
I , andb
=
b
= 0 , which transforms both (2) and ((3) into "  1 E ( k z (1)
z (2) k 22 ) + const.
Therefore,scaling the expression in (1) by 4 and using this modiﬁcation, we obtain
[Equation]
where we ignored the constant terms.
We refer to (4) as the stochastic CorInfoMax objective function.
In Section 4, we propose a SSL method based on the optimization of this objective function.
LDMI based self-supervised learning
This section introduces the proposed correlative information maximization (CorInfoMax) approachfor SSL.
We start by describing the presumed setting for the pretext task, which is matching the latentrepresentations of different augmentations of the same input in Section 4.1.
Section 4.2 is the mainsection where we propose the correlative information maximization algorithm for SSL.
Finally, wediscuss the implementation complexity of the proposed approach in Section 4.3.
4.1 Self-supervised learning setting
We start by describing the presumed self-supervised learning setup, which is illustrated in Figure 2: • The input is a sequence of tensors {X [ l ] 2 R M ⇥ K ⇥ N , l 2 Z } , where each index sample corresponds to a batch of N images with dimensions M ⇥ K .As discussed in Section 1, the pretext taskthat we consider is matching the latent representations of the two augmentations of the same input. Therefore, we ﬁrst apply the augmentation functions a 1 ( · ) and a 2 ( · ) to obtain the transformedversions of the input sequence, i.e., {X ( q ) [ l ] = a q ( X [ l ]) 2 R M ⇥ K ⇥ N , l 2 Z } , q 2 { 1 , 2 } . We
will illustrate particular choices of augmentations in the examples provided in Section 5.
We
note that the proposed setting can be potentially extended to multi-modal schemes where the
augmentations are deﬁned for multiple modalities of the same input, such as visual and sound.
• We represent Siamese networks to be trained by SSL with the mappings f q ( U ; W ( q ) DNN ) for q = 1 , 2 ,
where U 2 R M ⇥ K ⇥ N is the input, W ( q ) DNN represents the trainable parameters of the q th network.
[Figure] X N M
K a [Figure] Augmentationa [Figure]
AugmentationX (1)
X (2) f 1 ( · ; W (1) DNN ) f 2 ( · ; W (2) DNN )
DNNDNN-
p 1 ( · ; W (1) P ) p 2 ( · ; W (2) P )
ProjectorProjector-
Z (1)
Z (2)
P P
N
N
• The outputs of the DNN for different augmented inputs are represented by Y ( q )
[ l ] =
f q ( X ( q ) [ l ]; W
[Equation]
q
[Equation]
[Equation]
N , q 2 { 1 , 2 } , l 2 Z , where F is the output feature dimension.•
To be used in self-supervised training, we project the outputs of the enconders usingprojector networks p q ( Y ; W ( q ) P ) , q = 1 , 2 , where Y is its input, and W ( q ) P is the set of trainableparameters of the q th projector.
• The outputs of the projector network for different augmentations are represented by Z ( q ) [ l ] =
p q ( Y ( q ) [ l ]; W ( q ) P ) 2 R P ⇥ N , q 2 { 1 , 2 } , l 2 Z , where P is the projector head dimension.
For the numerical experiments in the article, we consider the weight sharing setup where bothencoders use the same network weights.
We provide more details in Appendix B.
4.2 Correlative mutual information maximization
We deﬁne the CorInfoMax approach for SSL through the optimization problemmaximize
[Equation]
where ˆ J ( " ) is the sample based estimate of (4) at the l th -batch which can be written as
[Equation]
[Equation]
where the rightmost term stands for the sample based estimate of the term 2 "  1 E ( k z (1)  z (2) k 22 ) in(4), and
ˆ R z (1)
[ l ] , ˆ R z (2) [ l ] and ˆ R z (1) z (2) [ l ] are the auto-covariance matrix and cross-covariance matrixestimates for the projector heads at the l th batch.
If the batch size is large enough, these estimatescan only be obtained from the current batch samples.
However, due to hardware limitations andconsidering the fact that intermediate batch sizes offer better accuracy [14, 15], sufﬁciently-largebatch sizes may not be possible to obtain reliable covariance estimates.
Therefore, we adopt recursivecovariance matrix estimation across batches [38].
The corresponding covariance update expressionstake the form
[Equation]
(7) where {+
+}   < 1 is the forgetting factor, and ˜ Z ( q ) [ l ] represents the batch of mean-centralizedprojector outputs deﬁned as ˜ Z ( q )
[ l ] =
Z ( q ) [ l ]  µ ( q ) [ l ] 1 TN , where { µ ( q ) [ l ] , q 2 { 1 , 2 }} representsthe mean estimates for the projector outputs, updated by µ ( q ) [ l ]
=  µ ( q ) [ l  1] + (1
[Equation]
for q 2 { 1 , 2 } .
We can write the update expression for the cross-covariance matrix as ˆ R z (1) z (2)
[ l ] = ˆ R z (1) z (2) [ l  1] + (1   ) 1 N ˜ Z
[ l ]˜ Z (2) [ l ] T .
(See Appendix C for CorInfoMax pseudocode.)
It is informative to inspect the terms in the objective function
ˆ J ( z (1) , z (2) ) in (6):
• Minimization of 2 "  1 N k Z
(1) [ l ]  Z (2) [ l ] k 2 F , acts as a force to pull the representations of alternativeaugmentations toward each other, which we refer to as the attraction factor .•
Maximization of log det( ˆ R z (1)
[ l ] + " I ) acts as a dispersion force causing non-degenerate (for" ⇡ 0 ) expansion of projection vectors in the P -dimensional space, which we informally referto as big-bang factor .
Therefore, the covariance determinant acts as a regularization or barrierfunction avoiding collapse in the latent space.
Consequently, it provides a convincing replacementfor negative samples in contrastive methods to prevent feature collapse.
Appendix K illustrates theeigenvalues of the projector covariance vector for the CIFAR-10 training, which demonstrates that
the CorInfoMax criterion spreads information throughout the feature space, avoiding dimensional
collapse.
Figure 1.(c) illustrates the learning dynamics of the CorInfoMax approach based on (6) with a toypicture.
In this ﬁgure, the ellipsoid is a representative surface for the level set of the quadratic function( z  µ ( q ) )
T R  1 z ( q ) ( z  µ ( q ) ) , which reﬂects the spreading pattern of the latent vectors around theirmean.
The black arrows represent the gradient of the log det( R z ( q ) ) regularization factor, which actsas a force to push latent vectors away from the center of the ellipsoid ( µ ( q ) ) and therefore correspondsto the expansion force of the "big-bang" factor.
In the same ﬁgure, the white arrows correspond to thegradients corresponding to the attraction factor, i.e., the Euclidian distance between two positive pairs.
In summary, we can view the learning dynamics of CorInfomax SSL as an expansion in the latentspace, while the representations of positive samples are attracted to each other.
More informationon the CorrInfoMax optimization setting is provided in the Appendix M. Furthermore, Appendix
J contains a visualization of the embeddings obtained using the CorInfoMax criterion.
4.3 Computational complexity
The main difference between our and related SSL methods is the extra log-determinant computation.
The log determinant is typically computed using a decomposition method such as the LUdecomposition, whose gradient requires a matrix inversion.
Both determinant and inversion havethe same complexity as matrix multiplication [39], i.e. O ( n ↵ ) with 2 < ↵  3 for multiplyingtwo n ⇥ n matrices.
In our experiments with GPUs, we observe an almost ﬂat runtime cost up ton = 1024 , which implies that the overhead due to the log det cost is insigniﬁcant in practice.
We ﬁndthe extra cost of log determinant and its gradient to be negligible compared to the rest of the modelcomputation, whose runtime is dominated by the encoder, as conﬁrmed by the runtime experimentsreported in Appendix G.
Experiments
5.1 Implementation details
Datasets: We perform experiments on CIFAR-10, CIFAR-100 [40], Tiny ImageNet [41], ImageNet100 and ImageNet-1K [42] datasets.
For ImageNet-100, we use the same subset of ImageNet-1K[42] as related work [43–45].
(See Appendix D for more details.)
Training Procedure: The experiments consist of two consecutive stages: pretraining and linearevaluation.
We ﬁrst perform the unsupervised pretraining of the encoder network f by applying theproposed CorInfoMax method described in Section 4.2 on the training dataset.
After completingpretraining, we perform the linear evaluation, a standardized protocol to evaluate the quality of thelearned representations [4, 8, 46].
For the linear evaluation stage, we ﬁrst perform supervised training of the linear classiﬁer using therepresentations obtained from the encoder network f with frozen coefﬁcients on the same trainingdataset.
Then, we obtain the test accuracy results for the trained linear classiﬁer based on thevalidation dataset.
Computational Resources: For ImageNet-100 and ImageNet-1K, we pretrain our model on {+up to 
+} A100 Cloud GPUs.
The remaining datasets are trained using a single T4 and V100 Cloud GPU.
Linear evaluations are performed using the same type and amount of computational resources. Thedetails related to batch sizes are described in 5.1.3.
5.1.1 Input augmentations
During the pretraining stage , two augmented versions of each input image are generated as shown in
Figure 2.
During this process, each image is cropped with random size, resized to the original resolution, followed by random applications of horizontal mirroring, color jittering, grayscale conversion,
Gaussian blurring, and solarization.
Since we use the same augmentation parameters as BYOL [4]and VicReg [6]: each augmentation branch uses the same probability values for these randomizedoperations except Gaussian blurring and solarization, which use different probabilities.
During the training phase of the linear evaluation stage, a single augmentation of each input imageis produced by random cropping and resizing followed by a random horizontal ﬂip.
For the test phaseof linear evaluation, we use resize and center crop augmentations, similar to [5, 6].
We provide moredetails about the augmentations in Appendix E.
5.1.2 Network architecture
Encoder Network : For CIFAR datasets, we use a modiﬁed form of ResNet-18 architecture [47] similarto [8, 25, 33].We use standard ResNet-50([47]) for Tiny ImageNet, ImageNet-100 and ImageNet-1K,and also standard ResNet-18([47]) for ImageNet-100.
In all cases, the last fully connected layeris removed.
Therefore, the encoder output size is 512 and 2048 for ResNet-18 and ResNet-50,respectively.
The encoder network f shares weights between augmented branches (see Appendix B).
Projector Network: The output of the encoder network is fed into the projector network, as in Figure2.
The projector network p is a 3 -layer MLP, with ReLU activation functions for the hidden layersand linear activation functions for the output layer.
The projector dimensions are 2048 - 2048 - 64for the CIFAR-10 dataset, 4096 - 4096 - 128 for the CIFAR-100, Tiny ImageNet and ImageNet-100,and 8192 - 8192 - 512 for the Imagenet-1K. Contrary to common practice [5, 6], we do not use batchnormalization [48] in the projection network p .
Finally, we perform L 2 -normalization on the projectoroutput.
Linear Classiﬁer: For the linear evaluation phase, we employ a standard linear classiﬁer whose inputis the weight-frozen encoder network’s output.
5.1.3 Optimization
For pretraining, we use 1000 epochs with a batch size of 512 for CIFAR, and 800 epochs with a batchsize of 1024 for Tiny ImageNet.
For ImageNet-100 experiments, we use 400 epochs for ResNet-18and 200 epochs for ResNet-50, with a batch size of 1024 for both. ImageNet-1K experiments areconducted as 100 epochs with a batch size of 1536.
We use the SGD optimizer with a momentumof 0 .
9 and a weight decay of 1 e  4 .
The initial learning rate is 0 . 5 for CIFAR datasets and Tiny ImageNet, 1 . 0 for ImageNet-100, and 0 .
2 for ImageNet-1K. These learning rates follow the cosinedecay with a linear warmup schedule.
We use the modiﬁed form of (6) as our objective function, where we replace 2 "  1 N with ↵ , theattraction coefﬁcient, which we consider as a separate hyper-parameter.
In our experiments, thediagonal perturbation is " = 1e  8 , while ↵ = 250 for CIFAR-10, ↵ = 1000 for CIFAR-100,↵ = 2000 for ImageNet-1K, and ↵ = 500 for Tiny ImageNet and ImageNet-100. The forgettingfactor  = 0 . 01 for all datasets except Tiny ImageNet and ImageNet-1K, which have  = 0 .
Details with coefﬁcients of our loss function are provided in Appendix F.
For linear evaluation, the linear classiﬁer is trained for 100 epochs with a batch size of 256 for alldatasets.
We used the SGD optimizer with a momentum of 0.9, cosine decay schedule with initial andminimum learning rates of 0 . 2 and 2e  3 respectively, without no weight decay.
For ImageNet-1K,these values are changed to 2 and 0 .
5.2 Results
We evaluate the learned representations from the CorInfoMax pretraining by following the linearevaluation protocol explained in Sec. 5.1.
Table 1 shows that CorInfoMax achieves state-of-the-artperformance in linear classiﬁcation after pretraining.
Due to the limited size of the validation sets(5-10K), differences less than ⇡ 0 .
5% are not statistically signiﬁcant.
The full comparison is providedin Table 10 in Appendix H. It is also interesting to observe the progress of the LDMI measure duringthe CorInfoMax training process, which is illustrated in Appendix L
[Table]
[Table] We compare the semi-supervised learning performance of CorInfoMax with VICReg [6].
For a faircomparison, VICReg is pretrained for 100 epochs on ImageNet-1K, then semi-supervised learningperformances of both models are evaluated by ﬁne-tuning the encoders with 1(%) and 10(%) of thelabeled ImageNet-1K dataset. The results are presented in Table 2, while details are provided in Appendix F.5.
[Table]
We experimented with transfer learning on object detection and instance segmentation with the sameprocedure and code of MoCo[52], and reproduce the results with the MoCo V2 [53] checkpoint. Our algorithm shows competitive performance with MoCo V2 [53], related details are provided in
Appendix I.
Discussions and conclusions
We proposed a novel SSL framework based on the correlative information maximization, CorInfoMax,which provides a natural solution for obtaining organized representations that are scattered in thelatent space and avoiding collapse.
Our experiments demonstrate the state-of-the-art performanceof CorInfoMax in numerous downstream tasks.
While achieving so, CorInfoMax does not yieldany signiﬁcant impact on training time.
Future work can further improve CorInfoMax focusing onshortcomings like improved selection of hyperparameter and augmentation, a core problem in SSL
[Table]
overall.