MAST: Masked Augmentation Subspace Training for Generalizable Self-Supervised Priors
Authors: Chen Huang
Hanlin Goh
Jiatao Gu
Joshua M. Susskind
Venue: ICLR 2023 poster (2023)
Keywords: Self-Supervised Learning (SSL)
Generalization
Computer vision

Abstract: Recent Self-Supervised Learning (SSL) methods are able to learn feature representations that are invariant to different data augmentations, which can then be transferred to downstream tasks of interest. However, different downstream tasks require different invariances for their best performance, so the optimal choice of augmentations for SSL depends on the target task. In this paper, we aim to learn self-supervised features that generalize well across a variety of downstream tasks (e.g., object classification, detection and instance segmentation) without knowing any task information beforehand. We do so by Masked Augmentation Subspace Training (or MAST) to encode in the single feature space the priors from different data augmentations in a factorized way. Specifically, we disentangle the feature space into separate subspaces, each induced by a learnable mask that selects relevant feature dimensions to model invariance to a specific augmentation. We show the success of MAST in jointly capturing generalizable priors from different augmentations, using both unique and shared features across the subspaces. We further show that MAST benefits from uncertainty modeling to reweight ambiguous samples from strong augmentations that may cause similarity mismatch in each subspace. Experiments demonstrate that MAST consistently improves generalization on various downstream tasks, while being task-agnostic and efficient during SSL. We also provide interesting insights about how different augmentations are related and how uncertainty reflects learning difficulty.

Published as a conference paper at ICLR 2023
MAST: MASKED AUGMENTATION SUBSPACE TRAIN- ING FOR GENERALIZABLE SELF- PRIORS
Chen Huang, Hanlin Goh, Jiatao Gu & Josh Susskind
Apple Inc.
(chen-huang, hanling in, jgu32, j
ABSTRACT
Recent Self-Supervised Learning (SSL) methods are able learn feature repre- sentations that are invariant to different data augmentations, which can then be transferred to downstream tasks of interest. However, different downstream tasks require different invariances for their best performance, the optimal choice of augmentations for SSL depends on the target task. In this paper, we aim to learn self-superviséd features that generalize well across variety of downstream tasks (e.g., object classification, detection and instance segmentation) without knowing any task information beforehand. We do so by Masked Augmentation Subspace Training (or MAST) to encode in the single feature space the priors from different data augmentations in factorized way. Specifically, space into separate subspaces, each induced by a learnable mask that selects rele- vant feature dimensions model invariance to specific augmentation. We show the success of MAST in jointly capturing generalizable priors from different aug- mentations, using both unique and shared features across the subspaces. We fur- ther show that MAST benefits from uncertainty modeling reweight ambiguous samples from strong augmentations that may cause similarity mismatch each subspace. Experiments demonstrate that MAST consistently improves generaliza- tion on various downstream tasks, while being task- agnosag and efficient during SSL. We also provide interesting insights about how different augmentations are related and how uncertainty reflects learning difficulty.
disentangle the feature
INTRODUCTION
Self-Supervised Learning (SSL) for image representation has made significant progress over the past few years. The feature representations are typically learned to be invariant to different data augmentations (e.g., Random Flip and Color Jitter). For example, the popular contrastive SSL methods (Chen et | 2020aa He al. 2020) learn invariances by discriminating augmented views of the same image (positive pair) from those of different images (negative pair), while recent non- contrastive SSL methods (Chen & He 2021 Grill et al.) 2020: Bardes al. 2022) simply maximize the similarity between positive pairs. Such learned features are shown to generalize across many downstream tasks, including classification, object detection, instance segmentation, etc.
Despite achieving strong transfer performance, we lack a good theoretical understanding about why both contrastive and non-contrastive SSL methods generalize so well. Balestriero & LeCun (2022) recently proposed a unified framework that demonstrates the key for generalization lies in the align- ment between the pairwise relation in SSL (characterized by augmented inputs) and downstream task. This is also in line with other theories (Arora et al. 2019 HaoChen et al. (221)) that quantify how data augmentations implicitly encode the class distributions downstream tasks. Motivated by these theoretical analyses, aim at working SSL method that can directly capture mcaningful priors from data augmentations in order to encourage generalization for a range of tasks.
Invariance, achieved through augmentation, is useful mechanism to facilitate generalization For example, one can imagine that invariance to Random Flip will boost generalization on many vision tasks. However, shown in Fig. the optimal set of augmentations (thus invariances) highly de- pends on the downstream task, which has been similarly observed in (Tian et al.) 2020). Sometimes, invariances (e.g., to Color Jitter) that prove helpful for one downstream task (e.g., object detection)

---

Published as a conference paper at ICLR 2023
ImageNet classity (Aco)
CUB-200 birds classity (Acc)
VOC detection (APw)
COCO segmentation (AP)
0
Palaihe performace
-20
-1.0
-0.6
-15
Augmemattinns ramoved
Augmentations removed
Augmentations removed
Augrientatione removed
Hemove Color Jftter
Pemove Goussian Blu
Remove Fladom Flip
Remove Grayscale
Remove Rendom
The change of downstream performance per task due to the selective removal of each augmentation during SSI., relative to one baseline SSL. method MoCo (He et al., 2020) trained with 5 standard data augmentations, including Color Jitter, Gaussian Blur, Random Flip, Random Grayscale, RandomResizedCrop. Performance drop after removing one augmentation indicates that it helps the corresponding task, while performance gain indicates harm. We observe that Random- ResizedCrop strongly benefits all tasks. Other augmentations encode their specific invariances that can be more helpful for some tasks than others. Some augmentations (e.g., Color Jitter) even hurt one task (CUB-200 bird classification) despite being quite helpful for other tasks.
Figure 1:
may even hurt generalization on another (e.g., birds classification which requires accurate color information of similarly shaped bird species). Hence it is impossible to maximize generalization by finding good augmentations for SSL in a task-dependent way, not only because the invariances learned across augmentations may contradict each other, but also because we often do not know the target task priori during SSL. Furthermore, manually finding suitable augmentations for general- izing to new task is quickly cumbersome.
In this paper, we propose new SSL method to learn generalizable features without presuming any downstream task information. Our method called Masked Augmentation Subspace Training (MAST) achieves this goal by learning single but disentangled feature space to encode a set of potentially contradicting augmentation invariances. For each augmentation, to it in a specialized feature subspace, which is induced from the full space with a learned masking operation. This allows us learn unique features for each augmentation subspace well as features shared between them. There are two main benefits of such subspace training: 1) we explicitly avoid feature suppression (Li et al. 2020) by jointly training separate feature subspaces, such that the learning of one subspace will not compromise that of another; 2) we obtain a disentangled, full feature space that is pre-trained in task-agnostic way, but does not discard the diverse information (i.e., invariances) required by all possible downstream tasks. We further model uncertainties in the feature representations to reduce harm of ambiguous samples from strong augmentations.
learn invariance
In order to examine how representation effectiveness scales with augmentation diversity, we run experiments with different numbers of augmentations, starting from standard augmentations used typically in SSL (Chen et al., 2020a)), and extending to an additional 10 augmentations (totaling 15). Note Tian et al. (2020) and Wang & Qi (2021) also learn from stronger augmentations, but not in our factorized and imcertainty-awar fashion. When it comes to transfer learning down- stream tasks, we simply drop our subspace masks and finetune the full feature representations for high efficiency. This is in contrast to LooC (Xiao et al., 2021 which needs to combine multiple feature "heads", with one head being invariant to all augmentations and other heads being sensitive to particular augmentation but invariant to others. Both the "leave-one-out" training and feature ensembling strategies in LooC lead to redundancy of parameters and high cost (thus only allowing a few augmentation-speen
heads).
Experiments show that MAST, while being efficient and task-agnostic, achieves state-of-the-ard transfer performance diverse downstream vision tasks. Investigations of the subspace masks and uncertainties also provide interesting insights in how different augmentations are related, and in how uncertainty reflects learning difficulty to avoid similarity mismatch during invariance learning.
To summarize, here are our main contributions:
We introduce MAST to make SSL. representations disentangled and uncertainty-awant to effectively encode different augmentation invariances for good generalization.
We show MAST is efficient, is resistant to feature suppression, and achieves state-of tfetth--rt downstream performance on diverse vision tasks without presuming any task information during pre-training.
2

---

Published as a conference paper at ICLR 2023
We provide interesting insights about how different augmentations related in SSL and how uncertainty reflects learning difficulty and impacts learning adaptively.
2 RELATED WORK
Self-Supervised Learning. Contrastive learning is one popular form of SSL, which contrasts differ- ent views of input (positive pairs) against other inputs (negative pairs). Recent contrastive meth- ods mainly differ in how they sample negative pairs, e.g., from a large batch as in SimCLR (Chen et al., 2020a) or from memory bank as in MoCo (v2) (He al., 2020) Chen al.] 2020b). One downside of contrastive methods is that they require a multitude of negative pairs to work well, leading to memory inefficiency. Recent non-contrastive SSL methods such as BYOL (Grill al.) 2020). SimSiam (Chen & He| 202T) and OBoW (Gidaris et al.) (02T) improve by not using nega- tive pairs. They only maximize the consistency between positive pairs teacher-student frame- work, and some variants like SwAV (Caron et al. 2020). Barlow Twins (Zbontar et al.) 202T) and VICReg (Bardes et al. 2022) achieve state-of-the-arr performance on several downstream tasks. Note et al. (2020) and Wang & (20TT) point out the important role of using strong data aug- mentations on generalization, but none of the aforementioned works explore the connection between learning factorized augmentation invariances and generalization.
Theoretical analysis on generalization. There has been significant interest in theoretically explain- ing the success of contrastive SSL from various perspectives, e.g., mutual information maximiza- tion (Bachman et al. 2019) or dimension collapse (Jing et al. 2022). Most relevant to our work theoretical studies (Arora et al. 2019 HaoChen et al. 202 that show the augmentation dis- tributions of same-class inputs have sufficient overlap when the costrastive loss is small, in which case data augmentations can implicitly encode the class structure for downstream tasks. [allestriero & LeCun (2022) further prove that generalization depends on the alignment between pairwise data relations with augmentations and downstream task. Inspired by these theories based on data and aug- mentation distributions, we propose to improve the generalization of SSL by factorized learning of augmentation properties well as uncertainty modeling. The factorization allows each downstream task to choose which underlying symmetries in the data require specific invariances.
Empirical works on generalization. Designing better augmentation strategies is one effective way to create task-relevant inductive biases for generalization. TTia et al. (2020) propose to learn the augmented views, but require labeled data from the target task, which is not scalable. Another line of works study the feature suppression phenomenon in SSL that can negatively impact downstream performance via "shortcuts" (i.e., by suppressing useful predictive features). To avoid shortcut so- lutions, Robinson et al. (2021) propose implicit feature modification to encourage instance discrim- ination using a wider variety of features; while Li et al. (2020) force contrastive learning to learn features that also predict the input, without discarding important information. Our method au- tomatically avoids shorteuts by jointly training multiple masked feature subspaces within the full space. Similar our subspace training idea LooC (Xiao et al.) 2021] where both augmentation- variant and -invariant feature spaces are learned and then concatenated to support different tasks. By contrast, our MAST is parameter-efficient by learning single feature space, can capture a lot more augmentation invariances efficiently in masked subspaces, and is aware of feature uncertainties.
Uncertainty modeling. Uncertainty-awet aware methods (Gal, 2016) have been extensively studied in supervised learning to handle ambiguous or noisy inputs. For SSL, only few studies estimate un- certainties helpful in specific applications, such as self-supervised depth estimation (Poggi et al. 2020) and multiview stereo (Xu et al., 2021). General methods for uncertainty-
SSI are lack- ing in the literature. One exception is variational inference SimSiam (Nakamura et al. 2022), a probabilistic extension of the non-contrastive SSL method SimSiam (Chen & He [222] I that models uncertainty using spherical posterior distributions. Here we propose a simple alternative for uncer- tainty modeling in our SSL framework, where a Gaussian feature assumption suffices in each of our mugmentation-speaae specific subspaces.
MAST: MASKED AUGMENTATION SUBSPACE TRAINING
We propose new SSL method to learn generalizable priors without presuming any downstream task information. To achieve this goal, we identify three desiderata for our method: 1) effective in captur- ing distinct priors (about invariances) from different augmentations. This will ease transfer learning
3

---

Published as a conference paper at ICLR 2023
Gaussian embedding
Augmentation- specific mask
Masked subspace
Augmentation |T|= =
Representation
K
R
View
encoder
projector
m
Ai 2
T
Image.x
Loss
f- T
View V
encoder
projector
Masks
Masked Augmentation Subspace Training for SSL. Given input image x, two augmented views and v' are produced by composing K image augmentations sampled from distribution T. Then different views are encoded into their feature representations, followed by projection into Gaussian embeddings ((, Σ) encoding uncertainty. For a disentangled learning of K types of augmentation invariances, we jointly learn mask my for each augmentation attend to the corresponding subspace where a similarity loss is computed. After pretraining, everything but the encoder discarded. The representations of the encoder are used for downstream tasks, without requiring overhead from ensembling subspace features or masks during transfer learning.
Figure 2: MAST:
on different downstream tasks, where each task automatically selects useful priors in the pre- trained representations via finetuning; 2) free of feature suppression (Li et al. 2020). i.e., learning priors in some features comes no cost of suppressing others. Particularly, we should accommodate learn- ing different invariances that can be either universally helpful for all downstream tasks or that would otherwise contradict each other (as shown in Fig. [T: 3) efficient enough to learn a large number of priors without too much cost (time or memory), and incurs overhead for transfer learning.
3.1 DESCRIPTION OF MAST
We achieve the desiderata above by MAST, see the schematic in Fig. Assume we have sam- pled batch of n images, and an augmentation set TT consisting of K atomic augmentation operators o1,..., OK such as RandomResizedCrrg and Color Jitter. For each image x, we produce two aug- mented views e = t(x) and o' t'(x) by sampling random compositions of image augmentations t T(or,..,xx) and t' T(o......x)) with their specific augmentation parameters. We encode the views and v by an encoder I into their representations = f(v) and y = [(v'), which are then g(v) and E' = =gyo))
mapped by projector y onto embeddings
Typically, a similarity loss computed between and z' learn augmentation invariances. How- ever, this makes it hard to learn distinct invariances because learning is unfactorized in the presence of the compositions of K augmentations applied on æ. Even worse, unfactorized learning from contradicting augmentation compositions could cause feature suppression that hurts generalization.
Disentangled learning is the main idea behind our MAST method to learn different invariances without feature suppression. Let Z [zi,...,22, 2n,] and Z' [=]...., 211 be the two batches of d-dimensional embeddings of the augmented views V = [u.....22 Un] and V = [of..... r. We build MAST top of the VICReg framework (Bardes al. 2022) which minimizes the distance between the embeddings of different views, with two regularization terms:
((Z, 7') AD Z') + alvar (7, Z') + (2, 2'),
where D') enforces augmentation invariance in form of the mean-sop meca--quared euclidean distance E, || </11B/n. without any feature normalization. The variance term Lver (2, Z') keeps the information content of each embedding (in Z or Z') above certain level. While the covariance term Leoe (Z, Z') is introduced to prevent informational collapse by decorrelating different dimensions in each embedding. X, a and 3 are the loss coefficients.
The above variance and covariance terms result in maximum information embeddings that are well suited for disentangled learning. As such, we propose to disentangle the embedding space into dis- tinct subspaces, each of which captures invariance to a specific augmentation. Instead of using bulky MLP heads to generate subspaces, we introduce K masks M E IR K over the embedding space at a much lower cost. Specifically, each mask m = M. is associated with a fixed augmentation operator og from set T. and m has the same dimensionality d as the embedding vector 2;-
To attend to the kth subspace, we only need to apply an element-wise masking operation @ using the corresponding mask my select relevant embedding dimensions (i.e, 2; m.). Therefore, to learn
4

---

Published as a conference paper at ICLR 2023
(a)
Mosk for Chsor -Ntue
(b)
my Gousskan Str
Misk for
Stakk 1i for Random Grayssate
Mi, CescooCopp
2000
Fintwiriin Giimer
Subepace invariant to Rendom Grayacale
Subspace invariant to Random/ieeseedCCpp
Figure 3: (a) Visualizing the masks learned for augmentation subspaces on ImageNet. The masks disentangle the full embedding space but allow shared dimensions between different subspaces. (b) 2D visualizations of two example subspaces on CUB-200 dataset (Wah et al. 2011). The subspaces are obtained by masking the ferent colors correspond to different classes. We observe meaningful subspaces with desired invari- ances: in the subspace learned be invariant to the RandomResizedCopp augmentation, classes are reasonably separated probably because random image patches can exploit color cues that are im- portant for fine-grained birds classification; while the subspace invariant the Random Grayscale augmentation is less discriminative (only for this particular task), because such invariance discards color information so one has to use other cues, e.g., birds' pose and shape for classification.
feature embeddings using the learned masks. Dif-
disentangled augmentation invariances from the paired views V and V', we replace the unfactorized distance function in Eq. ① with the sum of K masked embedding distances:
Dm (2
1 n1
Σ
2 mm 2 m&
(M) M|1
where we mask K times for the original embedding x; of an augmented view with K augmentations. The term (M) helps to encourage sparse masks. Note the embeddings and masks are jointly trained (see Section for more training details). Fig. shoows example masks that induce 5 distinct but overlapping subspaces. This not only avoids potential feature suppression separate subspaces, but also encourages collaborative learning in the shared embedding dimensions. F (Appendix reaffirms that different augmentation subspaces benefit different tasks.
Eventually, we obtain a single but disentangled embedding space that stores different invariance priors for generalization. Note the disentanglement at the embedding level will have a disentangle- ment effect at the representation level. Thus, only use the feature representations for testing on downstream tasks, dropping all the pre-trained embeddings and subspace masks. As a result, we introduce no computation overhead from subspace ensembling during transfer learning.
Uncertainty modeling plays an important role when the input is ambiguous (e.g., due to strong augmentations), but still remains largely unexplored problem in SSL. We address this issue with a simple method that explicitly models uncertainty using Gaussian embeddings. Concretely, instead of having projector g y z to map representations y to deterministic embeddings = E Rd. we extend to the projector with mean j E Rd and covariance Σ E d We consider the uncertainty independently for every dimension of the mean embedding μ.
+ (n, Σ) that produces a multivariate Gaussian distribution N(u, Σ) be diagonal for efficiency, modeling
With this independence assumption, we can use masking to easily induce embedding subspaces that are still Gaussian, i.e., .' (ig,, Ex) where ã p O my and Σi Σ @ diag(m)). We find the Gaussian embeddings suffice to model the uncertainties of augmented views in each subspace where only one type of augmentation is involved. Then we further improve Eq. (2) learn factorized augmentation invariances in an uncertainty-awart manner:
2 1 tn tr + tr(
Dmg (V. M)
(4)
where pi and nC are the mean embeddings of views v; and of in the ktth subspace, and their dis- tance is re-weighted by their uncertainty measures (trace tr(-) of covariance matrix) to determine
5

---

Published as a conference paper at ICLR 2023
the closeness of the two Gaussians. Intuitively, when the embeddingg have high uncertainties, the similarity loss should be downweighted to prevent similarity mismatch between ambiguous views.
To stablize the learning of uncertainties, we introduce an extra regularization term to encourage consistency between the predicted Gaussian embeddings Ni 2s) and ( 2,) for views v and o/. This term is based symmetric KL. divergence, with closed-form expression for Gaussians:
1 Vv(V.V")
DRI(NGILN)) DxLL(V)I
(5)
i=1
tr (E) (RI(N( 2
+ (14 - " m.) + log
(6)
[
Combining Eqs. (3-6) and substituting them into Eq. Q arrive final loss function:
{(V, :) mD (V,V:; V" M) + (p)) (V, V*) alear (V, VV Bleov (V,V). (7)
3.2 IMPLEMENTATION DETAILS
Architecture and hyper-parameters. Our pretraining is mostly performed the unlabeled Ima- geNet dataset (Deng et al. 2009) with 1000 classes unless otherwise stated. For fair comparison, the encoder f is the widely-adopted ResNet-50 backbone (He et al., 2016) which outputs 2048- dimensional feature representations. The projector has two fully-connected layers of size 8192, followed by separate fully-con culle-connected heads that respectively output the mean μ and diagonal covari- ance Σ (both of d = 4096 dimensions) of Gaussian embeddings. Loss coefficients in Eq. 7 are set as a = 25, 3 = 1 following VICReg (Bardes et al. 2022), and we set À = 25d/ K, A1 = (dK) and X2 25 to generate comparable loss magnitudes. Fig. g in Appendix D shows the pretraining performance on different datasets is not very sensitive to the loss coefficients.
Data augmentations. Our MAST-5 method uses K 5 augmentations adopted by most recent SSL methods for fair comparison purposes. The augmentations include Color Jitter (of brightness, contrast, saturation, hue), Gaussian Blur, Random Flip, Random Grayscale, RandomResizedCrop While MAST-15 uses 10 more augmentations (hence K = 15). The added augmentations are ShearX(Y)) TranslateX(Y), Rotate, Invert borrowed from RandAugment Cubukk al. (2019), well as Sharpness, Gaussian Noise, Sobel Filter and Cutout (Devries & Taylor 2017) for diversity.
The training protocol follows that in (Bardes al. 2022 with batch size 2048, the LARS optimizer (You et al. 2017) runs for 1000 epochs with weight decay 10-6 and learning rate 1.6. The learning rate follows cosine annealing schedule (Loshchilov & Hutter (017). For effective learning of disentangled subspaces, we found it better to schedule the available augmen- tations in two stages. In the first 500 epochs, we sample one random augmentation og at a time to train the corresponding subspace and mask my. setting K = 1 in Eq. (4). This ensures all the factorized subspaces are trained well in this stage. For the remaining 500 epochs, random augmen- tation compositions (..., .k....... are allowed to learn disentangled invariances from progressively stronger augmentations (K linearly increases from to its maximum value). For each composing augmentation, we use its default augmentation probability but random magnitude for robust training.
Training details:
EXPERIMENTS
Main results. We start with evaluating our MAST method for SSL on ImageNet. Table [ssows the results from both linear and semi- supi-sspeeeieed evaluations, using the same optimization procedures of our baseline VICReg (Bardes et al. 2022). The numbers reported for our MAST-5 and MAST-15 variants are the mean scores from 3 runs with diffeernnt random initialization.
can be seen from Table T that when standard augmentations are used for SSL, MAST-5 achieves strong performance on-par with or better than that of state-of-the-a contrastive and non-contrastive methods. MAST-5 also consistently outperforms our baseline method VICReg, indicating the power of disentangled and uncertainty- tive contrastive learning (PCL) is related to us in avoiding feature suppression, but by strengthening contrastive learning with predictive inpainting task. MAST-5, on the other hand, achieves much
learning of different augmentation invariances. Note predic-

---

Published as a conference paper at ICLR 2023
Table Evaluation on ImageNet using representations trained with ResNet-50) backbone in terms of Top-1 and Top-5 accuracies (%). We evaluate linear classification using the frozen representations from ImageNet, and semi-supervised classification using the fine-tuned representations from 1% and 10% ImageNet samples. Most methods in the middle cell (including our MAST-5) use the same 5 augmentations (Chen et al. 2020a), while the bottom cell compares methods using much more augmentations, e.g., from RandAugment (Cubuk et al. 2019).
Linear
Semi-supervised
Method
Top-1
Top-5
Top-1
Top-5
1%
10%
1%
10%
Supervised
25.4
56.4
48.4
80.4
MoCo (He et al. 2020
60.6
SimCLR (C Chen
2020a
69.3
89.0
48.3
65.6
75.5
87.8
PCL (Li et al. 2020
71.0
MoCo v2 Chen et al.
2020b
SimSiam Chen & He
2021
71.3
SwAV (Crron et al. 2020
71.8
OBoW Gidans al.
|02|1
2021
73.8
82.9
90.7
BYOL Grill et al. 2020
74.3
91.6
53.2
68.8
78.4
89.0
SwAV (w/ mulli-crop) (Caron el al.
2020
75.3
53.9
70.2
78.5
Barlow Twins (Zbontar of al. 2021
73.2
91.0
55.0
69.7
79.2
89.3
VICReg (Bardes et al. 2022
73.2
91.1
54.8
69.5
79.4
89.5
MAST-5 (ours)
74.9
91.8
55.2
70.2
80.1
90.2
InfoMin Aug Tian 2020
73.0
91.1
CLSA (Wang & Qi. 2021
72.2
MAST-15 (ours)
75.8
92.1
55.8
71.4
81.0
90.9
Table 2: Transfer performance downstream tasks using the pretrained representations Im- ageNet. We evaluate linear classification using frozen representations on datasets Places205 (Top-1 accuracy %), VOCO7 (mAP %) and iNat18 (Top-1 accuracy %). We also evaluate fine-tuning per-
formance for object detection on VOCO7+12 (APs0) and for instance segmentation on COCO (AP).
Linear Classification
Detection
Segmentation
Method
Places205
VOCO7
iNat18
VOCO7+12
COCO COCO
det
seg
Supervised
53.2
87.5
46.7
81.3
39.0
35.4
MoCo (He et al.
2020
46.9
79.8
31.5
SimCLK Chen al.
2020a
52.5
85.5
37.2
MoCo v2
Chen
20220
51.8
86.4
38.6
82.5
39.8
36.1
SimSiam Chen
2021
82.4
BYOL Grill ct 2020)
54.0
86.6
47.6
40.4
37.0
SwAV (w/ mullicrrp)) (Caron et al.
2020
56.7
88.9
48.6
82.6
41.6
37.8
OBoW (Gidrri et al.
2021
56.8
89.3
82.9
Barlow wins (Zbonaar el al.
2021
54.1
86.2
46.5
82.6
40.0
36.7
VICReg (Brrdss al.
2027
54.3
86.6
47.0
82.4
39.4
36.4
MAST-5 (ours)
58.9
49.5
83.1
41.9
37.2
InfoMin Aug (Tian al.
2020
82.7
40.6
36.7
CI.SA (Wagg & QL 2021
93.6
83.2
42.3
MAST-15 (ours)
59.6
92.8
50.2
84.0
43.1
38.3
stronger performance by explicitly training different feature subspaces that can encode contradicting invariances without undesired trade-offs. Fig. supports this hypothesis with one visual example of bird classification where one subspace sensitive to color cues, and the other captures invariance to color, exploiting pose and shape instead. Our method can be further extended to MAST-15 learning with 15 types of augmentations, which outperforms InfoMin Aug and CLSA that similarly learn from stronger augmentation compositions but in a suboptimal unfactorized way.
Next, following the setup in (Bardes et al. 2022). we evaluate the transfer performance of the AmageNet-pretrainet representations various downstream tasks: linear classification on the scene dataset Places205 (Zhou et al., 2004)), multi-label image classification dataset VOCO7 (Everingham al. 2010) and fine-grained image classification dataset iNaturalist2018 (Van Horn et al. 2018).

---

Published as a conference paper at ICLR 2023
Table 4: Downstream classification accuracies (%) of ImageNet- 100 pretrained representations CUB-200 and F10wers-102. indicates en- hanced baselines by AugSelf authors. We com- pare accuracy gains in subscripts over each baseline.
Table 3: Evaluation of linear classification on ImageNet-100 using representations pre-trained with a ResNet-50 backhone on the same dataset in terms of Top-1 and Top-5 accuracies (%). En- semble: feature ensembling. Repr: effective rep- resentation dimensionality used for classification.
CUB-200
Flowers- 102
Method
Top-1
Top-5
5-sbot
10-shol
Method
Top-1
Top-5
Repr.
MoCo x2 (Chen et al 2020bb
36.7
61.7
67.9
77.3
MoCo v2 (Chen et al. 2020b
81.0
95.2
2048
LooC (Xintr
39,6
69.2
70.9 so
80.8
IFM-MoCo V2 (Kobinson al. 2021
81.4
2048
McCo v2"
Chen et al
2020b
32.2
78.5
81.2
LooC (Xioo et ar 2021
79.2
94.7
2048
AugSelf
Lee
47.0
81.7
84.5
LooC++ Xtao 2021 (ensemble)
81.2
95.2
2048x3
SimSiam"
Chen & He
2021
38.4
83.6
85,9
MAST-5 (defauitt
81,9
95.4
2048
AugSelf
Lee et
45.3
86.4
88.3
MAST-5 (ensemble)
82.1
95.4
2048x5
VICReg MAST:S
(Badess
2022
37.5
65.2
68.3
77.5
MAST-15
82,5
95.5
2048
20.3
69.8 46
70.1 18
81.2 k3
MAST-15
41.1
70.0
72.4 41
81.9
We also transfer the representations to object detection on VOCO7+12 using Faster R-CNN (Rnn
al. 2015) with a R50-C4 backbone, and instance segmentation on COCO (Lin et al. 2014) using
Mask-R R-CNN (He et al. 2017) with R50-FPN backbone. Table 2 shows larger performance gap
between MAST-5/15 and prior works for almost all tasks, demonstrating the superior generalization
capabilities of our learned invariance priors despite being task-agnostic.
Comparing with more related methods: Implicit feature modification (IFM) (Robinson et al, 2021) and Leave-one-out Contrastive Learning (LooC) (Xiao al. (02I). All comparing methods use the same 100-class ImageNet dataset and 500 epochs for pretraining. IFM is a method that avoids shortcut solutions by altering sample features during contrastive learning. Table 3 shows IFM improves the MoCo v2 baseline by a small margin, but not as much as the gain offered by our disentangled method (MAST-5) that explicitly reduces feature suppression.
LooC combines inggeenaaiioninvariin and -variant feature subspaces to reduce the harm of partic- ular invariances for some downstream tasks. Table 3 shows that LooC, when regularized with such subspace training, actually hurts performance on the pretraining dataset ImageNet- 100 (in compar- ison to baseline MoCo v2). This suggests that the "disentanglement" learned by LooC may still suppress certain useful features. By contrast, our disentangled learning with MAST-5 seems more effective, giving notable gains over MoCo v2. Note the advanced LooC++ variant can recover the loss of LooC by ensembling feature subspaces. However, LooC++ cannot generate solid gains over the MoCo v2 baseline, and at the same time suffers from a much larger cost from feature ensem- bling. When we apply the idea of ensembling (masked) feature subspaces to MAST, we observe slightly improved performance at the cost of increased feature dimensionality. Hence our default MAST-5 approach just uses the original unmasked representations without any ensembling, which is efficient and performant already. MAST- 15 further improves performance by learning from more augmentations. Table validates that MAST-5/15 remain competitive with LooC and recent strong method AugSelf (Lee et al. 2021) on two downstream tasks: fine-grai fineigrineed classification on CUB-200 dataset (Wah al. 201) 1), and few-shot classification on Flowers- 102 dataset (Nilsback & Zisser- man 2008). Table in Appendix D provides more fine-grained classification results.
Ablation study is conducted in Ta evaluate different components of MAST. We use the setup where representations are pretrained on the ImageNet- kk dataset and tested on 4 representative tasks.
Recall that MAST-15 gains a lot over MAST-5 by using more augmentations. We start with asking: "is it more about disentangled learning or just richer augmentations?" Row 2 answers this question by applying the same K 15 augmentations our baseline VICReg. We can see that this variant falls far behind MAST-15 (row 1), and even behind the baseline (row 4) on classification tasks. This indicates that naively composing more augmentations could hurt downstream performance without disentangled learning to avoid feature suppression. Then for MAST-15 with such capability, will even more augmentations help? We experimented with adding common augmentations Solarize, Equalize, Posterize, Motion Blur in row 3 (so K = 19). No large improvements are observed probably because the added augmentations are redundant (more on augmentation relations later). We leave for future work the investigation of advanced or learned augmentation techniques.
Next, we ablate other components of our method based on MAST-5 for efficiency. Recall that MAST-5 benefits from two-staged scheduling of augmentations (from one at time to linearly increasing number of augmentations). We find such augmentation scheduling can improve VICReg as well (row 6), but not enough to compete with MAST-5 (row 5). Lastly, by comparing row 5 to
8

---

Published as a conference paper at ICLR 2023
Table 5: Ablation study of different MAST components (ImageNet pretraining + 4 testing tasks).
Linear classify
Linear classify
Detection
Segmentation
Comment
Row
Method
ImageNet (Top-1)
Places205 (Top-1)
VOCO7+ (APsa)
COCO
Disentangled learning or richer aug.?
1
MAST-15
75.8
59.6
84.0
38.3
2
VICReg 15)
82.6
36.9
Will aug. help?
3
MAST-19
75.8
59.4
84.0
38.4
4
VICReg (K 5)
73.2
54.3
82.4
36.4
Baseline
5
MAST-5
74.9
37.2
Staged aug. schedule
6
VICReg
73.4
82.4
36.5
Handcrafted masks
7
MAST-5
74.3
58.1
82.5
36,9
No uncertainty
8
MAST-5
74.1
57.5
82.7
36.8
Crier
Coder
Grogedadee
Auga
ghoteg
MAST- 15
Fanstase
Framanio
Equaliv
Auga
Plyerrco A
755
Auga
7
Drop
75.7
0,
Top-c aacomnsy
Grascakk
752
2.85
used
ShaarX
75.2
Trandow
0.50
75.1
Sobul
74.8
Equillo
Bix
Adding Auge scrted intreeting mean commistion with others
Figure 4: (a) Correlation between augmentations measured by their masks' cosine similarity. (b) ImageNet linear classification Top-1 accuracy vs. increasing number of augmentations during SSL. Augmentations with low correlation with others provide diverse priors and large performance gains.
rows 7-8, we can easily find the benefits of learned subspace masks (vs. handcrafted disjoint masks: all I's for each subspace of d/K dimensions, and 0 elsewhere) and uncertainty modeling.
Visualizations. We exploit the learned subspace masks quantify augmentation correlations unsu- pervisedly by computing the cosine similarity between masks. Fig. 4aa) shows the example pairwise correlation between augmentations on ImageNet. We can observe weakly correlated augmenta- tions such as (Gaussian Blur, Rotate) and (Gaussian Blur, ShearY). They define diverse augmen- tations that could introduce complementary invariance priors for good generalization (if potential contradictions can be reduced, e.g., by our factorized learning). There are also strongly correlated augmentations like (SheurX, ShearY) that offer redundant invariance information. The 4 optional augmentations Solarize. Equalize, Posterize, Motion Blur highly correlated with others too (e.g., with Color Jitter and Gaussian Blur), therefore we exclude them from MAST-15. Fig. ((bb) demon- strates the positive role of augmentation diversity generalization. We gradually add augmenta- tions to MAST-5 for SSL, such that distinct augmentations with low mean correlation with others observed improve the downstream performance quickly,
added first. Those augmentations and then performance plateaus with redundant, correlated augmentations. This suggests that our correlation measure can indeed be a rough guidance on composing generalizable augmentations.
Appendix B further shows how MSAT correctly models uncertainties for strongly augmented sam- ples order improve SSL performance (comparison with a related uncertainty model (Nakamura et al. 2022 is also provided), and how uncertainty reflects learning difficulty.
5 CONCLUSION
This paper studies the important role of data augmentation in creating useful invariance priors dur- ing SSL, which is the key inductive bias for good generalization on downstream tasks. Instead of enumerating all the possible augmentation compositions and selecting the best set of invariances for a target task, we propose a unified framework to embed all the distinct priors of invariance into feature representations that can be readily used for different tasks. Our MAST method, featuring disentangled representation learning and uncertainty modeling, proves competitive when evaluated on variety of downstream vision tasks.

---

Published as a conference paper at ICLR 2023
REFERENCES
Sanjeev Arora, Hrishikesh Khandeparkar, Mikhail Khodak, Orestis Plevrakis, and Nikunj Saunshi. A theoretical analysis of contrastive unsupervised representation learning. In Proceedings of the 36th International Conference on Machine Learning, volume 97, pp. 5628-56333 2019.
Philip Bachman, R Devon Hjelm, and William Buchwalter. Learning representations by maximiz- ing mutual information across views. In Advances in Neural Information Processing Systems, volume 32, 2019.
Randall Balestriero and Yann LeCun. Contrastive and non-contrastive self-supervised learning re- cover global and local spectral embedding methods. arXiv preprint arXiv:2205//1508, 2022.
Adrien Bardes, Jean Ponce, and Yann LeCun. VICReg: Variance- ietianc--nvaiinn--ovviinna regular- ization for supe--ud learning. In International Conference Learning Representations, 2022.
Mathilde Caron, Ishan Misra, Julien Mairal, Priya Goyal, Piotr Bojanowski, and Armand Joulin. Unsupervised learning of visual features by contrasting cluster assignments. Proceedings of Advances in Neural Information Processing Systems (NeurIPS), 2020.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. P Proceedings the 37th International Conference on Machine Learning, volume 119, pp. 1597-1607, 2020a.
Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Conference on Computer Vision and Pattern Recognition (CVPR), 2021.
Xinlei Chen, Haoqi Fan, Ross Girshick, and Kaiming He. Improved baselines with momentum contrastive learning. arXiv preprint arXiv:2003.0422 2020b.
Adam Coates, Andrew Ng, and Honglak Lee. An analysis of single-layer networks in unsuper- vised feature learning. In Proceedings of the Fourteenth International Conference Artificial Intelligence and Statistics, volume 15, pp. 215-223, 2011.
Ekin D. Cubuk, Barret Zoph. Jonathon Shlens, and Quoc V. Le. Randaugment: Practical automated data augmentation with reduced search space. arXiv preprint 1rXiv 2019.
Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei. Imagenet: A large-scale hierarchical image database. In CVPR, 2009.
Terrance De Vries and Graham W Taylor. Improved regularization convolutional neural networks with cutout. arXiv preprint arXiv: 040524455. 2017.
Mark Everingham, Luc Van Gool, Christopher I. Williams, John M. Winn, and Andrew Zisser- man. The pascal visual object classes (VOC) challenge. Int. J. Comput. Vis., 88(2):303-338, 2010.
Yarin Gal. Uncertainty in deep learning. Phd Thesis, University of Cambridge, 2016.
Spyros Gidaris, Andrei Bursuc, Gilles Puy, Nikos Komodakis, Matthieu Cord, and Patrick Pérez. Obow: Online bag-o--visulllwords generation for self-supervised learning. In IEEE Conference on Computer Vision and Pattern Recognition, CVPR, pp. 6830-6840, 2021.
Bean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Beraaro Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, Bilal Piot, koray kavukcuoglu, Remi Munos, and Michal Valko. Bootstrap your own latent new approach self-supervised learning. In Advances in Neural Information Processing Systems, volume 33, pp. 21271-21284 2020.
Jeff Z. HaoChen, Colin Wei, Adrien Gaidon, and Tengyu Ma. Provable guarantees for self- supervised deep learning with spectral contrastive loss. In Advances Neural Information Pro- cessing Systems, 2021.
10

---

Published as a conference paper at ICLR 2023
Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. Deep residual learning for image recog-
nition. CVPR, 2016.
Kaiming He, Georgia Gkioxari, Piotr Dollr, and Ross Girshick. Mask R-CNN. In 2017 IEEE International Conference on Computer Vision (ICCV), pp. 2980-2988, 2017.
Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In 2020 JEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 9726-9735, 2020.
Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross B. Girshick. Masked autoencoders scalable vision leamnrr. In IEEE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 15979-15988, 2022.
Li Jing, Pascal Vincent, Yann LeCun, and Yuandong Tian. Understanding dimensional collapse in contrastive self-supervised learning. In International Conference on Learning Representations, 2022.
Jonathan Krause, Michael Stark, Jia Deng, and Li Fei-Fei. 3d object representations for fine-grained categorization. In IEEE International Conference on Computer Vision Workshops, pp. 554-561, 2013.
Hankook Lee, Kibok Lee, Kimin Lee, Honglak Lee, and Jinwoo Shin. Improving transferability of representations via sugmentation-awat self-supervision. In Advances in Neural Information Processing Systems, 2021.
Tianhong Li, Lijie Fan, Yuan Yuan, Hao He, Yonglong Tian, Rogerio Feris, Piotr Indyk. and Dina Katabi. Addressing feature suppression in unsupervised visual representations. arXiv preprint arXiv:2012.09962 2020.
Tsung-Y Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C. Lawrence Zitnick. Microsoft COCO: Common objects in context. In ECCV, pp- 740-755, 2014.
Ilya Loshchilov and Frank Hutter. SGDR: Stochastic gradient descent with warm restarts. In Inter- national Conference on Learning Representations, 2017.
Hiroki Nakamura, Masashi Okada, and Tadahiro Taniguchi. Self-supervised representation learning as multimodal variational inference. arXiv preprint arXiv:2203. /1437, 2022.
Maria-Elena Nilsback and Andrew Zisserman. A visual vocabulary for flower classification. In IEEE Conference on Computer Vision and Pattern Recognition, volume 2, pp. 1447-1454, 2006.
Maria-Elena Nilsback and Andrew Zisserman. Automated flower classification over a large number of classes. In Sixth Indian Conference on Computer Vision, Graphics & Image Processing, pp. 722-729, 2008.
Matteo Poggi, Filippo Alcotti, Fabio Tosi, and Stefano Mattoccia. On the uncertainty of self- supervised monocular depth estimation. In IEEE/CVF Conference Computer Vision and Pat- tern Recognition (CVPR), 2020.
Filip Radenovi, Giorgos Tolias, and Ondej Chum. Fine-tuning cnn image retrieval with no human annotation. IEEE Transactions on Pattern Analysis and Machine Intelligence, 4157)1165--168, 2019.
Shaoqing Ren, Kaiming He, Ross Girshick, and Jian Sun. Faster R-CNN: Towards real-time object detection with region proposal networks. In Advances in Neural Information Processing Systems, volume 28, 2015.
Joshua David Robinson, Li Sun, Ke Yu, kayhan Batmanghelich, Stefanie Jegelka, and Suvrit Sra. Can contrastive learning avoid shortcut solutions? In Advances Neural Information Processing Systems, 2021.
11

---

Published as a conference paper at ICLR 2023
Yonglong Tian, Chen Sun, Ben Poole, Dilip Krishnan, Cordelia Schmid, and Phillip Isola. What makes for good views for contrastive learning? In Advances in Neural Information Processing Systems, 2020.
Grant Van Horn Oisin Mac Aodha, Yang Song, Yin Cui, Chen Sun, Alex Shepard, Hartwig Adam, Pietro Perona, and Serge Belongie. The inaturalist species classification and detection dataset. In The IEEE Conference on Computer Vision and Pattern Recognition (CVPR), June 2018.
C. Wah, S. Branson, P. Welinder, Perona, and S. Belongie. The caltech-uesd birds-200- 201-2011 dataset. Technical report, 2011.
Xiao Wang and Guo-Jun Qi. Contrastive learning with stronger augmentations. arXiv preprint arXiv22104.07713, 2021.
Ross Wightman, Hugo Touvron, and Herve Jegou. ResNet strikes back: An improved training procedure timm. In NeurIPS 2021 Workshop on ImageNer: Past, Present, and Future, 2021.
Jianxiong Xiao, James Hays, Krista A. Ehinger, Aude Oliva, and Antonio Torralba. Sun database: Large-scale scene recognition from abbey to 700. In IEEE Computer Society Conference on Computer Vision and Pattern Recognition, pp. 3485-3492, 2010.
Tete Xiao, Xiaolong Wang, Alexei Efros, and Trevor Darrell. What should not contrastive in contrastive learning. In International Conference Learning Representations, 2021.
Hongbin Xu, Zhipeng Zhou, Yali Wang, Wenxiong Kang, Baigui Sun, Hao Li, and Yu Qiao. Digging into uncertainty in self-supervised multi-view stereo. In Proceedings of the IEEE/CVF Interna- tional Conference on Computer Vision, pp. 6078 7778-6077-21 2021.
Yang You, Igor Gitman, and Boris Ginsburg. Large batch training of convolutional networks. arXiv preprint arXiv: 1708. 00888888, 2017.
Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stphane Deny. Barlow twins: learning via redundancy reduction. In JCML, 2021.
Self-supervised
Bolei Zhou, Agata Lapedriza, Jianxiong Xiao, Antonio Torralba, and Aude Oliva. Learning deep features for scene recognition using places database. In Advances Neural Information Process- ing Systems, volume 27, 2014.
12

---

Published as a conference paper at ICLR 2023
(a) Input
Augmented views
Meen uncertainty
(b)
0.6
Uncenainity
D.13
0.21
0.23
0.17
0.15
0.18 (aw)
0.68
0.46
0.74
0.39
0.61
0.58 (high)
class
class 'Panayy
mean Top-1 acc: 63.6%
mean Top-1 8OC: 18.
0.15
0.17
0.18
0.23
0.28
0.20 (aw)
0.42
0.79
0.74
0.54
0.85
0.67 (high)
Figure 5: (a) Example uncertainty estimations for augmented views of images on Flowers-1 17 dataset. Views with strong augmentations have high uncertainties in their feature representations. (b) Mean uncertainty estimations for two flower classes vs. mean linear classification accuracies (Top-1 %) of the two classes. The lower-perforning class 'Pansy' has higher uncertainty and larger variance in uncertainty estimates, due to large variations in the flower color, size and shape. This indicates our uncertainty measure can reflect the learning difficulty of the classification task.
A ADDITIONAL IMPLEMENTATION DETAILS
Uncertainty modeling. We have separate heads in the projector g to predict the mean and diagonal covariance of Gaussian embeddings. The mean head is composed of Sieeralize-Meean (GeM) pooling layer (Radenovi et al. 2019) followed by a fully connected layer that outputs f E R2 with d = 4096. The covariance head is composed of GeM layer and fully connected layer with ReLU activation function (to ensure positivity) to output the d-dimensional diagonal of a covariance matrix Σ. The separate GeM layers in the mean and covariance heads beneficial, since they help to learn different p-norms for the mean and covariance.
Mask learning. Our learned subspace masks act as element-wise gating functions to select relevant embedding dimensions. We ensure the positivity of masks M d by learning them as M max(0, U) using a ReLU activation function where U E Rdx During training, we found useful initialize the masks random rather than fixed ones (e.g., when initial mask values all the same). More specifically, we initialize each mask (i.e., each column of U) such that the corresponding subspace of d/ K dimensions follows Gaussian distribution and Gaussian noise is added for all dimensions. The Gaussian prior for each subspace mask cases mask training in the beginning and expedites the search of disentangled solutions. While the added Gaussian noise is large enough (with mean 0.2 and variance 0.01) to help find meaningful mask patterns in all dimensions.
B UNCERTAINTY VISUALIZATION AND COMPARISON
Uncertainty visualization. In MAST, uncertainty is estimated as a vector, i.e., diagonal covariance of Gaussian embeddings encode uncertainty in every dimension of the embeddings. For the ease of visualization, we need a scalar uncertainty measure to quantify how uncertain we are about the entire image input and its representations. Here we compute the scalar uncertainty based on a volumetric measure (trace of covariance matrix), which is then linearly rescaled to the range of [0, 1]. The higher the value is, the larger uncertainty the representations have. Fig. 5 shows some example uncertainty estimations on Flowers-17 dataset (Nilsback & Zisserman 2006). We observe high uncertainties estimated for strongly augmented images, while low uncertainties for the weakly augmented. This shows our uncertainty measure data augmentations. This a desirable behavior since when the augmentations are strong enough to produce ambiguous views uncertainty to downweight the distance between ambiguous views. As result, similarity mismatch avoided during invariance learning.
reasonably capture the strength of
change the identity of input images,
the estimated high
13

---

Published as a conference paper at ICLR 2023
Table 6: Evaluation of linear classification on Flowers-17 and Cars using representations pre- trained on the respective datasets in terms of Top-1 and Top-5 accuracies (%).
Flowers-17
Cars
Method
Top-1
Top-5
Top-1 1
Top-5
SimSiam (Chen & He 2021)
24.6
51.5
12.7
33.2
VI-SimSiam (Nakamura el al. 2022
19.7
43.9
14.9
37.2
MAST-5 (uncertanty only)
26.5
55.7
16.4
40.3
MAST-5 (uncertainty disentangled learning)
28.9
59.1
18.7
44.1
Fig. ())) further connects uncertainty with the classification performance (linear probing). It can be seen that the 'Pansy" class exhibits high uncertainty and large variance in uncertainty estimates due to the large variations of flower appearances. Accordingly, the 'Pansy' class has much lower classification accuracy than that of 'Sunflower' class with low uncertainty. This validates that our uncertainty measure can indeed reflect the learning difficulty. For the difficult 'Pansy' class, the large uncertainty variance actually offers the chance for us to adapt invariance learning and avoid similarity mismatch for improvement.
Comparison. We compare with VI-SimSiam (Nakamura al. 2022) that models uncertainty for SSL using spherical posterior distributions and variational inference. We follow VI-SimSiam to pre- train for 100 epochs and use ResNet- 18 as the encoder. Pre-training is performed on two datasets spectively: Flowers-17 (Nilsback & Zisserman 2006) with 17 flower categories (80 images for each category), and Cars (Krause et al. 2013) with 197 models (16, 185 images). After pre-training, we train a linear classifier on the frozen representations using the training set of the corresponding dataset, and evaluate linear classification accuracy in its test set.
Table summarizes the results in Top-1 and Top-5 accuracies. We observe that the uncertainty- aware VI-SimSiam method is not always helpful when compared to its baseline SimSiam (see the notable performance drop on Flowers-17). For MAST-5, we experiment with a special variant 'un- certainty only' to not benefit from disentangled subspace learning (i.e., no masks involved). Note MAST-5 (uncertainty only) is not directly comparable to VI-SimSiam since we have different base- lines (VICReg vs. SimSiam) and implementation details. The goal here is only to test how our uncertainty modeling mechanism works in general, when using simple Gaussian assumption about feature embeddings. Results show that MAST-5 (uncertainty only) obtains quite competitive perfor- mance on both datasets, and is further improved by disentangled learning consistently.
C MORE ANALYSIS OF THE AUGMENTATION SUBSPACES
Optimal augmentations (subspaces) are task-dependent. Fig. I supports the claim by demon- strating the difference between the optimal augmentations for different downstream tasks. This motivates us develop MAST, by learning disentangled augmentation invariances in masked fea- ture subspaces that are ingmentation-specife the relationship between augmentations and tasks, using the learned augmentation-speeiit masks as a bridge. The end task can freely select from the factorized latent space which feature masks or augmentations are relevant to the task. This can help us identify optimal augmentations for one task in cheaper way than the leave-one-ore out strategy in Fig.
Now with MAST, we find it much casier to quantify
We take the classification task as one example. For correctly classified samples, we examine the class predictions (only for the ground-truth class, using the corresponding weights of the classifier) by differently masked feature embeddings júj of each sample. The higher the class prediction, the more the corresponding masked augmentation subspace (and its captured invariance) benefits the given task. Fig. 6 shows the class predictions across 5 augmentation subspaces for ImageNet and CUB-200 classification, respectively. We observe that, again, different tasks benefit from a different set of augmentations (e.g., Color Jitter, Random Grayscale, RandomBesizedCCop for ImageNet, and Gaussian Blur; Random Flip, RandomResizedCrop for CUB-200).
Augmentation invariance vs. preserving augmentation-sseeeffi information. Since our MAST method focuses on disentangled learning of different augmentation invariances, now verify how
14

---

Published as a conference paper at ICLR 2023
Cass peciedoo
dasspeddit on
Clu prodition
Crean Geea foantoon
Recidor Goosetone
Handendiation comeriaatiniiii Frrd
Glou Ferston Pi
Crtir Ooee fonsoo Crosnon
Rostion cogpeau
Pop
Class prodetion
Clas predeceon
Clus prodction Straset
Gue ane the flansen Pir
stoa
Randon Pi
Randon chepeade
Gaes aor thr Raceen Pizz
siee
fostrm Dioy wooe
Randon Pib
Randon chepaesh
Dasten thay
Cate
date
(a) ImageNet
(b) CUB-200
Figure 6: Class predictions by different augmentation-speeiiii subspaces of correctly classified im- ages on (a) ImageNet and (b) CUB-200 datasets. MAST-5 is used here. Note we attend to different embedding subspaces via masking based on the learned masks for augmentations. Class predic- tions are obtained by passing the masked embeddings through the classifier weights associated with the ground ground-truth label. The higher the class prediction, the more beneficial the corresponding aug- mentation subspace (and its captured invariance) is to the classification task. As can be observed, the optimal set of augmentation invariances are different for the two classification tasks.
(a) Color Jitter
(b) Gaussian Blur
(c) RandomFeeszzeCrap
1,000
1.000
metric
0.975
Inrariance mettic
0.975
Iwertorce mettic
0.975
0.850
0.850
0.800
5ABT or subapace
MABT Bu subopace
*MASF-S Crip subepace
0.325
ViCheg
0.025
+VCRug
0.525
*VICRug
0.25
1.5
0.800
0.9
1,2
1.5
0.800
0.6
Strength factor
Standard dev lation of Goussian filter
Cropping size ratio
1.000
1.000
Invriincce metrr
Invartunce metrr
Invartance metrr
0.500
0.950
MASTS Crop
MAAT-5 Crop subspace
*MAST-S Grop subopace
0.550
Color subspace
0.525
*MAT subspace
0.700
MMSTS G. Balr sudepace
*MAST- Gur subspace
MAATT G. Bur sutspace
0.800
0
025 0.5
1.5
0,800
0.3
8.6
0.8
1.2
1.5
0.680
0.5
0.6
0.7
0.5
0.8
Strength factor
Standard deviation of Gaussian filter
Cropping bice ratio
Figure 7: Invariance metric for ImageNet-pretrained feature embeddings. The invariance met- ric is defined as the cosine similarity between masked embeddings from augmented and original samples. The higher the metric, the more invariant the subspace is to the tested augmentation. The invariance metric is computed against 3 augmentations with continuous augmentation parameters: (a) Color Jitter with the strength factor multiplied over the default configuration (larger value in- dicates stronger augmentation), (b) Gaussian Blur with the standard deviation of Gaussian filter (larger value indicates stronger augmentation). (c) RandomeRssizedCruy with the cropping size ratio (smaller value indicates stronger augmentation). Top row: augmentation invariances are well cap- tured by our masked subspaces (better than baseline VICReg features). Bottom row: our subspaces are disentangled since each is specialized to the corresponding augmentation.
well the invariance property is captured in each of our augmentation subspaces. The invariance metric is used by computing the cosine similarity between masked embeddings j from augmented and original samples. Higher metric indicates higher level of augmentation invariance. Fig 7 shows this metric on ImageNet when we vary the augmentation parameters of 3 standard augmentations
Color Jitter, Gaussian Blur, RandomResizeddCrp.. We exclude other (Random Flip, Random Grayscale) whose augmentation parameters are discrete.
Two observations from the figure: 1) in top row, our subspaces show great robustness under various augmentation intensities, especially when compared the VICReg baseline that learns all augmen-
15

---

Published as a conference paper at ICLR 2023
Table Linear classification accuracy (%) of rotation using prerto-retrriiedd representations.
Method
SimSiam (Chen & Hc 2021
SimSiam+AugSelf (Lee al. 2021
MAST-5
Accuracy
59.11
64.61
63.47
Table 8: ImageNet performance function of pretraining epochs. Performance is measured as the Top-1 accuracy (%) of linear classification. Our default number of pretraining epochs is 1000.
Epochs
200
400
800
1000
VICReg (Bardes et al. 2022 MAST-5
70.2
72.3
73.0
73.2
70.9
73.5
74.5
74.9
A
0.7
1.2
1.5
1.7
tation invariances in an unfactorized way. 2) the bottom row shows the disentanglement of different subspaces, since each subspace demonstrates specialization to the corresponding augmentation.
One might ask: will we loose all eegmenntaioonssesitivy information when we learn factorized in- variances from those augmentations? Short answer is no, as evidenced by our strong performance on various downstream tasks that rely specific augmentation characteristies (e.g., color-sensitive bird classification on CUB-200). We hypothesize that different subspaces may offer complemen- tary information about augmentations for different tasks to choose from. For example, in Fig. (b),, the color information suppressed in the subspace invariant to Random Grayscale is available in the subspace invariant to RandomResizedCrop:
To further support the above hypothesis, we follow the AugSelf work (Lee et al. 2021 to solve pretext task of 4-way rotation classification (0*,00 .0",, 180", 270") using pretrained representa- tions. The same pretraining protocol followed: pretraining ResNet-18 for 200 epochs on STL-10 dataset (Coates et al. 2011), and standard augmentations are used (no Rotate). To evaluate rotation classification accuracy, linear classifier is trained on top of the frozen representations. Table 7 shows competitive performance of MAST-5 (63.47%) vs. AugSelf (64.61%) which is explicitly trained to be augmentation prediicivee.... This validates that we can learn invariances to standard aug- mentations while maintaining sensitivity the unseen rotation. other words, MAST is able to preserve sugmentation-inaaria and -specific information implicitly in separate subspaces.
D ADDITIONAL RESULTS
Impact of pretraining epochs. Recall our MAST models are pre-trained in two stages, with in- creasing number of augmentations (i.e. stronger compositions). As result, MAST typically ben- efits significantly from longer training to take full advantage of disentangled learning with strong augmentation compositions. Table exemplifies the ImageNet performance as a function of pre- training epochs, from 200 to default 1000. We can observe that 1) our MAST-5 outperforms its baseline model VICReg across all epoch numbers, more epochs lead larger gains, but the gain is still decent (0.7%) at low epoch number 200.
Sensitivity of MAST to loss coefficients. Recall the loss function of MAST in Eq. 7 consists of 2 terms borrowed from the VICReg baseline with their default loss coefficients. We introduce 25d/ K, A = 600/(dK) and X 25
more loss terms, and only set their loss coefficients to have comparable loss magnitude with that of other terms. In practice, we use the exact same loss coefficients across different augmentation distributions and different pretraining datasets. Fig B shows the pretraining performance is not very sensitive to X, > especially when they are scaled by factor between 0.5 and 2. Note for comprehensive evaluations, we monitor the pretraining performance the large dataset ImageNet and two relatively small datasets Flowers- 17 and Cars, using their representative evaluation tasks. We could have casily tuned X, λ₁ X2 per dataset for best performance, but we choose use the same loss coefficients for simplicity.
16

---

Published as a conference paper at ICLR 2023
ImageNot pretraining
(b) Flowers-17 pretrainiing
(c) Cars pretraining
ImageNet Inear classifcaation
VOC07-12 detection
Flowers-17 inear classlication
Cars linerr ciassilicaion
(K)
75.2
61.4
Top- AO0S 1
29.0
accuracy (8)
10.
28.0
63.0
18.5
28.5
10.5
6--2
5--4g
5-0.5 (default)
5-3
Figure 8: Sensitivity of our MAST-5 to the loss coefficients X, Åi, A2. We sweep over the loss coefficients by varying multiplier (default pretraining performance monitored with (a) the large ImageNet dataset (evaluated on 2 tasks), as well relatively small datasets (b) Flowers-17 17 and (c) Cars. Performance is not very sensitive to the loss coefficients, especially when they are scaled by E [0.5,2]
1) applied to each of them. The impact on
Table Linear classification accuracy (%) on more fine-grained classification datasets, using ResNet-50 pretrained on ImageNet- Both the LooC/LoCC++ and AugSelf methods are built on top of the MoCo v2 baseline.
ImageNet-100
iNat-lk
Cars
SUN397
Method
Top-1
Top-5
Top-1
Top-5
Top-1
Top-1
MoCo v2 (Chen et al.
2020b
81.0
95.2
36.2
62.0
33.86
46.50
LooC (Xiao et al. 2021
79.2
94.7
44.0
69.3
LooC++ Xiao et al.
2021
81.2
95.2
46.1
71.5
AugSelf
Lee al.
2021
82.4
37.35
48.52
MAST-5
81.9
95.4
46.9
72.1
39.07
49.87
MAST-15
82.5
95.5
48.1
72.8
39.86
50.74
More results of fine-grained classification. As supplement to Tables 3 and
Table 9 pro-
vides more results on additional fine- aine-griiedd classification datasets iNat- Van Horn al. 2018),
Cars (Krause et al. 2013) and SUN397 (Xiao et al.) 2010 with relatively large number of classes. The same linear evaluation protocol is adopted, based on the ResNer-50 pretrained on ImageNet-100 for 500 epochs. We compare with both LooC and new method AugSelf (Lee et al. 2021) AugSelf combines contrastive learning with an auxiliary loss that predicts augmentation parameters to im- prove SSL generalization. Note both LooC and AugSelf are built top of the MoCo v2 baseline. It can be observed from the table that our MAST-5 is quite competitive on all datasets, and MAST-15 further improves and consistently achieves the best performance.
E DISCUSSION AND COMPARISON WITH MAE- BAE-BASED RECONSTRUCTION METHOD
The recent MAE method (He et al. 2022) has achieved great success without explicit learning of augmentation invariances. Hence, our MAST and MAE belong different categories of SSL meth- ods: MAST falls into the category of discriminative methods that heavily rely on data augmenta- tion, and MAST further improves by disentangled and uncertainty-awart learning of augmentation invariances; while MAE is an autoencoding method that reconstructs the masked out patches of in- put image. MAE found to perform decently even without augmentations, partially because the randomness of masking already provides strong regularization on pretraining.
Although the two methods MAST and MAE will inevitably exhibit different behaviors, we can still get general sense of how they compare in performance under the same experimental setup. For that, we conduct an apples-to apples-to-apple to-apples comparison using the same network architecture (ResNet-50). Fast pretraining for 100 epochs is performed evaluated by both linear probing and deep fine-tune tenintunng protocols, in order to thoroughly compare the discriminative and reconstruction methods. For linear probing, follow MoCo (He et al. 2020
ImageNet. The learned representations
then
17

---

Published as a conference paper at ICLR 2023
Table 10: MAST vs. MAE ImageNet. Top-1 accuracy (%) of ResNet-50 is measured by linear probing and deep fine-tuning after 100 epochs of pretraining.
Method
Linear
Fine-tuning
MAE (He et al. 2022 MAST-S
37.8
77.1
70.8
78.8
train a linear classifier on frozen features with an initial learning rate of 30 and batch size of 256. For deep fine-tuning, we follow the A3 training recipe in (Wightman et al. 2021) with its specific optimizer and learning rate scheduling.
Table 10 summarizes the comparing results. We can see that, when using ResNet for pretraining, the MAE-based reconstruction method seem unable learn as strong representations our MAST- based discriminative method, since MAE yields worse results using both of our evaluation protocols. One future work is to investigate whether our advantage will translate from CNN Vision Trans- former models.
Besides empirical evaluations, we note the main sesttecturr-idepenneenn advantages of MAST over MAE are twofold: disentangled augmentation learning and uncertainty learning, which are missing the latter family of methods to our knowledge. would be quite interesting to evaluate the two components the MAE framework and go beyond discriminative learning. Intuitively, we treat random masking as one type of "augmentation" in MAE, then disentangled learning from e.g., different augmentation strengths (masking ratios) could lead local and global features learned. On the other hand, uncertainty could be important for ambiguous patch reconstruction in MAE, egg., when the masking ratio too high, or when there barely visual cues around the masked patch location.
18


Reviews:
[{"review_id": "uimEZzuilUR", "reviewer": "Reviewer_oMUf", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "In this paper, the authors propose a method that tend to learn task-independent representations. Specifically, the proposed method uses multiple masks to generate different subspace representations to avoid feature suppression. The proposed method also employs the uncertainty modeling technique against variance from strong augmentations. The experiments demonstrate that the proposed method outperforms SOTA baselines in various downstream tasks.", "summary_of_the_review": "Although the proposed method is complex and more experiments should be added, it still shows a new way to learn task-independent representations. I recommend this work.", "strength_and_weaknesses": "Strength:\n\n1. How to learn task-independent representations is a key research problem in SSL.\n2. The proposed method shows remarkable improvements in various settings.\n3. The paper is well organized and easy to follow. \n\nWeakness:\n\n1. The proposed method is too complex, with 5 hyper-parameters in the loss function. Although authors claim all hyper-parameters do not need to be careful turn, only results on ImageNet-100/1K are provided. How about on other datasets, e.g., CIFAR-10/100?\n2. The proposed method shows significant improvements on ImageNet-1K compared with MoCo V2, over 3%, but on ImageNet-100, the improvements drop to below 1%. Why does it happen? \n3. The authors do not provide the results of the short-term training, such as 200 epochs. Does the proposed method require long-term training to fully explore the strong augmentations? \n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "lntcxoCFojq", "reviewer": "Reviewer_Psus", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper illustrates method (MAST) on how to transfer data augmentation to a downstream task for Self-Supervised Learning (SSL) methods.\n\nFrom their paper, their conclusion is that their method can show important role of data augmentation in creating useful invariance priors during SSL.\n\nHere are their contribution:\n- introduce MAST to make SSL representations disentangled and uncertainty-aware to\neffectively encode different augmentation invariances for good generalization.\n•show MAST is efficient, is resistant to feature suppression, and achieves state-of-the-art\ndownstream performance on diverse vision tasks without presuming any task information\nduring pre-training.\n• provide interesting insights about how different augmentations are related in SSL and\nhow uncertainty reflects learning difficulty and impacts learning adaptively.\n", "summary_of_the_review": "This paper mainly contributes a MAST by proposing a unified framework to embed all the distinct priors of invariance into\nfeature representations that can be readily used for different tasks. Generally, it is trying to solve a problem of generalization and downstream task. Generally, their method has shown a good direction on self supervised learning.", "strength_and_weaknesses": "Strength:\n- Their paper has pointed out the problem of current deep learning field is that current model or data augmentation lacks of generalization.\n\n- Weakness:\n\nIn their data augmentation part of 3.2, they set K=5 and K=15. Is there any performance trend of the different number of K?\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "JpcehufRAI0", "reviewer": "Reviewer_7wWs", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "The authors propose learning embedding in augmentation subspaces via masking, which specifies which dimensions are relevant for which augmentation.  In the loss function, the mask is applied to the embeddings before comparison.  Also, the loss function includes a L1 norm term that encourages  masks to be sparse (with fewer features used).   To model uncertainly in the embeddings, the approach learns the mean (replacing the embedding previously) and covariance matrix.  The loss function is normalized by the covariance.   To stablize the learning of the uncertainty model, they use a regularization term to consistency between the Gaussian embeddings before masking via symmetric KL divergence.\n\nThey compare with 13 algorithms on ImageNet via linear evaluation and semi-supervised learning.  They also evaluate representations learned on ImageNet for downstream tasks on 5 different datasets via linear classification, object detection, and instance segmentation.  Their results indicate that their approach generally compares favorably.  In their ablation study, they look at the effect of disentangled embedding, scheduling, learned mask, and Gaussian uncertainty.  They analyzed the correlation between masks and add more uncorrelated masks than the more correlated ones.   They found \n\n\n\n", "summary_of_the_review": "Learning embeddings in augmentation subspaces via masking is interesting.   Their results indicate that their approach generally compares favorably. ", "strength_and_weaknesses": "Strengths:\n\n1.  Learning embeddings in augmentation subspaces via masking is interesting.\n\n2.  Their results indicate that their approach generally compares favorably. \n\n3.  The paper is generally well written.\n\nWeaknesses:\n\n1.  While performance improves, further discussion/analysis on which subset of the masked embeddings are beneficial to different tasks would be interesting.   ", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "HnmDRJxpPk7", "reviewer": "Reviewer_JyME", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "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": "", "main_review": "", "summary_of_the_paper": "This paper proposes a self-supervised learning method, Masked Augmentation Subspace Training (MAST), which learns more generalizable representations by modeling augmentation-specific invariance properties. This paper shows that MASK outperforms existing SSL methods under various downstream tasks, including semi-supervised learning, detection/segmentation, and fine-grained classification.\n", "summary_of_the_review": "Although this paper successfully achieves strong empirical results in various downstream tasks, the method descriptions are unclear, and some important comparisons and analyses are missing. I think they are critical, so I vote for rejection.\n", "strength_and_weaknesses": "Strengths\n- This paper aims at learning more generalizable representations, which is an important problem in the SSL literature.\n- The proposed method outperforms previous SSL methods under various downstream tasks.\n\nWeaknesses\n- Method explanation is unclear: how can Eq (2) learn augmentation-specific masks $m_i$? For example, how to know $m_1$ is related to color jittering, and $m_2$ is related to Gaussian blur, as shown in Figure 3a? Eq (2) does not force any relationship between a mask $m_i$ and an augmentation operation $o_j$.\n  - I'm wondering whether MAST uses $K$ forward passes for each training iteration as LooC [1] did. If not, how to learn augmentation-specific embedding spaces? This could be critical because Eq (2) may cause misunderstanding.\n- Recent MAE-based approaches do not learn augmentation-invariant representations. What is the advantage of this paper over them? This paper should be compared with the MAE methods.\n- The hyperparameter choices of λ = 25d/K, λ1 = 600/(dK) and λ2 = 25 seem to be carefully chosen. So, I think ``work well for all experiments without the need of careful tuning.'' is not true. If the authors want to claim that the sentence is true, they should conduct experiments with varying λ, λ1, and λ2 and also should verify that performance is not sensitive to the hyperparameter choices.\n- There exist many fine-grained classification tasks, but Table 4 only provides two benchmarks, CUB and Flowers. More experiments with the tasks should be provided (see SimCLR [2], BYOL [3], or AugSelf [4] papers).\n  - This paper should be compared with [4]: [4] achieves 45% top1 accuracy on CUB-200 and 88% on Flowers-102.\n- There is no analysis of the correlation between each augmentation-specific embedding space and augmentation-related information (e.g., see Table 1 & Figure 3 in LooC [1] or see Table 8 & Figure 5 in AugSelf [4]).\n- There is no analysis of uncertainty. Such an accuracy comparison in Table 5 does not explain how the uncertainty affects learning representations. Also, this paper should provide some empirical evidence verifying that the uncertainty is successfully learned.\n\n[1] Xiao et al., What should not be contrastive in contrastive learning, ICLR 2021 \\\n[2] Chen et al., A Simple Framework for Contrastive Learning of Visual Representations, ICML 2020 \\\n[3] Griall et al., Bootstrap Your Own Latent: A New Approach to Self-supervised Learning, 2020 \\\n[4] Lee et al., Improving Transferability of Representations via Augmentation-Aware Self-Supervision, NeurIPS 2021\n", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Decision: Accept: poster