In-Context Symmetries: Self-Supervised Learning through Contextual World Models
Authors: Sharut Gupta
Chenyu Wang
Yifei Wang
Tommi Jaakkola
Stefanie Jegelka
Venue: NeurIPS 2024 poster (2024)
Keywords: Self-Supervised Learning; Context; Equivariance
Primary area: machine_vision
TL;DR: Learning general representations in self-supervised learning while having the versatility to tail down to task-wise symmetries when given a few examples as the context.

Abstract: At the core of self-supervised learning for vision is the idea of learning invariant or equivariant representations with respect to a set of data transformations. This approach, however, introduces strong inductive biases, which can render the representations fragile in downstream tasks that do not conform to these symmetries. In this work, drawing insights from world models, we propose to instead learn a general representation that can adapt to be invariant or equivariant to different transformations by paying attention to context --- a memory module that tracks task-specific states, actions and future states. Here, the action is the transformation, while the current and future states respectively represent the input's representation before and after the transformation.  Our proposed algorithm, Contextual Self Supervised Learning (ContextSSL), learns equivariance to all transformations (as opposed to invariance). In this way, the model can learn to encode all relevant features as general representations while having the versatility to tail down to task-wise symmetries when given a few examples as the context. Empirically, we demonstrate significant performance gains over existing methods on equivariance-related tasks, supported by both qualitative and quantitative evaluations.

In-Context Symmetries: Self-Supervised Learning through Contextual World Models
Sharut Gupta', Chenyu Wang', Yifei Wang', Tommi Jaakkola
Stefanie Jegelka TU Munich, MIT CSAIL stefje@mit. edu
MIT CSAIL
(sharut, wangchy, yifei_w, jaakkola) Omitt. edu
Abstract
At the core of self-supervised learning for vision the idea of learning invariant equivariant representations with respect to set of data transformations. This approach, however, introduces strong inductive biases, which can render the repre- sentations fragile in downstream tasks that do not conform these symmetries. this work, drawing insights from world models, we propose to instead learn general representation that adapt to be invariant equivariant to different transformations by paying attention to context - a memory module that tracks task-specific states, actions, and future states. Here, the action the transformation, while the current and future states respectively represent the input's representation before and after the transformation. Our proposed algorithm, Contextual Self- Supervised Learning (CONTEXTSSL), learns equivariance to all transformations (as opposed to invariance). As result, the model learns encode all relevant features as general representations while, importantly, it can adapt to restrict to task-wise symmetries when given a few examples as the context. Empirically, we demonstrate significant performance gains over existing methods on equivariance- related tasks, supported by both qualitative and quantitative evaluations. Code available at https:/ thub. com/ Saam//haut//////e CaerruttInnCCoiii
Introduction
Self-supervised learning (SSL) of image representations has made remarkable progress in recent years [9, 6, 47, 31, 18, 3, 19, 20, 38, 34, 10, 26, 11, 46, 39, 48), achieving competitive performance its supervised counterparts various downstream tasks, such image classification.
Most of these works are based on the joint-embedding architecture (as shown Figure 2(a)) which encourages the representations of semantically similar (positive) pairs to be close, and those of dissimilar (negative) pairs to be more orthogonal. Typically, positive pairs are generated by classic data augmentation techniques that correspond to common pretext tasks, c.g., randomizing color, texture, orientation, and cropping. The alignment of representations for positive pairs can be guided by either invariance [9, 6, 11, 26, 46, 20), which promotes insensitivity to these augmentations, equivariance [22, 13, 12, 16, 1, 17], which maintains sensitivity to them. However, enforcing invari- ance equivariance to pre-defined set of augmentations introduces strong inductive priors which are far from universal across a range of downstream tasks. For example, invariance to image flipping useful for image classification but significantly hurt performance retaining sensitivity to flipping is crucial. This often results in brittle representations that necessitate retraining the model with different augmentations tailored to each downstream task [42, 12].
image segmentation, where
Equal contribution
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

Context (Rotation)
Source
I-NN 2-NN 3-NN
Target
Contestall Self Superviuppeiee Learning
(Rg. R..
Coneext (Color)
(O b))
(0
Figure I: We apply transformation (rotation color) source image in latent space and retrieve the nearest neighbor (NN) of the predicted representation when the context contains pairs of data transformed by (top row) 3D rotation (R" R" RRt): (bottom row) color transformation (0, Φ). In the top row, we see that CONTEXTSSL learns cquivariance to rotation and invariance to color as the NN representations match the target's s angle but its color. the bottom row, adapts the color context and enforces the reverse, be equivariant to color and invariant to rotation.
This rigidity of traditional SSL methodologies contrasts sharply with human perceptual abilities, which are highly adaptive, tuning into relevant features based on the context the environment or task at hand. For example, humans focus more on color details when identifying flowers, and on spatial orientation such as rotation angle when determining time on analog clocks. suggests that the required feature invariances or equivariances should also vary across different tasks or contexts, which motivates our central question.
Can incorporating context into self-supervised vision algorithuns eliminate augmeentation-based inductive priors and enable dynamic adaptation varying task symmetries?
This work suggests a positive answer to this question by proposing to enhance the current joint embedding architecture with a finite context few demonstrations that inform about task-specific symmetries, as shown in Figure 2(c). Based contrastive
abstract representation of task, containing
this idea, propose Contextual Self-Supervised Learning (CONTEXTSSL.), learning framework that uses a transformer module to adapt to selective invariance or equivariance to transformations by paying attention to context representing task. Unlike previous approaches with built-in symmetries, the ability of CONTEXTSSL to adapt to varying data symmetries-an without undergoing any parameter updates- it to learn general representation across tasks, devoid of specific inductive priors.
This unique prospect makes model promising approach to building world models [24, 25, 28, 37, 45] for vision. World models are essential for building representations the world based on past experiences, akin to how humans form their internal world representations. Recently, efforts have been made to adapt world modeling into vision through Image World Models (IWM) [17] ( (Figure 2(b)), that consider transformations actions and the input and its transformed counterpart world states at different time steps. However, these approaches also enforce equivariance to a predefined set of actions, such as color jitter. CONTEXTSSL addresses this challenge by enhancing traditional IWMs with context, a model we refer to as Contextual World Models. We demonstrate that in the absence of context, CONTEXTSSL. learns a general representation by encoding all relevant features and data transformations. As the context increases, the model tailors its symmetries to task, encouraging cquivariance to a subset of transformations and invariance the rest (as shown in Figure 1). This approach promotes learning a general representation that can flexibly adapt to the symmetries relevant various downstream tasks, eliminating the need to learn separate representations for each task. We empirically validate our approach on the MIMIC-III, UCI Adult, 3D Invariant Equivariant Benchmark (3DIEBench) and CIFAR- 10 datasets, extending naturally occurring sensitive features such as gender and transformations such rotations, cropping, and blurring.
To summarize, the main contributions of work are:
We propose CONTEXTSSI. a self-supervised learning algorithm that adapts to task-specific symmetries by paying attention to context. Our method resolves the long-standing challenge of

---

Predictor
Contextual World Models
Context
14+1
Y4+1
Ap-1
Yi+1
X1
3/1
Ae+1
N141
(b)
Figure 2: Family of approaches in self-supervised learning (a) Joint Embedding methods [9, 8] encode invariances input transformations a by aligning representations across views of the same image; (b) Image World Models [17. I] train a world model in the latent space and encode equivariance to input transformations; (c) Contextual World Models (ours) selectively enforce equivariance or invariance to a subset of input transformations based on context {(z;, as Di) }
enforcing fixed invariances and equivariances to handcrafted data augmentations, enabling adaptive and task-sensitive representations without parameter updates.
We show that learning with context is prone to identifying shortcuts and subsequently propose two key modules to address it: context mask and an auxiliary predictor.
We demonstrate the efficacy of our approach on MIMIC-III, UCI Adult, 3DIEBench and CIFARIO. showing its ability to selectively learn invariance or equivariance to naturally occurring features and sensitive attributed like gender and transformations such as color and rotation while maintaining similar performance on invariant benchmarks.
2
Augmentation-basee Inductive Bias in Self-Supervised Learning
The goal of self-supervised learning (SSL) is derive meaningful data representations without relying on human-labeled data. Given an unlabeled dataset D, SSL methods learn representation function fo X' >> Z that maps input data X' a latent space Z.
2.1 Role of data augmentations in Self- Seff-Supervised Learning
Data augmentations are arguably the most important component modern SSL methods, where the representation function is learned map the augmented views of data into latent space. The choice of data augmentations plays a crucial role in the quality of the learned representations. Formally, define augmentation A random variable distributed over set of N data transformations
with domain A {ai....,,a, aN), where a; Rd + R2d denotes an input mapping, and d, d' its input and output dimensions, respectively. Among existing SSL methods, there are generally two ways utilize augmentations, either through invariant learning or equivariant learning. In invariant learning, two random augmentations of the example are drawn and their representations are pulled together during feature learning be invariant to the data augmentations as shown in Figure 2(a). Instead, in equivariant learning, the features are learned to be sensitive to data augmentations. Formally, for a representation Z, one can use H H(A))) as measure of the degree of feature invariance equivariance: if H(AZ)) is relatively small, the representation Z is nearly equivariant the augmentation A; otherwise, if H (A)Z) very large (close to H(A)).
invariant to A. Recent SSL methods [22, 16, 35, 13, 12] have shown that enforcing equivariance can often lead to better representations compared to enforcing invariance, for two key reasons: 1) Invariance restricts the expressive power of the features learned it removes information about features or transformations that may be relevant in fine-grained tasks [32, 43]; 2) contrastive learning benefits from partial invariance through implicit equivariance of the projection head [29].
Herr, the concept of equivariance is used a loose sense, meaning that the learned features are sensitive to data augmentations. Note that since some augmentations are non-invertible (e.g., grayscale). they do not form group, and exact equivariance is not well-defined.
3

---

2.2 Drawbacks of Hardcoding Symmetries in Self-Supervised Pretraining
As discussed above, a common theme in existing SSL methods to enforce invariance or equivariance to specific set of augmentations A. For instance, in SimCLR, A is chosen to be a manually selected set of random augmentations such as random cropping, flipping, and color jitter. Therefore, the learned representations, either invariant or equivariant to these augmentations, are tailored to the specific symmetry imposed during pretraining. However, real world scenarios, there is no single symmetry that is universally applicable across all tasks. For example, object recognition (e.g., chair) often requires invariance to image color, while certain tasks, e.g., flower recognition, need sensitivity to color information instead. Either to include or not to include color information as part of the augmentations can lead to suboptimal performance certain tasks, causing a fundamental dilemma in existing SSL. This leads to brittle representations over a range of downstream tasks, as the model needs to be retrained on different augmentations depending on the downstream tasks, as consistently observed in previous works [42, 12].
3 Beyond Built-in Symmetry: Contextual Self-Supervised Learning
Recognizing the limitations of existing augmentation-speen SSL methods. we propose a new paradigm: Contextual Self-Supervised Learning (CONTEXTSSL). Unlike traditional methods, this approach learns a single model that adapts to be either invariant or equivariant based on context- specific augmentations, tailored to the needs of the task or data at hand. Instead of enforcing fixed set of symmetries, CONTEXTSSI. learns these symmetries from contextual cues, thus capturing the unique set of features of downstream tasks. This adaptability allows it to serve SSL framework, capable of learning from a diverse array of pretraining tasks with varying symmetry priors and seamlessly adapting to different downstream tasks.
general-purpoon
To design CONTEXTSSL, we draw inspiration from world modeling [24, 25, 37, 45], widely used framework in reinforcement learning (RL). World modeling aims to build representations of the world from past experience by predict the next state 2z+1 from the current state æ, and action a,. This next state prediction task captures the inherent mechanisms of the system and facilitates decision making. Traditionally applied RI., the benefits of world modeling in vision have been largely unexplored. Recently. Image World Models (IWM) [15] established parallel between world models and the image-based SSL by considering data transformations as actions, the representation of input data as world state at time and that of the transformed input next world state. However, IWMs have two key drawbacks: 1) similar to previous SSL approaches, they rely on a predefined set of data augmentations, such as color, which are not tailored to specific downstream tasks and influence the learned features; 2) they lack the memory module of world models that tracks previous experience in terms of past states, actions and corresponding next states and provides context to fully define the current state.
In light of these ideas and challenges, we model CONTEXTSSI. in vision self self-supervised learning as Contextual World Models. In this way, CONTEXTSSI. addresses the key drawbacks of IWMs by 1) encouraging the model to preserve all meaningful features to be able to adapt to symmetry from context, and 2) incorporating context to adapt to different task-specific symmetries, removing the need to re-train separate representations for each downstream task. This general ability is akin to human perception that captures versatile aspects of the input, while focusing on specific details depending on the context at hand. For instance, humans focus more color details when identifying flowers, and on spatial orientation such as rotation angle when determining the time on analog clocks.
3.1 Contextual World Models
Drawing inspiration from the in-context learning [7] of foundation models in natural language processing, a natural way to incorporate the memory capabilities of world models is by encoding these abilities as contextual information. In this work, we propose an expressive and efficient implementation of CONTEXTSSL through Contextual World Models, where we design a transformer- based module to encode the context and extract contextually equivariant or invariant representations. We begin by baking symmetries in the context - (x,, a, y) using positive pairs x and y transformed by series of different augmentations. The key intuition behind our approach is selective inclusion of augmentation parameters for specific transformation groups: excluding parameters enforces invariance, while including them enforces equivariance. This is because providing augmentation

---

parameters allows the model to learn the impact of transformations (equivariance), whereas excluding them during alignment enforces invariance, akin invariant versus equivariant learning in SSL. We elaborate on these ideas below.
Symmetries as Context. Given set of groups of input transformations (9,....; G.M ). the goal of CONTEXTSSL is to build a general representation that is adaptive to a set of multiple symmetries corresponding to these different groups. For example, each data augmentation, e.g., rotation, trans- lation, as well as their compositions, can serve as different transformation groups. Each group Go can be represented through the joint distribution P(z, a, 3/Ge), where I is the input sample (sampled from an unlabeled dataset), a represents the parameters of the transformation drawn from Ge and applied to z, and y is the transformed input. In principle, can be transformed by composition of augmentations drawn from multiple transformation groups. For instance, in self-supervised learning, is common to enrich the learning process by transforming an input image through rotations, crops, and blurring. In such a case, a represents a subset of the transformation parameters belonging to the group 9c. applied to x to produce y. We approximate this probability distribution by drawing K' samples from the joint distribution and form a context C(G.) [[<] d1, v1 v1),... (IK, KK yk)], where Fú, a₁ 4/1 ~ P(x, a, y Go), E [K]. Therefore, the goal of ContextSSL is to learn data rep- resentations 2 = f(x, a/C') and
[xx(C)) that are adaptive to the data symmetries informed by the context C. Specifically, our goal is to train representations that become more equivariant to the underlying transformation group Ge with increasing context. Further, if 2 and y are transformed by augmentations from groups apart from Gc, we aim to learn more invariance to these groups with increase in context C(G.). The degree of equivariance of representation can be quantified by the error in maintaining consistent transformations. Based on this, a representation Z is considered "more equivariant (invariant)" if it has lower (higher) error in predicting the transformation parameters i.e. H(A Z))
Contextual World Models. To implement this broad goal, we propose to adaptively learn the symmetries represented by Go by training the model:
yi h((2i, a:); (Ir1 (1; v1), (Ii- 1, 04-1,, 9;-K),
(1)
While the requested prediction yi concerns only the inputs z; and ag the model can now pay attention to the experience so far, enforcing relevant symmetries for the augmentation group Ge. The predictor
is updated by minimizing the loss at each context length Σx1 ((h((z, a.); Ct-1), y1) where C = [[(1,, a1, 3/11 J1),.... (#;-1, a; 11- represents the context before index
natural way to facilitate such context- based training is through attention mechanisms in transformer- based autoregressive models. Large language models exhibit a remarkable capability of in-context learning - the ability to generalize to unseen tasks on-the-fly fly merely by paying attention to a few demonstrative examples of the task. Gupta et al. [21] among others, have leveraged this capability to generalize to different distributions merely by paying attention to unlabeled examples from a domain. Inspired by this, we train a decoder-only transformer model in-context by conditioning on the relevant context CG..) representing the transformation group Ge.
3.2 Contextual Supervieed Learning (CONTEXTSSL)
Motivated by the above ideas, we begin by constructing pairs of points ((2 } by either 1) sampling a transformation group 9 and transforming 2; by augmentation from G to ya or 2) if available, sampling a meta-latent and its transformation parameters difference between their individual latent parameters. We use the former construction in datasets such as CIFAR10 but use meta-latents such as 3D pose, lighting etc. for datasets such as 3DIEBench [16]. Note that pairs of data can also be transformed by a scries of augmentations sampled from other transformation groups However, as previously discussed, the transformation parameters used in the context C(G) of group G are solely those of the augmentations belonging to the group.
Following this, as illustrated in Figure 2, each input sample (zz,,,:)} from the context is indepen- dently transformed by the encoder into its corresponding latent representation. Next, representations of the input samples I; are concatenated with their corresponding transformation action a; This concatenated vector (I,; az) and the representation of the corresponding transformed input y/s col- lectively form the context corresponding to the symmetry G. The corresponding output embeddings are then aligned using the InfoNCE loss, which minimized each context length. If a, set to zero for all tokens in a sequence, CONTEXTSSL enforces invariance to G, since it aligns T; and y,
5

---

without conditioning on the transformation parameters, Overall, we optimize the following loss:
exp h((2,, ((1,2,,,aa((((()))100000000)
LCONTENISSI = (0I.. BC()))
Bo(g)
log
j== exp (22
where transformed data tokens yj (i # i) form the negatives. We use a similar symmetric loss term using y as the anchor, (Ir4, (4) and (z)) (i + i) the positive and negatives respectively.
At inference, we tailor the extraction of representations to match the specific requirements of the downstream task, whether it benefits from equivariance or invariance to transformation group G. In particular, if the task benefits from equivariance, we extract the representations of the test data at the maximum context length used during training K, by constructing { (22; Qg. y/) ys)),= its preceding context. Here a, belongs to the group 9 and is used to transform other unlabelled data from the test set z; into ye. On the contrary, if the downstream task benefits from invariance to the group, use {(x,, O, yı H=) as the preceding context. This notion can be generalized to enforce equivariance to a subset of groups and invariance another. Specifically, including the augmentation parameters for transformations in a group & in the context enforces equivariance, while excluding them enforces invariance. In both cases, the data are still transformed using augmentations, regardless of the type of symmetry desired. This flexibility of context creation in CONTEXTSSL allows us to tailor the representations to different symmetries and optimize for the model's performance across a range of tasks. However, this implementation bears two key challenges, as detailed below.
Context Masking, Given that (æ,, a,) precedes y/ in the context sequence, trivial solution to minimizing the alignment loss arises where the model treats the embeddings of (z;, a,) identical to Y: due its access -4-- This phenomenon, often referred to as shorteut learning, poses significant challenge as it leads the model to collapse to constant representations for each pair (I:; yi). all while perfectly minimizing the loss. We address this challenge by masking each token in the context. As a consequence, when encoding the token yt. the transformer only has access to past context C; = [[#1, 01, (11,91).................. BI....., (Zi-1, a;- 1, 9;)1),, excluding its corresponding positive sample (z).
the input token (I,, 0:) for
This masking approach ensures that both the anchor and its corresponding positive share the same context, thus promoting the alignment of positive samples based on semantic relationships rather than mere replication. However, shown in Figure 4 for p 0, a residual challenge of shortcut learning persists when distinguishing the positives from the negatives. Since the context corresponding to each negative is different from that of the anchor and the positive, the model could employ trivial solutions, such as using the mean of the context vector to differentiate between positives and negatives. To mitigate this issue, we introduce an additional layer of randomness to our masking strategy. Specifically, for each token in the context vector, we implement random masking with a probability p for tokens preceding it. This ensures that for given anchor token, both the positive and the negatives have different contexts from the anchor, thereby necessitating a deeper, semantic understanding to effectively distinguish the positives from the negatives.
Avoiding collapse to Invariance. A trivial but undesirable solution that minimizes our optimization objective is invariance to the input transformations i.e. the trained model can ignore the transfor- mation parameters and collapse back to behaviors associated with invariance-based methods. As illustrated in Figure 5, naively training CONTEXTSSL leads to poor equivariance with respect to the transformations. Previous works [16] have also identified this concern and proposed specialized architectures that incorporate transformation parameters directly into the model, thereby outputting the predictor's weights and ensuring effective utilization of these parameters. For our setting, we introduce a rather simple approach that involves jointly training an auxiliary predictor. This predictor is designed to predict the latent transformations of the target sample y, from the concatenated input vector (I,, a:).
Experimental Results
To evaluate the efficacy of our proposed algorithm CONTEXTSSI, our experiments are designed to address the following questions:

---

i) How does CONTEXTSSL fare against competitive invariant and equivariant self sefe--seeeeised learning approaches in terms of performance across varying context sizes and different sets of data transformations?
ii) How effectively can CONTEXTSSL. identify task-specific symmetries, both within the scope of self-supervised learning and beyond?
iii What roles do specific components such as selective masking and the auxiliary latent transformation predictor play in facilitating the learning of general and context-adaptable representations?
4.1 Quantitative Assessment of Adaptation to Task-Specific Symmetries
We use the 3D Invariant Equivariant Benchmark (3DIEBench) [16] and CIFARTO to test our approach. We compare CONTEXTSSL with 1) VICReg [6] and SimCLR [9] among the invariant self-supervised approaches: 2) EquiMOD [13], SEN [35] and SIE [16] amongst the equivariant baselines. To discard the performance gains potentially arising from CONTEXTSSL 's transformer architecture, for each approach .N, we replaces the original projection head or predictor with our transformer model, denoted as N We further test this at For all our equivariant baselines on 3DIEBench, we train equivariant approaches to be equívariant to either only 3D rotation, color transformations, or both. We report the test performance invariant representations, we employ linear classification over frozen features. For the equivariant counterpart, report R2 on the task of predicting the corresponding transformation. Additionally, we use Mean Reciprocal Rank (MRR) and Hit Rate at (H@k) to evaluate the performance of our context predictor. More details about pretraining algorithms and training setup are provided in Appendix B.
context lengths 0, 2. 14, 30, and 126. To assess the quality of the
Rotation context
Color context
Method
Rotation peediction (R*)
Color prodiction (89)
Classilication (top-1)
Retation prodiction
Insuriant
0.6
SimCLR
0,506
0.148
85.3
SimCLR*((-0)
0.478
0.070
81.4
0.5
SimCLR
0.247
42.3
VICReg
0.371
0.023
76.3
VICReg ((=()
0.356
0.062
73.3
Equivarion!
Higher betler
Lower fetler
14
30
126
Rolation
SIE
EquiMOD
0.671
0,512
0
0.097
77.3
82.4
1.0
0.633
0.055
81.5
0.8
CONTEXTSSL, rot context
0.744
0.023
80.4
Lower is better
Higher better
EquiMOD
0.429
0.859
82.1
SIE
0,304
0.975
70.3
SEN
0.386
0.949
77.6
CONTEXTSSL, color context
0.344
0
80.4
30 Context Length
14
Figure 3: (Left) Quantitative evaluation of learned representations on invariant (classification) and equivariant (rotation prediction, color prediction) tasks; (Right) Performance of CONTEXTSSL on equivariant (top) rotation prediction; (bottom) color prediction tasks with varying context length for context corresponding
rotation and color. The algorithm increasingly demonstrates equivariance to rotation (color) as the rotation (color) context length increases while simultancously becoming more invariant to color (rotation).
Invariant Classification and Equivariant transformation prediction task. As shown in Figure invariant self-supervised learning methods such as SimCLR and VICReg achieve high downstream classification accuracies but underperform in equivariant augmentation prediction tasks. Among the equivariant baselines, EquiMOD persistently maintains its downstream classification accuracy but exhibits improvements in augmentation prediction tasks only when trained to be equivariant to color. SIE and SEN exhibit sensitivity to the trained transformations and remain less sensitive to the others. However, their degree of invariance or equivariance is much worse compared to CONTEXTSSI. Besides, aligning them with different targeted symmetry groups requires retraining the entire model. contrast, CONTEXTSSL. exhibits equivariance to both rotation and color in the absence of context. As seen from the two rows corresponding to CONTEXTSSL in Section 4.1,
7

---

when the context corresponds to pairs of data with transformations sampled from the rotation (color) group, the model adaptively learns be invariant color (rotation) while improving equivariance rotation (color). Appendix C.8 shows that CONTEXTSSI. learns equivariance or invariance to the same transformation based on the context.
Table 1: Quantitative evaluation of learned predictors equivariant only rotation based on Mean Reciprocal Rank (MRR) and Hit Rate H@k on the validation dataset. CONTEXTSSL. learns to be more equivariant rotation with context.
Method
MRR (?)
(t)
2
126
30
126
2
14
30
126
EqaMMD
0.16
0.05
0.222
SEN CONTEXTSSL
0.05
0.22
0,240 0.270
0.373
0.396
0.402
0,108
0.129
0.223
0.245
0.292
0.366
0,412
0.541
0.561
0.568
Equivariant Measures Based on Nearest Neighbours Retrieval, Table illustrates the performance of CONTEXTSSL MRR and H@ compared to baseline methods with trained equivariance to rotation. CONTEXTSSL outperforms the baseline models, and its performance on all the metrics consistently improves with increasing context length, showing adaptation to rotation-specific features.
4.2 Role of Context Mask and Auxiliary Predictor
Role of Context Mask. To illustrate how context masking effectively eliminates shorteuts, we conduct ablation study with varying masking probabilities, detailed in Figure 4. We observed that as masking probability increases, performance on both classification prediction tasks initially improves but later declines, reaching optimal performance masking probability of 90%.
Classifiation accuracy (9)
80
0.74
1.0
Rotation prediction
Color prediction R2
0.8
0-0
60
0.72
0.6
0-0.50
40
0.70
0.4
0.30
20
0.2
0.68
0.0
0.00 0.20 0.50 0.75 0.90 0.98 Context Probability
02
14
30
126
02
14
126
Context Length (rotation)
Context Length (rotation)
(a)
(b)
(c)
Figure 4: Role of context mask to avoid context based shortcuts in CONTEXTSSL
We demonstrate that the auxiliary predictor erucial for the model to achieve equivariance. In its absence, as depicted in Figure 5, while the model retains its performance on the invariant classification task, it fails to learn equivariance, and cannot effectively adapt to different contexts.
Role of Auxiliary Predictor.
SimCLR
VICReg
SEN
SIB
EquiMOD
ConnEOSSL
Rolation predicton
0.6
0.8
Sounce
L-N
06
0.4
04
2NNN
0.2
SmCLR
ENN
Context (rotation)
Context (colar)
(a)
(b)
Figure 5: Role of auxiliary predic- tor to avoid the trivial solution of invariance.
Figure 6: Nearest neighbors of different methods taking as input the source image and rotation angle. CONTEXTSSL aligns best with the rotation angle of the target image.

---

4.3 Qualitative Assessment of Adaptation to Task-Specific Symmetries
We conduct qualitative assessment of model performance by taking the nearest neighbors of the predictor output when imputting a source image and transformation variable, as shown in Figure 6. The nearest neighbors of invariance models (SimCLR and VICReg) have random rotation angles. Equivariance baselines (SEN, SIE, EquiMOD) correctly generate the target rotation angle for some of the 3-nearest neighbors but fail in others. CONTEXTŠSL. outperforms by successfully identifying the correct angle in all 3-nearest neighbors while remaining invariant to color variations. Additional qualitative assessments for CONTEXTSSL with varying context are provided Appendix C.3.
4.4 Expanding to Diverse Data Transformations
Unlike 3DIEBench where meta-latents for each data are available, we manually construct positives by applying augmentations like crop and blur on CIFARIO. The results for the combinations of crop and blur are reported in Table 2. Consistent with our previous results, while almost retaining the clas- sification performance as SimCLR, CONTEXTSSL learns to adaptively enforce equivariance to crop (blur) and invariance blur (crop) depending upon the context. Note that the invariance performance initially improves with increasing context length but then diminishes. This occurs due to the 90% random masking ratio during training, which necessitates out-of-distribution generalization when the context length is large. Results on additional transformation pairs are provided in Appendix C.4.
Table 2: Performance of CONTEXTSSL on invariant (classification) and equivariant (crop prediction, blur prediction) tasks in CIFAR-10 under the environment of crop, i.e. CONTEXTSSL. (crop), and blur, i.e. CONTEXTSSL (blur).
Method
Crop prediction (XP)
Blur prediction (R*)
Classification (top-1)
0
2
14
30
126
14
30
126
Represcntation
SimCLR SimC1.R (c=0) SunCLR*
0.459
0.371
89.1
0.448
0.362
0.444
59.9
CONTENTSSI. (cmp)
0.608
0.607
0.607
0.608
0.608
0.920
0.854
0.624
0.667
0.694
88.5
CONTENTSSL (blur)
0.609
0,482
0.434
0.417
0.465
0.920
0.923
0.925
0.925
0.925
88.5
4.5
Context World Models for -eturrl Symmetries
So far, we have evaluated CONTEXTSSL on synthetic transformations such as color shifts and rotations using datasets like 3DIEBench, as well as on hand-crafted augmentations like cropping and blurring in 10. However, CONTEXT SSL is much more versatile and can be applied to selectively learn sensitivity or insensitivity to naturally occurring symmetries and sensitive features. To demonstrate this, we test CONTEXTSSL on a widely used physiological dataset, MIMIC-III [30] and a standard fairness benchmark, UCI Adult [2].
MIMIC-III is a large collection of medical records, which includes crucial identifiers like medications, patient demographics, hospital length of stay (LOS), and survival data. We investigate if CON- TEXTSSL. could learn representations that are dynamically invariant or equivariant to the sensitive attribute, gender. Equivariance is necessary for gender-specific medical conditions, where different medication dosages may be required, while invariance is critical for ensuring fairness in tasks such as predicting hospital stay duration or medical costs. To construct the context, for each data point æ,, we create the transformed data y; by flipping the value of gender. As shown in Table 3, when the model equivariant to gender, both gender prediction accuracy and medical treatment prediction improve with context. On the contrary. when the environment is invariant, embedding fairness of hospital stay (LOS) prediction as measured by equalized odds (EO) and equality of opportunity (EOPP), improves with context. We observe similar results for fairness of income prediction in the UCI Adult dataset, as shown in Appendix C.5.
4.6 Context World Models Beyond Self-Supervieed Learning
While our analysis has primarily focused on self-supervised learning, the concept of context is versatile and extends beyond representation learning. In principle, irrespective of the task at hand, paying attention to context can learn and identify features defined by it. To validate this and explore broader

---

Table 3: Performance of CONTEXTSSL on invariant (length of stay (LOS) prediction) and equivariant
(treatment prediction accuracy) tasks in the MIMIC III dataset.
g
Gender prediction Acc
LOS prodicction Ace
Equalizee adds
Equality of oppoetunity
Treatment protiction Ace
Context Lengch
0
126
126
126
0
126
126
Equivariant
0.969
0.944
D.35
0. 0.344
Invariant
0.969
0.626
0.942
0.943
0.028
0.023
0.023
0.004
0.335
0.316
applications of our algorithm, we consider supervised learning task where our transformer model trained to directly predict the labels corresponding to an input image. We further corrupt the labels to be directly influenced by the augmentation group transforming the data. Specifically, for 3DIEBench dataset, we add a constant value of 10 to each label if the context corresponds to the rotation group and leave it unchanged otherwise. We report classification performance along with rotation and color prediction equivariant measures. As shown in Table 4, CONTEXTSSL: classification accuracy improves with context, demonstrating its ability to better identify the underlying symmetry group with increase in context. Additional results are provided in Appendix C.6. Further, CONTEXTSSL serves as a general framework that can adapt to different training regimes such as supervised learning.
Table 4: Performance of CONTEXTSSL on equivariant tasks (including classificaion) for context- dependent labels. CONTEXTSSI. adapts context-dependent lahels with varying context.
Method
Rotation prediction (R')
Color prediction (R))
Classification (top-1)
2
14
126
SimCLL (color)
0.537
0.056
72.0
SimL1.8 (rotation)
0.537
(1056
SimCLR (c=0) (color)
0.427
-0.007
80.4
SimCLR (c=0) (rocation)
0.427
-0.007
5.2
SimC3 R (oolor)
0.424
(1,243
15.1
14.8
14.00
SinCLR* (rotation)
0.424
0,243
56.1
58.2
58.4
58.4
59.1
CONTEXTSSL (color)
0.556
0.542
0.538
0,540
0.539
0.913
0.973
0.981
0.982
0.982
8.9
$2.4
82.7
82.8
83.0
CONTEXIŠSL, (rotation)
0.556
0.624
0.661
0.665
0.666
0.913
0.379
0,111
0.095
73.5
82.7
82.6
83.0
5 Conclusion and Future Perspectives
The field of language modeling has witnessed a significant paradigm shift over the past decade, moving towards foundation models that generalize across a variety of tasks either directly through distillation. However, this shift toward gencralization has been conspicuously absent in the vision domain. This is largely because self-supervised approaches for vision still heavily rely on inductive priors strongly introduced by enforcing either invariance or equivariance to data augmentations. This renders representations brittle in downstream tasks that do not conform to these priors and necessitates retraining the representation separately for each task. This work forgoes any notion of pre-defined symmetries and instead trains model to infer the task-relevant symmetries directly from the context through what we term Contextual Self-Supervised Learning (CONTEXTSSL) The ability of our model learn selective equivariances and invariances based on mere context opens up new avenues for effectively handling a broader range of tasks, particularly in dynamic environments where the relevance of specific features may change over time. However, to -and-cd transformations in the data and do not explore naturally occurring symmetries. Nonetheless, CONTEXTSSL lays the groundwork for models that can potentially discern and adapt to the underlying patterns of tasks, recognize shortcuts, and more effectively generalize across unseen scenarios. Through this work, we hope contribute a broader understanding of how machines can learn more like humans - contextually, adaptively, and with an eye toward the infinite variability of the real world.
limit scope of symmetries
10

---

Acknowledgement
This research was supported part by Office of Naval Research grant N00014-20-1-202 (MURI ML-SCOPE), NSF Award CCF-2112265 (TILOS AI Institute), and the Alexander von Humboldt Foundation. CW and TJ acknowledge support from NSF Expeditions grant (award 191 8839: Collabo- rative Research: Understanding the World Through Code) and Machine Learning for Pharmaceutical Discovery and Synthesis (MLPDS) consortium. We acknowledge MIT SuperCloud and Lincoln Laboratory Supercomputing Center for providing HPC resources that have contributed to this work.
References
[1] Mahmoud Assran, Quentin Duval, Ishan Mísra, Piotr Bojanowski, Pascal Vincent, Michael Rabbat, Yann LeCun, and Nicolas Ballas. Self-supervised learning from images with a joint- embedding predictive architecture. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pages 15619-15629, 2023. URL https: arxiv. org/ abs/ 2301. 08243.
[2] Arthur Asuncion, David Newman, al. Uci machine learning repository, 2007.
[3] Philip Bachman, R Devon Hjelm, and William Buchwalter. Learning representations by maximizing mutual information across views. Advances in neural information processing systems, 32, 2019.
[4] Hangbo Bao, Li Dong, Songhao Piao, and Furu Wei. Beit: Bert pre-training of image trans- formers. International Conference on Learning Representations, 2022.
[5] Amir Bar, Yossi Gandelsman, Trevor Darrell, Amir Globerson, and Alexei Efros. Visual prompting via image inpainting. Advances in Neural Information Processing Systems, 35: 25005-25017, 2022.
[6] Adrien Bardes, Jean Ponce, and Yann LeCun. Vicreg: Bariane-invarinee-ccvaaaa regular- ization for self-supervised learning. International Conference Learning Representations, 2022. URL bittps://arxie
htggaba/21000 04900.9000.0......
[7] Tom Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared D Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, et al. Language models are few-shot learners. Advances in Neural Information Processing Systems, 2020.
[8] Mathilde Caron, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. Emerging properties in self-supervised vision transformers. In Proceedings of the IEEE/CVF international conference on computer vision, pages 9650-9660, 2021. URL bttts:://rxxs org/ /204.. 14294.
[9] Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A simple framework for contrastive learning of visual representations. In International conference on machine learning, pages 1597-1607. PMLR, 2020. URL http://pproceedings. mlr press/v119/
pdf.
[10] Ting Chen, Simon Korbllith, Kevin Swersky, Mohammad Norouzi, and Geoffre Hinton. Big self-supervised models are strong semi-supervised learners. Advances neural information processing systems, 33:22243-22255, 2020.
(11) Xinlei Chen and Kaiming He. Exploring simple siamese representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 15750-15758, 2021.
[12] Rumen Dangovski, Li Jing, Charlotte Loh, Seungwook Han, Akash Srivastava, Brian Che- ung, Pulkit Agrawal, and Marin Soljačié. Equivariant contrastive learning. In International Conference on Learning Representations, 2022.
[13] Alexandre Devillers and Mathieu Lefort. Equimod: An equivariance module to improve self-supervised learning. In International Conference Learning Representations, 2023.
11

---

[14] Alaaeldin El-Nouby, Michal Klein, Shuangfei Zhai, Miguel Angel Bautista, Alexander Toshev, Vaishaal Shankar, Joshua M Susskind, and Armand Joulin. Scalable pre-training large autoregressive image models. arXiv preprint arXiv:2401.08541, 2024.
[15] Quentin Garrido, Laurent Najman, and Yann Lecun. Self-supervised learning of split invariant equivariant representations. preprint arXiv:2302. 10283, 2023.
[16] Quentin Garrido, Laurent Najman, and Yann Lecun. Self-supervised learning of split invariant equivariant representations. International Conference on Machine Learning, 2023. URL
https://arxiv. org/ http/ddf/232222 2302. 10283. pdf.
[17] Quentin Garrido, Mahmoud Assran, Nicolas Ballas, Adrien Bardes, Laurent Najman, and Yann LeCun. Learning and leveraging world models in visual representation learning. arXiv preprins arXiv:2433.00004 2024.
[18] Spyros Gidaris, Praveer Singh, and Nikos Komodakis. Unsupervised representation learning by predicting image rotations. International Conference on Learning Representations, 2018.
[19] Spyros Gidaris, Andrei Bursuc, Gilles Puy, Nikos Komodakis, Matthieu Cord, and Patrick Pérez. Obow: Online vis-of--vsuulwwoooe generation for self-supervised learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 6830-6840, 2021.
[20] Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, al. Bootstrap your own latent-a new approach to self-supervised learning. Advances in neural 2020.
information processing systems,
[21] Sharut Gupta, Stefanie Jegelka, David Lopez-Paz, and Kartik Ahuja. Context is environment. In The Twelfth International Conference on Learning Representations, 2023.
[22] Sharut Gupta, Joshua Robinson, Derek Lim, Soledad Villar, and Stefanie Jegelka. Structur- ing representation geometry with rotationally equivariant contrastive learning. International Conference on Learning Representations, 2023.
[23] David Ha and Jürgen Schmidhuber. World models. Advances in Neural Information Processing Systems, 2018
[24] Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. International Conference on Learning Representa- lions, 2020.
[25] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse domains through world models. arXiv preprint arXiv:2301 303004444444..... 2023.
[26] Kaiming He, Haoqi Fan, Yuxin Wu, Saining Xie, and Ross Girshick. Momentum contrast for unsupervised visual representation learning. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 9729 9738, 2020.
[27] Kaiming He, Xinlei Chen, Saining Xie, Yanghao Li, Piotr Dollár, and Ross Girshick. Masked autoencoders are scalable vision learners. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 16000-16009, 2022.
[28] Anthony Hu. Lloyd Russell, Hudson Yeo, Zak Murez, George Fedoseev, Alex Kendall, Jamie Shotton, and Gianluca Corrado. Gaia-1: A generative world model for autonomous driving. arXiv preprint arXiv:2309. 17080, 2023.
[29] Li Jing, Pascal Vincent, Yann LeCun, and Yuandong Tian. Understanding dimensional collapse in contrastive self-supervised learning. International Conference on Learning Representa- tions, 2022.
[30] Johnson, T Pollard, and Mark III. Mimic-iii clinical database (version 1.4). physionet. 2016, 2016.
12

---

[31] Gustav Larsson, Michael Maire, and Gregory Shakhnarovich. Learning representations for au- tomatic colorization. In Computer Asion-ECCV 2016: 14th European Conference, Amsterdam, The Netherlands, October 11-14, 2016, Proceedings, Part IV 14, pages 577-593. Springer, 2016.
[32] Hankook Lee, Kibok Lee, Kimin Lee, Honglak Lee, and Jinwoo Shin. Improving transferability of representations via augmentation-awarr sulf--upervision.. Advances in Neural Information Processing Systems, 34:17710- 2021.
[33] Xiang Lisa Li and Percy Liang. Prefix-tuning: Optimizing continuous prompts for generation. Annual Meeting of the Association for Computational Linguistics, 2021.
[34] Ishan Misra and Laurens van der Maaten. Self-supervised learning of pretext-invarians rep- resentations. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 6707-6717 17, 2020.
[35] Jung Yeon Park, Ondrej Biza, Linfeng Zhao, Jan Willem van de Meent, and Robin Walters. Learning symmetric embeddings for equivariant world models. International Conference on Machine Learning. 2022.
[36] Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, al. Language models are unsupervised multitask learners. OpenAl blog, 1(8):9, 2019.
[37] Ramanan Sekar, Olch Rybkin, Kostas Daniilidis, Picter Abbeel, Danijar Hafner, and Deepak Pathak. Planning explore via self-supervised world models. In International conference on machine learning, pages 8583-8592. PMLR, 2020.
[38] Ravid Shwartz-Ziv, Micah Goldblum, Hossein Souri, Sanyam Kapoor, Chen Zhu Yann LeCun, and Andrew G Wilson. Pre-train your loss: Easy bayesian transfer learning with informative priors. Advances in Neural Information Processing Systems, 35:27706-27765 2022.
[39] Nenad Tomasev, Ioana Bica, Brian McWilliams, Lars Buesing, Razvan Pascanu, Charles Blundell, and Jovana Mitrovic. Pushing the limits of self resnets: Can we outperform supervised learning without labels on imagenet? arXiv preprini arXiv:2201.05112 2022.
[40] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez. Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in neural information processing systems, 30, 2017.
[41] Xinlong Wang, Wen Wang, Yue Cao, Chunhua Shen, and Tiejun Huang. Images speak in images: A generalist painter for in-context visual learning. In Proceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pages 6830-6839, 2023.
[42] Tete Xiao, Xiaolong Wang, Alexei A Efros, and Trevor Darrell. What should not be contrastive in contrastive learning. In International Conference on Learning Representations, 2021.
[43] Yuyang Xie, Jianhong Wen, Kin Wai Lau, Yasar Abbas Ur Rehman, and Jiajun Shen. What should be equivariant in self-supervised learning. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 41 11 4-200 2022.
[44] Zhenda Xie, Zheng Zhang, Yue Cao, Yutong Lin, Jianmin Bao, Zhuliang Yao, Qi Dai, and Han Hu. Simmim: A simple framework for masked image modeling. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pages 9653-9663, 2022.
[45] Mengjiao Yang, Yilun Du, Kamyar Ghasemipour, Jonathan Tompson, Dale Schuurmans, and Pieter Abbeel. Learning interactive real-world simulators. International Conference on Learning Representations, 2024.
[46] Jure Zbontar, Li Jing, Ishan Misra, Yann LeCun, and Stéphane Deny. Barlow twins: Self- supervised learning via redundancy reduction. In International Conference on Machine Learn- ing, pages 12310-12333 123200 PMLR, 2021.
[47] Jinghao Zhou, Chen Wei, Huiyu Wang, Wei Shen, Cihang Xie, Alan Yuille, and Tao Kong. ibot: Image bert pre-training with online tokenizer. International Conference on Learning Representations, 2022. URL nattps://ett //tpxv/axvvvv//////////ee/aaaa org/oos///////ooooooo/ooooee ab//a//s//11..// 0121110078822.
13

---

[48] Pan Zhou, Yichen Zhou, Chenyang Si, Weihao Yu, Teck Khim Ng, and Shuicheng Yan. Mugs: multi-granulaa self-supervised learning framework. arXiv preprint arXiv:2203 14415, 2022.
14

---

Appendix
Related Work
15
B
Supplementary experimental details and assets disclosure
16
B.1
Assets
16
B.2 Hardware and setup
16
B.3 Datasets
16
B.4
Baseline Algorithms
16
B.5
Training Protocol
17
B.6
Evaluation metrics
17
Additional Experiments
18
C.1
Quantitative Assessment of Adaptation to Task-Specific Symmetries
18
C.1.1 Invariant Classification and Equivariant transformation prediction task
18
C.1.2 Equivariant Measures Based on Nearest Neighbours Retrieval
19
C.2 Role of Context Mask and Auxiliary Predictor
19
C.2.1 Role of Context Mask
19
C.2.2 Role of Auxiliary Predictor
19
C.3 Qualitative Assessment of Adaptation Task-Specific Symmetries
20
C.3.1 Comparison with Baseline Approaches
20
C.3.2 Nearest Neighbour Retrieval with Varying Context
20
C.4
Expanding to Diverse Data Transformations
21
C.5 Context World Models For Naturally Occurring Symmetries
22
C.6 Context World Models Beyond Self-Supervised Learning
23
C.7 Performance Encoder Representations and Predictor Embedding
24
C.8 Enforcing Invariance or Equivariance to the Same Transformation Using Context
24
A Related Work
Self-Supervised Learning. Existing SSL methods generally belong to two categories: invariant learning [9, 6, 11, 26, 46, 20] and equivariant learning. The representative method for invariant learning is contrastive learning, which draws the representations of positive samples together in the latent space such that the representations are invariant data augmentation. Contrastive learning can learn highly discriminative features at the cost of losing certain image information due to the invariance constraint [42]. Motivated by this limitation, recent works explore merging contrastive learning with equivariant learning tasks by separate embedding [42, 16), nugmentaiion-condiioneee predictor [13, 17], and explicit equivariant transformation [22]. However, existing works still inherit the limitations of contrastive learning: its symmetry prior is built
given set of manual augmentations and is not adaptive to downstream tasks. In contrast, our method enables the contextual world model to adapt its symmetry the contextual data, which is more flexible and generalizable to various tasks.
World Models. World modeling has achieved notable success in reinforcement learning (RL) for -baeed planning [23, 37, 24] and vision [25, 28, 45], where involves predicting future states based on current observations and actions. This concept, however, has not yet been fully leveraged in visual representation learning. Nevertheless, Garrido et [17] shows that several families of self- supervised learning approaches can be reformulated through the lens of world modeling. Equivariant
15

---

self-supervised learning methods. Specifically, Masked Image Modeling approaches [27, 4, 14, 44] consider masked pixels and target pixel reconstruction their action and next state. Other equivariant learning approaches [13, 35, 16] consider data transformations and representation of the target image as their action and next state pair. However, unlike true world modeling, these approaches do not track past experiences, component critical for generalization. Our method instead leverages context to track past experiences in terms of state, action, and next-state triplets, enabling it to adapt and generalize to varying environments,
In-context Learning. Our work is inspired by and extends the concept of in-context learning (ICL) [7] to training. Initially studied in the context of language, in-context learning has recently been adapted for vision tasks [21, 41, 33], allowing models to infer environmental features or tasks directly from input prompts without predefined notions. For example, Visual Prompting [41, 5] uses task input/output example pair and a query image at test time, and uses inpainting to generate the desired output. Gupta et al. [21] propose using unlabeled data as context at training to extract environment-speer specific signals and address domain generalization. ICL has been extensively explored in various domains, including vision, language, and multimodal tasks. However, our work is the first to apply ICL to vision self-supervised representation learning.
B Supplementary experimental details and assets disclosure
B.1
Assets
We do not introduce new data in the course of this work. Instead, we use publicly available widely used image datasets for the purposes of benchmarking and comparison.
B.2 Hardware and setup
Each experiment was conducted on NVIDIA Tesla V100 GPUs, each with 32GB of accelerator RAM. The CPUs used were Intel Xeon E5-2698 v4 processors with 20 cores and 384GB RAM. All experiments were implemented using the PyTorch deep learning framework.
B.3 Datasets
3D Invariant Equivariant Benchmark (3DIEBench). To test equivariance and invariance to multiple data transformations, we use the 3D Invariant Equivariant Benchmark (3DIEBench) [16] which has been specifically designed to address the limitations of existing datasets in evaluating invariant and equivariant representations. It contains images of 3D objects along with their latent parameters such object rotation, lighting color, and floor color. Since we have access available to individual meta latent parameters, transformation parameters between two views of an object are calculated as the difference between their individual latents. We test our approach under two settings 1) Considering two transformation groups: rotation and color with the aim of learning invariance to one and equivariance to another after conditioning on context; Considering one transformation group, say rotation and learning to enforce invariance or equivariance to rotation with context. As previously mentioned, all methods are trained for 1000 epochs using batch size of 512 on 128 x 128 resolution images. We use the standard training, validation and test splits, made publicly available by the authors [16].
3DEEBench
CIFAR10. 3DIEBench dataset is limited to only rotations and color as transformation groups. We extend our approach to include more common self-supervised benchmarks, such as CIFAR- 10, incorporating transformations like blurring, color jitter, and cropping. Unlike 3DIEBench, manually construct positive pairs by applying compositions of these handcrafted augmentations. We consider three transformation groups: crop, blur and color. Similar 3DIEBench, consider combinations of two groups for each training run. We use the standard training, validation and test splits.
B.4 Baseline Algorithms
Among the invariant self-supervised approached, we compare our approach to VICReg [6] and and SimCLR [9]. For each method, comparisons are drawn using their originally proposed architectures.
16

---

For the equivariant baselines, we consider EquiMOD [13], SIE [16] and SEN [35]. Similar to Garrido et al. [16]. For SEN, we use the InfoNCE loss instead the original triplet loss. To discard the performance gains potentially arising from CONTEXTSSL'S transformer architecture, for each approach, we consider an additional baseline that replaces the original projection heads or predictor with our transformer model. Given an algorithm name N, we refer to this baseline as V Amongst these, we report the best performing variant in our results. For N we conduct analysis in two distinct settings: 1) 'no context' or c = 0 invariant condition, and 2) a fully contextualized setting with a context length of 126.
B.5 Training Protocol
To ensure a fair comparison across different algorithms for each dataset, we use a standardized neural network backbone. Precisely, for our encoder, we use ResNet- backbone pre-trained on ImageNet. For CONTEXTSSL, output features from the encoder are transformed into the context sequence, which is then processed by the decoder-only Transformer [40] from the GPT-2 Transformer family [36]. Our model configuration includes 3 layers, 4 attention heads, and a 2048-dimensional embedding space, consistently applied across all datasets. Linear layers are utilized to convert the input sequence into the transformer's ' latent embedding of dimension 2048 and to map the predicted output vectors to the output space of dimension 512.
We fix the maximum training context length to 128. Since for every y, the corresponding token (ara, a.) is masked out, context length I. corresponds effective context length L- 2. Thus, we report CONTEXTSSL'S performance over varying test context length of 0, 2, 14, 30 and 126. On all datasets, we train CONTEXTSSL with the Adam optimizer with a learning rate of 5e and weight decay 1e-3. For baseline self-supervised approaches, in their original architecture, we use learning rate of le- with no weight decay. However, when tested using the transformer architecture, we choose one of the above two optimizer hyperpameters. Consequently, performance of the best performing model is reported among the two baselines. Similar to Garrido et al. [16], we report hyper-parameteer and architectures specific to each method:
SimCLR [9] We train using 2048-2048-222 dimensional multi-layered perceptron (MLP) based projection head with a temperature of 0.5.
VICReg [6] We train using 2048-2048 204-2204-2048 MLP for the projection head and use weight of 10 for both the invariance loss and variance loss and for covariance loss.
SEN [35] Similar to other approaches
a projection head of dimension 2048-2048-
2048 and temperature 0.1.
Equi [13] We use the standatd projection head of dimensions 1024-1024- 128 and use
equal weighing of the invariance and the equivariance loss.
SIE [16] We use two 1024-10241000 1024 projection heads, one for invariant latent space and other for equivariant. When trained to learn equivariance to only rotation or only color, we use weight of 10 for both the invariance loss and variance loss, I for the covariance loss and 4.5 for the equivariant loss. However, when trained to be equivariant to both rotation and color jointly, we use 10 as the equivariant weight.
B.6 Evaluation metrics
In line with established self-supervised learning methodologies, we begin by assessing the quality of the learned representations through downstream tasks. For evaluating invariant representations, we employ linear classification over frozen features. To evaluate equivariant representations, we predict the corresponding data transformation. This prediction takes representations from two differently transformed views of the same object and regresses on the applied transformation between them. Further, we use Mean Reciprocal Rank (MRR) and Hit Rate at & (H@k) evaluate the performance for our context predictor. Given the source data and the transformation action, we identify the k nearest neighbors in the embedding space. MRR calculated as the average reciprocal rank of the target embedding within these nearest neighbors. Hit rate-k (H@k) assigns score of 1 if the target embedding within the k-nearest neighbors of the predicted embedding and 0 otherwise. Similar to Garrido et al. [16], we restrict the search for nearest neighbors to different views of the same object, thus ensuring that the predictor is not penalized for retrieving an incorrect object in a pose similar to the correct one.
17

---

C Additional Experiments
C.1 Quantitative Assessment of Adaptation to Task-Specifie Symmetries
In this section, present additional results the quantitative assessment of model performance 3DIEBench, including the evaluation of learned representations on equivariant tasks (rotation and color prediction) to predict individual latent values. In contrast, the results in Section 4.1 focus predicting relative latent values between pairs image embeddings as inputs.
Table 5: Quantitative evaluation of learned representations on equivariant (rotation prediction, color prediction) tasks to predict individual latent values.
G
Method
Rotation prediction (/t²)
Color prediction (A?2
2
14
30
126
0
14
30
126
Invoriant
SimCLR
0.791
SimCLR (c=0)
0.773
SimCLR
0.544
0.498
VICReg
0.660
VICReg (c=0)
0.615
0.061
Rotation
Equivariant
Color
EquiMOD
0.712
0.221
SIE
0.760
0.972
SEN
0.617
0.888
Rotation
EquiMOD
0.707
0.033
SIE
0.790
0.001
SEN
0.723
0).437
CONTEXTSSL
0.838 0.839
0.840
0.840
0.840
0.895
0.620
0.021
0.014
0.021
EquiMOD
0.660
0.855
Color
SIE
0.560
0.974
SEN
0.713
0.876
CONTENTSSL
0.838 0.800
0.699
0.666
0.685
0.895
0.981
0.985
0.985
0.986
C.1.1 Invariant Classification and Equivariant transformation prediction task
As shown in Table 5, invariant self-supervised learning methods such as SimCLR and VICReg underperform in equivariant augmentation prediction tasks. The equivariant baselines, EquiMOD, SIE, and SEN, exhibit improvements compared to the invariant baselines in some of the augmentation prediction tasks. However, their degree of equivariance much worse compared CONTEXTSSL. Besides, aligning them with different targeted symmetry groups requires retraining the entire model. In contrast, CONTEXTSSL employs single model capable of learning equivariance to rotation and invariance to color (or vice versa) based on the given context. As seen from the two rows corresponding to CONTEXTSSL in Section 4.1, when the context corresponds to pairs of data with transformations sampled from the rotation (color) group, the model adaptively learns to be invariant color (rotation) while retaining equivariance rotation (color).
Results in Section 4.1 are the average value over three random seeds. We provide the standard deviation for rotation and color prediction of CONTEXTSSL in Table 6 and Table 7.
Table Performance of CONTEXTSSL 3DIEBench in rotation prediction under the environment of rotation, i.e. CONTEXTSSL (rotation), and color, i.e. CONTEXTSSL (color), with standard deviations over three random seeds.
Method
Rotation prediction (R²)
30
126
CONTEXTSSI SSL (rotation)
0.734 + 0.002
0.740 +: 0.004
0.743 1: 0.001
0.743 4: 0.001
0.744 1 0.001
CONTEXTSSL (color)
0.735 = 0.001
0.614 4 0.108
0.389 + 0.054
0.345 + 0.040
0.344 1 0.003
18

---

Table 7: Performance of CONTEXTSSL in 3DIEBench in color prediction under the environment of rotation, L.e. CONTEXTSSL (rotation), and color, i.e. CONTEXTSSL (color), with standard deviations three random seeds.
Method
Color prediction (R²)
30
126
CONTEXTSSL (rotation)
0.908 = 0.002
0.664 0.166
0.037 + 0.010
0.023 + 0.001
0.046 + 0.007
CONTEXTSSL (colur) 0.908 = 0.002
0.981 0.002
0.985 + 0.001
0,986 + 0.001
0.986 + 0.001
C.1.2 Equivariant Measures Based on Nearest Neighbours Retrieval
Similar to Table 1, we provide the performance of CONTEXTSSL on MRR and H@k compared to baseline methods with trained equivariance to rotation. While Table I uses the validation set data as the retrieval library, Table provides the results using the training set data. CONTEXTSSL outperforms the bascline models, and its performance all the metrics consistently improves with increasing context length, showing adaptation to rotation -spccinn features.
Table 8: Quantitative evaluation of learned predictors equivariant only rotation based on Mean Reciprocal Rank (MRR) and Hit Rate H@k on training dataset. CONTEXTSSL learns to be more cquivariant rotation with context.
Method
MRR (?)
126
30
126
14
126
EquiMOD
0.17
0.06
0.24
SEN CONTEXTSSL
0.24
0.650
0.669
C.2 Role of Context Mask and Auxiliary Predictor
In this section, we provide additional results for the role of context mask and auxiliary predictor.
C.2.1
Role of Context Mask
Classification accuracy 186
80
Rotelion prediction R2
0.98
p-0.00
60
0.6
Color prediction
p=000
0.96
p-0.50
40
0.5
00700
0.94
p 0-0.0
20
0.4
p=0.88
0.3
0.50 0.75
14
30
128
02
14
128
Context Mask Probabiity
Context Length (color)
Context Length (color)
(b)
(c)
Figure 7: Role of context mask to avoid context based shortcuts in CONTEXTSSL under color context
In addition to Figure 4, we provide the performance of the rotation and color prediction tasks with varying masking probabilities under the environment of color in Figure 7. We observed that as masking probability increases, performance on both classification and prediction tasks initially improves but later declines, reaching optimal performance a masking probability of 90%.
Results in Figure 4 and Figure are the average value over three random seeds. We provide the standard deviation for rotation and color prediction of CONTEXTSSL in Table 9 and Table 10.
C.2.2 Role Auxiliary Predictor
We provide the complete results corresponding to Figure 5 in Table to demonstrate that the auxiliary predictor is crucial for the model to achieve equivariance. In its absence, while the model retains performance on the invariant classification task, it fails learn equivariance, performs similarly to the invariant models, and cannot effectively adapt to different contexts.
19

---

Table 9: Performance of CONTEXTSSL rotation prediction tasks in 3DIEBench under different random masking probabilities, with standard deviations over three random seeds.
Context
Probahility
Rotation prediction (R*)
2
14
30
126
0.00
0.677 0.004
0.677 = 0.002
0.673
0.009
0.682
0.003
0.683 +: 0.003
0.20
0.710 + 0.002
0.721 = 0.006
0.727 = 0.002
0.729
0.001
0.729 = 0.001
0.50
0.725 +
0.738 = 0.005
0.743 = 0.001
0.743 +: 0.001
0.744 + 0.001
Rotation
0.75
0.734 0.002
0.738 0.006
0.742 0.004
0.741
0.004
0.741 + 0.002
0.90
0.734 0.002
0.740
0.743
0.001
0.743
0.000
0.744 0.001
0.98
0.726 0.002
0.725
0.003
0.002 0.726
0.726 0.000
0.726 1 0.003
0.00
0.677 0.004
0.676
0.005
0.620
0.019
0.569
0.019
0.655 = 0.010
0.20
0.710 + 0.002
0.689 = 0.013
0.427 0.031
0.336 0.007
0,282 = 0.022
0.50
0.725 + 0.001
0.683 + 0.006
0.390 + 0.031
0.282 + 0.013
0.287 + 0.002
Color
0.75
0.734 0.002
0.718 + 0.002
0.499 +: 0.035
0.378 0.054
0.472 + 0.015
0.735
0.614 0.8
0.389 0.054
0.345 0.040
0.344 0.003
0.98
0.726 0.002
0.508 0.127
0.529 0.141
0.571 0.125
0.665 + 0.023
Table 10: Performance of CONTEXTSSL color prediction tasks in 3DIEBench under different random masking probabilities, with standard deviations over three random seeds.
Context
Probahility
Color prediction (R*)
2
14
126
0.00
0,981
0.002
0.940 = 0.033
0.613
0.123
0.406
0.125
0.807 + 0.080
0.20
0.975 1 0.001
0.866 1 0.171
0.465
0.113
0.194
0.057
0.124 1: 0.027
0.50
0,971 + 0.002
0.904 = 0.086
0.699
1 0.028
0.205
0.054
0.091 1 0.016
Rotation
0.75
0,980 + 0.001
0.727 = 0.351
0.358 = 0.233
0.162 0.021
0.076 - 0.009
0.90
0.908 0.002
0.664 = 0.166
0.037
0.010
0.023 - 0,001
0.046 = 0.007
0.98
0,982 0.001
0.309
0.139
0.303 0.118
0.253 + 0.033
0.00
0,981 0.002
0.989
0.001
0.989
0.001
0.20
0,975 0.001
0.984
0.002
0.987
0.001
0.987 0.001
0,987 1 0.001
Color
0.50
0,971 + 0.002
0.982 = 0.002
0.986 - 0.002
0.987
- 0.002
0,988 : 0.001
0.75
0,980 + 0.001
0.983 = 0.001
0.987 = 0.001
0.987 0.001
0,983 - 0.001
0.90
0.908 0.002
0.981 = 0.002
0.985
0.001
0.986 + 0.001
0.986 + 0.001
0.98
0.982 0.001
0.982 0.001
0.981 0.001
0.981 + 0.001
0,981 + 0.001
Table 11: Performance of CONTEXTSSL on classification, rotation and color prediction tasks in 3DIEBennh with and without the auxiliary predictor.
Method
Rotation prediction (J2*)
Color prediction (II*)
Classification (top-))
126
14
30
126
Representation
SimCLR SinCLR* (c=0)
0.227
-0.004
85.3
-0.004
SimCLR*
0.245
0.028
42.3
CONTEXTSSL (wfo) (rotation)
0.227
0.227
0.226
0.226
0.227
-0.003
-0.003
-0.003
-0,004
-0.004
80.8
CONTEXTSSL (color)
0.227
0.227
.0.003
-0.004
CONTEXTŠSL (rotation)
0.734
0.740
0.743
0.743
0.744
0.908
0.664
0.037
0.023
0.046
80.4
CONTEXTSSI, (color)
0.735
0.614
0.389
0.345
0.344
0.908
0.981
0.985
0.986
0.986
80.4
C,3 Qualitative Assessment of Adaptation to Task-Specific Symmetries
C.3.1
Comparison with Baseline Approaches
We provide additional results to the qualitative assessment comparing with different models in Figure 8. The nearest neighbors of invariance models (SimCLR and VICReg) have random rotation angles. Equivariance baselines (SEN, SIE, EquiMOD) correctly generate the target rotation angle for some of the 3-nearest neighbors but fail in others. CONTEXTSSL outperforms by successfully identifying the correct angle in all 3-nearest neighbors while remaining invariant to color variations.
C.3.2 Nearest Neighbour Retrieval with Varying Context
In this section, we conduct a qualitative assessment of model performance by taking the nearest neighbors the predictor output when inputting a source image and transformation variable, and show the change in retrieving quality in Figure 9, Figure 10, and Figure 11. We observe that the
20

---

SimCLR
ViCReg
SEN
SIE
EquiMOD
CONTEXISSL
Source
I-NN
2-NN
Target
3-NN
Figure 8: Nearest neighbors of different methods taking as input the source image and rotation angle. CONTEXTSSL aligns best with the rotation angle of the target image.
nearest neighbors have closer rotation angle (color) the target image under rotation (color) context as context length increases, indicating CONTEXTSSL 's ability to adapt to the given context as context length increases.
I-NN
2-NN
3-NN
4-NN
2-NN
3-NN
4-NN
Source
Target
126
Increasing
context length
Context (Rotation)
Context (Color)
Figure 9: Nearest neighbors of CONTEXTSSL. taking as input the source image and rotation angle at different context lengths. As context increases, CONTEXTSSL. aligns better with the rotation angle (color) of the target image when the context is based on rotation (color).
C.4 Expanding to Diverse Data Transformations
Unlike 3DIEBench where meta-latents for each data available, manually construct positives by applying augmentations like crop and blur on CIFAR10. The results for the combinations of crop and blur reported in Table 2. We additionally provide the results for the combinations of crop and color in Table 13 and crop and blur in Table 2. Consistent with our previous results, while almost retaining the classification performance as SimCLR, CONTEXTSSL learns to adaptively enforce cquivariance and invariance to different environments depending upon the context.
In addition to the results for predicting relative latent values between pairs of image embeddings as input in Table 2, Table 13, and Table 12, we provide the evaluation of learned representations equivariant tasks (rotation and color prediction) to predict individual latent values, shown Table 14, Table 16, and Table 15 respectively. Both results lead to the same conclusion, that CONTEXTSSL. is able adaptively enforce equivariance and invariance different environments depending upon the context.
21

---

I-NN
2-NN
3-NN
4-NN
I-NN
2-NN
3-NN
4-NN
Source
Target
32
126
Increasing
context length
Context (Rotation)
Context (Color)
Figure 10: Nearest neighbors of CONTEXTSSL taking as input the source image and rotation angle at different context lengths. As context increases, CONTEXTSSL aligns better with the rotation angle (color) of the target image when the context is based rotation (color).
I-NN
Source
Target
126
Increasing
context length
(Rotation)
Context (Color)
Figure 11: Nearest neighbors of CONTEXTSSL taking as input the source image and rotation angle different context lengths. As context increases, CONTEXTSSL aligns better with the rotation angle (color) of the target image when the context is based on rotation (color).
Table CIFAR-10 Color-Blur, Performance CONTEXTSSL on invariant (classification) and equivariant (color prediction, blur prediction) tasks in CIFAR-10 10 under the environment of color, i.e. CONTEXTSSL (color), and blur, i.e. CONTEXTSSL (blur).
Method
Color prediction (R²)
Blur prediction (R²)
Classification (top-1)
126
126
Representation
SimCLR SimCLR* (c=0)
0.154
0.371
89.1
0.054
0.361
88.9
SimCLR+
0.318
0.444
59.9
CONTEXISSL (color)
0.518
0.519
0.519
0.519
0.519
0916
0.793
0.699
0.735
0.823
88.9
CONTENTSSL. (blur)
0.518
0.353
0.241
0.259
0.333
0.916
0.916
0.916
0.916
0.917
88.8
C.5
Context World Models For Naturally Occurring Symmetries
We demonstrate how CONTEXTSSL can dynamically learn equivariance or invariance to naturally occurring features such gender in the UCI Adult dataset [2]. For each data point J,, we create the transformed data y; by flipping the value of gender. As showm Table 17, when the environment equivariant gender, both gender prediction and income prediction improve with context. When the
22

---

Table 13: CIFAR-10 Crop-Color. Performance of CONTEXTSSL invariant (classification) and equivariant (crop prediction, color prediction) tasks in CIFAR-10 under the environment of crop, i.e. CONTEXTSSL (crop), and color, i.e. CONTEXTSSL (color).
Method
Crop prodiction (R²)
Color prediction (R*)
Classification (top-
14
30
126
2
14
30
126
Representation
SimCLR SimCLR* (c=0)
0.459
0.154
89.1
0.448
0.054
88.9
SimCLR+
0.318
CONTEXISSL (crop)
0.606
0.606
0.607
0.607
0.607
0.522
0.378
0.253
0,264
0.301
87.5
CONTEXTSSL. (color)
0.525
0.527
0.527
0.527
Table 14: CIFAR-10 Crop-Blur. Performance of CONTEXTSSL. on equivariant (crop prediction, blur prediction) tasks in CIFAR-10 under the environment of crop, i.e. CONTEXTSSL (crop), and blur, ie. CONTEXTSSL (blur), to predict individual latent values.
Method
Crop prediction (R²)
Blur prediction (R²)
2
14
30
126
SimCLR SimCLR F (c=0)
0.382
0.122
0.375
0.111
SimCLR+
0,202
0.322
CONTEXTSSL (crop) 0.576
0.575
0,576
0.576
0.576
0.835
0.795
0.630
0.644
0.663
CONTEXTSSL (blur) 0.575
0.504
0.463
0.443
0.474
0.835
0.835
0.836
0.837
0.837
Table 15: CIFAR-10 Color-Blur. Performance of CONTEXTSSL on equivariant (color prediction, blur prediction) tasks in CIFAR-IO under the environment of color, i.c. CONTEXISSL (color), and blur, i.c CONTEXTSSL (blur), to predict individual latent values.
Method
Color prediction (Rt²)
Blur prediction (Rt2)
2
14
30
126
2
14
30
126
SimCLR SimCLR (c=0)
0.121
0.122
0,039
0.111
S
0.242
0.322
CONTEXTŠSL (color) 0.488
0.488
0.488
0.488
0.488
0.837
0.711
0.628
0.672
0.730
CONTEXTSSL (blur) 0.488
0.376
0,286
0.309
0.362
0.837
0.838
0.838
0.838
0.837
Table 16: CIFAR-10 Crop-Blur. Performance of CONTEXTSSL on equivariant (crop prediction. color prediction) tasks in CIFAR-10 10 under the environment of crop, i.e. CONTEXTSSL (crop), and color, i.e. CONTEXTSSL. (color), predict individual latent values.
Method
Crop prediction (!t²)
Color prediction (22²)
2
14
30
126
2
14
30
126
SimCLR SimCLR (c=0)
0.382
0.121
0.375
0.039
SimCLR
0,202
0.242
CONTEXTSSI. (crop)
0.570
0.572
0.572
0.572
0.572
0.495
0.417
0.342
0.356
0.373
CONTEXISSL (color) 0.570
0.490
0.447
0.492
0.515
0.495
0.496
0.497
0.497
0.497
environment is invariant, embedding fairness of income prediction measured by equalized odds (EO) and equality of opportunity (EOPP), improves with context.
C.6 Context World Models Beyond Self-Supervised Learning
We report classification performance along with rotation and color prediction equivariant measures. The results for predicting relative values are shown in Table 4 and the results for predicting individual latent values are shown in Table 18. The equivariance (invariance) performance CONTEXTSSI. improves with increased context.
23

---

Table 17: Results UCI Adult dataset. Performance of CONTEXTSSL invariant (income pre-
diction accuracy as measured in terms of equalized odds and equality of opportunity) and equivariant (gender prediction accuracy) tasks on the UCI Adult dataset.
G
Gender prediction Ace 1
Income prediction AUC
Equalized odds
Equality opportunity
Context Length
126
126
0
126
0
126
Equivariant
0.985
0.999
0.900
0.900
0.114
0.130
0.061
0.101
Invariant 0.985
0.605
0,900
0.899
0.114
0.066
0.061
0.047
Table 18: Context-Dependent Labels Classification Task. Performance of CONTEXTSSL on equivariant (rotation prediction, color prediction) tasks for context-dependent labels to predict individual latent values. As context length increases, CONTEXTSSI. becomes more equivariant to color (or rotation) and more invariant rotation (or color) within the respective environment.
Method
Rotation prediction (R²²)
Color prediction (R²)
14
126
14
126
SimCLR SimCLR (c=0)
0,781
0.058
0.478
-0.003
SimCLR+
0.695
0.267
CONTEXISSL (color)
0.751
0.751
0.750
0.750
0.749
0.915
0.973
0.980
0.981
0.981
CONTEXTSSL (rolation)
0.750
0.778
0.797
0.795
0.795
0.915
0.375
0.104
0.091
0.090
C.7 Performance on Encoder Representations and Predictor Embedding
We analyze the difference between the performance representation and the performance on predictor embedding for both the invariance (classification) task and equivariance (rotation prediction) task in Table 19 and Table 20. CONTEXTSSL maintains almost the same performance for rotation prediction using either representations or embeddings, while the performance of all other baselines drops significantly when using the embeddings. Similar conclusions apply to the classification case, except for SimCLR+, for which the classification accuracy for both representations and embeddings low.
Table 19: Model performance in rotation prediction task, within the rotation equivariant environment. The R2 values are calculated for both the representations and the embeddings (output of projection head for invariant models (VICReg, SimCLR) or predictor for equivariant models (SEN, EquiMod, SIE, CONTEXTSSL). Unlike other models, which experience significant performance drop between representations and embeddings, CONTEXTSSL maintains consistent performance.
Method
Rotation prediction (R*)
Representations
Embeddings
Change
VICReg
0.37
0.23
-0.14
SimCLR
0.51
0.23
-0.28
SEN
0.63
0.39
-0.24
EquiMod
0.51
0.39
-0.12
SIE
0.67
0.60
-0.07
CONTEXTSSL (rotation)
0.74
0.74
-0.00
C.8 Enforcing Invariance or Equivariance to the Same Transformation Using Context
Apart from adaptively learning equivariance to subset of transformation groups and invariance the rest as shown in Section 4.1, we extend CONTEXTSSL to operate within environments characterized by single transformation. Motivated by this, we ask the question: Can CONTEXTSSL adapt to learn equivariance or invariance to the same transformation depending on the context?. At training, we randomly sample one of these environments. the enviromment corresponds to enforcing equiv- ariance, we construct our context in the same way before i.e. pairs of positives transformed using augmentations sampled from the transformation group. However, the environment corresponds to enforcing invariance, maximize alignment between positives transformed by augmentation
24

---

Table 20: Performance of CONTEXTSSL accuracy of predictor embeddings for context-dependent labels.
Method
Classification (top-1)
2
14
30
126
Representation
Change
SimCLR
52.7
85.3
-32.6
SimCLR (c=0)
72.4
83.4
-11.0
SimCLR
41.8
42.3
-0.5
CONTEXTSSL (rotation)
76.6
76.9
75.6
76.9
77.5
80.4
-2.9
CONTEXTSSL (color)
75.3
71.7
72.6
76.5
80.4
-3.9
sampled from the transformation group without conditioning on that augmentation. Take rotation example. As shown in Table 21, similar to results in two transformation setting (rotation and color) in Section 4.1, CONTEXTSSL effectively adapts to enforce invariance and equivarance rotation depending on the context. Results for predicting individual latents provided in Table 22.
in 3DIEBench
Table 21: Single Transformation Setting. Performance of CONTEXTSSL in 3DIEBench under the equivariant environment, i.e. CONTEXTSSL (rotation), and the invariant environment, i.c. CONTEXISSL (none), with respect to rotation.
Method
Rotation prediction (Jl*)
Classificatio (top-1)
14
30
126
Representation
SinCLR SimCLR (cw0)
0.506
85.3
0.478
83.4
SinCLR*
0.247
42.3
CONTEXISSL (rotation)
0.736
0.737
80.6
CONTEXTSSL (none)
0.737
0.717
0.477
0.377
0.473
80.6
Table 22: Single Transformation Setting. Performance of CONTEXTSSL in 3DIEBench under the equivariant environment, i.e. CONTEXTSSL (rotation), and the invariant environment, i.e. CONTEXTSSL (none), with respect to rotation, predict the individual latent values.
Method
Rotation prediction (R²")
14
30
126
SimCLR SimCLR (c=0)
0.791
0.773
SimCLR
0.544
CONTEXISSL. (rotation)
0.778
0.777
0.767
0.768
0.777
CONTEXTSSL (none) 0.839
0.829
0.721
0.667
0.698
25

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification:
We support all our claims through rigorous experiments on multiple benchmark
datasets.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
ii) Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Limitations our work are discussed in Section 5.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g.. independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
iii) Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
26

---

Justification:
Our work does not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proof' can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
iv) Experimental Result Reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification: All experimental details to reproduce the main results are provided in Ap- pendix B.
Code is available at http:///ggggaaccccce///ttttttoot Syymeeerreetetttttterrrrree.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution a specific model and empirical evaluation, it may be necessary either make it possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that appropriate to the research performed.
While NeurlPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how reproduce that algorithm. (b) If the contribution is primarily new model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility. In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should possible for other researchers to have some path to reproducing or verifying the results.
v) Open access to data and code
27

---

The authors should provide scripts to reproduce all experimental results for the new reproducible, they Does the paper provide open access to the data and code, with sufficient instruc- tions faithfully reproduce the main experimental results, as described in supplemental material?
Question:
Answer: [No]
Justification: We are working towards organizing the code base and will make it available by the rebuttal.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.. cc/
Serriidee/CodeSuumiissee
for more details.
While we encourage the release of code and data, we understand that this might be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
/ httpssscccppbccccceeeeceeecc guides/ CoodeSubmis ssi aesionPolicyy for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
proposed method and baselines. only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code is permitted.
vi) Experimental Setting/Details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We mention all experimental details in Appendix B.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can provided either with the code, in appendix, or as supplemental material.
vii Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: All experimental runs are averaged over three independent seeds. Further, Appendix C.2 and Appendix C.I report the mean and standard deviation across these seeds.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
28

---

The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation or the standard error of the mean.
It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out of range (c.g. negative emror rates).
If error bars reported in tables plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
viii) Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: Compute resources and requirements are mentioned in Appendix B.2.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that NeurIPS Code of Ethics https: : /ttp///eerrpp.eeeee//ee/eeeeeeet cc httpbliic//tiiccuuddelleeee...ee
didn't 't make it into the paper).
ix) Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the
Answer: [Yes]
Justification:
We obey all aspects of the Code of Ethics.
Guidelines:
The answer NA means that the authors have reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations in their jurisdiction).
x) Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This is a method oriented paper and not related to societal impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
29

---

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
xi) Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification:
This paper does not involve such models.
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
xii) Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We cite the authors for their models and dataset.
Guidelines:
The answer NA means that the paper does not existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
30

---

If this information is not available online, the authors are encouraged to reach out to the asset's 's creators.
xiii) New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: This paper does not release new assets.
Guidelines:
The answer NA means that the paper does not release new assets,
Researchers should communicate the details of the dataset/code/moodl as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
xiv)
Crowdsourcing and Research with Human Subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification:
This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
xv) Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: This paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
31


Reviews:
[{"review_id": "KEMcBL81d4", "reviewer": "Reviewer_WVFD", "summary": "This paper focuses on the problem of symmetry discovery in self-supervised learning. In particular, the goal is to learn models that are either sensitive to certain features like rotations and lightning or invariant to them, depending on the task. The authors propose to learn a world model that models transformations of the input images as a sequence of state, action, next state tuples. The major contribution is learning to adapt the representation of the world model based on the provided context of the task.", "questions": "What is the path towards making this method discover or adapt to naturally occurring symmetries?", "limitations": "Despite stating “Limitations of our work are discussed in Section 5” in the paper checklist, the discussion of the limitations in Section 5 is insufficient.", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 3, "strengths": "1. The authors present a novel combination of in-context learning and symmetry discovery. Their method successfully adapts its representation to be equivariant or invariant to different transformations \n\n2. The experimental evaluation uses a large number of strong baselines in self-supervised representation learning and learning of symmetric representations. The proposed method is superior in its ability to be sensitive to or invariant to certain features based on the context of the task.\n\n3. The paper contains an extensive ablation study to justify all components of the method.", "weaknesses": "1. I am not sure if it is meaningful to call this property equivariance: “if H(A|Z) is relatively small, the representation Z is nearly equivariant to the augmentation A”. Equivariance is specifically defined as the transformation of the input to the model having a predictable effect on the transformation of the output. This is different from simply having features that are predictive of a particular features (hence low entropy H(A|Z)). Would it be better to call this property something like sensitivity to a transformation?\n\n2. Section 4.1 does not make a strong case for ContextSSL outperforming the baselines. The results in Table 1 are somewhat mixed. Table 2 is not explained well.\n\n3. The paper does not make a clear case for the application of the proposed method outside of synthetic tasks. The 3DIEBench is artificially created to test equivariance and invariance to specific properties and the CIFAR-10 experiments do not actually demonstrate an improvement in classification accuracy. It is unclear how this method could be applied to more general and practical visual pre-training settings, such as CLIP [1] or DINO [2] self-supervised pre-training.\n\nReferences:\n\n[1] https://arxiv.org/abs/2103.00020\n\n[2] https://arxiv.org/abs/2304.07193 \n\n## Comments:\n\n* Table 1 is difficult to read. It is not immediately clear why ContextSSL is missing from the Rotation + Color section and why the context length != 14 fields for other methods are empty. Moreover, depending on the place in the table, a high or a low R^2 score could be the best result. That is very non-intuitive.\n* Clipped sentence: “We further test this at For all our equivariant baselines on 3DIEBench”.\n* The text “Contextual Self-Supervised Learning” in Figure 1 should be rotated by 180 degrees.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful to the reviewer for the time they put in to review our work. We are glad to see that they recognize several strengths in our work, including the novelty of our approach, comprehensive empirical evaluation using many baselines, and conducting thorough ablations. Below, we share our thoughts on the questions asked. \n\n\n> H(A|Z) as a definition of equivariance\n\nWe fully concur with the reviewer's observations. As noted in the footnote of Section 2.1, we use the term \"equivariance\" in a somewhat relaxed sense to denote that learned features are sensitive to data augmentations. However, it is common practice in equivariant self-supervised learning [1, 2, 3, 4] to use this definition to enforce equivariance. \n\n[1] Dangovski, Rumen, et al. \"Equivariant contrastive learning.\" arXiv preprint arXiv:2111.00899 (2021).\n\n[2] Lee, Hankook, et al. \"Improving transferability of representations via augmentation-aware self-supervision.\" NeurIPS (2021): 17710-17722.\n\n[3] Xie, Yuyang, et al. \"What should be equivariant in self-supervised learning.\" CVPR. 2022.\n\n[4] Scherr, Franz, Qinghai Guo, and Timoleon Moraitis. \"Self-supervised learning through efference copies.\" NeurIPS (2022): 4543-4557.\n\n> Confusion regarding Table 1 and explanation about the benefits of ContextSSL\n\nWe acknowledge Table 1 may be confusing as presented in the paper, and will improve the presentation. Table 1 compares ContextSSL with baselines and shows the effect of context length. We believe the empirical success of ContextSSL is significant, and to show that, we show a different version of Table 1 in the attached rebuttal document (through Table 1 and Figure 1), that is hopefully clearer; it separates context length from the comparison (the baseline methods are independent of context length). \n\n**Please refer to the detailed response about key observations from this table in the consolidated review above.** \n\n- On the 3DIEBench dataset, ContextSSL achieves higher $R^2$ scores for rotation (0.744) and color (0.986) compared to all baselines (the highest being 0.671 for rotation and 0.975 for color). It also matches or exceeds other equivariant models in linear probe performance. Unlike equivariant baselines, ContextSSL does this *without* training separate models for each equivariance. \n- ContextSSL seamlessly enforces equivariance or invariance to rotation or color by merely paying attention to different contexts, as shown in Figure 1 of the attached document. Thus *one* model can align the learned representation to priors that are beneficial for different downstream tasks. \n- Additional compelling evidence of ContextSSL's strong performance over baselines is presented in Table 2, Table 3, Table 4, and Figure 5 of the original manuscript.\n\n---\n\n> Confusion regarding Table 2\n\n\nWe provide more details regarding Table 2 here and will add the corresponding discussion in the paper.\n\nTable 2 shows that ContextSSL outperforms baseline approaches on two key metrics for equivariance: Mean Reciprocal Rank (MRR) and Hit Rate at k (H@k) [1]. ContextSSL's performance on these metrics consistently improves with increasing context length, demonstrating adaptation to rotation-specific features. To put these numbers into perspective, a H@1 score of 0.29 for ContextSSL signifies that the first nearest neighbor is the target embedding 29% of the time. In contrast, this occurs only 5% of the time for EquiMod and SEN, which is marginally better than the 2% expected by random chance. Notably, ContextSSL surpasses the baseline performances even with zero context, demonstrating its ability to learn equivariance without any contextual information.\n\n[1] Garrido, Quentin, Laurent Najman, and Yann Lecun. \"Self-supervised learning of split invariant equivariant representations.\" arXiv preprint arXiv:2302.10283 (2023).\n\n---\n\n\n> Extension beyond synthetic augmentations and towards adapting to naturally occurring symmetries\n\nWe thank the reviewer for raising this excellent question. It is indeed critical to evaluate our approach beyond self-supervised datasets, which use synthetic augmentations to enforce priors on the representation. To address this, we show that ContextSSL extends to naturally occurring symmetries and sensitive features in fairness and physiological datasets such as the MIMIC III and UCI Adult. To demonstrate this, we train ContextSSL to be selectively equivariant or invariant to gender by merely attending to different contexts. This is crucial; for instance, equivariance is needed for gender-specific medical diagnoses where different medicine dosages are required, while invariance is essential for fairness in tasks such as predicting hospital stay duration or medical cost. We present these results in Table 2 and Table 3 of the attached rebuttal document, with details in the caption. \n\n**Please refer to the detailed response about the extension to naturally occurring features and symmetries in the consolidated review above and to Table 2 and Table 3 in the attached document.** \n\n---\n\n> Additional Limitations\n\nWe would like to highlight some additional limitations of our work. \n- So far, ContextSSL has been evaluated on medium-sized datasets such as 3DIEBench and CIFAR. Its scaling law to massive datasets and more diverse environments is left to be explored in the future with more available compute.\n- Using the transformer network to learn a contextual world model increases training and memory costs, though these are relatively small compared to the encoding process.\n- So far, we have tested ContextSSL on two environments such as rotation and color or crop and blur. As a future work, we aim to expand our testing to continuous environments, moving beyond the constraints of finite settings.\n---\n\n\n> Minor errors \n\nWe thank the reviewer for their attention to detail. We will make these corrections in our revised manuscript.\n\n---", "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": "", "extra_scores": {}}, {"review_id": "GGbCvHmXdX", "reviewer": "Reviewer_y5eJ", "summary": "This work proposes to employ context modules to learn general representations such that invariance and equivariance to specific augmentations do not bias the representations. The method utilizes a module to learn to be both invariant or equivariant based on the context of the input augmentations, thus producing highly generalized features from learnt symmetries that are capable of preserving or disregarding a variety of transformations for the downstream tasks. The resulting performance leads to improved downstream evaluations in both invariant and equivariant settings, excelling above state-of-the-art in some settings.", "questions": "⁃\tI’m intrigued to understand how does “out of context” augmentations for downstream tasks impact performance? I assume given the automatically learnt symmetries that applying a full context would result in better performance to downstream cases where the context would be vastly different to the training data. \n\n⁃\tSee weaknesses", "limitations": "The limitations are appropriately addressed", "rating": 7, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "⁃\tThe paper is generally well presented and written, describing a clear and rationale problem statement supported by appropriate examples. \n\n⁃\tThe resulting framework is original and highly significant in the field of SSL. Notably, the addition of the context module could allow for a significant shift in the real-world application of SSL. \n\n⁃\tEmpirical results show significant improvement in equivariant downstream tasks further justifying this works significance in the field. Additionally, an extensive ablation and sensitivity analysis is performed guiding the reader into a greater understanding of the behavior of the method and the rationale behind implementation decisions. \n\n⁃\tExtensive details to support replication are provided.", "weaknesses": "⁃\tHow does the method handle more complex augmentation strategies, the proof of concept in the setting of 3DIE and CIFAR demonstrate strong performance yet these transformations, especially equivariance are highly controlled and unique to these datasets. It therefore would have been good to see more generalized augmentations including multiple combinations in the downstream context that better adapt to the real world setting and thus support the generalization claim of the work.\n\n⁃\tDoes not the choice of context length to extract representations for inference require significant supervision? Such an implementation requires the practitioner to construct augmentations per context and then provide such information for the downstream. Please correct my understanding if incorrect.\n\n⁃\tFrom the latter point, this method claims to not hard code symmetries. However, from my understanding this method is still hard coding symmetries to some extent as the practitioner is selecting which augmentations to select. In this case the context is enabling separability between representations that belong to each group symmetry. They are still hand-coded just conditioned on the learnt context for determining whether to be invariant or equivariant. \n\nMinor: \n\n⁃\tFigure 1 and 2 could perhaps be made clearer or positioned differently in the paper. It is not overly clear or useful in supporting the written explanation.\n\n⁃\tContext length should be headed in the tables.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer's thorough and insightful review, along with their positive feedback regarding the significance of our work for the SSL community, extensive evaluations and ablations, and our writing. In response to their review, have endeavored to address your concerns as concretely as possible\n\n\n> Extension to complex augmentation strategies\n\nWe thank all the reviewers for raising this excellent question. It is indeed critical to evaluate our approach beyond self-supervised datasets, which use synthetic augmentations to enforce priors on the representation. To address this, we show that ContextSSL extends to naturally occurring symmetries and sensitive features in fairness and physiological datasets such as the MIMIC III [1] and UCI Adult [2]. To demonstrate this, we train ContextSSL to be selectively equivariant or invariant to gender by merely attending to different contexts. This is crucial; for instance, equivariance is needed for gender-specific medical diagnoses where different medicine dosages are required, while invariance is essential for fairness in tasks such as predicting hospital stay duration or medical cost. We present these results in Table 2 and Table 3 of the attached rebuttal document, with details in the caption. From Table 2, we can observe that ContextSSL learns equivariance to gender in one context, improving gender and medical diagnosis prediction for MIMIC-III. In another context, ContextSSL achieves higher invariance to gender, resulting in superior performance on fairness metrics like equalized odds (EO) and equality of opportunity (EOPP) for hospital stay (LOS) prediction. We observe similar results for fairness of income prediction in the UCI Adult dataset, as shown in Table 3 of the attached document. \n\n[1] Johnson, A., T. Pollard, and R. Mark III. \"MIMIC-III Clinical Database (version 1.4). PhysioNet. 2016.\" (2016).\n\n[2] Arthur Asuncion and David Newman. UCI machine learning repository, 2007.\n\n \n---\n\n> Does the choice of context length require supervision at inference?\n\nThis is indeed a critical question. As shown in Table 1 and Table 4 of the attached rebuttal document, ContextSSL is robust to varying context lengths and generalizes well to longer contexts, eliminating the need for explicit supervision during inference. During training, we use random masking and subsequently test without masking, which enhances the model's robustness to varying context lengths. For example, although trained with an average context length of 9 under 90% data masking, the model extrapolates well to context lengths up to 128 during testing, as demonstrated in Table 1. \n\nFurthermore, depending on the useful priors of different downstream tasks, one only needs to construct the corresponding context and use the *maximum* context length. The degree of equivariance or invariance in ContextSSL increases with context length and is highest at the maximum context length, as observed empirically through all our experiments.\n\n---\n\n> Does our work still hard coding symmetries to some extent?\n\nIndeed, we still need to know the set of symmetries at training time. However, our work is the first to move beyond these fixed symmetries, training a representation that can dynamically adapt to be invariant or equivariant to a subset of these transformations. This enables learning a single representation that performs well across various downstream tasks, eliminating the need to retrain a new model for each task. So far, we have tested this on environments like rotation and color or crop and blur. We believe that it is an important stepping stone towards learning from diverse contexts. In practice, ContextSSL could be trained to handle a larger set of transformations, covering the entire set of commonly used augmentations.\n\nTo demonstrate how ContextSSL extends beyond these synthetic transformations to naturally occurring symmetries, we conduct experiments on the fairness and physiological datasets such as the MIMIC III [1] and UCI Adult [2]. We present these results in Table 2 and Table 3 of the attached rebuttal document, with details in the caption. \n\n**Please refer to the detailed response about the extension to naturally occurring features and symmetries in the consolidated review above and to Table 2 and Table 3 in the attached document.** \n\n[1] Johnson, A., T. Pollard, and R. Mark III. \"MIMIC-III Clinical Database (version 1.4). PhysioNet. 2016.\" (2016).\n\n[2] Arthur Asuncion and David Newman. UCI machine learning repository, 2007.\n\n---\n\n\n> Minor issues and typographical errors \n\nWe thank the reviewer for their attention to detail. We will make the following corrections in our revised manuscript:\n\n- *Positioning of Figure 1*: We will add more discussion on Figure 1 to highlight our approach clearly.\n\n- *Header for context length in Table 1*: We agree with the reviewer and have improved Table 1 with supporting plots and tables, as shown in the attached rebuttal document (Table 1 and Figure 1).", "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": "", "extra_scores": {}}, {"review_id": "yu9DDBCw2L", "reviewer": "Reviewer_iSU1", "summary": "This paper proposes ContextSSL, a novel self-supervised learning framework designed to enhance the existing joint embedding architecture by incorporating task-specific context. The main idea is to dynamically adapt symmetries by leveraging context in SSL. Consequently, ContextSSL can adapt to varying task symmetries without requiring parameter updates. The authors demonstrate the efficacy of ContextSSL on 3DIEBench and CIFAR10, showing that ContextSSL can selectively learn invariance or equivariance to transformations while maintaining general representations.", "questions": "**[Q]** Regarding the [W2], can ContextSSL benefit from few-shot classification, e.g., ImageNet accuracy of models trained with 1% of labels [1]? I believe few-shot can serve as context, and in this setup, ContextSSL might outperform other baselines.\n\n[1] Chen, Ting, et al. \"A simple framework for contrastive learning of visual representations.\" International conference on machine learning. PMLR, 2020.", "limitations": "They addressed the limitations.", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "**[S1]** This paper suggests an interesting direction for SSL, proposing that self-supervised representation incorporating context can enable dynamic adaptation to varying task symmetries.\n\n**[S2]** The overall writing is smooth and easy to follow.", "weaknesses": "**[W1]** It seems possible for invariant-based approaches to achieve context lengths of 0 to 126 by training a linear classifier. Why are these results ignored? More shots could also improve the performance of SimCLR and VICReg.\n\n**[W2]** Although ContextSSL performs well on the augmentation prediction task, it underperforms compared to other important baselines in linear classification, which is the most common task.\n\n**[W3]** It seems that ContextSSL can be trained on a single augmentation type, while other Equivariant-based approaches benefit from multiple augmentations.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for their time and feedback that helps us improve the work.  We believe that there are a few confusions that have resulted in the given rating. We have endeavored to address your concerns as concretely as possible and ask for your careful consideration of our clarifications. All of the discussions below shall be added to the paper to improve its clarity.\n\n---\n\n> Reporting results for invariant baselines for context length 0 to 128\n\nThe invariant and equivariant baselines do not utilize context in their training, and as a result, **their performance remains constant regardless of context length**. The reported performance for these baselines in Table 1 corresponds to zero-shot evaluations. We recognize the confusion caused by the centered values in the table. To clarify further, we have included supporting figures and tables in the attachment pdf that will also be added to the revised manuscript. \n\n---\n\n> Lower performance of ContextSSL on linear probe compared to baselines\n\nCore to ContextSSL is its ability to selectively enforce invariance or equivariance based on context. Thus, our experiments test if ContextSSL can do so without sacrificing performance on standard benchmarks like linear probe accuracy. Achieving high linear probe accuracy is not the primary goal of this work.\n\n- On the 3DIEBench dataset, ContextSSL achieves higher $R^2$ scores for rotation (0.744) and color (0.986) compared to all baselines (the highest being 0.671 for rotation and 0.975 for color). It also matches or exceeds other equivariant models in linear probe performance, as shown in Table 1. Unlike equivariant baselines, ContextSSL does this *without* training separate models for each equivariance. Moreover, in 3DIEBench, augmentations like rotation or color are independent of classification labels, indicating that equivariant models generally aren't expected to outperform invariant ones. To clarify this further, we will add additional clarification in the revised manuscript. \n\n- In datasets where equivariance to transformations is crucial and correlated with labels, ContextSSL achieves superior performance, with an 83% linear probe classification accuracy compared to 72% for baselines like SimCLR (as shown in Table 4).\n\n- As shown in Table 3, ContextSSL achieves an $R^2$ of 0.608 on rotation and 0.925 on color, significantly surpassing SimCLR's 0.459 and 0.371. Further, ContextSSL achieves this while also achieving a linear probe accuracy of 88.5%, comparable to SimCLR's 89.1%.\n\n- Additional compelling evidence of ContextSSL's strong performance over baselines is presented in Table 2, Table 3, Table 4, and Figure 5 of the original manuscript.\n\n**Please refer to the detailed response about key observations from Table 1 in the consolidated review above.** \n\n---\n\n> ContextSSL can be trained on a single augmentation type, while other Equivariant-based approaches benefit from multiple augmentations.\n\nSimilar to other invariant and equivariant baselines, ContextSSL is indeed trained with multiple augmentations by using multiple augmentations for generating positive samples in the context. \nAs shown in Table 1, all equivariant baselines are trained to be equivariant to either 1) rotation, 2) color, 3) both rotation and color. This requires training a separate model for each setting. ContextSSL, on the contrary, trains a single model using two contexts—one corresponding to rotation and the other to color—thus using multiple augmentations. Depending on which context is used, the model dynamically enforces either invariance/equivariance to rotation or color.  To avoid this confusion around Table 1, we show a different version of Table 1 in the attached rebuttal document (through Table 1 and Figure 1) that is hopefully clearer; it separates context length from the comparison (the baseline methods are independent of context length). \n\n**Please refer to the detailed response about key observations from this table in the consolidated review above.** \n\n---\n\n> ContextSSL benefits in few shot classification setting\n\nWe are afraid that there seems to be some misunderstanding. The linear probing metrics in Table 1 show zero-shot results for both ContextSSL and other methods, which, indeed, form a fair comparison. ContextSSL outperforms both the invariant and equivariant baselines in terms of quantitative equivariance measures such as $R^2$ in Table 1, Mean Reciprocal Rank (MRR), and Hit Rate in Table 2. While achieving the highest linear probe accuracy is not the goal of this work, ContextSSL still demonstrates competitive performance in Table 1 and surpasses other baselines in Table 4. \n\nTo further emphasize this, we compare the linear probe accuracy of the *predictor* of ContextSSL across different context lengths with that of SimCLR. As shown in Table 4 in the attached rebuttal pdf and Table 19 in the Appendix of our manuscript, ContextSSL outperforms SimCLR in linear probe accuracy across all context lengths. Note that the invariant and equivariant baselines do not operate on context, and, as a result, the performance of SimCLR in this table remains constant regardless of context length. \n\n---", "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": "", "extra_scores": {}}]

Author rebuttal:
We thank all the reviewers for their time and expertise in evaluating our paper. Their perceptive remarks and constructive feedback have been valuable in improving our work. In response, we have made several key revisions to address their concerns and have conducted additional experiments to enhance the support for our claims. Below is a brief summary of the key revisions: 

---

> Testing ContextSSL beyond synthetic augmentations and towards naturally occurring symmetries?

We thank all the reviewers for raising this excellent question. It is indeed critical to evaluate our approach beyond self-supervised datasets, which use synthetic augmentations to enforce priors on the representation. To address this, we show that ContextSSL extends to naturally occurring symmetries and sensitive features in fairness and physiological datasets such as the MIMIC III [1] and UCI Adult [2]. To demonstrate this, we train ContextSSL to be selectively equivariant or invariant to gender by merely attending to different contexts. This is crucial; for instance, equivariance is needed for gender-specific medical diagnoses where different medicine dosages are required, while invariance is essential for fairness in tasks such as predicting hospital stay duration or medical cost. We present these results in Table 2 and Table 3 of the attached rebuttal document, with details in the caption. From Table 2, we can observe that ContextSSL learns equivariance to gender in one context, improving gender and medical diagnosis prediction for MIMIC-III. In another context, ContextSSL achieves higher invariance to gender, resulting in superior performance on fairness metrics like equalized odds (EO) and equality of opportunity (EOPP) for hospital stay (LOS) prediction. We observe similar results for fairness of income prediction in the UCI Adult dataset, as shown in Table 3 of the attached document.

[1] Johnson, A., T. Pollard, and R. Mark III. "MIMIC-III Clinical Database (version 1.4). PhysioNet. 2016." (2016).

[2] Arthur Asuncion and David Newman. UCI machine learning repository, 2007.


---

> Clarifications about key results of ContextSSL and Table 1 of the manuscript

We acknowledge that in the initial manuscript, there was some confusion surrounding Table 1, which may not have effectively communicated the strengths of our approach. To address this, we present an improved version of Table 1 in the attached rebuttal document (through Table 1 and Figure 1). The new table demonstrates the empirical success of ContextSSL over invariant and equivariant baselines, and the new figure highlights ContextSSL's dynamic adaptability by paying attention to different contexts. Key observations from the Table are as follows:
- With context corresponding to rotation and color, respectively, ContextSSL achieves higher $R^2$ scores for rotation (0.744) and color (0.986) compared to all baselines (the highest being 0.671 for rotation and 0.975 for color). This indicates that it enforces equivariance to rotation and color in their respective contexts. ContextSSL also matches or exceeds other equivariant models in linear probe performance. Unlike equivariant baselines, ContextSSL does this \emph{without} training separate models for each augmentation group. 
- With contexts of rotation and color, ContextSSL achieves invariance to the other transformation, i.e., color ($R^2$ of 0.023) and rotation ($R^2$ of 0.344), respectively, comparable to SIE's $R^2$ values of 0.011 for color and 0.304 for rotation. However, ContextSSL achieves higher linear probe classification accuracy while training a single model, unlike SIE and other equivariant baselines that require two trained models, one for rotation and one for color. 
- ContextSSL seamlessly enforces equivariance or invariance to rotation or color by merely paying attention to different contexts, as shown in Figure 1 of the attached document. Thus *one* model can align the learned representation to priors that are beneficial for different downstream tasks. 
- Additional compelling evidence of ContextSSL's strong performance over baselines is presented in Table 2, Table 3, Table 4, and Figure 5 of the original manuscript.

---

Meta-review: This paper focuses on the transformations that self-supervised learning algorithms learn to be invariant or equivariant to. The Authors are concerned with the inductive biases brought by this paradigm, and propose to adapt invariance (or equivariance) of the learned representation to the context of the current task.

After the rebuttal, all Reviews are leaning towards acceptance (BA, WA and A), and they praised the main contribution of this paper, its writing and the experimental analysis. I recommend acceptance, and encourage the Authors to improve their paper by taking into account the feedback and discussions with Reviewers
Decision: Accept (poster)