MAST: M ASKED A UGMENTATION S UBSPACE T RAINING FOR G ENERALIZABLE S ELF -S UPERVISED P RIORS
Anonymous authors
Paper under double-blind review
A BSTRACT
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 classiﬁcation, 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.
Speciﬁcally, we disentangle the feature space into separate subspaces, each induced by a learnable mask that selects relevant feature dimensions to model invariance to a speciﬁc 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 beneﬁts 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 efﬁcient during
SSL.
We also provide interesting insights about how different augmentations are related and how uncertainty reﬂects learning difﬁculty.
I NTRODUCTION
Self-Supervised Learning (SSL) for image representation has made signiﬁcant 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 al., 2020a; He et al., 2020) learn invariances by discriminating augmented views of the same image (positive pair) from those of different images (negative pair), while recent noncontrastive SSL methods (Chen & He, 2021;
Grill et al., 2020; Bardes et al., 2022) simply maximize the similarity between positive pairs.
Such learned features are shown to generalize across many downstream tasks, including classiﬁcation, 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 uniﬁed framework that demonstrates the key for generalization lies in the alignment 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., 2021) that quantify how data augmentations implicitly encode the class distributions in downstream tasks.
Motivated by these theoretical analyses, we aim at a working SSL method that can directly capture meaningful priors from data augmentations in order to encourage generalization for a range of tasks.
Invariance, as achieved through augmentation, is a useful mechanism to facilitate generalization.
For example, one can imagine that invariance to Random Flip will boost generalization on many vision tasks.
However, as shown in Fig.
1, the optimal set of augmentations (thus invariances) highly depends 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) may even hurt generalization on another ( e.g. , birds classiﬁcation which requires accurate color
R e l a t i v e p e r f o r m a n ce c h a n g eAugmentations removed - -   -
ImageNet classify (Acc) -
CUB-200 birds classify (Acc) -2.
Augmentations removed -1.
Augmentations removed - -   -
VOC detection (AP 50 ) -
Crop
[Figure]
Color even hurts
COCO segmentation (AP) -1.
Augmentations removed -0.
Remove Color Jitter
Remove Gaussian Blur
Remove Random Flip
Remove Grayscale
Remove RandomResizedCrop information of similarly shaped bird species).
Hence it is impossible to maximize generalization by ﬁnding 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 a priori during SSL.
Furthermore, manually ﬁnding suitable augmentations for generalizing to a new task is quickly cumbersome.
In this paper, we propose a 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 a single but disentangled feature space to encode a set of potentially contradicting augmentation invariances.
For each augmentation, we learn invariance to it in a specialized feature subspace, which is induced from the full space with a learned masking operation.
This allows us to learn unique features for each augmentation subspace as well as features shared between them.
There are two main beneﬁts 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 a 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.
In order to examine how representation effectiveness scales with augmentation diversity, we run experiments with different numbers of augmentations, starting from 5 standard augmentations used typically in SSL (Chen et al., 2020a), and extending to an additional 10 augmentations (totaling 15).
Note Tian et al.
and Wang & Qi (2021) also learn from stronger augmentations, but not in our factorized and uncertainty-aware fashion.
When it comes to transfer learning on downstream tasks, we simply drop our subspace masks and ﬁnetune the full feature representations for high efﬁciency.
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 a 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-speciﬁc heads).
Experiments show that MAST, while being efﬁcient and task-agnostic, achieves state-of-the-art transfer performance on 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 reﬂects learning difﬁculty to avoid similarity mismatch during invariance learning.
To summarize, here are our main contributions :
• We introduce MAST to make SSL representations disentangled and uncertainty-aware to effectively encode different augmentation invariances for good generalization.
• We show MAST is efﬁcient, is resistant to feature suppression, and achieves state-of-the-art downstream performance on diverse vision tasks without presuming any task information during pre-training.
• We provide interesting insights about how different augmentations are related in SSL and how uncertainty reﬂects learning difﬁculty and impacts learning adaptively.
R ELATED W ORK
Self-Supervised Learning.
Contrastive learning is one popular form of SSL, which contrasts different views of an input (positive pairs) against other inputs (negative pairs).
Recent contrastive methods mainly differ in how they sample negative pairs, e.g. , from a large batch as in SimCLR (Chen et al., 2020a) or from a memory bank as in MoCo (v2) (He et al., 2020; Chen et al., 2020b).
One downside of contrastive methods is that they require a multitude of negative pairs to work well, leading to memory inefﬁciency.
Recent non-contrastive SSL methods such as BYOL (Grill et al., 2020), SimSiam (Chen & He, 2021) and OBoW (Gidaris et al., 2021) improve by not using negative pairs.
They only maximize the consistency between positive pairs in a teacher-student framework, and some variants like SwAV (Caron et al., 2020), Barlow Twins (Zbontar et al., 2021) and
VICReg (Bardes et al., 2022) achieve state-of-the-art performance on several downstream tasks.
Note Tian et al.
and Wang & Qi (2021) point out the important role of using strong data augmentations 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 signiﬁcant interest in theoretically explaining the success of contrastive SSL from various perspectives, e.g. , mutual information maximization (Bachman et al., 2019) or dimension collapse (Jing et al., 2022).
Most relevant to our work are theoretical studies (Arora et al., 2019;
HaoChen et al., 2021) that show the augmentation distributions of same-class inputs have sufﬁcient overlap when the costrastive loss is small, in which case data augmentations can implicitly encode the class structure for downstream tasks.
Balestriero & 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 augmentation distributions, we propose to improve the generalization of SSL by a factorized learning of augmentation properties as well as uncertainty modeling.
The factorization allows each downstream task to choose which underlying symmetries in the data require speciﬁc invariances.
Empirical works on generalization.
Designing better augmentation strategies is one effective way to create task-relevant inductive biases for generalization.
Tian 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 solutions, Robinson et al.
(2021) propose implicit feature modiﬁcation to encourage instance discrimination using a wider variety of features; while Li et al.
(2020) force contrastive learning to learn features that can also predict the input, without discarding important information.
Our method automatically avoids shortcuts by jointly training multiple masked feature subspaces within the full space.
Similar to our subspace training idea is LooC (Xiao et al., 2021) where both augmentationvariant and -invariant feature spaces are learned and then concatenated to support different tasks.
By contrast, our MAST is parameter-efﬁcient by learning a single feature space, can capture a lot more augmentation invariances efﬁciently in masked subspaces, and is aware of feature uncertainties.
Uncertainty modeling.
Uncertainty-aware methods (Gal, 2016) have been extensively studied in supervised learning to handle ambiguous or noisy inputs.
For SSL, only a few studies estimate uncertainties helpful in speciﬁc applications, such as self-supervised depth estimation (Poggi et al., 2020) and multiview stereo
(Xu et al., 2021).
General methods for uncertainty-aware SSL are lacking 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, 2021) that models uncertainty using spherical posterior distributions.
Here we propose a simple alternative for uncertainty modeling in our SSL framework, where a Gaussian feature assumption sufﬁces in each of our augmentation-speciﬁc subspaces.
MAST: M ASKED A UGMENTATION S UBSPACE T RAINING
We propose a 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 capturing distinct priors (about invariances) from different augmentations.
This will ease transfer learning
˜ μ k , ˜Σ k
Masked subspace encoder f View v
View v ′
Representation y
Image x encoder f
y ′ projector g
Gaussian embedding
Masks μ Σ μ ′ Σ′ m k ⊙ ˜ μ ′ k , ˜Σ′ k
Loss t ∼ T
Augmentation set | T | =
K t ′ ∼ T projector g
Augmentation-speciﬁc mask on different downstream tasks, where each task automatically selects useful priors in the pre-trained representations via ﬁnetuning; 2) free of feature suppression (Li et al., 2020) i.e. , learning priors in some features comes at no cost of suppressing others.
Particularly, we should accommodate learning different invariances that can be either universally helpful for all downstream tasks or that would otherwise contradict each other (as shown in Fig. 1); 3) efﬁcient enough to learn a large number of priors without too much cost (time or memory), and incurs no overhead for transfer learning.
3. D ESCRIPTION OF MAST
We achieve the desiderata above by MAST, see the schematic in Fig.
Assume we have a sampled batch of n images, and an augmentation set T consisting of K atomic augmentation operators o 1 , . . .
, o K such as RandomResizedCrop and Color Jitter .
For each image x , we produce two augmented views v = t ( x ) and v (cid:48)
=
t (cid:48)
( x ) by sampling random compositions of image augmentations t ∼ T ( o 1 ,...,o K ) and t (cid:48) ∼ T ( o 1 ,...,o K ) with their speciﬁc augmentation parameters.
We encode the views v and v (cid:48) by an encoder f into their representations y =
f ( v ) and y (cid:48) =
f ( v (cid:48) ) , which are then mapped by a projector g onto embeddings z = g
( y ) and z (cid:48)
=
g ( y (cid:48) ) .
Typically, a similarity loss is computed between z and z (cid:48) to learn augmentation invariances.
However, this makes it hard to learn distinct invariances because learning is unfactorized in the presence of the compositions of K augmentations applied on x .
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 =
[ z 1 , . . .
, z n ] and Z (cid:48) =
[ z (cid:48) 1 , . . .
, z (cid:48) n ] be the two batches of d -dimensional embeddings of the augmented views V
=
[ v 1 , . .
. , v n ] and V (cid:48)
=
[ v (cid:48) 1 , . . .
, v (cid:48) n ] .
We build MAST on top of the VICReg framework (Bardes et al., 2022) which minimizes the distance between the embeddings of different views, with two regularization terms:
[Equation]
, Z (cid:48) ) ,
[Equation]
where D ( Z, Z (cid:48) )
enforces augmentation invariance in form of the mean-squared euclidean distance (cid:80)
i ||
z i − z (cid:48)
i || 22 /n , without any feature normalization.
The variance term (cid:96) var ( Z, Z (cid:48) ) keeps the information content of each embedding (in Z or Z (cid:48) ) above a certain level.
While the covariance term (cid:96) cov ( Z, Z (cid:48) ) is introduced to prevent informational collapse by decorrelating different dimensions in each embedding.
λ, α and β are the loss coefﬁcients.
The above variance and covariance terms make the resulting maximum information embeddings well suited for disentangled learning.
As such, we propose to disentangle the embedding space into distinct subspaces, each of which captures invariance to a speciﬁc augmentation.
Instead of using a set of bulky MLP heads to generate subspaces, we introduce masks
M ∈ R d × K ≥over the embedding space at a much lower cost.
Each mask m k
= M : ,k has the same dimensionality d as the embedding vector z i , and there are K masks to learn for the same amount of augmentation operators in T .
To attend to the k th subspace, we only need to apply an element-wise masking operation (cid:12) using the corresponding mask m k to select relevant embedding dimensions ( i.e. , z i (cid:12) m k ).
Therefore, to learn
[Figure]
disentangled augmentation invariances from the paired views V and V (cid:48) , we replace the unfactorized distance function in Eq.
(1) with the sum of K masked embedding distances:
[Equation]
where (cid:96) sp ( M ) is a constraint to encourage sparse masks.
The disentangled embeddings and masks
M are trained jointly (see Section 3.2 for the two-staged training details for better disentanglement).
Fig.
3 shows example masks that induce distinct but overlapping subspaces.
This helps to avoid feature suppression by encoding contradicting invariances in separate subspaces, but at the same time encourages collaborative learning in the shared embedding dimensions among subspaces.
Fig.(Appendix C) reafﬁrms that different augmentation subspaces beneﬁt 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 disentanglement effect at the representation level.
Thus, we 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 a largely unexplored problem in SSL.
We address this issue with a simple method that explicitly models uncertainty using Gaussian embeddings.
Concretely, instead of having a projector g : y
→ z to map representations y to deterministic embeddings z ∈ R d , we extend to the projector g :
y → ( µ, Σ) that produces a multivariate Gaussian distribution N ( µ, Σ) with mean µ ∈ R d and covariance Σ ∈ R d × d .
We consider Σ to be diagonal for efﬁciency, modeling the uncertainty independently for every dimension of the mean embedding µ .
With this independence assumption, we can use masking to easily induce embedding subspaces that are still Gaussian, i.e. , N (˜ µ k , ˜Σ k )
where ˜ µ
k = µ (cid:12) m k and ˜Σ k = Σ (cid:12) diag ( m k ) .
We ﬁnd the Gaussian embeddings sufﬁce to model the uncertainties of augmented views in each subspace where only one type of augmentation is involved.
Then we further improve Eq.
(2) to learn factorized augmentation invariances in an uncertainty-aware manner:
[Equation]
where ˜ µ ik and ˜ µ (cid:48)
ik are the mean embeddings of views v i and v (cid:48)
i in the k th subspace, and their distance is re-weighted by their uncertainty measures (trace tr ( · ) of covariance matrix) to determine the closeness of the two Gaussians.
Intuitively, when the embeddings have high uncertainties, the similarity loss should be downweighted to prevent similarity mismatch between ambiguous views.
To stablize the learning of uncertainties
[Equation]
we introduce an extra regularization term to encourage consistency between the predicted Gaussian embeddings N i ( µ i , Σ i ) and N (cid:48) i ( µ (cid:48)
i , Σ (cid:48) i ) for views v i and v (cid:48)
i .
This term is based on symmetric KL divergence, with a closed-form expression for Gaussians:
[Equation]
Combining Eqs.
(3-6) and substituting them into Eq.
[Equation]
we arrive at our ﬁnal loss function:
[Equation]
3. I MPLEMENTATION DETAILS
Architecture and hyper-parameters.
Our pretraining is mostly performed on the unlabeled ImageNet 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 2048dimensional feature representations.
The projector g has two fully-connected layers of size 8192, followed by separate fully-connected heads that respectively output the mean µ and diagonal covariance Σ (both of d = 4096 dimensions) of Gaussian embeddings.
Loss coefﬁcients in Eq.
(7) are set as α = 25 , β
= 1 following VICReg (Bardes et al., 2022), and we set λ
= 25 d/K , λ 1 = 600 / ( dK ) and λ 2 = 25 to generate comparable loss magnitudes.
Fig.
8 in Appendix D shows the pretraining performance on different datasets is not very sensitive to the loss coefﬁcients.
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 Cubuk et al.
(2019), as well as Sharpness, Gaussian Noise, Sobel Filter and Cutout (DeVries & Taylor, 2017)
for diversity.
Training details : The training protocol follows that in (Bardes et 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 a cosine annealing schedule (Loshchilov & Hutter, 2017).
For effective learning of disentangled subspaces, we found it is better to schedule the available augmentations in two stages.
In the ﬁrst 500 epochs, we sample one random augmentation at a time (with its default parameters) to train the corresponding subspace and mask, setting K = 1 in Eq.
This ensures all the factorized subspaces are trained well in this stage.
For the remaining 500 epochs, random augmentation compositions are allowed to learn disentangled invariances from progressively stronger augmentations ( K linearly increases from 1 to its maximum value).
For each composing augmentation, we use its default augmentation probability but random magnitude for robust training.
E XPERIMENTS
Main results.
We start with evaluating our MAST method for SSL on ImageNet.
Table 1 shows the results from both linear and semi-supervised evaluations, using the same optimization procedures of our baseline VICReg (Bardes et al., 2022).
The numbers reported for our MAST-5 and MASTvariants are the mean scores from 3 runs with different random initialization.
It can be seen from Table 1 that when 5 standard augmentations are used for SSL, MAST-5 achieves strong performance on-par with or better than that of state-of-the-art contrastive and non-contrastive methods.
MAST-5 also consistently outperforms our baseline method VICReg, indicating the power of disentangled and uncertainty-aware learning of different augmentation invariances.
Note predictive contrastive learning (PCL) is related to us in avoiding feature suppression, but by strengthening contrastive learning with a predictive inpainting task.
MAST-5, on the other hand, achieves much
[Table] [Table]
stronger performance by explicitly training different feature subspaces that can encode contradicting invariances without undesired trade-offs.
Fig.
3 supports this hypothesis with one visual example of bird classiﬁcation where one subspace is 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
[Table]
which outperforms InfoMin Aug and CLSA that similarly learn from stronger augmentation compositions but in a suboptimal unfactorized way
[Table]
Next
[Table]
following the setup in (Bardes et al., 2022)
[Table]
we evaluate the transfer performance of the
ImageNet-pretrained representations on various downstream tasks: linear classiﬁcation on the scene dataset Places205 (Zhou et al., 2014), multi-label image classiﬁcation dataset VOC07 (Everingham et al., 2010) and ﬁne-grained image classiﬁcation dataset iNaturalist2018 (Van Horn et al., 2018).
[Table]
We evaluate linear classiﬁcation on the ﬁne-grained classiﬁcation dataset CUB-200 (in Top-1 and
Top-5 accuracies %) and Flowers-102 dataset (in 5-shot and 10-shot classiﬁcation accuracies %)
[Table]
We also transfer the representations to object detection on VOC07+12 using Faster R-CNN (Ren et al., 2015) with a R50-C4 backbone, and instance segmentation on COCO (Lin et al., 2014) using
Mask-R-CNN (He et al., 2017) with a 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 modiﬁcation (IFM) (Robinson et al., 2021) and Leave-one-out Contrastive Learning (LooC) (Xiao et al., 2021).
All comparing methods use the same 100-class ImageNet dataset and 500 epochs for pretraining.
IFM is a method that avoids shortcut solutions (caused by feature suppression) 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) that explicitly reduces feature suppression.
LooC is built on the idea of combining both augmentation-invariant and -variant features (similarly in different subspaces) to reduce the negative effects of particular invariances that can be harmful for some downstream tasks.
Table 3 shows that the original feature embeddings of LooC when regularized with such subspace training actually hurt performance on the source dataset, in comparison to baseline MoCo v2.
This suggests that the “disentanglement” learned by LooC may still suppress certain useful features.
By contrast, our disentangled and uncertainty-aware learning with MASTseems feature suppression-free, giving notable gains over MoCo v2.
Note LooC++ can recover the loss of LooC by ensembling feature subspaces.
However, LooC++ still remains on-par with MoCo v2, while suffering from a much larger cost of ensembling those redundantly trained subspaces.
When we apply the subspace ensembling idea to MAST, we also observe (slightly) improved performance at the cost of increased feature dimensionality.
Hence our default MAST-5 approach uses the shared representations before subspace masking, which is efﬁcient and performant already.
MAST-15 further improves performance by learning from more augmentations.
Table 4 validate our beneﬁts again on two downstream tasks: ﬁne-grained classiﬁcation on CUB-200 dataset (Wah et al., 2011) with 200 bird classes, and few-shot classiﬁcation on Flowers-102 dataset (Nilsback &
Zisserman, 2008) with 102 ﬂower categories.
Ablation study is conducted in Table 5 to evaluate different components of MAST.
We use the setup where representations are pretrained on the ImageNet-1k 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 to our baseline VICReg.
We can see that this variant falls far behind MAST-15 (row 1), and even behind the baseline (row 4) on classiﬁcation 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 4 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 efﬁciency.
Recall that
MAST-5 beneﬁts from a two-staged scheduling of augmentations (from one at a time to linearly increasing number of augmentations).
We ﬁnd 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
[Table]
GrayscaleFlipRotateShearXShearYTranslateXTranslateYInvertSharpnessSobelG
[Table]
NoiseCutoutSolarizeEqualizePosterizeM
[Table]
[Figure]
Augmentations with low correlation with others provide diverse and large performance rows 7-8, we can easily ﬁnd the beneﬁts of our learned subspace masks (vs. handcrafted disjoint masks
1’s
Visualizations.
We exploit the learned subspace masks to quantify augmentation correlations unsupervisedly by computing the cosine similarity between masks.
Fig.
4(a) shows the example pairwise correlation between 19 augmentations on ImageNet.
We can observe weakly correlated augmentations such as ( Gaussian Blur , Rotate ) and ( Gaussian Blur , ShearY ).
They deﬁne diverse augmentations 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 ( ShearX , ShearY ) that offer redundant invariance information.
The 4 optional augmentations Solarize, Equalize, Posterize, Motion Blur are highly correlated with others too ( e.g. , with Color Jitter and Gaussian Blur ), therefore we exclude them from MAST-15.
Fig. 4(b) demonstrates the positive role of augmentation diversity on generalization.
We gradually add augmentations to MAST-5 for SSL, such that distinct augmentations with low mean correlation with others are added ﬁrst.
Those augmentations are observed to improve the downstream performance quickly, 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 samples in order to improve SSL performance (comparison with a related uncertainty model (Nakamura et al., 2022) is also provided), and how uncertainty reﬂects learning difﬁculty.
C ONCLUSION
This paper studies the important role of data augmentation in creating useful invariance priors during 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 uniﬁed 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 a variety of downstream vision tasks.
A A DDITIONAL I MPLEMENTATION D ETAILS
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 a Generalized-Mean (GeM) pooling layer (Radenovi et al., 2019) followed by a fully connected layer that outputs µ ∈ R d with d = 4096 .
The covariance head is composed of a GeM layer and a 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 are beneﬁcial, 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 ∈ R d × K ≥by learning them as M = max(0 , U ) using a ReLU activation function where U ∈ R d × K .
During training, we found it useful to initialize the masks as random rather than ﬁxed ones ( e.g. , when initial mask values are all the same).
More speciﬁcally, we initialize each mask ( i.e. , each column of U ) such that the corresponding subspace of d/K dimensions follows a Gaussian distribution and Gaussian noise is added for all dimensions.
The Gaussian prior for each subspace mask eases 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 ﬁnd meaningful mask patterns in all dimensions.
B U NCERTAINTY V ISUALIZATION AND C OMPARISON
Uncertainty visualization.
In MAST, uncertainty is estimated as a vector, i.e. , diagonal covariance of Gaussian embeddings to 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(a) 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 can reasonably capture the strength of data augmentations.
This is a desirable behavior since when the augmentations are strong enough to produce ambiguous views or even change the identity of input images, we can use the estimated high uncertainty to downweight the distance between ambiguous views.
As a result, similarity mismatch is avoided during invariance learning.
[Table]
Fig.
5(b) further connects uncertainty with the classiﬁcation 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 ﬂower appearances.
Accordingly, the ‘Pansy’ class has a much lower classiﬁcation accuracy than that of ‘Sunﬂower’ class with low uncertainty.
This validates that our uncertainty measure can indeed reﬂect the learning difﬁculty.
For the difﬁcult ‘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 et al., 2022) that models uncertainty for
SSL using spherical posterior distributions and variational inference.
We follow VI-SimSiam to pretrain for 100 epochs and use ResNet-18 as the encoder.
Pre-training is performed on two datasets respectively: Flowers-17 (Nilsback & Zisserman, 2006) with 17 ﬂower categories (80 images for each category), and Cars (Krause et al., 2013) with 197 car models (16,185 images).
After pre-training, we train a linear classiﬁer on the frozen representations using the training set of the corresponding dataset
[Table]
and evaluate linear classiﬁcation accuracy in its test set.
Table 6 summarizes the results in Top-1 and Top-5 accuracies.
We observe that the uncertaintyaware 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 ‘uncertainty only’ to not beneﬁt 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 baselines (VICReg vs. SimSiam) and implementation details.
The goal here is only to test how our uncertainty modeling mechanism works in general, when using a simple Gaussian assumption about feature embeddings.
Results show that MAST-5 (uncertainty only) obtains quite competitive performance on both datasets, and is further improved by disentangled learning consistently.
C M ORE A NALYSIS OF THE A UGMENTATION S UBSPACES
Optimal augmentations (subspaces) are task-dependent.
Fig.
1 supports the claim by demonstrating the difference between the optimal augmentations for different downstream tasks.
This motivates us to develop MAST, by learning disentangled augmentation invariances in masked feature subspaces that are augmentation-speciﬁc.
Now with MAST, we ﬁnd it much easier to quantify the relationship between augmentations and tasks, using the learned augmentation-speciﬁc 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 to identify optimal augmentations for one task in a cheaper way than the leave-one-out strategy in Fig.
We take the classiﬁcation task as one example.
For correctly classiﬁed samples, we examine the class predictions (only for the ground-truth class, using the corresponding weights of the classiﬁer) by differently masked feature embeddings ˜ µ ik of each sample.
The higher the class prediction, the more the corresponding masked augmentation subspace (and its captured invariance) beneﬁts the given task.
Fig.
6 shows the class predictions across augmentation subspaces for ImageNet and CUB-200 classiﬁcation, respectively.
We observe that, again, different tasks beneﬁt from a different set of augmentations ( e.g. , Color Jitter, Random Grayscale, RandomResizedCrop for ImageNet, and Gaussian Blur, Random Flip, RandomResizedCrop for CUB-200).
Augmentation invariance vs. preserving augmentation-speciﬁc information.
Since our MAST method focuses on disentangled learning of different augmentation invariances, we now verify how
[Figure]
[Table]
[Figure]
[Table]
[Figure] [Figure]
[Figure]
a n ce m e t r i c
[Figure]
I n v a r i a n ce m e t r i c
[Figure]
I n v a r i a n ce m e t r i c 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 ˜ µ ik 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, RandomResizedCrop .
We exclude other 2 ( Random Flip, Random Grayscale ) whose augmentation parameters are discrete.
We can see from the ﬁgure that our subspaces indeed show great robustness under various augmentation intensities, especially when compared to the VICReg baseline that learns all augmentation invariances in an unfactorized way.
One might ask: will we loose all augmentation-sensitive information when we learn factorized invariances from those augmentations?
Short answer is no, as evidenced by our strong performance on various downstream tasks that rely on speciﬁc augmentation characteristics ( e.g. , color-sensitive bird classiﬁcation on CUB-200).
We hypothesize that different subspaces may offer complementary information about augmentations for different tasks to choose from.
For example, in Fig. 3(b), the color information suppressed in the subspace invariant to Random Grayscale is available in the subspace invariant to RandomResizedCrop .
[Table] [Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
[Table]
[Figure]
by varying a multiplier s (default s = 1 ) applied to each of them.
The impact on pretraining performance is monitored with (a) the large ImageNet dataset (evaluated on 2 tasks), as well as relatively small datasets (b) Flowers-17 and (c) Cars.
Performance is not very sensitive to the loss coefﬁcients, especially when they are scaled by s ∈ [0 . 5 , 2] .
To further support the above hypothesis, we follow the AugSelf work (Lee et al., 2021) to solve a pretext task of 4-way rotation classiﬁcation (0 ◦ , 90 ◦ , 180 ◦ , 270 ◦ ) using pretrained representations.
The same pretraining protocol is followed: pretraining ResNet-18 for 200 epochs on STLdataset (Coates et al., 2011), and standard augmentations are used (no Rotate ).
To evaluate rotation classiﬁcation accuracy, a linear classiﬁer is trained on top of the frozen representations.
Tableshows competitive performance of MAST-5 (63.47%) vs. AugSelf (64.61%) which is explicitly trained to be augmentation-predictive.
This validates that we can learn invariances to standard augmentations while maintaining sensitivity to the unseen rotation.
In other words, MAST is able to preserve augmentation-invariant and -speciﬁc information implicitly in separate subspaces.
D A DDITIONAL R ESULTS
Impact of pretraining epochs.
Recall our MAST models are pre-trained in two stages, with increasing number of augmentations (i.e. stronger compositions).
As a result, MAST typically beneﬁts signiﬁcantly from longer training to take full advantage of disentangled learning with strong augmentation compositions.
Table 8 exempliﬁes the ImageNet performance as a function of pretraining epochs, from 200 to default 1000.
We can observe that 1) our MAST-5 outperforms its baseline model VICReg across all epoch numbers, 2) more epochs lead to larger gains, but the gain is still decent (0.7%) at low epoch number 200.
Sensitivity of MAST to loss coefﬁcients.
Recall the loss function of MAST in Eq.
(7) consists of 2 terms borrowed from the VICReg baseline with their default loss coefﬁcients.
We introducemore loss terms, and only set their loss coefﬁcients λ = 25 d/K , λ 1 = 600 / ( dK ) and λ 2 =to have a comparable loss magnitude with that of other terms.
In practice, we use the exact same loss coefﬁcients across different augmentation distributions and different pretraining datasets.
Fig.shows the pretraining performance is not very sensitive to λ, λ 1 , λ 2 , especially when they are scaled
[Table]
by a factor between 0.5 and 2.
Note for comprehensive evaluations, we monitor the pretraining performance on the large dataset ImageNet and two relatively small datasets Flowers-17 and Cars, using their representative evaluation tasks.
We could have easily tuned λ, λ 1 , λ 2 per dataset for best performance, but we choose to use the same loss coefﬁcients for simplicity.
More results of ﬁne-grained classiﬁcation.
As a supplement to Tables 3 and 4, Table 9 provides more results on additional ﬁne-grained classiﬁcation datasets iNat-1k (Van Horn et 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 ResNet-50 pretrained on ImageNetfor 500 epochs.
We compare with both LooC and a new method AugSelf (Lee et al., 2021).
AugSelf combines contrastive learning with an auxiliary loss that predicts augmentation parameters to improve SSL generalization.
Note both LooC and AugSelf are built on top of the MoCo v2 baseline.
It can be observed from the table that our MAST-5 is quite competitive on all datasets, and MASTfurther improves consistently and achieves the best performance.