Unsupervised Visual Representation Learning via Mutual Information Regularized Assignment
Anonymous Author(s)
AffiliationAddress email
Abstract
This paper proposes Mutual Information Regularized Assignment (MIRA), apseudo-labeling algorithm for unsupervised representation learning inspired byinformation maximization.
We formulate online pseudo-labeling as an optimizationproblem to find pseudo-labels that maximize the mutual information between the label and data while being close to a given model probability.
We derive a fixed-pointiteration method and prove its convergence to the optimal solution.
aIn contrast tobaselines, MIRA combined with pseudo-label prediction enables a simple yet effective clustering-based representation learning without incorporating extra trainingtechniques or artificial constraints such as sampling strategy, equipartition constraints, etc.
With relatively small training epochs, representation learned by MIRAachieves state-of-the-art performance on various downstream tasks, including thelinear/ k -NN evaluation and transfer learning.
Especially, with only 400 epochs, ourmethod applied to ImageNet dataset with ResNet-50 architecture achieves 75.5%linear evaluation accuracy.
Introduction
There has been a growing interest in using a large-scale dataset to build powerful machine learningmodels [43].
Self-supervised learning (SSL), which aims to learn a useful representation withoutlabels, is suitable for this trend; is actively studied in the fields of natural language processing [19, 20]and computer vision [10, 29].
In the vision domain, recent SSL methods are commonly designedto use augmented views and train visual representation to be augmentation-invariant.
They haveachieved state-of-the-art performance surpassing supervised representation in a variety of visual tasks,including semi-supervised learning [8, 50], transfer learning [21] and object detection [13].
Meanwhile, a line of works use clustering for un-/self-supervised representation learning.
Theyexplicitly assign pseudo-labels to embedded representation via clustering and the model is therebytrained to predict such labels.
These clustering-based methods can account for inter-data similarity;representations are encouraged to encode the semantic structure of data.
Prior works [48, 46, 4, 31]have shown encouraging results in small-scaled settings; Caron et al.
[6] show that it can be alsoapplied to the large-scaled dataset or even to a non-curated dataset [7].
Recently, several works [2, 8,37] have adapted the philosophy of augmentation invariance and achieved strong empirical results.
They typically assign pseudo-labels using augmented views while predicting the labels looking atother differently augmented views.
Despite its conceptual simplicity, a naive application of clustering to representation learning is hardto achieve, especially in large-scale dataset.
This is because clustering-based methods are prone tocollapse, i.e., all samples are assigned to a single cluster.
To address this, recent methods heavily relyon extra training techniques or artificial constraints, such as pre-training [47], sampling strategy [6], equipartition constraints [2, 8], etc.
However, it is unclear if these additions are appropriate or howsuch components will affect the representation quality.
In this paper, we propose Mutual Information Regularized Assignment (MIRA), a pseudo-labelingalgorithm that enables clustering-based SSL without any artificial constraints or extra trainingtechniques.
MIRA is designed to follow the infomax principle [38] and the intuition that goodlabels are something that can reduce most of the uncertainty about the data .
Our method assigns apseudo-label in a principled way by constructing an optimization problem.
For a given training modelthat predicts pseudo-labels, the optimization problem seeks a solution that maximizes the mutualinformation (MI) between the pseudo-labels and data while taking the model probability into account.
We formulate the problem as a convex optimization problem and derive the necessary and sufficientcondition of solution with the Karush-Kuhn-Tucker (KKT) condition.
The solution can be achievedby fixed-point iteration that we prove the convergence.
We remark that MIRA does not require anyform of extra training techniques or artificial constraints, e.g., equipartition constraints.
We apply MIRA to clustering-based representation learning and verify the representation quality onseveral standard self-supervised learning benchmarks.
We demonstrate its state-of-the-art performanceon linear/k-NN evaluation, semi-supervised learning, and transfer learning benchmark.
We furtherexperiment with convergence speed, scalability, and different components of our method.
Our contributions are summarized as follows:
• We propose MIRA, a simple and principled pseudo-label assignment strategy based on mutual
information.
Our method does not require extra training techniques or artificial constraints.
• We apply MIRA to clustering-based representation learning and it shows comparable performanceagainst the state-of-the-art methods with half of the training epochs.
Especially it achieves75.5%
top-1 accuracy on ImageNet linear evaluation with only 400 epochs of training and bestperformance in 9 out of 11 datasets in transfer learning.
• Representation by MIRA also consistently improves over other information-based SSL methods [22, 50].
Especially, our method without a multi- crop augmentation strategy achieves 73.8%top-1 accuracy and outperforms BarlowTwins [50], an information maximization-based selfsupervised method.
Related works
Self-supervised learning SSL methods are designed to learn the representation by solving pretexttasks.
Recent state-of-the-art SSL methods train their representation to be augmentation invariant.
They are based on various pretext tasks: instance discrimination [10, 11, 13, 14], metric learning [27,12], self-training [51, 9], and clustering [2, 6, 8]; our method belongs to the clustering-based SSLmethod.
Meanwhile, these methods are prone to collapsing into a trivial solution where everyrepresentation is map into a constant vector.
To address this, a variety of schemes and mechanisms aresuggested, e.g., the asymmetric structure, redundancy reduction, etc.
We will review more relevantworks in detail below.
Collapse preventing Many SSL approaches rely on extra training techniques and artificial assumptions to prevent collapsing.
In clustering-based methods, DeepCluster [6] adapts a sample strategy tosample elements uniformly across pseudo-labels to deal with empty clusters; SeLa [2] and SwAV [8]impose equipartition constraints to balance the cluster distribution.
Similarly, SelfClassifier [1] uses auniform pseudo-label prior, and PCL [37] employs concentration scaling.
DINO [9] and ReSSL [51]address collapsing by specific combinations of implementation details, i.e., centering and scaling withan exponential moving average network; while their mechanism for preventing collapse is unclear.
In this work, we show our method can naturally avoid collapsing without any of these assumptionsor training techniques.
We achieve results better than baselines with a simple but novel informationregularization algorithm.
We take a more detailed comparison with SeLa and SwAV in Sec. 3.3.
Information maximization Information maximization is a principal approach to learn representation and to avoid collapse.
DeepInfoMax [30] propose the MI maximization between the local andglobal views for representation learning; the existence of negative pairs prevents training toward thetrivial solution.
BarlowTwins [50] and W-MSE [22] addresses collapsing with redundancy reduction,
[Figure]
①
[Figure]
② model
[Figure]
indirectly maximizing the information content of the embedding vectors [3]. Among clustering-basedapproaches, IIC [33] maximizes the MI between the embedding codes for representation learning;most similarly to ours, TWIST [25] proposes to combine the mutual information between the dataand class prediction as a negative loss term with a consistency loss.
Both IIC and TWIST use the
MI as a loss function and directly optimize their model parameters with gradient descent of the loss.
However, direct optimization of MI terms by updating model parameters often leads to a sub-optimalsolution [25]; TWIST copes with this issue by appending the normalization layer before softmaxand introducing an additional self-labeling stage.
In contrast, MIRA addresses the difficulty of MImaximization in a principled way via explicit optimization.
Method
In this section, we explain our pseudo-labeling algorithm–MIRA.
When applying MIRA to representation learning, we follow the basic framework of clustering-based representation learning thatalternates between pseudo-labeling , i.e., cluster assignments, and model training to predict such labels.
Figure 1 illustrates our representation training cycle.
We will first explain our main contribution,
MIRA ( pseudo-labeling ) and then explain how it applies to model training .
Our idea is to employ the information maximization principle into pseudo-labeling.
We formulatean optimization problem for online clustering that assigns soft pseudo-labels to mini-batch samples(Sec.
The problem accounts for the model probability and mutual information between thepseudo-labels and data.
We propose an iterative method to solve the optimization problem (Sec. 3.2).
For the model training, we use the swapped prediction loss as in [8] (Sec. 3.3).
3.1 MI regularized cluster assignment
We have a model 1 f θ parametrized by θ that outputs K -dimensional logit f θ ( x ) ∈ R K for an imagex , where K is a predefined number of clusters.
The model probability p of an image x is then givenby the temperature τ t scaled output of the model— p := softmax ( f θ ( x ) /τ t ) —as in [8, 9].
For a minibatch of input images X = { x } Bi =1 , we denote the model probability P = { p i } Bi =1 ⊂ R K .
In ourpseudo-labeling, for the given model probability P , we want to assign pseudo-labels W ∗ =
{ w ∗ } Bi
=1that will be used for training the model by predicting them.
We argue that such pseudo-labels should maximize the mutual information between themselves anddata while accounting for the model probability P .
Let B ∈ { 1 , ..., B } and Y W ∈ { 1 , ..., K } be the random variables associated with the data index in mini-batch and label by probability distributions
W = { w } Bi =1 , respectively.
Our online pseudo-label (cluster) assignment is determined by solvingthe following optimization problem:
[Equation]
where ∆ K := { x ∈ R K + | x ⊺ 1 K = 1 } , ˆ I indicates an empirical (Monte Carlo) estimates of MI,and β is a trade-off parameter.
The problem consists of the (1) KL divergence term that makespseudo-labels to be based on the model probability p and (2) MI term between the pseudo-labels anddata to induce more information into the pseudo-labels.
By combining these two terms, we provide arefined pseudo-label that take account of both the model probability and MI.
To make the optimization problem tractable, we substitute the MI term ˆ I with the mini-batch estimatesof the entropy ˆ H
( Y W |B ) and marginal entropy ˆ H ( Y W ) in Eq.
We get:
[Equation]
where w j = 1 B (cid:80) Bi
=1 w ij is the marginal probability of a cluster j with W .
In practice, we find theoptimal point W ∗ of the optimization problem Eq. 3.
3.2 Solving strategy
To solve efficiently, we propose a fixed-point iteration that guarantees convergence to the uniqueoptimal solution W ∗ of our optimization problem.
The method is based on the following proposition.
Proposition 1.
For β ∈ [0 , 1) , the problem Eq.
3 is a strictly convex optimization problem; has aunique optimal point W ∗ that satisfies the following necessary and sufficient condition.
[Equation]
The proposition is driven by proving the strict convexity and then applying the Karush–Kuhn–Tucker(KKT) condition.
By substituting the necessary and sufficient condition (Eq. 4) of proposition 1 intow j = 1 B (cid:80) Bi =1 w ij , we get the necessary and sufficient condition with w ∗ : [Equation] Based on Eq. 5, we propose the following update rule for { u ( n ) j } Kj =1 ⊂ R + :
[Equation]
where u ( n ) j converges to w ∗ j as n → ∞ .
We can easily get w ∗ ij by Eq.
4 when the marginalprobability w ∗ j is given.
The proof of the proposition and convergence is in the Appendix.
By using the iterative updates of Eq. 6, we get our desirable pseudo-labels.
This requires a fewlines of code that are simple to implement.
We find that a few steps of iterations are enough fortraining.
This is supported by the convergence analysis in Sec. 4.3.
We use this fixed point iterationfor pseudo-labeling and name the method– M utual I nformation R egularized A ssignment ( MIRA )since it finds the pseudo-labels that are regularized by the mutual information.
3.3 Representation learning with MIRA
We explain how our pseudo-labeling algorithm is applied to representation learning.
We integrate thecomputed pseudo-labels with swapped prediction loss [8].
Specifically, given the two mini-batches of differently augmented views X (1) , X (2) , MIRA outputs the pseudo-labels U (1) , U (2) for each minibatch independently.
In parallel, model f θ provides the temperature τ s scaled softmax predictions
Q (1) , Q (2) of each mini-batch.
The swapped prediction loss is given as follows:
[Equation]
This loss function (Eq. 7) is minimized with respect to the parameters θ of the model f
θ used toproduce the predictions Q (1) , Q (2) .
For more detailed information about swapped prediction loss,please refer to [8].
In this paper, we verify our pseudo-labeling algorithm MIRA for a representation learning purposewith Eq.
For convenience, in the rest of this paper, we call the representation learning with
MIRA also as MIRA.
We note that MIRA can integrate recently suggested SSL components such asexponential moving average (EMA) or multi-crop augmentation strategy following the baselines [14,8, 9].
The pseudo-code for MIRA is provided in the Appendix.
We discuss some further details asfollows:
Preventing collapse The MI term in Eq.
3 takes a minimum value when collapsing happens.
MIRAnaturally avoids collapsed solution via penalizing assignment that exhibits low MI.
To be morespecific, unless starting from the collapsed state, MIRA finds MI-maximizing points around themodel prediction; will not choose collapsed pseudo-labels.
Hence, the iterative training to predictsuch labels will not lead to collapsing whenever the prediction of pseudo-labels is achievable.
Ourempirical results verify that MIRA doesn’t require extra training techniques or artificial constraints toaddress collapsing.
Comparison to SwAV and SeLa SeLa [2] and SwAV [8] assume the equipartition of data intoclusters.
They formulate their pseudo-labeling process into optimal transport (OT) problem; solvingit with the iterative Sinkhorn-Knopp (SK) algorithm
Mathematically, the difference to MIRAis in how to deal with the marginal entropy .
SeLa and SwAV constrain the marginal entropy tomaximum value–equipartition while MIRA decides marginal entropy by MI regularization 2 .
Asanoet al.
[2] argue that their pseudo-labels with OT problem maximize the information between labelsand data indices under the equipartition constraints.
However, it more resembles assuming MImaximization and finding the assignments that are OT to the model probability.
In contrast, MIRAdirectly maximizes the MI without artificial constraints.
While SwAV performs better than SeLa inmost self-supervised benchmarks, we empirically verify that MIRA improves over SwAV in variousdownstream tasks.
Experiments
In this section, we evaluate the representation quality learned via MIRA.
We first provide theimplementation details of our representation learning with MIRA (Sec. 4.1).
We present our mainresults on linear, k -NN, semi-supervised learning, and transfer learning benchmarks in comparison toother self-supervised baselines (Sec. 4.2).
Finally, we conduct an analysis of MIRA (Sec. 4.3).
4.1 Implementation details
We mostly follow the implementation details from our baselines [8, 9, 50].
More training detailsabout evaluation procedures and analysis are described in the Appendix.
Architecture The training model (network) consists of an encoder, projection head, and classification head.
We use a widely used ResNet50 [28] as our base encoder and use the output ofaverage-pooled 2048 d embedding as our representation for both representation training and downstream evaluations.
The projection head is a 3 -layer fully connected MLP of sizes [2048 , 2048 , d ] ;hidden layers are followed by batch normalization [32] and ReLU.
The classification head is used topredict the pseudo-labels; is composed of an L2-normalization layer and a weight-normalized layerof the size d × K as in [8, 9].
We use d = 256 and K = 3000 .
[Table]
[Table]
Training details We train our model on the training set of the ImageNet-1k ILSVRC-2012 dataset[18] without using class labels.
We use the same data augmentation scheme (color jittering, Gaussianblur, and solarization) and multi-crop strategy (two 224 × 224 and six 96 × 96 ) used in [9].
We usea batch size of 4096 and employ the LARS optimizer [49] with a weight decay of 10 − 6 .
We uselinearly scaled learning rate of lr × batch size / 256 [26] with a base learning rate of 0 . 3 .
3 We adjustthe learning rate with 10 epochs of a linear warmup followed by cosine scheduling.
We also use EMAnetwork by default.
When the EMA is used, we set the momentum update parameter to start from0.99 and increase to 1 by cosine scheduling.
We use temperature scales of τ s = 0 . 1 , τ t = 0 .
225 withtrade-off coefficient β
= 2 / 3 .
We assign soft pseudo-labels after 30 steps of the fixed point iteration.
We further discuss this choice in Sec. 4.3.
Otherwise stated, we use the encoder model trained by
MIRA with 400 epochs training and multi-crop augmentations for the evaluations in this section.
4.2 Main results
Linear evaluation Tables 1 and 2 report linear evaluation results.
We follow the linear evaluationsettings in [27, 10].
We train a linear classifier on the top of the frozen trained backbone with thelabeled training set of ImageNet.
We train for 100 epochs using a LARS optimizer with a batchsize of 1024.
We use a base learning rate of 0 . 1 and adjust the learning rate by cosine annealingschedule.
We apply random-resized-crop and horizontal flip augmentations for training.
We evaluatethe representation quality by the linear classifier’s performance on the validation set of ImageNet.
Table 1 shows linear evaluation performance in top-1 accuracy for different un-/self-supervisedrepresentation training epochs.
We train and evaluate MIRA with and without multi-crop augmentations.
With multi-crop augmentations, MIRA consistently outperforms baselines while achieving75.5% top-1 accuracy with only 400 epochs of training.
We also report that 200 epochs of trainingwith MIRA can outperform the 800 epochs results of other baselines that don’t use multi-crops.
Without multi-crop augmentations, MIRA is comparable to MoCo-v3 [14] and performs slightlyworse than BYOL [27].
However, MIRA performs the best among the clustering-based [6, 8] andinformation-driven [50, 25] methods.
In Table 2, we compare MIRA to other self-supervised methods with the final performance.
MIRAachieves the state-of-the-art performance on linear evaluation of ImageNet with only 400 epochs oftraining.
While TWIST can achieve similar performance to MIRA within 450 epochs, they require anextra training stage with self-labeling ; without it, they achieve 74.1% accuracy with 800 epochs oftraining.
In contrast, MIRA doesn’t require additional training.
[Table]
we evaluate the
[Table] [Table]
[Table]
Semi-supervised learning In Table 4, we evaluate the trained model on the semi-supervisedlearning benchmark of ImageNet.
Following the evaluation protocol in [27, 10], we add a linearclassifier on top of the trained backbone and fine-tune the model with ImageNet 1% and 10% subsets.
We report top-1 and top-5 accuracies on the validation set of ImageNet.
For the 1% subset, MIRAoutperforms the baselines; both the top-1 and top-5 accuracies achieve the best.
For the 10% subset,
MIRA is comparable to SwAV [8].k -NN evaluation We further evaluate the quality of learned representation via the nearest neighborclassifier.
We follow the procedures of [9].
First, representations of the labeled training data are stored.
Then, the label of the new validation data is predicted with the majority vote of k -nearest storedrepresentations.
We use the same evaluation settings in [9] with 20 nearest neighbors, temperaturescaling 4 of 0 .
07 , and cosine distance metric.
Table 3 shows the k-NN classification accuracies on the validation set of ImageNet.
We use 1/10/100%subsets of ImageNet training dataset to produce labeled representations.
For ImageNet 1%
and 10%subsets, we use the same subsets of semi-supervised learning evaluation.
The results show that ourmethod achieves state-of-the-art k -NN evaluation performance with ResNet50.
To be more specific,our method outperforms the previous state-of-the-art DINO [9] on 100% and 10% subset evaluationby 1.2 ∼ 1.4%.
We note that BarlowTwins [50], a method also motivated by information-maximization,shows a strong performance of 47.7% in the 1% subset evaluation.
Transfer learning We further evaluate the representation learned by MIRA on the transfer learningbenchmark following [21] that includes FGVC aircraft [39], Caltech-101 [24], Standford Cars [34],
CIFAR-10/100 [35], DTD [15], Oxford 102 Flowers [40], Food-101 [5], Oxford-IIIT Pets [41],
SUN397 [45], and Pascal VOC2007 [23] datasets.
We follow the linear evaluation procedure in [21]that fits a multinomial logistic regression model on the extracted representations of 2048d from thetrained backbone.
First, we perform a hyperparameter search on the L2-normalization coefficient ofthe logistic regression model; then the final performance is evaluated on the model that is retrainedon all training and validation sets with the found coefficient.
[Figure]
Table 5 shows the performance of our algorithm compared to other baselines in 11 datasets.
MIRAoutperforms supervised representation on 10 out of 11 datasets.
Compared to the other self-supervisedmethods, representation learned by MIRA achieves the best performance in 9 out of 11 datasetsand improves 0.9% over the second-best baseline method on average.
The results confirm that therepresentation trained with MIRA has a strong generalization ability for classification.
4.3 Analysis
Convergence of pseudo-label assignment We study the speed of convergence of the proposedfixed-point iteration in MIRA.
We also experiment with the Sinkhorn-Knopp (SK) algorithm
[16]used in SwAV [8] as a baseline.
We experiment with both methods on the ImageNet with a batchsize of 512.
We observe the converging behavior with the pre-trained models from MIRA and SwAV.
Results are averaged over 1000 randomly sampled batches.
Figure 2 shows the result of the converging behavior of our method (blue) and SK algorithm
(yellow)on trained models of MIRA (left) and SwAV (right) .
Our fixed-point iteration converges faster thanthe SK algorithm in both pre-trained models.
Especially our default setting of 30 steps of updates aresufficent for our fixed point iteration.
Multi-crop and EMA Table 6 reports an ablation study on how EMA and multi-crop augmentationaffects our representation quality.
We train a model for 200 epochs in the settings with and without
EMA or Multi-crop.
Both EMA and Multi-crop augmentations greatly improve the linear evaluationperformance as in [8, 9].
We take a further comparison with baselines that are in the same setups.
Withthe only difference in the pseudo-labeling algorithm, our method outperforms SwAV [8] by 1.3% intop-1 accuracy.
While DINO [9] also uses both multi-crop and EMA, our method outperforms DINOwith fewer training epochs.
The results validate the effectiveness of our pseudo-labeling algorithm. These results validate the effectiveness of our pseudo-labeling algorithm.
[Table]
Scalability We further validate MIRA’s scalability on the small-and medium-scaled datasets.
ResNet-18 is used as a base encoder throughout the experiments.
While changing the base encoder,other architectural details remain the same as in ImageNet-1k.
We do not apply multi-crop augmentations while using the EMA.
We use image sizes of 32 × 32 and 256 × 256 for small and mediumdatasets, respectively.
Following the procedures in [17], we report the linear evaluation performance on the validation set.
More experimental details about the optimizer, batch size, augmentations, etc.,are provided in the Appendix.
[Table] [Table]
The results are in Table 7.
In CIFAR-10 and ImageNet-100, our method outperforms other selfsupervised baselines by 0.4% and 0.7% in top-1 accuracy, respectively.
For CIFAR-100, our methodis comparable to the best performing baseline–BarlowTwins; MIRA performs better in top-5 accuracy.
Training with small batch Throughout the experiments in Sec.
4.2, we use a batch size of 4096.
While such batch size is commonly used in self-supervised methods, large amounts of GPU memoryare required; hence limiting the accessibility.
In Table 8, we test our method with a smaller batchsize of 512 that can be used in an 8 GPU machine with 96GB memory.
In this setting, we use the
SGD optimizer with a weight decay of 10 − 4 .
We also test the robustness of pseudo-labeling with the
Sinkhorn-Knopp algorithm in SwAV [8] reproduced by us and compare the results.
We report a top-1 linear evaluation performance of both methods after 100 epochs of training.
In theresult, the performance gap between our method and SwAV is amplified from 2.9% to 6% in thereduced batch size of 512.
One possible explanation is that since SwAV is based on the equipartitionconstraint, the performance of SwAV harshly degrades when the batch size is not enough to matchthe number of clusters.
[Table] [Table]
[Table]
Discussion
Conclusion This paper proposes the mutual information maximization inspired pseudo-labelingalgorithm MIRA.
We formulate pseudo-labeling into an optimization problem and solve it in aprincipled way.
We apply MIRA to representation learning and demonstrate its effectiveness inself-supervised learning benchmarks.
We hope that our simple yet theoretically guaranteed approachto information maximization will guide many future applications
[Table]
Limitation and negative social impact
Our information maximization perspective pseudo-labelingseems applicable to various tasks and domains, e.g., semi-supervised training [36].
We validatethe effectiveness only in self-supervised visual representation learning.
Furthermore, despite ourimproved training efficiency, the self-supervised learning methods still require a huge amount ofcomputations compared to supervised learning.
Such computational requirements may accelerate theenvironmental problems of global warming.