MOCA: Self-supervised Representation Learning by Predicting Masked Online Codebook Assignments
Authors: Spyros Gidaris
Andrei Bursuc
Oriane Siméoni
Antonín Vobecký
Nikos Komodakis
Matthieu Cord
Patrick Perez
Venue: Accepted by TMLR ()

Abstract: Self-supervised learning can be used for mitigating the greedy needs of Vision Transformer networks for very large fully-annotated datasets. Different classes of self-supervised learning offer representations with either good contextual reasoning properties, e.g., using masked image modeling strategies, or invariance to image perturbations, e.g., with contrastive methods. In this work, we propose a single-stage and standalone method, MOCA, which unifies both desired properties using novel mask-and-predict objectives defined with high-level features (instead of pixel-level details). Moreover, we show how to effectively employ both learning paradigms in a synergistic and computation-efficient way. Doing so, we achieve new state-of-the-art results on low-shot settings and strong experimental results in various evaluation protocols with a training that is at least 3 times faster than prior methods. We provide the implementation code at https://github.com/valeoai/MOCA.

Published in Transactions on Machine Learning Research (02/2024)
MOCA Self-supervised Representation Learning by Predicting Masked Online Codebook Assignments
Spyros Gidarist, Andrei Bursuc', Oriane Siméoni', Antonin Vobecky Nikos Komodakis Matthieu Cord', Patrick Pérez
Valeo.ai ai
Czech Institute of Informatics, Robotics and Cybernetics at the Czech Technicol University in Progue Czech Technical University in TrggueFacculle of Electricol Engineering
University of Crete
IACM-Forth
Archimedes/Atiee Athena RC
Correspondance: spyros. spyros.gidaris watssggiaaissssvll..coom
Reviewed on OpenReview: http:://pppppee..s..///////t//ta
Abstract
Self-supervised learning be used for mitigating the greedy needs of Vision Transformer networks for very large fully-annotated datasets. Different classes of self-supervised learning offer representations with either good contextual reasoning propertics, e.g,, using masked image modeling strategies, or invariance to image perturbations, e.g., with contrastive methods. In this work, we propose a single-stage and standalone method, MOCA, which unifies both desired properties using novel mask-and-predict objectives defined with high-level features (instead of pixel-level details). Moreover, we show how to effectively employ both learning paradigms in synergistic and computation--eficien way. Doing so, we achieve new state-of-the- results low-shot settings and strong experimental results in various evaluation protocols with a training that is at least 3 times faster than prior methods. We provide the implementation code at
1 Introduction
Sele-super representation learning for Vision Transformers (ViT) (Devlin et al., 2018; Vaswani al., 2017) has attracted significant amount of attention over the last years (Bao ct al., 2022; Li ct al., 2021; He et al., 2022; Xie al., 2022; Zhou et al., 2022; Zhai et al., 2022a; Assran et al., 2022; Peng et al., 2022). Indeed, contrary to convolutional neural networks, ViTs lack image-specific inductive bias -e.g., spatial locality or weight sharing as they treat an image simply as a sequence of tokens, each token corresponding to a local image patch. As a result, although ViTs exhibit a tremendous capacity to learn powerful image representations, they require much larger amounts annotated data to do so. In this context, self-supervised learning aims to help transformers overcome such large training requirements and deliver their full potential by utilizing unlabeled image data that can often be readily available even large quantities.
When it comes to applying self-supervised learning to ViT architectures, the most prominent paradigm follows hide-and-predict approach, which consists in hiding tokens on the input and training the model to predict information for the missing lokens, e.g., predicting token ids (Bao et al., 2022) or reconstructing pixels (He et al., 2022). This paradigm originated in the natural language processing (NLP) domain (Devlin al., 2018; Radford et al., 2019) and enforces the learning of contextual reasoning and generative skills. Most applications in the image domain, usually named Masked Image Modeling (MIM), define low-level image information, such as pixels (He et al., 2022; Xie et al., 2022), as prediction targets. Although they
1

---

Published in Transactions on Machine Learning Research (02/2024)
MAE
DINO
iBOTi
MSNi
NOCA
80
70
60
MOCA
IBOTY
60
Top 1 (%)
76
* DINO
50
mloU (%)
50
40
60
30
40
50
74
MSNT
20
10
40
50
100
150
200
5
13
100
372
2975
100
372
2975
Pro-training time (hours)
# Labeled images per class
# Training images
# Training images
(a) K-NN classification ImageNet-1k
(b) Low-shot classification on ImageNet-Ik
(c) Linear probing
(d) Fine-tuning
Cityscapes segmentation
Cityscapes segmentation
Figure 1: Comparison of MOCA with statc-of-thc-art methods using ViT-B/16. (a) K-NN ImageNet classification accuracy pre-training time; (b) One-, (wo-, five- and m13-shot (the last correspouding to 1% training data) ImageNet classification accuracy; (c) Semantic segmentation on Cityscapes using linear probing and (d) fine-tuning with 100, 372, and 2975 training images. MOCA achieves superior results whilst. requiring times less training time. 't" denotes usage of multiple crops (Caron et al., 2021).
reach strong results when fine-tuning on a downstream task with sufficient training data (Bao et al., 2022; He et al., 2022), they do not offer high-level "ready-lo-use" representations, typically evaluated in linear probing and k-NN classification settings. This is partly because predicting low-level pixel details not conducive to lcarning representations that are "readily" available for downstream semantic tasks
Additionally, they do not promote the learning of invariances to image perturbations, which is important characteristic of good image representations. This is actually what several methods which are discriminative have been designed to learn via contrastive- (Chen et al., 2020a; He et al., 2020; Misra & Maaten, 2020), teacher-student- (Grill et al., 2020; Gidaris et. al., 2021; Assran et al., 2022; Caron et al., 2021) or clustering- style (Caron et al., 2018; Asano et al., 2020; Caron et al., 2020) objectives. However, such methods do not necessarily promote contextual reasoning and generative skills, crucial for effective visual representations, as the hide-and-predice methods more explicitly do.
Motivated by the above observations, we propose dual self-supervised learning strategy that exhibits the learning principles of both discriminative and hide-and-predics paradigms. To accomplish this non-trivial goal, we propose novel masking- masseng-baseed based method, named MOCA, in which the two learning paradigms are defined in the same space of high-level features. In particular, our masking-based based method enforces the good "reconstruction" of patch-wist codebook assignments that encode high-level and perturbation invariant features. Those assignment vectors are produced an online and standalone fashion (i.e., without requiring pre-trained models or multiple training stages) using teacher-student scheme and an online generated codebook: Indeed, the teacher produces spatially-dense codebook assignments from the unmasked image views while the student is trained to predict these assignments with masked views. Our masked- based training framework integrates in unified way both dense patch-wise (local) loss that encourages detailed feature generation and global image-wise loss that enforces consistent image representation across different views masked or not of the same original image. As a result, our method converges quickly, reaching state-of-the-ar performances with a training at least. three times faster than the one of prior methods (see Fig. 1a))).
To summarize, our contributions are: (1) We propose a self-supervised representation learning approach called MOCA (for Masked Online Codebook Assignments prediction) that unifies both perturbation invariance and dense contextual reasoning. As will show, combining these two learning principles in effective way is not a trivial task. (2) We demonstrate that MOCA achieves statc-of-the-art results in low-shot settings (e.g,, see Fig. 1(b-d)) and strong results when using the learned representation as an initialization for full fine-tuning, as well as for linear or k-NN classification settings, for which pixel reconstruction approaches under-pent perform. Finally (3), our approach consists of a single end-to-end training stage that does not require prior pre-trained models. At the same time, MOCA able to learn good representations with a much smaller training computation budget than competing approaches (see Fig. 1(a)).
2

---

Published in Transactions on Machine Learning Research (02/2024)
2
Related Work
Vision transformers. Stato-of-tho-art in NLP tasks (Devlin et al., 2018), Transformers (Vaswani et al., 2017) have been recently adapted to vision (Dosovitskiy et al., 2021; Touvron et al., 2021a; Zhao el al., 2020). The cadence to which tricks (Liu et al., 2021; Touvron et al., 2021b; Yuan et al., 2021) and architecture improvements (Touvron et al., 2021a; Zhai et al., 2022b) have been published has allowed for Transformer models to now also be high-ranked in most vision tasks (Carion et al., 2020; Liu et al., 2021; Touvron et al., 2022). However powerful the architecture is, fully-supervised training of Transformer models has proven to be non-trivial (Liu et al., 2020; Huang al., 2020; Xiong et al., 2020), requiring extremely large amounts of annotated data (Dosovitskiy et al., 2021) and luning, more than for the equivalent in size CNNs, thus making the learning of downstream tasks difficult. At the same time, unlike CNNs, ViTs do not easily saturate with more training data, making them particularly interesting for training on large collections of unlabelled images in a self-supervised manner (Bommasani et al., 2021).
Self-supervised learning. Self-supervised learning has emerged as an effective strategy for learning visual representations from unlabelled data without any human manual annotations. In its most common form, a model is trained on an annotation-free pretext task, e.g., (Doersch et al., 2015; Gidaris et al., 2018; Larsson ct al., 2016; Misra et al., 2016; Noroozi & Favaro, 2016; Pathak et al., 2016), and then fine-tuned on a downstream lask of interest with fewer annotated samples. The rise of contrustive objectives (Bachman el al., 2019; Dosovitskiy et al., 2014; Oord et al., 2018; Wu et al., 2018) has finally shown that self-supervised models can surpass the performance of supervised models on downstream tasks. These methods contrast representations between similar and dissimilar views, relying on well-thought data augmentations to create positive samples (Chen et al., 2020a;b; He et al., 2020; Misra & Maaten, 2020; Tian et al., 2020). Other forms of contrastive learning rely on clustering (Asano et al., 2020; Caron et al., 2018; 2019; 2020; Gidaris et al., 2020) or on various forms of diff-distillatinn (Bardes et al., 2022: Caron et al., 2021; Chen & He, 2021; Gidaris et al., 2021; Grill et al., 2020; Zbontar et al., 2021) by learning to predict the similar representations for different views of the same image. Recently, a number of these methods have been successfully adapted to ViT architectures (Assran el al., 2022; Caron et al., 2021; Chen et al., 2021). Such discriminative methods do not perform spatial reasoning. In this work, we propose to facilitate learning by also enforcing local reasoning with an objective inspired by masked image modeling methods which we describe below.
Masked image modeling. The success of transformers in NLP has also been enabled by effective self- supervisod pre-training tasks such as mask autoencoding in BERT (Devlin et al., 2018) or hinguage-modeling in GPT (Radford & Narasimhan, 2018; Radford et al., 2019). Translated in the vision paradigm as Masked Image Modeling (MIM), such tasks have rapidly gained popularity through their simplicity and their direct compatibility with ViTs procossing images as sequences of patch tokens (Bao et al., 2022; Chen et al., 2022; He et al., 2022; Xie et al., 2022; Li et al., 2021; Zhou et al., 2022; Fang et al., 2023). MIM approaches come with different reconstruction targets for the masked input tokens: RGB pixels (He et al., 2022; Xie et al., 2022), hand-crafted HOG descriptors (Wei et al., 2022a) or token features computed by a teacher network (Assran et al., 2022; Baevski et al., 2022; Bao et al., 2022: El-Nouby et al., 2021; Li et al., 2021; Dong et al., 2023; Zhou et al., 2022). For the latter, most of the methods rely on static vocabulary of tokens generated by a pre-trained generative model, e.g., VQ-VAE (Ramesh et al., 2021), where the MIM task consists in reconstructing the tokens corresponding to the masked patches (Bao et al., 2022; Dong et al., 2023; Li et al., 2022; Peng et al., 2022). Recent methods resort to higher capacity teachers like CLIP (Radford et al., 2021) leading to significant performance boosts (Hou et al., 2022; Peng et al., 2022; Wei et al., 2022b;c; Xue et al., 2023). Although convenient, this strategy involves multiple training stages, including the computationally expensive training of the visual tokenizer. Alternatively, the target tokens can be computed online in a self- distillatiion-ikk manner (Assran et al., 2022; Baevski et al., 2022; Kakogeorgiou et al., 2022; Zhou et al., 2022). However, since these ViT-based models essentially start from scratch, they may require many epochs to bootstrap the learning of token targets and of rich representations. MOCA elficiently follows this approach, being at least three times faster than current methods. Our scheme trains the student to densely predict teacher token assignments of masked patches over an online generated codebook. Other works quantize tokens into static pre-computed codebooks (Peng et al., 2022) or predict assignments of global image features over online codebook (Assran et al., 2022; Caron et al., 2020), though potentially losing spatial information.
3

---

Published in Transactions on Machine Learning Research (02/2024)
3 Our approach
(s(x²)
S(-) Student enooder
D(-)
T(} Teacher encoder
D(.): Decoder
GAP
Pred(.)
deac
9r (x'
DpAG
LLOC
MT(("))
GAP
generate
assignements
K
T())
gr(x')
original image
ZK
Figure 2: Overview of MOCA The teacher (bottom) lakes as an input two unnasked random views x(122) of the same image and generates dense token-wise code assignments @r(x) for them (i.e., soft-assigns codebook items to the patch tokens). The student (top) receives as an input a randomly masked image version x of view x! and trained to minimize two types of self-supervised losses: (1) A masked same-view token assignmend prediction loss LLoC, which requires predicting the teacher-produced assignment vectors of view x² from the corresponding masked image x This is spatially dense loss that enables learning representations with dense contextual reasoning. (2) A masked cross-view average assignment prediction loss LIMG, which is to predict with the global image embedding of the first view x the average assignment vector of the opposite view x² ('GAP' stands for Global Average Pooling). This is an image-wise loss that promotes learning image representations that are invariant with respect to different augmentations of the input. The same objectives are applied a symmetric way when the student gets as input the masked image version x² of view x² (not shown). We implement the TLoC objective with a condenser-based decoder that gets as input patch-token embeddings from an intermediate layer of the student encoder and the global image embeddings from its last layer (see See. 3.3). The teacher is an exponential moving average ("EMA') of student.
Our goal to build cole-and-predictt self-supervised approach that (a) is defined over high-level concepts and (b) standalone, i.e, does not require any pre-trained model or off-line training stage. To this end, we follow teacher-student setup where the teacher transformer T(-) is momentum-updated version (He et al., 2020; Richemond et al., 2020) of the student transformer S(.) and generates largels for ils student (Figure 2).
For defining these targets, we use what we will refer to hereafter as token assignment vectors (over an online codebook), which constitute the high-level visual concepts over which our hido-and-predict framework is defined. We detail the definition of these assignment vectors Sec. 3.1. We then describe in Sec. 3.2 how we use these teacher-prouuced assignment. vectors for defining the hide-and-prodict objectives for the student. training. In Sec. 3.3 we discuss important design choices and in Sec. 3.4 techniques for faster training.
3.1 Teacher-produced token assignment vectors
Given input an image token sequence x {x,) the teacher transformer first extracts dene -inennsini token embeddings T(x) = (M and then soft-quantizes them over a codebook C [cI,.....cp cx] of K embeddings of dimension dene- This produces the soft-assignment vectors @r(x) {rr (x) where qr (x) is a K -ddimensional vector computed as the cosine similarity the token embedding T(x), and each codebook embedding Ck. followed by a softmax with temperature Tr, which controls the softness of the assignment. In essence, (x(x); encodes the assignment of T(x); to its closest (in terms of cosine distance) codebook embeddings. We these leacher-produced assignment vectors for defining the hide-and-predicl objectives for the student training, which is described in Sec. 3.2. For the softmax temperature where PMSD is the exponential moving average (with momentum 0.99) of per-token difference between the maximum and average cosine similarity over the codebook embeddings. For the codebook, we use K 4096.
use 7T 10-PMSn
4

---

Published in Transactions on Machine Learning Research (02/2024)
Simple online codebook construction. Following Gidaris et al. (2021), we build the codebook C as queue of size K containing teacher-token embeddings randomly sampled from previous mini-batches. During each training step, after computing assignment vectors, we randomly sample Knew < K token embeddings from the current mini-batch, each from a different image. Specifically, we uniformly select images and then randomly choose (again with uniform distribution) one teacher patch-token embedding from each. These selected embeddings replace the oldest ones the queue, keeping only those from the past K/Knew mini-batches in memory. Codebook embeddings are initially randomly initialized from a normal distribution at training step O. Since the learning rate is very low during the initial training steps (due the warmup period), the random initialization of the embeddings has minimal impact on the learning process in the early stage.
3.2 Hide-and and-prediict token assignment vectors
Let x" with eE {1, 2} be the token sequences of two random views sampled from the same training image. On the teacher side, they are processed by the transformer T(-) to produce a sequence of image token assignment vectors 9r(x"). On the student side, the transformer S(-) receives as input randomly masked versions of the x" images, which are denoted by x" with e E (1,2), and is trained to minimise two types of masked-based self-supervised losses:
1. A masked some-view loken assignment prediction loss L.oC, which aims at predicting the sequence of teacher assignment vectors qr (x") from the corresponding masked image x" for E {1,2). This is a spatially dense loss whose role is to enable learning reprosentations with dense contextual reasoning skills.
2. A masked cross-view average assignment prediction loss LIMG, which amounts to predicting the teacher's average assignment vector of the first view, (x*), from the opposite masked view x² and vere-versa. This is an image-wisc loss whose role is to promote lcarning representations that are invariant with respect to different input augmentations.
The overall training loss for the student model is:
L = -LING + (1 X)L
(1)
where A E [0, 1]. Unless stated otherwise, we use x = 0.5. Next, we describe these two objectives in detail.
Masked same-view token assignment prediction. With this objective, given as input the masked view X", the student is trained to predict the assignment vectors of the tokens as generated by the teacher from x". In the following, for each (1,2), we denote by uh C {1,..., N} the set of indices of all tokens that remain unmasked (i.c., all tokens from view x" that are scen by the student) and by m. C N} the set of indices of all the masked tokens (it holds m. U up = {L....,N and m, Nu, = -)
For computational efficiency reasons, the student encoder Si.)) processes only the unmasked tokens. Therefore, for the prediction of the masked loken assignment vectors, we employ an additional decoder transformer D(.). For this prediction take place, in the defoult setup, we employ the following steps:
Encoding the masked view:
We first give the unmasked tokens x"= encoder transformer S(-) to produce the deue-dimensional student token embeddings (S(x");)eu..
= (x] heu. as input to the student
Decoder input embedding: Before we give the encoder embeddings as input to the decoder, we apply to them a linear projection layer to convert their dimensionality from dene to date, where datr is the dimensionality of the decoder embeddings. Then, we append the encoder embeddings with a learnable masked embedding for each masked token E m, and add (non-learnable sin-cos-based) positional embeddings to them. We denote the set of embeddings produced from this step as z".
Local token assignment prediction: Last, the above embeddings are given input to the decoder transformer D(-). The decoder computes the tokens embeddings D(z") = (D(z") ( which will be used
We always assume that the [CL.S] token is assigned the index 0.
5

---

Published in Transactions on Machine Learning Research (02/2024)
for predicting the token assignment vectors as follows:
qs (x",, softmax _ D(z"), wh). 7a
(2)
[w4. wi] are Lz-normalized daee-dimensional prototypes (one per codebook embedding in C¹) used for predicting the assignment vectors. As we explain in more detail in Sec. 3.2.1, these prototypes are not learnt but are dynamically generated by a weight generation module. Also, Ta is softmax temperature that controls the softness of the predicted probability.
for i E NJ, The weights =d=
The masked same-view token assignment prediction loss is
LLOC Lioc + Liocs where
(3)
Lioc = Σ ( Ve E (1,2},
(4)
and CE(a, b)(a.b)=--3 -- b(g//k)()))))) the cross-entropy loss between two distributions and b.
Masked cross-view average assignment, prediction. In this objective, first reduce the teacher assignments (T(x)= {x): to single K-dimensional vector yr(x) via global average pooling, i.e., we reduce (x) to a 'bag-of- wford (BoW) representation. Then, similar to BoW-prediction methods (Gidaris & Komodakis, 2019; Gidaris et al., 2021), for each view e the student must predict with its global image embedding" 5(x") the reduced assignment vector of the opposite view as follows:
(x") = softmax( S(x") T W
(5)
As in the same-view token assignment prediction, ть is the softmax temperature, and W = [w/..... w are Lg-normalized deme-dimensional prototypes that dynamically generated by weight generation module.
The cross-view image-wise loss that is minimized for the two masked image views x" with v e {1,2], is
(ys LIMG CFF ((()s(x))) (x') vx())) CF(()s("))
3.2.1 Assignment predictions using dynamic prototype generation modules
Our method updates the codebook at the teacher side at each training iteration with new randomly-sampled teacher-tokon emboddings. Due this, instead of lcarning the prototypes Wd and Wh that are usod on the student side for the two assignment prediction tasks defined by Eqs. (2) and (5), we dynamically generate them using two weight generation modules. Tn particular, we employ the generation networks G²(-) and G+)-) that at each training step take as input the current codebook C' = .c....., Ck] and produce for them and Wd = GOCC) = [G*( G
the prototypes weights 4hb = (C) = [G' respectively. The networks (-) and (.) are implemented with 2-layer perceptrons (i.e., each as a sequence of Linear, BN, ReLU, and Linear layers), whose input and output vectors are normalizd and have skip connection between them (more implementation details in Appendix C.1).
We 7b Ta = the softmax temperatures in Eqs. (2) and (5) respectively.
3.3 Learning high-level target assignment vectors
For our site-and-prediic approach to be effective, the teacher must learn assignment vectors encoding high-level visual concepts. In the context of our teacher-student scheme, where the teacher is a slower-moving version of the student, this depends the student first learning patch-token embeddings (from where the assignment vectors are extracted at the teacher side) that capture this type of visual features. Below we detail the critical role some design choices have in that.
1. Eieu, 8(R"), We found this to work better 2 As global image embedding we use the average token embedding 5(x") Taul than using the [CT.S] taken embedding S(x")o- )o-

---

Published in Transactions on Machine Learning Research (02/2024)
Table 1: Training efficiency. (a) Training savings and (b) k-NN results on ImageNet-Tk 1k with ViT-Base/ Rase//e models pre-trained for 100 epochs. 'Partial Dec.: partial decoding; '# #askss number of different masks per view; 'k-NN:: top-1 accuracy (%); "Time' and 'Memory:: per epoch training time and GPU memory footprint measured with a single 8-A100 node and batch size 2048.
(a) Training savings by switching full partial decoding.
(b) Partial decoding & 2nd masking round.
The
decoder depth is 2.
Decoder depth
Time (min)
Memory (GB)
Partial Dec.
#Masks
k-NN (%)
Time (min)
Memory (GB)
2
9.00 + 8.15
30.4 + 23.0
1
71.6
9.00
30.4
4
10.25 + 8.55
37.7 + 26.2
1
71.8
8.15
23.0
8
12.41 + 9.30
62.4 + 32.6
2
74.4
10.35
39.5
Condenser-based decoding. As in Gao & Callan (2021); Peng et al. (2022), we change the default setup described in Sec. 3.2 and create a bottleneck in the decoding task that promotes the student's global image representation to caplure more spatial structure from the input view &, crucial for predicting patch token assignments. The bottleneck is introduced by changing the way the decoder input embeddings 2" are formed.
In particular, instead of using as input to the decoder the token patch embeddings (S(*")))Ew. (coming from its last layer), we give as input (a) token embeddings (S"(x") from an intermediate transformer layer of together with (b) the last layer's global image embodding S(x"). Using lower-layer emboddings for the patch tokens introduces bottleneck. To compensate, the decoder must rely more on the global image embedding S(x") that it receives input. This forces the global image embedding to capture additional spatial structure information, potentially beneficial for downstream tasks such k-NN and linear probing.
Eieu, S(x"); (called [AVG] token for brevity) used for performing (a) the cross-view image- wise prediction task (see Eq. (5)), which is erucial for promoting perturbation invariances, as well as (b) the condenser-based decoding that was explained above. Therefore, using the [AVG] token (instead of the [CLS] token) results in the student's ss patch embeddings {S(**):)iem, (from which the [AVG] token is computed) to receive more direct supervisory signal from the two prediction losses. Consequently, these embeddings (i) are steered to encode the cnntuurbation-nnvarinn and context/ context/structurer /structure-aware information necessary for performing the two prediction objectives, and (ii) are typically of more semantic nature.
Hide-and-predict using the [AVG] token. The average patch embedding S(x") - T.. E..
The impact of these design choices is studied in Sec. 4.1.
3.4 Masking strategies for efficient training
Partial decoding for efficient training. We randomly select the token indices m. to be masked from each view, with a relatively high percentage of masked tokens (i.e., 65%). This results in significant computational and GPU memory savings from the student encoder side during training. However, some of these savings are lost if the dccoder processes all masked tokens. To mitigate this, we proposo inputting only a small subset m/, C my of masked lokens, along with the visible lokens, to the decoder. For example, instead of inputting all 65% of image tokens that we mask from the student, we input only random subset, e.g., 20% of the total image tokens. As shown in Tab. la, this results in noticcable savings in computation and GPU memory, especially for many-layer decoders, without any loss in the quality of learned representations (as we demonstrate in Sec. 4.1; see Tab. 1b). Note that although we only tested partial decoding in the context of our approach, we believe it could benefit other aide-and-predicc methods as well, e.g., MAE (He et al., 2022).
Faster training convergence with 2nd masking round. With our partial decoding, we can allocate our computation budget to enhance the model's performance with other techniques. Specifically, we apply a second round of random token masking to both views and then perform the assignment prodiction objectives with the newly generated masked views in the same manner as before. This technique artificially increases the batch size and saves time since we can reuse the teacher assignment vectors from the first masking round. Tt.
7

---

Published in Transactions on Machine Learning Research (02/2024)
NOCA ablation. k-NN classification (%) on Image.Net-lk with a ViT-B/16 trained for 100 pre- pre-training epochs and one masking round. Default setting (in gray): X = 0.5, mask ratio of 65%, [AVG] token used global image embedding, condenser- intermediate student layer. 'Time' and 'Memory': per epoch training time (min) and GPU memory footprint (GB) measured with single 8-A100 node and batch size 2048.
Table 2:
decoder with depth fed with patch-embeddings from the 8
(4) bonden decoding is more effective.
(b) Decoder depth. "Shallower"
(c) Masking ratio. The values 55% and 65% give the best results.
is more accurale and faster.
LIMG (A=1.0)
LIMG & LLoC
Depth
Masking
55%
65%
75%
80%
Objectives
(A=0.5)
k-NN
71.7
71.8
70.9
70.5
k-NN
72.6
71.8
69.8
68.4
Condenser
N/A
Time
8.08
8.15
8.55
9.3
Time
9.00
8.15
7.37
7.35
k-NN
66.8
67.8
71.8
Memory
21.5
23.0
26.2
32.6
Memory
28.6
23.0
18.2
15.7
(d) [CLS] vs. [AVG] as global image embedding.
(e) Weight MLMa + (1 - ALoc.. Both objectives are important.
(f) Intermediate layer in Condenser-baser decoding.
im. embed. [CLS]
[AVG]
A
1.0
0.75
0.5
0.25
0.00
Layer E
8
9
10
k-NN
60.2
71.8
k-NN
66.8
70.2
71.8
71.5
13.1
k-NN 71.4 71.8
71.5
70.9
can be seen as form of batch augmentation (Hoffer et al., 2020; Touvron et al., 2022) that reduces gradient variance by having in the same Aini-batch multiple randomly augmented copics of the same sample.
4 Experiments
Setup. We evaluate our MOCA method by training ViT-B/16 models on the ImageNet-1k (Russakovsky et al., 2015) dataset. We use the AdamW optimizer (Loshchilov & Hutter, 2019) with B1 0.9, 82 0.999 and weight decay 0.05. The batch size is 2048 split. over 8 A100 GPUs. For the learning rate Ir, we use linear warm-up from 0 to its peak value for 30 epochs and then decrease it over the remaining epochs with cosine annealing schedule. The peak Ir is 1.5 x 10-4 and the number of training epochs is 100 or 200. More implementation details are provided in the appendix.
4.1 Method analysis
We first study several aspects of our approach with ablative experiments, in which we train ViT-B/16 models for 100 epochs. Except stated otherwise, the models include the condenser- based decoder (Sec. 3.3), partial decoding, and a single token masking round. We evaluate the learned representations on the k-NN ImageNet. classification task.
Integrating the two masked prediction objectives. Here we show that a naive integration of the two masked-prediction objectives does not work effectively.
Ablating the condenser-based decoding in Tob. 2a. Compared to a model trained only with the image- objective LIMG, naive implementation of our dense hide-and-predict objective ILocc without the condenser design principle provides only small k-NN performance improvement (i.e., +1 point). Instead, when using the condenser, the k-NN performance increase is much larger (i.e., +5 points). We further compare the performance with and without condenser when pre-training for 200 epochs (instead of 100) in Tab. 3: condenser Ieads to significantly better results in all evaluation protocols, apart from ImageNet finc-tuning. On the other hand, in this longer pre-training selting the naive implementation of LLoC achieves worse results than only using LIMG in k-NN and linear probing. These results validate that the condenser design choice indeed leads a more effective hido-and- prodict training (sce discussion in Sec. 3.3). In Tab. 2f we see that using & 8 as the intermediate student layer for providing the patch-token embeddings that are given as input to the condenser-based decoder, gives the best results.
8

---

Published in Transactions on Machine Learning Research (02/2024)
Table 3: Ablating hide-and-predict variants. ImageNet- 1k classification results with ViT-B/16 models
pre-trained for 200 epochs with one masking round. The improvements are over the Livg-only objective.
Target
+hide-and-predicc variants
LIMG
Objectives
Assign.
Assign. (ours)
Pixels
Embeddings
Condenser
N/A
X
X
k-NN
72.0
70.0 (-2.0)
75.4 (+3.4)
72.3 (+0.3)
73.6 (+1.6)
Linear
76.1
75.7
(-0.4)
77.7 (+2.0)
76.4 (40.3)
76.9 (+1.2)
Fine-tune
82.7
83.5 (+0.8)
83.4 (+0.7)
83.2 (+0.5)
83.0 (+0.3)
[CLS] us. [AVG] global embedding tokens (Tab. 2d). Using the [CLS] token as the global embedding, instead of the average embedding token ([AVG]), significantly decreases the k-NN performance (i.e., -11.6 points). We postulate that this is because, as discussed in Sec. 3.3, the [AVG] token lcads to learning assignment vectors that encode more high-level concepts and thus makes our hide-and-predict method more effective.
Balancing the two prediction objectives. The À hyperparameter controls the importance of the two assignment prediction objectives (see Eq. (1)). We see in Tab. 2e that À = 0.5 gives the best results while switching off any of the two objectives (A= = 1 or À = 0) seriously deteriorates the k-NN performance. This suggests that the objectives work in a complementary and synergistic way. For instance, for X = 0 (no LING loss) the results are very poor. The reason is that without the image-wise loss, it is hard to bootstrap the learning of teacher-prood produucd assignment vectors that encode high-level concepts (see discussion in Sec. 3.3). In the appondix, we report more experiments, results, and discussions about the behavior when training without the image-wise loss LIMG-
Decoding and masking strategics.
Studying the impact of decoder depth in Tab. 2b. We observe that our method behaves better with "shallower" decoders, in contrast to what is the case in MAE (which uses a decoder with eight layers). This expected since our hide-and-predic, objectives require the decoder to predict high-level concepts. In that case, is better for the applied loss to be "closer" to the student (i.e., having a shallower decoder) the output student embeddings to capture such high-level information in a more "ready-to-use" way. On the other hand, MAE's decoder, which must reconstruct pixels, is better to be deeper and thus spare the MAE's encoder from capturing such low-level details. Furthermore, using a "shallower" decoder leads to significant. computation and GPU memory savings in our case.
to promote
Partial decoder and 2nd masking round in Tab. 1b. Switching off partial decoding (i.c., decoding all the masked tokens) leads to an increase in the per-epoch training time and GPU memory usage with zero benefits on the quality of the learned features as measured with k-NN. Instead, having second round of token masking, although increasing the training time and GPU memory usage, leads to a significant k-NN performance boost. Therefore, we employ it for the final model that we train for the results in Sec. 4.4.
Studying the masking ratio impact in Tab. 2c. The 55% and 65% ratios give the best results.
4.2 Comparison of hide-and predict objectives
In our work, propose hide-and-predde objective LLOC that is defined in the space of assignment vectors over an onlino-updated codebook. In Tab. 3, we compare this objective ('Assign.') against (a) a hide-and-prediet objective defined in the pixel space with per-patch normalization as in MAE (He el al., 2022) ('Pixels"); and (b) a tite--rrrittoojeettve in which the decoder must directly regress the dimenninni output teacher-token embeddings (before any code assignment) using a cosine-distance loss ('Embeddings)) We keep the imago-wise LIMG objective intact for all three models. The Embeddings model is implemented with condenser-based decoder of depth two, as with our Assignments model, while the Pixels model uses
9

---

Published in Transactions on Machine Learning Research (02/2024)
Table Ablations on Cityscapes semantic segmentation with linear probing. We report mloU (%) results with ViT-B/ 16 models pre-trained for 200 epochs. For training we use 100 372 images from the Cityscapes training set using random splits from (French al., 2020).
Linear probing
Objectives
Condenser
# Masks
100
372
LIMG
N/A
1
44.6
50.9
LIMG & LLOC
1
44.2
51.5
LIMG LLoC
1
46.8
52.4
LIMG LLOC
2
49.9
55.2
standard decoder and depth four (which gives better results). All other model siper-parameters remain the same for fair comparison. We see that the Assignments model produces superior results on all three ImageNet classification evaluation protocols. This demonstrates the advantage of defining hide-and-predict objective with online codebook assignment vectors targets.
4.3 Ablations on low-shot semantic segmentation
Here we conduct ablation studies evaluate (a) our hide-and-prodict objective LLoc, (b) the condenser decoder design, and (c) the second masking round, on the downstream task of low-shot Cityscapes semantic segmentation with linear probing: detailed information on this evaluation protocol can be found in Sec. 4.4 and Appendix C.3. The results confirming these design choices are presented in Tab, 4.
4.4 Comparative results
We compare our method against other iide-and-predict methods as well as self-supervised methods based on teacher-student or contrastive objectives. To this end, we train a ViT-B/16 /6 model for 200 epochs with a condenser-based decoder, two masking rounds (the 1st with 55% mask ratio and the 2nd with 75% mask ratio), and partial docoding only 20% of patch tokens per masking round.
Full and low-shot image classification. We first evaluate the learned representations on ImageNet classification with the k-NN, linear probing, and fine-tuning evaluation protocols and using the full training set. Moreover, we evaluate low-shot Tma classification with logistic regression (Assran et al., 2022).
Comparison with hide-and-predict approaches. In Tab. 5a, we compare our method MOCA against other self-supervised methods with MIM objectives. We observe that on linear probing MOCA achieves superior performance to all the other works, with more than 2 points. This demonstrates that MOCA is better at learning "ready-to-use" features. At the same time, its fine-tuning performance is on par with the state-of-the-art, falling only 0.4 points behind the top-performing method. This is despite MOCA using only 200 epochs and a smaller computational budget in general. For instance, as we in Tab. 5c, the total training time of MOCA is more than 3.5 times smaller compared to MAE, which is (one of) the most training- -effiinentf methods.
Comparison with teacher-student and contrastive methods. In Tab. 5b, we compare our method against other self-supervised methods based on contrastive objectives or teacher-student schemes. We observe that our method is superior to competing approaches that do not use multiple crops and on par with them when they employ such techniques. Again, we emphasize that our method is much more computationally efficient. than these techniques (soc Tab. 5c).
Low-shot ImogeNet-Ik Ik classification. Here we adopt the low-shot evaluation protocol of MSN (Assran et al., 2022) and use few as 1, 2, or 5 training images per class as well as using 1% of the ImageNet- 1k's training data, which corresponds to 113 images per class. In particular, for this low-shot training setting, we extract global image embeddings with the pre-trained transformer (average patch embeddings computed by the
We found that implementing the Pixels hide-and hiredee version with condenser based decoder has worse performance than the model trained with only the image-wise LIMG objective.
10

---

Published in Transactions on Machine Learning Research (02/2024)
Table 5: Comparisons with prior methods. ImageNet-1k 1k classification and training time results with ViT-B/16 models. MOCA, pre-trained for 200 pro-training epochs and with two token masking rounds, is compared with (a) the hide-and-predien methods MaskFeat (Wei et al., 2022a), CAE (Chen et al., 2022), BEiT (Bao et al., 2022); SimMIM (Xie et al., 2022), MAE (He et al., 2022), CAN (Mishra al., 2022), and (b) the discriminative methods MoCo-v3 (Chen et al., 2021), DINO (Caron et al., 2021), iBOT (Zhou et al., 2022) and MSN (Assran et al., 2022). In (c) we compare the per-epoch ('Epoch', in minutes) and total ('Total', in hours) training time. The per-epoch time measurements are based on a single 8 A100 GPU node. In (d) we evaluate using the low-shot image classification protocol. t: use of multiple crops (Caron et al., 2021). The k-NN and linear probing results of MSN were compuled by us.
(a) Hide-and-predict methods
(b) Discriminative methods
top-1 I accuracy (%)
top-1 accuracy (%)
Method
# Epoch
Linear
F-tune
Method
# #Ppoch
k-NN
Linear
F-tune
MaskFeat
1600
84.0
MoCo-v3
300
76.7
83.0
CAE
1600
70.4
83.9
DINO
400
68.9
72.8
BEIT
800
37.6
83.2
IBOT
400
71.2
76.0
SimMIM
800
56.7
83.8
DINO
400
76.1
78.2
82.8
MAE
1600
68.0
83.6
iBOT+
400
77.1
79.6
84.0
CAN
800
74.8
83.6
MSNf
600
73.4
77.2
83.4
MOCA (ours)
200
78.7
83.6
MOCA (ours)
200
77.2
78.7
83.6
(c) Training time comparisons
(d) Low-shot ImageNett classification evaluation
Method
# #pooch
Epoch (mn)
Total (hr)
# Tmaaees per class
MAE
1600
5.35
142
Method
# Epoch
2
113
DINO
400
16.42
109
MAE
1600
8.2+0.3
25.0+0.3
40.5+0.2
51.1
IBOT
400
17.40
116
DINOI
400
41.8:10.3
51.91006
61.4100.
67.2
DINO
400
23.16
154
iBOTI
400
46.1100.3
56.210.7
64.7.00.3
69.7
iBTTY
400
24.33
162
MSN
600
49.
58.9+0.4
65.5+0.3
69.4
MSNI
600
19.07
191
MOCA (ours)
200
55.4+0.4
63.9+0.2
69.9=0.1
72.4
NOCA (ours)
200
11.23
38
Table 6: Semantic segmentation Cityscapes. We report mloU (%) results using Segmenter (Strudel et al., 2021) model with ViT-B/16 backbonc. For training we use 100, 372 or 2975 (full training split) images from the Cityscapes training set using random splits from (French et al., 2020) for the 100- and 372-image settings. t: use of multi-crop augmentation during pre-training (Caron et al., 2021).
Linear probing
Finetuning
Method
100
372
2975
100
372
2975
MAE (Ie et al., 2022)
33.2
39.6
40.6
41.4
62.6
77.7
DINOF (Caron al., 2021)
45.8
51.0
55.5
47.7
65.7
75.7
IBOTI (Zhou et al., 2022)
47.8
54.0
59.3
48.6
66.8
77.8
MOCA (ours)
49.9
55.2
60.1
51.5
71.1
78.2
teacher) and train a lincar classifier on top of with the few available training data. As we see from the results in Tab. 5d (see also Fig. 1bb), MOCA outperforms prior methods with ViT-B/16 surpassing the second best method MSN (Assran et al., 2022) by a large margin.
Full and low-shot semantic segmentation. We further evaluate our method on the Cityscapes (Cordts et al., 2016) semantic segmentation dataset, by either linear probing or finetuning. We use the Seg- menter (Strudel et al., 2021) model with ViT-B/Ir B/6 backbone initialized either with our MOCA or other self-supervised methods. In the linear probing setup we append a lincar layer to the frozen encoder and finetuning setup we use a single mask transformer layer a decoder following the original setup from Strudel et al. (2021). We train using the full Cityscapes training set of 2975 images well as 100 374 training
11

---

Published in Transactions on Machine Learning Research (02/2024)
images, representing 1/30 and 1/8 of the full training set. For these 100 and 374 low-shot settings, use three different splits of 100 or 374 training images respectively following the protocol of French et al. (2020) and report the average mloU performance over the three splits. Learning rates were optimized for every method separately, with other hyperparameters kept as the default oncs used in the Segmenter (Strudel et al., 2021) codebase. We evaluate the trained segmentation models on the 500 Cityscapes validation images.
We report the linear probing results in Fig. 1(c) and fine-tuning results in Fig. 1(d) and provide detailed results in Tab. 6. They show that MOCA consistently achieves superior performance compared to prior methods in all the setups, with bigger improvements appearing in the 100 and 374 low-shot settings.
COCO detection and instance segmentation results. We present results on COCO detection and instance segmenta- tion in Tab. 7. MOCA demonstrates strong performance in these tasks. We use the COCO 2017 set consisting of 118K training images and 5k validation. For this experiment we follow the implementation of Liu et al. (2021). In detail, we adopt Cascade Mask R-CNN (He et al., 2017; Cai & Vasconcelos, 2019) as task layer and use a similar syper-paramenter configuration for ViT-B with (Zhou et al., 2022): multi-scale training (resizing image with shorter size between 480 and 800, with the longer side no larger than 1333), AdamW (Loshchilov & Hutter, 2019) optimizer with initial learning rate 2e the 1x schedule (12 epochs with the learning rate decayed by 10x at epochs and 11).
Table 7: COCO detection and instance segmentation with ViT-B/16. t: use of multiple crops (Caron et al., 2021).
Method DINOf
iBOTf
MOCA
Apbox 50.1
51.2
50.5
Apmaok
43.4
44.2
43.6
5 Conclusion
In this work we have proposed novel self-supervised learning framework that is able to exhibit the complementary advantages of both discriminative and hido-and-prodict approaches, thus promoting visual representations that are perturbation invariant, encode semantically meaningful features and also exhibit dense contextual reasoning skills. To that end, we make use of a novel masking- in the space of high-level assignment vectors that are produced in an online and standalone fashion (i.e., without the need for any pre-trained models or multiple training stages) using a teacher-student scheme and an online generated codebook. We have experimentally shown that our approach possesses significant. computational advantages compared prior methods, while exhibiting strong performance, not only in the case of full finetuning but also in the case of linear probing. k-NN, and low-shot settings. It thus allows to obtain high-quality image representations with much smaller training computational budget compared to prior statc-of-the-art approaches.
strategy that operates
Broader Impact Statement Self-supervised image representation learning can exacerbate existing biases within pre-training data, potentially leading to biased visual recognition systems that disproportionately affect individuals based on gender, ethnicity, skin tone, or geographical location. Additionally, the collection of large-scale pro-training image datasets raises substantial privacy and copyright concerns, particularly when images contain sensitive, protected personal information. For instance, unauthorized access misuse of such data can result in privacy breaches, with potential repercussions for individuals depicted in the dataset. Therefore, practitioners deploying self-superrvied pre-training in real-world applications must exercise caution, considering factors such as biases, privacy, and liconsing associated with the pro-training data.
Acknowledgements This work was performed using HPC resources from GENCI-IDRIS (Grants 2022- AD011012884R1 and 2022-AD0110134133) and from CINES under the allocation GDA2213 for the Grand Challenges AdAstra GPU made by GENCI This work supported by the Ministry of Education, Youth and Sports of the Czech Republic through the c-INRRA CZ (ID:90254) and by CTU Student Grant SGS211640HK33T337 This research received the support of EXA4MIND, a European Union' 's Horizon Europe Research and Innovation programme under grant agreement No 101092944. Views and opinions expressed are however those of the author(s) only and do not necessarily reflect those of the European Union or the European Commission. Neither the European Union nor the granting authority can be held responsible for them.
12

---

Published in Transactions on Machine Learning Research (02/2024)
References
Yuki Markus Asano, Christian Rupprecht, and Andrea Vedaldi. Self-labelling via simultaneous clustering and representation learning. In ICLR, 2020. 2, 3
Mahmoud Assran, Mathilde Caron, Ishan Misra, Piotr Bojanowski, Florian Bordes, Pascal Vincent, Armand Joulin, Mike Rabbat, and Nicolas Ballas. Masked Siamese networks for label-efficient learning. In ECCV, 2022. 1, 2, 3, 10, 11, 21
Philip Bachman, R Devon Hjelm, and William Buchwalter. Learning representations by maximizing mutual information across views. In NewrIPS, 2019. 3
Alexci Bacvski, Wei-Ning Hsu, Qiantong Xu, Arun Babu, Jiatao Gu, and Michael Auli. Data2vec: A general framework for self-supervised learning in speech, vision and language. In ICML, 2022.
Hangbo Bao, Li Dong, and Furu Wei. BETT: BERT pre-training of image transformers. In ICLR, 2022. 1, 2, 3, 11
Adrien Bardes, Jean Ponce, and Yann LeCun. Vicreg: Varfaace-imarraacc-coaaaaae regularization for self-supervised learning. In ICLR, 2022.
Rishi Bommasani, Drew A Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S Bernstein, Jeannctte Bohg, Antoine Bossclut, Emma Brunskill, et al. On the opportunitics and risks of foundation models. Xiviv preprint urXie:2108.0722 2021.
Zhaowei Cai and Nuno Vasconcelos. Cascade R-CNN: high quality object detection and instance segmentation. TPAMT, 2019. 12
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In ECCV, 2020.
Mathilde Caron, Piotr Bojanowski, Armand Joulin, and Matthijs Douze. Deep clustering for unsupervised
lcarning of visual features. In ECCV, 2018. 2, 3
Mathilde Caron, Piotr Bojanowski, Julien Mairal, and Armand Joulin. Unsupervised pre-training image
features on non-curated data. In ICCV, 2019. 3
Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. In NewrIPS, 2020. 2, 3
Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin, Emerging propertics in self-supervised vision transformers. In ICCV, 2021. 2, 3, 11, 12, 19, 21
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In ICML, 2020a. 2, 3, 20
Xiaokang Chen, Mingyu Ding, Xiaodi Wang, Ying Xin, Shentong Mo, Yunhao Wang, Shumin Han, Ping Luo, Gang Zeng, and Jingdong Wang. Context autoencoder for self-supervised representation learning. arXiv preprint XrXi:2200..0 03026, 2022. 3, 11
Xinlei Chen and Kaiming He. Exploring simple Siamese representation learning. In CVPR, 2021. 3
Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming He. Improved baselines with momentum contrastive learning. ar preprint arXiv:2003. N0x:2200300497, 2020b. 3
Xinlci Chen, Saining Xie, and Kaiming He. An empirical study of training self-supervised vision transformers.
In ICCV, 2021. 3, 11, 20
Kevin Clark, Minh- Minh-Thang Luong, Quoe Le, and Christopher D Manning. Electra: Pre-training text encoders discriminators rather than generators. In ICLR, 2020. 21
13

---

Published in Transactions on Machine Learning Research (02/2024)
Marius Cordts, Mohamed Omran, Sebastian Ramos, Timo Rehfeld, Markus Enzweiler, Rodrigo Benenson, Uwe Franke, Stefan Roth, and Bernt Schiele. The Cityscapes dataset for semantic urban scene understanding. In CVPR, 2016. 11, 21
Ekin D Cubuk, Barret Zoph, Jonathon Shlens, and Quoc V Le. Randaugment: Practical automated data augmentation with reduced search space. In CVPR, 2020. 21
Jacob Devlin, Ming- Wei Chang, Kenton Loe, and Kristina Toutanova. BERT: Pre-training of deep bidirectional transformers for language understanding. arXi preprint urXiv:1810.0458 2018. 1,
Carl Doersch, Abhinav Gupta, and Alexei A Efros. Unsupervised visual representation learning by context prediction. In ICCV, 2015. 3
Xiaoyi Dong, Jianmin Bao, Ting Zhang, Dongdong Chen, Weiming Zhang, Lu Yuan, Dong Chen, Fang Wen, and Nenghai Yu. PeCo: Perceptual codebook for BERT pre-training of vision transformers. In AAAI, 2023.
3
Alexcy Dosovitskiy, Jost Tobias Springenberg, Martin Riedmiller, and Thomas Brox. Discriminative unsuper- vised feature learning with convolutional neural networks. NewrIPS, 2014.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Jakob Uszkoreit, Mostafa Dehghani Neil Houlsby, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby. An image is worth 16x16 words: Transformers for image recognition at scale. In ICLR, 2021. 3, 20
Alaacldin El-Nouby, Gautier Izacard, Hugo Touvron, Ivan Laptev, Hervé Jegou, and Edouard Grave. Are large-scale datasets necessary for self-supervised pre-training? arXiv preprinl 2Xix::112.107/// 2021. 3
Yuxin Fang, Li Dong, Hangbo Bao, Xinggang Wang, and Furu Wei. Corrupted image modeling for self- supervised visual pre-training. In ICLR, 2023.
Geoff French, Samuli Laine, Timo Aila, Michal Mackiewicz, and Graham Finlayson. Semi-supervised semantic segmentation needs strong, varied perturbations. In BMVC, 2020. 10, 11, 12, 21
Luyu Gao and Jamie Callan. Condenser: a pre-training architecture for dense retrieval. In EMNLP, 2021.
Spyros Gidaris and Nikos Komodakis, Generating classification weights with GNN denoising autoencoders
for few-shot learning. In CVPR, 2019. 6
Spyros Gidaris, Praveer Singh, and Nikos Komodakis. Unsupervised representation learning by predicting image rotations. In ICLR, 2018. 3
Spyros Gidaris, Andrei Bursuc, Nikos Komodakis, Patrick Pérez, and Matthieu Cord. Learning representations by predicting bags of visual words. In CVPR, 2020. 3
Spyros Gidaris, Andrei Bursue, Gilles Puy, Nikos Komodakis, Matthieu Cord, and Patrick Perez. OBoW: Online ohg-oo-visualwwrrt generation for self-supervised learning. In CVPR, 2021. 2, 3, 5, 6, 19, 20
Jean-Bastion Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Koray Kavukeuoglu, Rémi Munos, and Michal Valko. Bootstrap your own latent: A new approach to self-supervised learning. In NewrIPS, 2020. 2, 3
Kaiming Hc, Georgia Gkioxari, Piotr Dollár, and Ross Girshick. Mask R-CNN. In ICCV, 2017. 12
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In CVPR, 2020. 2, 3,
Kaiming Hc, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders scalable vision learners. In CVPR, 2022. 1, 2, 3, 7, 9, 11, 18
14

---

Published in Transactions on Machine Learning Research (02/2024)
Elad Hoffer, Tal Ben-Nun, Itay Hubara, Niv Giladi, Torsten Hoefter, and Daniel Soudry. Augment your batch: Improving generalization through instance repetition. In CVPR, 2020. 8
Zejiang Hou, Fei Sun, Yen-Kuang Chen, Yuan Xie, and
Kung. MILAN: Masked image pretraining on language assisted representation. arXiv preprint arXim:2008.0000 2022.
Gao Huang, Yu Sun, Zhuang Liu, Daniel Sedra, and Kilian Q Weinberger. Deep networks with stochastic depth. In ECCV, 2016. 21
Xiao Shi Huang, Felipe Perez, Jimmy Ba, and Maksims Volkovs. Improving transformer optimization through better initialization. In ICML, 2020. 3
loannis Kakogeorgion, Spyros Gidaris, Bill Psomas, Yannis Avrithis, Andrei Bursuc, Konstantinos Karantzalos, and Nikos Komodakis. What to hide from your students: Attention-guided masked image modeling. In ECCV, 2022. 3
Gustav Larsson, Michael Maire, and Gregory Shakhnarovich. Learning representations for automatic colorization. In ECCV, 2016. 3
Xiaotong Li, Yixiao Ge, Kun Yi, Zixuan Hu, Ying Shan, and Ling- Yu Duan. MC-BEIT: Multi-choice discretization for image BERT pre-training. Tn ECCV, 2022. 3
Zhaowen Li, Zhiyang Chen, Fan Yang, Wei Li, Yousong Zhu, Chaoyang Zhao, Rui Deng, Liwei Wu, Rui Zhao, Ming Tang, and Jinqiao Wang. MST: Masked self-supervieed transformer for visual representation. In NewrIPS, 2021. 1, 3
Liyuan Liu, Xiaodong Liu, Jianfeng Gao, Weizhu Chen, and Jiawei Han. Understanding the difficulty of training transformers. EMNLP, 2020.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin transformer: Hierarchical vision transformer using shifted windows. In ICCV, 2021. 3, 12
Lya Loshchilov and Frank Hutter. SGDR: Stochastic gradient descent with warm restarts. In ICLR, 2017.
20, 21
Πya Loshchilov and Frank Hutter. Decoupled weight decay regularization. ICLR, 2019. 8, 12, 20, 21
Julien Mairal. Cyanure: An open-source toolbox for empirical risk minimization for Python, C++, and soon more. arXiu preprint arXio:1912.08166 2019. 21
Shlok Mishra, Joshua Robinson, Huiwen Chang, David Jacobs, Aaron Sarna, Aaron Maschinot, and Dilip Krishman. A simple, efficient and scalable contrastive masked autoencoder for learning visual representations. arXiv preprint (xxi:22110116711 2022. 11
Ishan Misra and Laurens van der Maaten. Self-supervised learning of pretext-invariant representations. Tn CVPR, 2020. 3
Ishan Misra, C Lawrence Zitnick, and Martial Hebert. Shuffle and learn: unsupervised learning using temporal order verification. In ECCV, 2016. 3
Mchdi Noroozi and Paolo Favaro. Unsupervised learning of visual representations by solving jigsaw puzzles. ECCV, 2016.
Aaron van den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprinl arXiv:1807.037: 807.03777, 2018. 3
Deepak Pathak, Philipp Krahenbuhl, Jeff Donahue, Trevor Darrell, and Alexei A Efros. Context encoders: Feature learning by inpainting. In CVPR, 2016. 3
Zhiliang Peng, Li Dong, Hangbo Bao, Qixiang Ye, and Furu Wei. BETT v2: Masked image modeling with vector-quantized visual vokenizers. arXiv preprint arXiv:2208.00:: 06366, 2022. 1, 3,
15

---

Published in Transactions on Machine Learning Research (02/2024)
Alec Radford and Karthik Narasimhan. Improving language understanding by generative pre-training. 2018.
3
Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Nlya Sutskever. Language models are unsupervised multitask learners. 2019. 1, 3
Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever. Learning transferable visual models from natural language supervision. In ICML, 2021. 3
Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Nly Sutskever. Zeero-hoo text-to-image generation. In ICML, 2021. 3
Pierre Н Richemond, Jean-Bastien Grill, Florent Altché, Corentin Tallec, Florian Strub, Andrew Brock, Samuel Smith, Soham De, Razvan Pascamu, Bilal Piot, and Michal Valko. BYOL works even without batch statisties. arXiv preprint ar Xiv:20000 10244, 2020. 4
Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander Berg, and Li Fei-Fei. Imagenct large scale visual recognition challenge. IJCV, 2015. 8
Robin Strudel, Ricardo Garcia, Ivan Laptev, and Cordelia Schmid. Segmenter: Transformer for semantic segmentation. In ICCV, 2021. 11, 12, 21, 22
Christian Szegedy, Vincent Vanhoucke, Sergey Loffe, Jon Shlens, and Zbigniew Wojna. Rethinking the Inception architecture for computer vision. In CVPR, 2016. 21
Yonglong Tian, Chen Sun, Ben Poole, Dilip Krishnan, Cordelia Schmid, and Phillip Isola. What makes for good views for contrastive learning. In NewrIPS, 2020. 3
Hugo Touvron, Matthieu Cord, Matthijs Douze, Francisco Massa, Alexandre Sablayrolles, and Hervé Jégou. Training data-efficient image transformers & distillation through attention. In ICML, 2021a. 3
Hugo Touvron, Matthieu Cord, Alexandre Sablayrolles, Gabriel Synnaeve, and Hervé Jégou. Going deeper with image transformers. In ICCV, 2021b. 3
Hugo Touvron, Matthieu Cord, and Hervé Jégou. DeiT III: Revenge of the VIT. In ECCV, 2022. 3, 8
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. In NewrIPS, 2017. 1, 3, 20
Chen Wei, Haoqi Fan, Saining Xie, Chao-Yuan Wu, Alan Yuille, and Christoph Feichtenhofer. Masked feature prediction for self-supervised visual pre-training. In CVPR, 2022a. 3, 11
Longhui Wei, Lingxi Xie, Wengang Zhou, Houqiang Li, and Qi Tian. MVP: Multimodality-guiieee visual pro-training. In ECCV, 2022b.
Yixuan Wei, Han Hu, Zhenda Xie, Zheng Zhang, Yue Cao, Jianmin Bao, Dong Chen, and Baining Guo. Contrastive learning rivals masked image modeling in fine-tuning via feature distillation. ar preprinl arXiv:220.. 14144, 2022c.
Zhirong Wu, Yuanjun Xiong, Stella Yu, and Dahua Lin. Unsupervised feature learning via non-parametric instance-level discrimination. In CVPR, 2018. 3, 21
Zhenda Xie, Zheng Zhang, Yue Cao, Yutong Lin, Jianmin Bao, Zhuliang Yao, Qi Dai, and Han Hu. SimMIM: A simple framework for masked image modeling. In CVPR, 2022. 1, 3, 11
Ruibin Xiong, Yunchang Yang, Di He, Kai Zheng, Shuxin Zheng, Chen Xing, Huishuai Zhang, Yanyan Lan, Liwei Wang, and Tieyan Liu. On layer normalization in the transformer architecture. In TCML, 2020. 3
16

---

Published in Transactions on Machine Learning Research (02/2024)
Hongwei Xue, Peng Gao, Hongyang Li, Yu Qiao, Hao Sun, Houqiang Li, and Jiebo Luo. Stare at what you see: Masked image modeling without reconstruction. In CVPR, 2023. 3
Li Yuan, Yunpeng Chen, Tao Wang, Weihao Yu, Yujun Shi, Zi- Zi-Hang Jiang, Francis E.H. Tay, Jiashi Feng, and Shuicheng Yan. Tokens-to-token ViT: Training vision transformers from scratch imagenet. In ICCV, 2021.
Sangdoo Yun, Dongyoon Han, Seong Joon Oh, Sanghyuk Chun, Junsuk Choe, and Youngjoon Yoo. CutMix: Regularization strategy train strong classifiers with localizable features. In ICCV, 2019. 21
Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny. Barlow twins: Self-supervised learning via redundancy reduction. In ICML, 2021. 3
Shuangfei Zhai, Navdeep Jaitly, Jason Ramapuram, Dan Busbridge, Tatiana Likhomanenko, Joseph Yitan Cheng, Walter Talbott, Chen Huang, Hanlin Goh, and Joshua Susskind. Position prediction as an effective pretraining strategy. In ICML, 2022a. 1
Xiaohua Zhai, Alexander Kolesnikov, Neil Houlsby, and Lucas Beyer. Scaling vision transformers. In CVPR, 2022b.
Hongyi Zhang, Moustapha Cisse, Yann N Dauphin, and David Lopez- Paz. mixup: Beyond empirical risk minimization. In JCLR, 2018. 21
Hengshuang Zhao, Jiaya Jia, and Vladlen Koltun. Exploring self-attention for image recognition. In CVPR, 2020.
Jinghao Zhou, Chen Wei, Huiyu Wang, Wei Shen, Cihang Xie, Alan Yuille, and Tao Kong. iBOT: Image BERT pre-training with online tokenizer. In ICLR, 2022. 3, 11, 12, 21
17

---

Published in Transactions on Machine Learning Research (02/2024)
Appendix
A Additional experimental results
A.1 Training without the LIMG objective
In Tab. 8, we report more results from MOCA trainings with only the Lroc loss and the LIMG loss deactivated (i.e., setting À = 0). The decoder of these models implemented using: (a) the default approach described in Sec. 3.2, or (b) the condenscr- based approach described in Sec. 3.3. We also provide results for when having two decoders (c), one using the default approach and the other using the condenser-based approach, each trained with a separate LLOC loss.
We observe that, with only the LLOC loss, the default decoder (model (a)) is better than the condenser-based decoder (model (b)); which is the opposite from the behavior when both the LLoC and the LIMG objectives are used. The fact that the optimal choice for the decoder is different depending on whether LIMG used or not hints that the two objectives work together in a synergistic way. Nevertheless, even with the default decoder (model (a)) the k-NN and Linear probing performances are still quite low. However, model (c), which combines the two decoders, fares much better. Indeed, in this case, the model better bootstraps its learning process and, thus, learns better image representations.
We also compare with variants where the hide-and-prediin objectives are defined the pizel space with per-patch normalization as MAE (He et al., 2022) ("Pixels'). We see that in this pixel-reconstructionn case, the decoder type has small impact. Also, our resigmment-prodictin models (a) and (c) achieve significantly better k-NN and Linear probing performance than these rixel-reconstruction models, which shows the interest. in exploiting higher-level features.
A.2 Alternative codebook updating approach
During our pre-training, we update the codebook C by randomly selecting teacher-token embeddings from previous mini-batches. However, in this exploration, tested alternative approach using farthest point sampling (FPS) to select teacher-token embeddings for the codebook updates. Following the pre-training setup outlined in Sec. 4.1, the use of FPS resulted in decline in k-NN performance from 71.8% to 70.3%. We hypothesize that the reason for this drop in performance is that this FPS strategy is more prone to inserting outlier tokens to the codebook and/or the codebook changes more drastically from one iteration to the next, making the learning process harder.
A.3 Longer pre-training
As we see in Table 9, longer pre-training in MOCA improves performance, but with diminishing returns. This behavior is similar to other pre-training methods, where performance converges after a certain number of pre-training epochs. MOCA's advantage lies in its faster convergence, effectively reducing pre-training time.
A.4 Bigger backbones
As we see Tables 9 and 10, MOCA can scale with bigger backbones and achieves strong results while has smaller pre-training time.
B Creating the condenser-based decoder inputs
Here we describe more detail how we produce the inputs for the condenser-basee based decoder of MOCA models. For each masked view x", we create the decoder's input embeddings from the following student-produced embeddings:
18

---

Published in Transactions on Machine Learning Research (02/2024)
Table 8: Hide-and-predice models trained without the LIMG loss (A = 0). ImageNet-1k classification
results with ViT-B/ /6 models trained for 100 pre-training epochs with a single token masking round. For these models, we use the default and/or the condenser decoder versions.
Decoder approach
Top-1 accuracy
Default Condenser
Targets
k-NN
Linear
Assign.
34.4
51.6
(b)
Assign.
13.1
33.7
(c)
V
Assign.
52.5
64.2
(d) V 4
Pixels 14.6
38.8
(e)
Pixels
15.6
42.7
(f)
Pixels
13.9
41.9
Table Scaling pre-training epochs or model size. Top-1 1 accuracy results on ImageNet with ViT-B/16 ViT-L/16.
Method
#Epoch
k-NN
Lincar
Finctun.
1-shot
2-shot
5-shot
NOCA ViT-B/16
200
77.2
78.7
83.6
55.4
63,9
69.9
NOCA VTT-B/16
400
77.6
79.3
83.7
55.5
64.6
69.9
NOCA VTT-L/16
200
78.4
80.3
84.9
56.3
65.2
70.7
Table 10: VIT-L/16 comparisons. Top-1 accuracy results ImageNet. The time measurements based single 8 A100 GPU node. t: use of multiple crops (Caron et al., 2021).
Method
#Epoch
Time
k-NN
Linear
Finetun.
MAE
1600
187h
75.1
85.9
IBOTI
300
320h
78.0
81.0
84.8
MOCA
200
78.4
80.3
1. The average patch-token embedding at the output of S (i.c., the last transformer layer of S): S(x") = |..| Eieu, 5(x"); Note that in ViTs, after the last transformer layer, there is a LayerNorm layer. The patch-token embeddings (S(*"); bieu,
the outputs of this LayerNorm layer.
2. The patch-token embeddings (S"))))) from intermediate transformer layer of S. As before, these intermediaate-layr embeddings have been passed from an additional LayerNorm layer.
Before feeding these embeddings to the decoder, we first apply a linear projection layer p(-) in order to convert their dimensionality from dene to daec. Then, we append a learnable masked embedding e for each masked token i m, and add the (non-learnable -in-cos-based) positional embeddings {ef Jimo, where e is the positional embedding for the [AVG] token (i.e., for S(x")). This produces the following input set of tokens embeddings:
z" 2== u (p(S' (x*)) He.. Ne +efhimm.
JiEm.
(7)
In the case of partial decoding, the set m, in the above Eq. (7) is replaced by the subset m/ C mp-
C Additional implementation details
C.1 Weight generation modules G*()) and G*(-)
As described in the main paper (Sec. 3.2.1), for masked same-view token assignment prediction (see Eq. (2)) and the masked cross-view average assignment prediction (see Eq. (5)), use the weight generation modules (d(-) and G+)-) respectively. Like Gidaris et al. (2021), we use MLP networks with the following configuration:
(imear(1536,,00 L2Norm Linear Linear(768,, 1566,1536) BatchNorm + ReLU Linear (1536,d) d) L2Norm,
19

---

Published in Transactions on Machine Learning Research (02/2024)
Table 11: MOCA 's settings. We define the masking percentage as 100 Imal and the percentage of predicted tokens (i.e., partial decoding) as 100 l
(a) NOCA 's model setting
(b) MOCA 's optimization setting
Hyperparameter
Value
Hyperparameter
Value/Method
Codebook size K
4096
Teacher momentum a
0.99
+ 1.0 (cosine annealing)
Codebook words Knew per training step
4
optimizer
AdamW (Lowhchilov & Hutter, 2019)
Masking mask percentage for 1st round
55%
base learning rate
1.5e-4
Masking mask percentage for 2nd round
75%
weight decay
0.05
Masking - the percentage of predicted tokens
20%
optimizer momentum
8i = 0.9, A2 = 0.999
Decoder- intermediate layer& E for input tokens
8
batch size
2048
Decoder depth
2
learning rate schedule
cosine decay Loshchitov Hutter (2017)
Decoder embedding size ddec
512
warmup epochs
30
Decoder self-attention heads
16
Loss weighting parameter A
0.5
Table 12: ImageNet-lk linear classification.
Hyperparameter
Value/Method
optimizer
SGD
base learning rate
0.04
weight decay
0.0
optimizer momentum
0.9
batch size
1024
training epochs
100
learning rate schedule
cosine decay Losheialov Hutser (2017)
augmentation
RandomResizedCrro
where d is 5512=512 for G+(-) and deno=768 for G²(-). Also, these MLPs have a skip connection that begins after the input L2Norm unit and concludes before the output L2Norm unit. Given as imput the codebook C [cI...... cx), G.) and G-) produce the prototypes weights Wd (C) (c)) (cx)) and 4h = (G ..... (cx)) respectively. The Zz-normalizat in the L2Norm layers is across the channels dimension and the normalization in the BatchNorm layers is across the K codebook items, which is the batch dimension in these MLP networks.
C.2 Pre-training MOCA
Here we share additional implementation details for the pre-training of our MOCA representations. In Tab. 11a we provide the implementation details for the ViT-B/16-base MOCA model that we used for producing the results of Sec. 4.4 in the main paper and in Tab. 11b the optimization setting for its training. For the ViT-B/16, we follow the standard VIT architecture (Dosovitskiy et al., 2021) and implement it with 12 transformer layers, 12 attention heads, and denc 768 channels. We use sine-cosine positional embeddings (Vaswani et al., 2017) for both the encoder and the decoder inputs. Also, as in OBoW (Gidaris et al., 2021), at the teacher side, during the average pooling operation for computing the reduced assignment vectors, ignore the token assignment vectors the edge border patch tokens (i.e., we ignore the rows/columns of border patch tokens from each image side).
In order to generate the two unmasked random views x² and x2, input for the teacher, we employed the widely used augmentation strategy found most prior works, such as SimCLR (Chen et al., 2020a) and MoCo- v3 (Chen et al., 2021). The PyTorch pscudo-code for this augmentation stategy provided in Appendix C.4. To generate the masked views that the student gets as input, we simply applied random token masking the x and x random views.
20

---

Published in Transactions on Machine Learning Research (02/2024)
Table 13: Full fine-tuning for ImageNet-lk classification.
Hyperparameter
Value/Metin
optimizer
AdamW (Loshchllov & Hutter, 2019)
base learning rate
1.0
weight decay
0.05
optimizer momentum
B1 = 0.9, 0g
0.999
layer-wise Ir decay (Clazk et nl., 2020; Buo al., 2022)
0.65
batch size
1024
learning rate schedule training epochs
cosine decay (Loshchilov & Hutter, 2017)
100
warmup epochs
5
drop path (Huanx ah., label smoothing Seexedy (2016)
0.2
0.1
mixup (Zhang al, 2018)
0.8
cutmix (Yon ah, 2019) augmentation
1.0
RandAug (9,0.5) (Cubux at., 2020)
C.3 Evaluation protocols
Here we report implementation details for the evaluation protocols that we used in our work. We note that in all of them we use the pre-trained teacher transformer.
K-NN ImageNet classification. For the k-NN evaluation protocol (Wu et al., 2018; Caron et al., 2021) we freeze the pre-trained transformer, extract features from the training images, and then for the test images use a k-nearest neighbor classifier with = 20 neighbors.
Linear ImageNett classification. In this case, we again freese the pre-trained transformer and train a linear classifier on the extracted training image features. We follow the linear classification setting of iBTT (Zhou et al., 2022). In Tab. 12 we provide implementation details and hyperparameter values.
End-to-end fine-tuning on ImageNet classification. We follow the end-to-end fine-tuning setting of iBOT (Zhou et al., 2022). In Tab. 13 we provide implementation details and hyperparameter values.
Low-shot ImageNet classification. Following MSN (Assran et al., 2022), we freeze the pre- tra-trained transformer, extract image features from the few available training images, and then train linear classifier with them using L2-regularized logistic regression. For the logistic regression, we use the cyanure package (Mairal, 2019).
Cityscapes semantic segmentation with linear probing. We study the quality of learned represen- tations using linear probing either with the entire Cityscapes (Cordts et al., 2016) dataset containing 2975 training images or in the few-shot setup. For all experiments we evaluate the 500 validation images and report the best mean intersection over union (mIoU). We use the Segmenter (Strudel et al., 2021) model with a frozen ViT-B/16 backbone initialized with one of the studied approaches. Then, we append single learnable linear layer on top.
Entire dalasel: When training linear probes using the entire dataset, optimize the network for 216 epochs following the protocol of Strudel et al. (2021). We optimize the learning rate for every studied method separately.
In the few-shol selup, we randomly sample three training subsets of the Cityscapes dataset that are shared across all the experiments with different methods. We use either splits of 100 or 374 training images (French al., 2020). For every method, search for the optimal learning rate on the first split and then use it for the remaining two training splits. We optimize every split for 100 epochs.
Cityscapes semantic segmentation with fine-tuning. Furthermore, we study different methods when used initialization of the Segmenter backbone when fine-tuning the entire network. In this setup,
21

---

Published in Transactions on Machine Learning Research (02/2024)
use single layer of mask transformer (Strudel et al., 2021) as decoder. As in the linear probe setup, we use either the entire dataset or 100-/372-large training splits and report the best mloU on the validation set.
Entire dataset: We train all methods for 216 epochs as done by Strudel et al. (2021) and optimize the learning rate for each method separately.
Few-shoo setup: We use the identical training splits as the few-shot linear probing and search for the optimal learning rate in the same way. We optimize every method for 216 epochs.
C.4 Image augmentations pseudo-code
Here we provide PyTTrrch pseudo-code for the image augmentations used in MOCA for generating the two unmasked random views x² and x².
import torchvision transforms T normalize = T. Normalize (memammmeam0a))00 0.456, 0.406), std= 229., 022 0.225)
aug_viewl viewl = T. Compoee
T. RandomResizeedCrop (224, scale=(0.2, 1.)),
T. (aallmmpppyy(T..CCooritterr (0.4, 0.4, 0.2, 0. 9)l p=0.8).
T. (andommrrayscalle(0002.))
(1 0.1, 2.0),
T. RandonHorizontaa111pp))
T. ToTensor
normalize])
aug_viewe = T. Compose
T. RandomResizedCrop (224, scale=(0.2, 1.)).
T. (nddompplly(T...oorr))ttee (0.4, 0.4, 0.2, 0.1)], p=0.8),
T. Randomdrayscaalepss0..)) (0ussianbluur 0.1, 2.0),
Solarization
T. RandomforizonntaFF1i...
T. ToTensor normalize])
22


Reviews:
[{"review_id": "T0olNwxoNq", "reviewer": "Reviewer_J994", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strengths:\n- Self-supervision is expensive to train. Being able to reduce training time while showing better performance is a good direction and helpful to the community. \n- The proposed approach picks up some interesting bits from prior works (SwAV, DINO, MAE..) with some novel ideas and combines them in a unified framework.\n- Overall, the paper is well written. Experiments and ablations are well thought of and executed. \n- Strategy for partial decoding and multiple masking rounds is quite interesting. \n- Promising results and comparison with prior works in this area. \n\nWeaknesses/Questions:\n- Unclear why student and teachers work with different 'prototypes'. More intuitions and details on this would be helpful. A short paragraph on how this relates to an approach like SwAV would also help with the motivation. \n- Currently the global loss operates on an augmented version of the image and a masked version of the image. While masking can be considered to be an augmentation, it is a lot more extreme. Did you also try to add the IMG loss on unmasked student inputs (like how discriminative methods did in the past) ?\n- More intuitions (and perhaps a figure) needed for the 'condenser' module since it significantly affects the performance. \n- I think the paper needs more transfer learning and few-shot learning experiments. You can perhaps use a benchmark like this [a] or a more recent benchmark to compare discriminative and MIM approaches. \n- Missing comparisons on MaskFeat ? \n\n[a] How Well Do Self-Supervised Models Transfer?", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes a self-supervised pre-training strategy for ViTs. The approach involves using a masking and predict strategy where the teacher produces code book assignments from unmasked views and the student is trained to predict them using masked views. Additionally, an image-wise loss on averaged features is also added to enforce global consistency hoping to learn good linearly separable features. \n\nAuthors show that the approach sees a fast convergence which leading to better results on selected benchmarks.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please refer to the weakness section for comments. Apart from some clarity fixes which might be easy to do, a more extensive transfer learning evaluation would help strengthen the paper.", "broader_impact_concerns": "The authors have provided a broader impact section - don't have any additional concerns here.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "tBr8SjImze", "reviewer": "Reviewer_JVWW", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "The paper proposed a novel framework that combines discriminative and generative pretraining so the advantage from both sides can be gained, i.e., the learned features are both ready-to-use (leading to high kNN and linear acc) and detailed (leading to high finetuned acc) features. As a result, the training efficiency and data efficiency are improved.\n\nMy concerns include:\n\n1. I find that the proposed method is more data-efficient than existing methods very interesting. However, I find little insight to help understand why it performs particularly better in the low-shot settings but is not so good in the data-sufficient settings.\n    \n2. From my point of view on self-supervised pretraining, the evaluation protocol of finetuning (including low-shot) and transfer learning is more important than kNN and linear probing, as they are more practical and show the value of pretraining. However, over half of this paper’s comparisons and conclusions are based on kNN and linear accuracies. For example, in Table 3, the model performs better without the condenser in finetuning. Plus, I hope to see more ablation studies on semantic segmentation experiments to learn the effect of each component of the framework. It is very likely that I am biased, and I am eager to hear from the authors about why they believe kNN and linear accuracies are so important.\n\n3. Some discussions may be questionable. For example, in Section 4.1, the authors say “On the other hand, MAE’s decoder, which must reconstruct pixels, is better to be deeper and thus spare the MAE’s encoder from capturing such low-level details.” However, in SimMIM, which also reconstructs pixels, only one linear layer is used in the decoder.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This work deals with the challenge of learning self-supervised visual representation. It proposes to unify the discriminative and generative pretraining frameworks so the learned features are both semantically meaningful and contextual reasoning. Experimental results show that the proposed method is more data-efficient than previous methods.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "Please see the weaknesses stated above.", "broader_impact_concerns": "I have no concerns about the ethical implications of the work.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "nSzzgT2JcI", "reviewer": "Reviewer_hvsU", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength:\n\n1. Self-supervised pre-training is a very important task in representation learning. An improved self-supervised learning method can potentially contribute to many downstream tasks.\n\nWeakness:\n\n1. It is not clear what data augmentation is applied for student and teacher views respectively. \n\n2. The specific design of updating the memory bank storing teacher-token deserves further investigation. Randomly sampling token embeddings to replace old tokens in the memory is sensible to non-i.i.d. sampling of training samples. A more stable way of updating the memory bank might be using farthest sampling to select most diverse tokens to update the memory bank.\n\n3. The fine-tuning performance between different self-supervised learning approaches are rather close according to Tab. 4 (a). Therefore, it begs the question whether the more complicated self-supervised pre-training is really necessary.\n\n4. Downstream tasks are only evaluated on two datasets, i.e. ImageNet classification and CityScapes segmentation. The evaluation is far from enough for the purpose of extensively evaluating the effectiveness of a new self-supervised pre-training method.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "No", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "This paper introduced a self-supervised learning approach that unifies masking-and-prediction and invariance to perturbation paradigms. In particular, a teacher-student framework is introduced and the student is forced to mimic the predictions made by the teacher network on unlabeled data for representation learning. Losses are defined at two levels where both global and local (dense) representations are aligned. Experiments on Knn classification, linear probing and fine-tuning on two datasets, including ImageNet-1k and CityScapes, suggest the superiority of the proposed method.", "claims_and_evidence_yn": "No", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "1. It is necessary to present details of data augmentation applied. It is also highly recommended to evaluate how different augmentation strategy could affect the results.\n\n2. The specific design of memory bank update needs to be further elaborated. I suggest evaluating update the memory bank with farthest sampling strategy.\n\n3. More diverse downstream tasks are for sure necessary. These could include more segmentation tasks, e.g. Pascal VOC, MS COCO, etc. as dense contextual reasoning is claimed to be the novelty.", "broader_impact_concerns": "There is no concern over ethical implications.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}, {"review_id": "I8xWhqllDg", "reviewer": "Reviewer_fwxy", "summary": "", "questions": "", "limitations": "", "rating": null, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "Strength\n\n1. The proposed idea is very interesting. Codebook models were very popular ten years ago in computer vision. The authors combine the idea of codebooks into self-supervised learning which I think is very smart and novel.\n2. The results are impressive. The proposed method basically outperform existing SOTA methods in SSL. This demonstrates the authors' hypothesis that a moderate level of perturbation invariance in feature representation helps with self supervised learning.\n3. The paper is nicely written with sufficient experimental validation and ablation study.\n\nWeakness\n\n1. Although I feel the paper is generally well written, the method intro can be a bit more clear. Some key procedures for codebook construction and assignments are not emphasized or highlighted, which I believe are core questions for codebook methods.\n2. Although the approach combines codebook idea and masking, related to weakness 1, I generally feel the proposed method is a bit complicated. Without sufficient differences, I am worried the proposed method may be not adopted by practitioners due to its addition complication.", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "Yes", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "paper_topic_and_main_contributions": "", "excitement": null, "reproducibility": null, "ethical_concerns": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "summary_of_contributions": "The paper proposes self-supervised learning techniques based on mask-and-predict and the codebook models. The idea is to mask the image randomly and then aggregate the feature representation using codebooks, and then use the masked embedding to predict the non-masked codebook representation (under augmentations). The idea is interesting and sort of combining new methods with old techniques. Codebook methods provide some level of perturbation invariance to the learning procedure. The paper performed experiments and demonstrated the effectiveness of the proposed method (MOCA). The results are impressive.\n\nIn general, I think the paper has novelty and shows promise, with some concerns on the complication. Overall positive.", "claims_and_evidence_yn": "Yes", "claims_explanation": "", "audience_yn": "Yes", "audience_explanation": "", "requested_changes": "I think the introduction of the methods can be improved by highlighting the following questions:\n\n1. How does the codebook construction look like? There are teachers and students. Do they share the same codebook? If the system is cold-started, then how does the codebook look like in the beginning? I suggest organize a separate section just about codebook construction.\n\n2. Codebook methods are usually resulting into a histogram but I don't see the proposed method goes that way. Do you only quantize the feature representations so that it becomes more stable?\n\n3. How does global average pooling contribute to the final performance?\n\n4. It is surprising to me the optimal lambda is 0.5, can you elaborate on this? Have you evaluated it on multiple datasets and how would its optimal value change across different datasets?", "broader_impact_concerns": "The authors have stated the concerns clearly in the paper which I tend to agree with. There will be bias and copyright issues coming up with these SSL approaches.", "additional_comments": "", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: Overall there seems to be agreement that this is a useful addition to the self-supervised learning literature, thanks to a few novel ideas as well as novel combination of previous ideas, that together lead to more efficient training of vision transformers.  Execution is also praised -- writing, ablation study -- and promising results.
Decision: Accept (poster)