Unsupervised Visual Representation Learning via Mutual Information Regularized Assignment
Abstract
This paper proposes Mutual Information Regularized Assignment (MIRA), a pseudo-labeling algorithm for unsupervised representation learning inspired by information maximization.
We formulate online pseudo-labeling as an optimization problem 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-point iteration method and prove its convergence to the optimal solution.
In contrast to baselines, MIRA combined with pseudo-label prediction enables a simple yet effective clustering-based representation learning without incorporating extra training techniques or artificial constraints such as sampling strategy, equipartition constraints, etc.
With relatively small training epochs, representation learned by MIRA achieves state-of-the-art performance on various downstream tasks, including the linear/ k -NN evaluation and transfer learning.
Especially, with only 400 epochs, our method applied to ImageNet dataset with ResNet-50 architecture achieves 75.6% linear evaluation accuracy.
Our implementation is available at https://github.com/movinghoon/mira .
Introduction
There has been a growing interest in using a large-scale dataset to build powerful machine learning models [45].
Self-supervised learning (SSL), which aims to learn a useful representation without labels, is suitable for this trend; it is actively studied in the fields of natural language processing [19, 20] and computer vision [10, 30].
In the vision domain, recent SSL methods commonly use data augmentations and induce their visual representation to be augmentation-invariant.
They have achieved state-of-the-art performance surpassing supervised representation in a variety of visual tasks, including semi-supervised learning [8, 53], transfer learning [22] and object detection [13].
Meanwhile, a line of work uses clustering for un-/self-supervised representation learning.
They explicitly assign pseudo-labels to embedded representation via clustering, and the model is thereby trained 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 [51, 49, 4, 32] have shown encouraging results in small-scaled settings; Caron et al.
[6] show that it can also be applied to the large-scaled dataset or even to a non-curated dataset [7].
Recently, several works [2, 8, 39] have adopted the philosophy of augmentation invariance and achieved strong empirical results.
They typically assign pseudo-labels using augmented views while predicting the labels by looking at other differently augmented views.
Despite its conceptual simplicity, a naive application of clustering to representation learning is hard to achieve, especially when training with large-scale datasets.
This is because clustering-based methods are prone to collapse, i.e., all samples are assigned to a single cluster; hence, recent methods heavily rely on extra training techniques or artificial constraints, such as pre-training [50], sampling strategy [6], equipartition constraints [2, 8], to avoid collapsing.
However, it is unclear if these additions are appropriate or how such components will affect the representation quality.
In this paper, we propose Mutual Information Regularized Assignment (MIRA), a pseudo-labeling algorithm that enables clustering-based SSL without any artificial constraints or extra training techniques.
MIRA is designed to follow the infomax principle [40] and the intuition that good labels are something that can reduce most of the uncertainty about the data .
Our method assigns a pseudo-label in a principled way by constructing an optimization problem.
For a given training model that predicts pseudo-labels, the optimization problem finds a solution that maximizes the mutual information (MI) between the pseudo-labels and data while considering the model probability.
We formulate the problem as a convex optimization problem and derive the necessary and sufficient condition of solution with the Karush-Kuhn-Tucker (KKT) condition.
This solution can be achieved by fixed-point iteration that we prove the convergence.
We remark that MIRA does not require any form of extra training techniques or artificial constraints, e.g., equipartition constraints.
We apply MIRA to clustering-based representation learning and verify the representation quality on several standard self-supervised learning benchmarks.
We demonstrate its state-of-the-art performance on linear/ k -NN evaluation, semi-supervised learning, and transfer learning benchmark.
We further experiment 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 algorithm based on mutual information.
Our method does not require extra training techniques or artificial constraints.
• We apply MIRA to clustering-based representation learning, showing comparable performance against the state-of-the-art methods with half of the training epochs.
Specifically, MIRA achieves 75.6% top-1 accuracy on ImageNet linear evaluation with only 400 epochs of training and the best performance in 9 out of 11 datasets in transfer learning.
• Representation by MIRA also consistently improves over other information-based SSL methods.
Especially our method without multi-crop augmentation achieves 74.1% top-1 accuracy and outperforms BarlowTwins [53], a baseline information maximization-based self-supervised method.
Related works
Self-supervised learning SSL methods are designed to learn the representation by solving pretext tasks, and recent state-of-the-art methods encourage their learned representations to be augmentation invariant.
They are based on various pretext tasks: instance discrimination [10, 11, 13, 14], metric learning [28, 12], self-training [54, 9], and clustering [2, 6, 8]; only a few account for encoding the semantic structure of data.
While some works [47, 21, 35] consider the nearest neighbors in the latent space, our method belongs to the clustering-based SSL method that flexibly accounts for inter-data similarity.
Meanwhile, many SSL methods are prone to collapsing into a trivial solution where every representation is mapped into a constant vector.
Various schemes and mechanisms are suggested to address this, e.g., the asymmetric structure, redundancy reduction, etc.
We will review more relevant works 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 sampling strategy to sample 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 a uniform pseudo-label prior, and PCL [39] employs concentration scaling.
DINO [9] and ReSSL [54] address collapsing by specific combinations of implementation details, i.e., centering and scaling with an exponential moving average network; their mechanism for preventing collapse is unclear.
In this work, we show our method can naturally avoid collapsing without any of these assumptions or training techniques.
We achieve results better than baselines with a simple but novel
[Figure]
[Figure]
fixed point iteration ① pseudo-labeling ② model learning pseudo-labels swapped prediction
[Figure]
information regularization algorithm.
We take a more detailed comparison with SeLa and SwAV after explaining our method in Sec. 3.3.
Information maximization Information maximization is a principal approach to learn representation and to avoid collapse.
DeepInfoMax [31] propose the MI maximization between the local and global views for representation learning; the existence of negative pairs prevents training toward the trivial solution.
BarlowTwins [53] and W-MSE [23] address the collapsing with redundancy reduction that indirectly maximizes the content information of embedding vectors.
Among clustering-based approaches, IIC [34] maximizes the MI between the embedding codes to enable representation learning; similar to ours, TWIST [26] proposes combining the MI between the data and class prediction as a negative loss term with an augmentation invariance 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, the direct optimization of MI terms by updating model parameters often leads to a sub-optimal solution [26]; TWIST copes with this issue by appending the normalization layer before softmax and introducing an additional self-labeling stage.
In contrast, MIRA addresses the difficulty of MI maximization 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 that alternates 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 formulate an optimization problem for online clustering that assigns soft pseudo-labels to mini-batch samples (Sec. 3.1).
The problem minimizes the KL divergence between the model prediction (probability) and pseudo-label while maximizing the mutual information between the data and pseudo-label.
We propose an iterative fixed point method to solve the optimization problem (Sec. 3.2).
For the model training, we use the swapped prediction loss [8] (Sec. 3.3).
3.1 Mutual information regularized pseudo-labeling
We have a classification model 1 f θ parametrized by θ that outputs K -dimensional logit f θ ( x ) ∈ R K for an image x , where K is a predefined number of clusters.
The model probability p of an image x is then given by the temperature τ t scaled output of the model— p := softmax ( f θ ( x ) /τ t ) —as in
Caron et al.
For a mini-batch of input images X
= { x i } Bi
=1 , we denote the set of model probabilities P = { p i } Bi =1 ⊂ R K .
In our pseudo-labeling, for the given model probabilities P , we want to assign pseudo-labels W ∗ =
{ w ∗ i } Bi =1 that will be used for training the model by predicting them.
We argue that such pseudo-labels should maximize the mutual information (MI) between themselves and data while accounting for the model probabilities P .
Let B ∈ { 1 , ..., B } and Y W ∈ { 1 , ..., K } be the random variables associated with the data index in mini-batch and labels by probability distributions W = { w i } Bi =1 , respectively.
Our online pseudo-label (cluster) assignment is determined by solving the following optimization problem:
[Equation]
where ∆ K := { w ∈ R K + | w ⊺ 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 makes pseudo-labels to be based on the model probability p and (2) MI term between the pseudo-labels and data to induce more information about data into the pseudo-labels.
By combining these two terms, we provide a refined 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 estimates of the entropy ˆ H ( Y W |B ) and marginal entropy ˆ H ( Y W ) in Eq.
We get:
ˆ I ( Y W ; B ) =
[Equation]
H
[Equation]
Y W )
[Equation]
where w j = 1 B (cid:80) Bi
=1 w ij is the marginal probability of a cluster j with W .
In practice
[Equation]
we find the optimal point W ∗ of the optimization problem Eq.
4 for pseudo-labeling
[Equation]
3.2 Solving strategy
To solve efficiently
[Equation]
we propose a fixed point iteration that guarantees convergence to the unique optimal solution W ∗ of our optimization problem.
The method is based on the following proposition.
Proposition 1.
For β ∈ [0 , 1) , the problem Eq.
4 is a strictly convex optimization problem with a unique optimal point W ∗ that satisfies the following necessary and sufficient condition:
[Equation]
The proposition 1 is driven by applying the Karush-Kuhn-Tucker (KKT) conditions to the optimization problem Eq. 4.
By substituting the necessary and sufficient condition (Eq. 5) into the marginal probability w j = 1 B (cid:80) Bi
=1 w ij , we get the necessary and sufficient condition for w ∗ :
[Equation]
Based on Eq.
6, we propose the update rule for { u ( n ) j } Kj =1 ⊂ R + using the fixed point iteration as follows:
[Equation]
where u ( n ) j converges to w ∗ j as n → ∞ since the necessary and sufficient condition (Eq. 6) is satisfied at the convergence.
We can easily get w ∗ ij by Eq.
5 when the optimal marginal probability w ∗ j is given.
The proof of the proposition and convergence is in the Appendix.
By using the iterative updates of Eq.
7, we get our desirable pseudo-labels W ∗ .
This requires a few lines of code that are simple to implement.
We observe that a few steps of iteration are enough for training.
This is supported by the convergence analysis in Sec. 4.3.
We use this fixed point iteration for 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 method is applied to representation learning.
We integrate the computed pseudo-labels with swapped prediction loss [8] to train the model; the model is trained to predict the pseudo-labels from the augmented views using the other views.
Specifically, given two mini-batches of differently augmented views X ( m ) =
{ x ( m ) i } Bi =1 , m ∈ { 1 , 2 } , MIRA independently assigns pseudo-labels U ( m ) =
{ u ( m ) i } Bi =1 for each mini-batch.
In parallel, model f θ provides the temperature τ s scaled softmax predictions 2 Q ( m ) =
{ q ( m ) i } Bi
=1 of each mini-batch.
The swapped prediction loss is given as follows:
[Equation]
This loss function (Eq.
8) is minimized with respect to the parameters θ of the model f θ used to produce the predictions Q (1) , Q (2) .
For more detailed information about swapped prediction loss, please refer to Caron et al.
The pseudo-code of MIRA for representation learning with Eq.
8 is provided in the Appendix.
In the following experiments, we verify the effectiveness of MIRA for a representation learning purpose.
We note that MIRA can integrate recently suggested self-supervised learning components, such as exponential moving average (EMA) or multi-crop (MC) augmentation strategy following the baselines [14, 8, 9].
For convenience, in the rest of this paper, we call the representation learning with MIRA also as MIRA.
We discuss some further details as follows:
Preventing collapse The MI term in Eq.
4 takes a minimum value when collapsing happens.
MIRA naturally avoids collapsed solution via penalizing assignment that exhibits low MI.
Specifically, unless starting from the collapsed state, MIRA finds MI-maximizing points around the model prediction; it will not choose collapsed pseudo-labels.
Hence, the iterative training to predict such labels will not collapse whenever the prediction of pseudo-labels is achievable.
Our empirical results verify that
MIRA does not require extra training techniques or artificial constraints to address collapsing.
Comparison to SwAV and SeLa SeLa [2] and SwAV [8] formulate their pseudo-labeling process into optimization problems, i.e., optimal transport (OT) problem, and solve it iteratively with SinkhornKnopp (SK) algorithm
To avoid collapse and apply the SK algorithm, they assume the equipartition of data into clusters.
Mathematically, the difference to MIRA is in how to deal with the marginal entropy .
SeLa and SwAV constrain the marginal entropy to maximum value–equipartition while MIRA decides it by MI regularization 3 .
Asano et al.
[2] argue that their pseudo-labels with the OT problem maximize the MI between labels and data indices under the equipartition constraint.
[Table]
[Table] [Table]
However, it more resembles assuming MI maximization and then finding the cluster assignments that are optimal transport to the model prediction.
In contrast, MIRA directly maximizes the MI by regularization without artificial constraints.
While SwAV performs better than SeLa in most self-supervised benchmarks, we verify that MIRA improves over SwAV in various downstream tasks.
Experiments
In this section, we evaluate the representation quality learned via MIRA.
We first provide the implementation details of our representation learning with MIRA (Sec. 4.1).
We present our main results on linear, k -NN, semi-supervised learning, and transfer learning benchmarks in comparison to other 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 of representation learning from our baselines [8, 9].
More training details about evaluation procedures and analysis are described in the Appendix.
Architecture The training model (network) consists of an encoder, projection head, and classification head as in Caron et al.
We employ a widely used ResNet50 [29] as our base encoder and use the output of average-pooled 2048 d embedding for representation training and downstream evaluations.
The projection head is a 3 -layer fully connected MLP of sizes [2048 , 2048 , d ] ; each hidden layer is followed by batch normalization [33] and ReLU.
The classification head is used to predict the pseudo-labels; it is composed of an L2-normalization layer and a weight-normalization layer of the size d × K .
We use d = 256 and K = 3000 .
Training details We train our model on the training set of the ILSVRC-2012 ImageNet-1k dataset [18] without using class labels.
We use the same data augmentation scheme (color jittering, Gaussian blur, and solarization) and multi-crop strategy (two 224 × 224 and six 96 × 96 ) used in Caron et al.
We use a batch size of 4096 and employ the LARS optimizer [52] with a weight decay of 10 − 6 .
We use linearly scaled learning rate of lr × batch size / 256 [27] with a base learning rate of 0 . 3 .
4 We adjust the learning rate with 10 epochs of a linear warmup followed by cosine scheduling.
We also use an exponential moving average (EMA) network by default.
When
[Table] [Table]
using EMA
[Table]
we set the momentum update parameter to start from 0.99 and increase it to 1 by cosine scheduling.
We use temperature scales of τ s = 0 . 1 , τ t = 0 .
225 with a trade-off coefficient β
= 2 / 3 .
We obtain soft pseudo-labels by 30 steps of the fixed point iteration.
We further discuss this choice in Sec. 4.3.
Otherwise stated
[Table]
we use the encoder model trained for 400 epochs with multi-crop augmentations for the downstream task evaluations in this section.
4.2 Main results
Linear evaluation Tables 1 and 2 report linear evaluation results.
We follow the linear evaluation settings in [28, 10].
We train a linear classifier on the top of the frozen trained backbone with the labeled training set of ImageNet.
We train for 100 epochs using an SGD optimizer with a batch size of 1024.
We choose a learning rate 5 with a local validation set in the ImageNet train dataset and adjust the learning rate by cosine annealing schedule.
We apply random-resized-crop and horizontal flip augmentations for training.
We evaluate the 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 training epochs.
We train and report the performance of MIRA in two settings, with and without multi-crop augmentations.
With multi-crop augmentations, MIRA consistently outperforms baselines while achieving 75.6% top1 accuracy with only 400 epochs of training.
We also report that 200 epochs of training with MIRA can outperform the 800 epochs results of other baselines that do not use multi-crops.
Without multicrop augmentations, MIRA performs slightly worse than BYOL [28].
However, MIRA performs the best among the clustering-based [6, 8] and information-driven [53, 26] methods.
In Table 2, we compare MIRA to other self-supervised methods with the final performance.
MIRA achieves state-of-the-art performance on linear evaluation of ImageNet with only 400 epochs of training.
While TWIST [26] can achieve similar performance to MIRA within 450 epochs, they require an extra training stage with self-labeling ; without it, they achieve 74.1% accuracy withepochs of training.
In contrast, MIRA does not require additional training.
Semi-supervised learning In Table 4, we evaluate the trained model on the semi-supervised learning benchmark of ImageNet.
Following the evaluation protocol in [28, 10], we add a linear classifier 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, MIRA outperforms 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 evaluate the quality of learned representation via the nearest neighbor classifier.
We follow the evaluation procedure of Caron et al.
We first store the representations of labeled training data; then, we predict the test data by the majority vote of the k -nearest stored representations.
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.
We
[Table]
[Figure]
use the same hyperparameter settings in Caron et al.
[9] with k = 20 nearest neighbors, temperature scaling 6 of 0 .
07 , and cosine distance metric.
Table 3 shows the classification accuracies on the validation set of ImageNet.
The results show that our method achieves the best evaluation performance.
Specifically, our method outperforms the previous state-of-the-art DINO [9] on 100% and 10% subset evaluation by 1.2 ∼ 1.4%.
We note that BarlowTwins [53], a method also motivated by information maximization, shows a strong performance of 47.7% in the 1% subset evaluation.
Transfer learning We evaluate the representation learned by MIRA on the transfer learning benchmark with FGVC aircraft [41], Caltech-101 [25], Standford Cars [36], CIFAR-10/100 [37],
DTD [15], Oxford 102 Flowers [42], Food-101 [5], Oxford-IIIT Pets [43], SUN397 [48], and Pascal
VOC2007 [24] datasets.
We follow the linear evaluation procedure from Ericsson et al.
[22] that fits a multinomial logistic regression classifier on the extracted representations of 2048d from the trained encoder.
We perform a hyperparameter search on the L2-normalization coefficient of the logistic regression model.
The final performance is evaluated on the model that is retrained on all training and validation sets with the found coefficient.
Table 5 shows the performance of our algorithm compared to other baselines in 11 datasets.
MIRA outperforms supervised representation on 10 out of 11 datasets.
Compared to other self-supervised methods, representations learned from MIRA achieved the best performance in 9 out of 11 datasets, with an average improvement of 0.9% over the second-best baseline method.
The results confirm that the representation trained with MIRA has a strong generalization ability for classification.
4.3 Analysis
Convergence of pseudo-label assignment We study the convergence speed of the proposed fixed point iteration in MIRA.
We also experiment with the Sinkhorn-Knopp (SK) algorithm
[16] used in
SwAV [8] as a baseline.
Both methods have experimented with a batch size 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 than the SK algorithm in both pre-trained models, and our default setting of 30 steps of updates is sufficient for convergence.
While we use 30 steps of updates to follow the theoretical motivation in our main experiments, we observe that choosing a small number of iterations is possible in practice.
Multi-crop and EMA Table 6 reports an ablation study on how EMA and multi-crop augmentations affect our representation quality.
We train a model for 200 epochs in the settings with and without
EMA or Multi-crop augmentations.
EMA and Multi-crop augmentations greatly improve the linear evaluation performance as in Caron et al.
We take a further comparison with baselines that are in the same setups.
While the only difference in the pseudo-labeling, our method outperforms
SwAV [8] by 1.3% in top-1 accuracy.
DINO [9] uses both multi-crop and EMA, our method outperforms DINO with fewer training epochs.
The results validate the effectiveness of MIRA.
[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 medium datasets, respectively.
Following the procedures in da Costa et al.
[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]
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 method is 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 learning, a large amount of GPU memory is required, limiting the accessibility.
In Table 8, we test our method with a smaller batch size 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] and compare the results.
We report a top-1 linear evaluation performance of both methods after 100 epochs of training.
In the result, the performance gap between our method and SwAV is amplified from 2.9% to 6% in the reduced batch size of 512.
One possible explanation is that since SwAV is based on the equipartition constraint, the performance of SwAV harshly degrades when the batch size is not enough to match the number of clusters.
[Table]
Discussion
Conclusion This paper proposes the mutual information maximization perspective pseudo-labeling algorithm MIRA.
We formulate online pseudo-labeling into a convex optimization problem with mutual information regularization and solve it in a principled way.
We apply MIRA for representation learning and demonstrate the effectiveness in various self-supervised learning benchmarks.
We hope our simple yet theoretically guaranteed approach to mutual information maximization will guide many future applications
[Table] Limitation and negative social impact Our mutual information regularization with optimization strategy seems applicable to various tasks and domains, e.g., semi-supervised training [38].
However, we validate the effectiveness only in self-supervised visual representation learning.
We note that the performance of MIRA for non-classification downstream tasks 8 , e.g., detection, is not as dominating as in the classification tasks.
In these tasks, methods that consider local or pixel-wise information achieve superior performance; incorporating such formulation into clustering-based methods seems to be an important future direction.
Furthermore, despite our improved training efficiency, the selfsupervised learning methods still require massive computation compared to supervised approaches.
Such computational requirements may accelerate the environmental problems of global warming.
Acknowledgments
This work was supported by Institute of Information & communications Technology Planning &
Evaluation (IITP) grant funded by the Korea government(MSIT)
(No.2022-0-00926).
We thank Suyoung Lee, Minguk Jang, Dong Geun Shin, and anonymous reviewers for constructive feedbacks and discussions to improve our paper.
Part of the experiments for this work was done at the Fundamental Research Lab of LG AI Research.
Checklist
For all authors... (a) Do the main claims made in the abstract and introduction accurately reflect the paper’s contributions and scope?
[Yes] (b) Did you describe the limitations of your work?
[Yes] See Section 5 DiscussionLimitation and negative social impact (c) Did you discuss any potential negative societal impacts of your work?
[Yes] See Section 5 Discussion - Limitation and negative social impact (d) Have you read the ethics review guidelines and ensured that your paper conforms to them?
[Yes] 2.
If you are including theoretical results... (a) Did you state the full set of assumptions of all theoretical results?
[Yes] in the Appendix (b) Did you include complete proofs of all theoretical results?
[Yes] in the Appendix 3.
If you ran experiments... (a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] in supplemental material (b) Did you specify all the training details (e.g., data splits, hyperparameters, how they were chosen)?
[Yes] in Appendix (c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)?
[Yes] in Appendix (d) Did you include the total amount of compute and the type of resources used (e.g., type of GPUs, internal cluster, or cloud provider)?
[Yes] 4.
If you are using existing assets (e.g., code, data, models) or curating/releasing new assets... (a) If your work uses existing assets, did you cite the creators?
[Yes] (b) Did you mention the license of the assets?
[Yes] in Appendix (c) Did you include any new assets either in the supplemental material or as a URL?
[Yes] code in supplemental material.
(d) Did you discuss whether and how consent was obtained from people whose data you’re using/curating?
[N/A] (e) Did you discuss whether the data you are using/curating contains personally identifiable information or offensive content?
[N/A] 5.
If you used crowdsourcing or conducted research with human subjects... (a) Did you include the full text of instructions given to participants and screenshots, if applicable?
[N/A] (b) Did you describe any potential participant risks, with links to Institutional Review Board (IRB) approvals, if applicable?
[N/A] (c) Did you include the estimated hourly wage paid to participants and the total amount spent on participant compensation?
[N/A]
A Appendix A.1 PyTorch pseudo-code for MIRA
[Table]
A.2 Proof of proposition
In this subsection, we derive the necessary and sufficient condition in proposition 1.
Denote B, K be some natural numbers.
We first prove the strict convexity of the optimization function f : R BK × 1 + → R :
[Equation]
log w ij + β k (cid:88) j = w j log w j , (9) where w j = 1 B (cid:80) Bi
=1 w ij .
Lemma 1.
For x ∈ R N × 1 + , s ( x ) = (cid:80)
Ni x i log x
i is a strictly convex function of x .
Proof.
Since the Hessian of s is a diagonal matrix with positive elements ∇ 2 x s ( x ) i,i =
1 /x i , s is a strictly convex function.
Corollary 1.
For W ∈ R BK × 1 + , f ( W ) is a strictly convex function of W .
Proof.
Note that
f 1 ( W ) = − 1 B (cid:80) Bi
=1 (cid:80)
Kj =1 w ij log p ij is a affine transformation of W , 2.
f 2 ( W ) = 1 − βB (cid:80) Bi
=1 (cid:80)
Kj =1 w ij log w ij is a strictly convex function of W by lemma 1, 3.
f 3 ( W ) =
β (cid:80) kj =1 w j log w j is a convex function of W since f 3 is the composition of the strictly convex
function s and affine
transformation
w
j
= 1 B (cid:80) Bi =1 w ij . The function f ( W ) = f 1 ( W ) + f 2 ( W ) + f 3 ( W ) that is the sum of the convex and strictly convex terms becomes strictly convex.
We show that the optimization function f is strictly convex.
The optimization problem (Eq. 4) is defined on the convex set, i.e., W ⊂ ∆ K :
= { x ∈ R K + | x ⊺ 1 K = 1 } ; hence, the problem is a strictly convex optimization problem 9 .
By the strict convexity, our optimization problem has a unique optimal point that has Karush–Kuhn–Tucker (KKT) conditions as the necessary and sufficient condition for optimality.
By using the KKT condition, we derive our necessary and sufficient condition Eq. 5.
Considering the constraint of the optimization domain, i.e., W ∈ ∆ K ⇔ ∀ i ∈ 1 : B, (cid:80) Kj =1 w ij = 1 , the Lagrangian function of our optimization problem becomes:
[Equation]
where λ is a Lagrange multipliers.
By the KKT conditions with (cid:80) Kj =1 w ij = 1 ,
[Equation]
proves the necessary and sufficient condition in proposition 1.
The Eq. 16 comes from (cid:80) j w ∗ ij = [Equation]
For w ∗ [Equation] this condition becomes:
[Equation]
A.3 Convergence of the fixed point iteration
In this subsection, we prove the convergence of our fixed point iteration (Eq. 7).
Denote B, K be some natural numbers.
Lemma 2.
For a , x ∈ R B × 1 + and β ∈ [0 , 1) , h a ( x ) := ( 1 B (cid:80) Bi
=1 a i x − 1 i ) − β is a concave function of x .
Proof.
Denote r [Equation]
[Equation]
for a simplicity of expression; hence h a ( x ) =
r a ( x ) − β .
The second order partial derivatives of h a ( x ) w.r.t.
x becomes:
[Equation]
For ω ∈ R B × 1 ,
[Equation]
The inequality in Eq.
26 comes from Cauchy–Schwartz inequality as bellows:
[Equation]
Since ∀ ω ∈ R B × 1 , ω ⊺ ∇ 2 x h a ( x ) ω < 0 , the Hessian of h a ( x ) is negative definite, h a ( x ) is concave.
Corollary 2.
For v ∈ R K × 1 +
[Equation]
a concave function of v .
Proof.
Since g j is a composition of the concave function h ( Q ⊺ ) j (by Lem. 2) and affine transformation with Q ; g j becomes concave.
We introduce the proposition from [44] that proves geometrical convergence of positive concave mapping.
Proposition 2 (Piotrowski and Cavalcante [44], Proposition 3) .
Let f : R M + → int ( R M + ) be a continuous and concave mapping w.r.t cone order with a fixed point x ∗ ∈ int ( R M + ) .
Then, the fixed point iteration of f , i.e., x n +1 = f
( x n ) , with x ∈ R M + converges geometrically to x ∗ with a factor c ∈ [0 , 1) .
By the proposition 2 and corollary 2, we prove the convergence of the following fixed point iteration:
Proposition 3.
For β ∈ [0 , 1) , Q ∈ R B × K + , and for any initialization v (0)
∈ R K × 1 + , the fixed point iteration:
[Equation]
converges to the fixed point.
Proof.
We can see that output of the fixed point iteration g j ( v ( n ) ; Q ) is always positive.
By corollary 2, g ( v ( n ) ; Q ) is a concave mapping.
Therefore, by proposition 2, the fixed point iteration converges to the fixed point.
Finally, by bijective change of variables
[Equation]
we derive our fixed point iteration in Eq.
Throughout updates, Eq.
32 continuously satisfies the relation v ( n ) j = (cid:0) u ( n ) j (cid:1) − β/ (1 − β ) to the fixed point iteration by Eq. 29; hence, the fixed point iteration by Eq.
32 also converges to the fixed point.
Furthermore, the fixed point by Eq.
32 satisfies the necessary and sufficient condition (Eq.19); is equal to w ∗ .
A.4 Implementation & evaluation details
For the transfer learning evaluation, we use the open-source implementation of Ericsson et al.
[22] at https://github.com/linusericsson/ssl-transfer .
For the k-NN evaluation, we refer to the official implementation of Caron et al.
[9] at https://github.com/facebookresearch/dino .
For the CIFAR-10/100 and ImageNet-100 training & evaluations, we refer to the implementations of da Costa et al.
[17] at https://github.com/vturrisi/solo-learn .
A.4.1 ImageNet-1k
Representation learning We describe more detailed implementation details here.
When we pretrain with EMA, we use the EMA backbone for downstream evaluations.
We exclude BN parameters and biases from weight decay in the LARS optimizer.
For 100/200 epochs training, we reduce the warm-up epochs to 5.
Similar to temperature annealing of Caron et al.
[9], we apply cosine annealing on β from 0 .
7 to 2 / 3 .
For hyper-parameters tuning, we use ImageNet-100 tuned hyper-parameters ( τ t = 0 .
225 , β = 2 / 3 ) accross all ImageNet-1k, ImageNet-100, and CIFAR-10/100 datasets.
For other settings, including the projector design, number of clusters, and classification (prototype) head, we follow SwAV implementation [8, 12].
Semi-supervised learning We explain our semi-supervised learning evaluation training.
We train 20 epochs with a batch size of 256 .
We employ an SGD optimizer with a cosine learning rate schedule.
For 1 % training, we find that freezing the backbone encoder (with 0 .
learning rate) performs the best; we use learning rates of 0 .
4 for the linear classifier.
For 10 % training, we use learning rates of 0 .and 0 .
1 for the backbone encoder and linear classifier.
More details about the linear evaluation As the local validation set (in the training dataset), we use the 1% training split from SimCLR in the ImageNet train set as our local validation set.
We do not use regularization methods such as weight decay, gradient clipping, etc.
While we follow Grill et al.
[28] for linear evaluation in the main results (Sec.4.2), it requires much computation.
Thus, for the additional linear evaluations in the analysis section (Sec. 4.3) and Appendix, we fix to use a LARS optimizer with a learning rate of 0.1 since it performs well regardless of different models and matches the performance with the SGD optimizer. For reference, we report MIRA’s linear evaluation performance with the LARS optimizer
[Table] [Table] A.4.2 CIFAR-10/100 and ImageNet-  Representation learning For experiments on CIFAR-10/100 and ImageNet-100, we follow the settings in da Costa et al.
We use the same settings of ImageNet-1k experiments except for the base encoder, optimizer, augmentation scheme, and batch size.
We employ ResNet18 as a base encoder for experiments on CIFAR-10/100 and ImageNet-100.
For CIFAR-10/100 datasets, we change the first 7 × 7 convolution layer with stride 2 into 3 × 3 convolution layer with stride 1 .
We use the SGD optimizer with a weight decay of 10 − 4 .
We employ a linearly scaled learning rate with a base learning rate of 0 . 3 as in ImageNet-1k and scheduled the learning rate with 10 epochs of a linear warmup followed by cosine scheduling.
For CIFAR-10/100 datasets, we remove the GaussianBlur and adjust the minimum scale of RandomResizedCrop to 0 .
We use batch sizes of 256 and 512 for
CIFAR10/100 and ImageNet-100, respectively.
We train for 1000 and 400 epochs for CIFAR-10/ and ImageNet-100, respectively.
Linear evaluation Following the da Costa et al.
[17], we report online and offline linear evaluation results for CIFAR-10/100 and ImageNet-100, respectively.
For online linear evaluation on CIFAR10/100, we use the SGD optimizer with a learning rate of 0 . 1 .
We do not apply weight decay and use cosine scheduled the learning rate.
For the linear evaluation on ImageNet-100, we use the LARS optimizer with learning rate of 0 . 1 and batch size of 1024 .
A.5 Standard deviations
We report the standard deviations for our main results.
The linear, semi-supervised, and transfer learning evaluations are conducted four times with four different random seeds.
[Table] [Table]
A.6 Results on 800 epochs training
We report MIRA 800 epochs training results on the linear and k -NN evaluations.
To account for longer epoch training, we use the initial EMA momentum value of 0.996 as in Chen et al.
For the linear evaluation of 800 epochs training, we follow the evaluation protocols in the main results [28].
[Table]
A.7 Experiments on the detection and segmentation task
We test our method on detectionsegmentation of the COCO 2017 dataset with Masked R-CNN,
R50-C4 on a 2x scheduled setting.
We use the configuration from the MoCo official implementation.
MIRA performs better than the supervised baseline and is comparable to MoCo; it is not as dominating as in the classification tasks.
[Table]
A.8 Ablation studies A.8.1 Ablation study on the number of clusters
In table
[Table]
we study the effect of the number of clusters on the performance of linear and k -NN evaluation (with 1%/10%/100% labels).
We train MIRA on ImageNet-1k for 100 epochs without multi-crop augmentations while varying the number of clusters.
When the number of clusters is sufficiently large ( ≥ 3000 ), we observe no particular gain by varying the number of clusters.
This is consistent with the observation in SwAV [8].
[Table]
A.8.2 Ablation study on the number of the fixed point iteration
We report the 100 and 400 epochs pre-training linear evaluation results with 1 and 3 fixed point iterations.
In the 100 epochs training, the models with a smaller number of fixed point iterations, 1it and 3it, perform slightly better (+0.2%); in 400 epochs training, the model with more iterations, 30it, performs better (+0.2 ∼ 0.3%).
The convergence to the fixed point is not the primary factor for learning; however, choosing a sufficiently large number of fixed point iterations in longer epoch training seems reasonable.
[Table]