In-Context Symmetries: Self-Supervised Learning through Contextual World Models

Authors: sharut  gupta
chenyu  wang
yifei  wang
tommi  jaakkola
stefanie  jegelka
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-15 17:01:19

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.

Main paper: 
 introduction: 
 Self-supervised learning (SSL) of image representations has made remarkable progress in recent years (Chen et al., 2020a; Bardes et al., 2022; Zhou et al., 2022a; Larsson et al., 2016; Gidaris et al., 2018; Bachman et al., 2019; Gidaris et al., 2021; Grill et al., 2020; Shwartz-Ziv et al., 2022; Misra and Maaten, 2020; Chen et al., 2020b; He et al., 2020; Chen and He, 2021; Zbontar et al., 2021; Tomasev et al., 2022; Zhou et al., 2022b) , achieving competitive performance to its supervised counterparts on various downstream tasks, such as image classification.Most of these works are based on the joint-embedding architecture (as shown in 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, e.g., randomizing color, texture, orientation, and cropping. The alignment of representations for positive pairs can be guided by either invariance (Chen et al., 2020a; Bardes et al., 2022; Chen and He, 2021; He et al., 2020; Zbontar et al., 2021; Grill et al., 2020) , which promotes insensitivity to these augmentations, or equivariance (Gupta et al., 2023b; Devillers and Lefort, 2023; Dangovski et al., 2022; Garrido et al., 2023b; Assran et al., 2023; Garrido et al., 2024) , which maintains sensitivity to them. However, enforcing invariance or equivariance to a 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 is useful for image classification but can significantly hurt performance on image segmentation, where 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 (Xiao et al., 2021; Dangovski et al., 2022) . This rigidity of traditional SSL methodologies contrasts sharply with human perceptual abilities, which are highly adaptive, tuning into relevant features based on the context of 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 the time on analog clocks. It 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 algorithms eliminate augmentation-based inductive priors and enable dynamic adaptation to varying task symmetries?Source (R x 1 , R y 1 , R z 1 )(θ 1 , ϕ 1 )Figure 1 : We apply a transformation (rotation or color) on a 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 x , R y , R z ); (bottom row) color transformation (θ, ϕ). In the top row, we see that CONTEXTSSL learns equivariance to rotation and invariance to color as the NN representations match the target's angle but not its color. In the bottom row, it adapts to the color context and enforces the reverse, be equivariant to color and invariant to rotation.This work suggests a positive answer to this question by proposing to enhance the current joint embedding architecture with a finite context -an abstract representation of task, containing a few demonstrations that inform about task-specific symmetries, as shown in Figure 2 (c). Based on this idea, we propose Contextual Self-Supervised Learning (CONTEXTSSL), a contrastive learning framework that uses a transformer module to adapt to selective invariance or equivariance to transformations by paying attention to context representing a task. Unlike previous approaches with built-in symmetries, the ability of CONTEXTSSL to adapt to varying data symmetries-all without undergoing any parameter updates-enables it to learn a general representation across tasks, devoid of specific inductive priors.This unique prospect makes our model a promising approach to building world models (Hafner et al., 2020; 2023; Hu et al., 2023; Sekar et al., 2020; Yang et al., 2024) for vision. World models are essential for building representations of 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) (Garrido et al., 2024) ( Figure 2(b) ), that consider transformations as actions and the input and its transformed counterpart as 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 a task, encouraging equivariance to a subset of transformations and invariance to the rest (as shown in Figure 1 ).This approach promotes learning a general representation that can flexibly adapt to the symmetries relevant to various downstream tasks, eliminating the need to learn separate representations for each task. We empirically validate our approach on the 3D Invariant Equivariant Benchmark (3DIEBench) and CIFAR-10, extending to transformations such as rotations, cropping, and blurring.To summarize, the main contributions of our work are:• We propose CONTEXTSSL, a self-supervised learning algorithm that adapts to task-specific symmetries by paying attention to context.• We show that learning with context is prone to identifying shortcuts and subsequently propose two key modules to address it: a context mask and an auxiliary predictor.• We demonstrate the efficacy of our approach on 3DIEBench and CIFAR10, showing its ability to selectively learn invariance or equivariance to transformations such as color and rotation while maintaining similar performance on invariant (classification) benchmarks. We extend CONTEXTSSL to supervised learning, demonstrating its ability to effectively leverage context to identify features defining a task. 
augmentation-based inductive bias in self-supervised learning: 
 The goal of self-supervised learning (SSL) is to derive meaningful data representations without relying on humanlabeled data. Given an unlabeled dataset D, SSL methods learn a representation function f θ : X → Z that maps input data x ∈ X to a latent space Z.In-Context Symmetries: Self-Supervised Learning through Contextual World Models(Proposed) Architecture for World Models21 y k x k y 1 x 1 a 1 a k y k+1 x k+1 a k+1 y k+1 x k+1 a Predictor (b) (c) y k+1 x k+1 a (a)Figure 2 : Family of approaches in self-supervised learning (a) Joint Embedding methods (Chen et al., 2020a; Bardes et al., 2022; Caron et al., 2021) encode invariances to input transformations a by aligning representations across views of the same image; (b) Image World Models (Garrido et al., 2024; Assran et al., 2023 ) 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 {(x i , a i , y i )} k i=1 2.1. Role of data augmentations in Self-Supervised LearningData augmentations are arguably the most important component in modern SSL methods, where the representation function is learned to 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, we define an augmentation A as a random variable distributed over a set of N data transformations with domain A = {a 1 , . . . , a N }, where a i : R d → R d ′ denotes an input mapping, and d, d ′ are its input and output dimensions, respectively. Among existing SSL methods, there are generally two ways to 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 to 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.1 Formally, for a representation Z, one can use H(A|Z) as a measure of the degree of feature invariance or equivariance: if H(A|Z) is relatively small, the representation Z is nearly equivariant to the augmentation A; otherwise, if H(A|Z) is very large (close to H(A)), Z is invariant to A. Recent SSL methods (Gupta et al., 2023b; Garrido et al., 2023b; Park et al., 2022; Devillers and Lefort, 2023; Dangovski et al., 2022) 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 as it removes information about features or transformations that may be relevant in fine-grained tasks (Lee et al., 2021; Xie et al., 2022a) ; 2) contrastive learning benefits from partial invariance through implicit equivariance of the projection head (Jing et al., 2022) .Self-Supervised PretrainingAs discussed above, a common theme in existing SSL methods is to enforce invariance or equivariance to a 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, in real world scenarios, there is no single symmetry that is universally applicable across all tasks. For example, object recognition (e.g., a 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 in 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 (Xiao et al., 2021; Dangovski et al., 2022) . 
beyond built-in symmetry: contextual self-supervised learning: 
 Recognizing the limitations of existing augmentationspecific 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 a fixed set of symmetries, CONTEXTSSL learns these symmetries from contextual cues, thus capturing the unique set of features of downstream tasks. This adaptability allows it to serve as a general-purpose SSL framework, capable of learning from a diverse array of pretraining tasks with varying symmetry priors and seamlessly adapting to different downstream tasks.To design CONTEXTSSL, we draw inspiration from world modeling (Hafner et al., 2020; 2023; Sekar et al., 2020; Yang et al., 2024) , a widely used framework in reinforcement learning (RL). World modeling aims to build representations of the world from past experience by predict the next state x t+1 from the current state x t and action a t . This next state prediction task captures the inherent mechanisms of the system and facilitates decision making. Traditionally applied in RL, the benefits of world modeling in vision have been largely unexplored. Recently, Image World Models (IWM) (Garrido et al., 2023a) established a 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 t and that of the transformed input as 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 CON-TEXTSSL in vision self-supervised learning as Contextual World Models. In this way, CONTEXTSSL 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 on color details when identifying flowers, and on spatial orientation such as rotation angle when determining the time on analog clocks.Drawing inspiration from the in-context learning (Brown et al., 2020) 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 transformerbased 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 a 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 to invariant versus equivariant learning in SSL. We elaborate on these ideas below.Symmetries as Context. Given a set of groups of input transformations {G 1 , . . . , 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, translation, as well as their compositions, can serve as different transformation groups. Each group G c can be represented through the joint distribution P (x, a, y|G c ), where x is the input sample (sampled from an unlabeled dataset), a represents the parameters of the transformation drawn from G c and applied to x, and y is the transformed input. In principle, x can be transformed by a composition of augmentations drawn from multiple transformation groups. For instance, in self-supervised learning, it 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 G c , 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 c ) = [(x 1 , a 1 , y 1 ), . . . , (x K , a K , y K )], where x i , a i , y i ∼ P (x, a, y|G c ), i ∈ [K].y i ≈ h((x i , a i ); (x 1 , a 1 , y 1 ), . . . , (x i-1 , a i-1 , y i-1 )). (1)While the requested prediction y i concerns only the inputs x i and a i , the model can now pay attention to the experience so far, enforcing relevant symmetries for the augmentation group G c . The predictor h is updated by minimizing the loss at each context length K i=1 ℓ(h((x, a i ); C i-1 ), y i ) where C i = {(x 1 , a 1 , y 1 ), . . . , (x i-1 , a i-1 , y i-1 )} represents the context before index i.Motivated by the above ideas, we begin by constructing pairs of points {(x i , y i ) K i=1 } by either 1) sampling a transformation group G and transforming x i by augmentation from G to y i ; or 2) if available, sampling a meta-latent and its transformation parameters as 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 (Garrido et al., 2023b) . Note that pairs of data can also be transformed by a series 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 {(x i , y i )} K i=1 from the context is independently transformed by the encoder into its corresponding latent representation. Next, representations of the input samples x i are concatenated with their corresponding transformation action a i . This concatenated vector (x i , a i ) and the representation of the corresponding transformed input y i collectively form the context corresponding to the symmetry G. The correspond-ing output embeddings are then aligned using the InfoNCE loss, which is minimized at each context length. If a i is set to zero for all tokens in a sequence, CONTEXTSSL enforces invariance to G, since it aligns x i and y i without conditioning on the transformation parameters. Overall, we optimize the following loss:L algo (h) = E G∼{G1,...,G M } E C(G) K i=1 -log exp h((x i , a i )|C i (G)) ⊤ h(y i |C i (G))/τ K j=1 exp h((x i , a i )|C i (G)) ⊤ h(y j |C j (G))/τwhere transformed data tokens y j (j ̸ = i) form the negatives. We use a similar symmetric loss term using y i as the anchor, (x i , a i ) and (x j , a j ) (j ̸ = i) as 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 a 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 {(x i , a i , y i )} K i=1 as its preceding context. Here a i belongs to the group G and is used to transform other unlabelled data from the test set x i into y i . On the contrary, if the downstream task benefits from invariance to the group, we use {(x i , 0, y i )} K i=1 as the preceding context. This notion can be generalized to enforce equivariance to a subset of groups and invariance to another. Specifically, including the augmentation parameters for transformations in a group G 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 (x i , a i ) precedes y i in the context sequence, a trivial solution to minimizing the alignment loss arises where the model treats the embeddings of (x i , a i ) identical to y i due to its access to x i . This phenomenon, often referred to as shortcut learning, poses a significant challenge as it leads the model to collapse to constant representations for each pair (x i , y i ), all while perfectly minimizing the loss. We address this challenge by masking out the input token (x i , a i ) for each token y i in the context. As a consequence, when encoding the token y i , the transformer only has access to past context C i = {(x 1 , a 1 , y 1 ), . . . , (x i-1 , a i-1 , y i-1 )}, excluding its corresponding positive sample (x i , a i ).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, as shown in Figure 3 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 a 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 transformation parameters and collapse back to behaviors associated with invariance-based methods. As illustrated in Figure 4 , naively training CONTEXTSSL leads to poor equivariance with respect to the transformations. Previous works (Garrido et al., 2023b) have also identi-fied 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 i from the concatenated input vector (x i , a i ). 
quantitative assessment of adaptation to task-specific symmetries: 
 We use the 3D Invariant Equivariant Benchmark (3DIEBench) (Garrido et al., 2023b ) and CIFAR10 to test our approach. We compare CONTEXTSSL with 1) VI-CReg (Bardes et al., 2022) and SimCLR (Chen et al., 2020a) among the invariant self-supervised approaches; 2) EquiMOD (Devillers and Lefort, 2023), SEN (Park et al., 2022) and SIE (Garrido et al., 2023b) 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 equivariant to either only 3D rotation, color transformations, or both.We report the test performance on context lengths 0, 2, 14, 30, and 126. To assess the quality of the invariant representations, we employ linear classification over frozen features.For the equivariant counterpart, we report R 2 on the task of predicting the corresponding transformation. Additionally, we use Mean Reciprocal Rank (MRR) and Hit Rate at k (H@k) to evaluate the performance of our context predictor. More details about pretraining algorithms and training setup are provided in Appendix B.Invariant Classification and Equivariant transformation prediction task. As shown in Table 1 , invariant selfsupervised 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 CONTEXTSSL. In contrast, CONTEXTSSL exhibits equivariance to both rotation and color in the absence of context. As seen from the two rows corresponding to CONTEXTSSL in Table 1 , when the context corresponds to pairs of data with transformations sampled from the rotation (color) group, the model adaptively learns to be invariant to color (rotation) while improving equivariance to rotation (color). Appendix C.7 shows that CONTEXTSSL learns equivariance or invariance to the same transformation based on the context. We observed that as masking probability increases, performance on both classification and prediction tasks initially improves but later declines, reaching optimal performance at a masking probability of 90%.Role of Auxiliary Predictor. We demonstrate that the auxiliary predictor is crucial for the model to achieve equivariance. In its absence, as depicted in Figure 4 , while the model retains its performance on the invariant classification task, it fails to learn equivariance, and cannot effectively Unlike 3DIEBench where meta-latents for each data are available, we manually construct positives by applying augmentations like crop and blur on CIFAR10. The results for the combinations of crop and blur are reported in Table 2 . Consistent with our previous results, while almost retaining the classification performance as SimCLR, CONTEXTSSLIn-Context Symmetries: Self-Supervised Learning through Contextual World Models learns to adaptively enforce equivariance to crop (blur) and invariance to 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.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, we consider a supervised learning task where our transformer model is 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, 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  
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 or through distillation. However, this shift toward generalization 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 a 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 to 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.In-Context Symmetries: Self-Supervised Learning through Contextual World Models Appendix A. Related WorkSelf-Supervised Learning. Existing SSL methods generally belong to two categories: invariant learning (Chen et al., 2020a; Bardes et al., 2022; Chen and He, 2021; He et al., 2020; Zbontar et al., 2021; Grill et al., 2020 ) 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 to data augmentation. Contrastive learning can learn highly discriminative features at the cost of losing certain image information due to the invariance constraint (Xiao et al., 2021) . Motivated by this limitation, recent works explore merging contrastive learning with equivariant learning tasks by separate embedding (Xiao et al., 2021; Garrido et al., 2023b) , augmentationconditioned predictor (Devillers and Lefort, 2023; Garrido et al., 2024) , and explicit equivariant transformation (Gupta et al., 2023b) . However, existing works still inherit the limitations of contrastive learning: its symmetry prior is built on a 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 to 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 model-based planning (Ha and Schmidhuber, 2018; Sekar et al., 2020; Hafner et al., 2020) and vision (Hafner et al., 2023; Hu et al., 2023; Yang et al., 2024) In-context Learning. Our work is inspired by and extends the concept of in-context learning (ICL) (Brown et al., 2020) to training. Initially studied in the context of language, in-context learning has recently been adapted for vision tasks (Gupta et al., 2023a; Wang et al., 2023; Bar et al., 2022; Li and Liang, 2021) , allowing models to infer environmental features or tasks directly from input prompts without predefined notions. For example, Visual Prompting (Wang et al., 2023; Bar et al., 2022)  
 
 Appendix: To evaluate the efficacy of our proposed algorithm CON-TEXTSSL, our experiments are designed to address the following questions:1. How does CONTEXTSSL fare against competitive invariant and equivariant self-supervised learning approaches in terms of performance across varying context sizes and different sets of data transformations?2. How effectively can CONTEXTSSL identify taskspecific symmetries, both within the scope of selfsupervised learning and beyond?3. 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?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.Each experiment was conducted on 1 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 of RAM. All experiments were implemented using the PyTorch deep learning framework.To test equivariance and invariance to multiple data transformations, we use the 3D Invariant Equivariant Benchmark (3DIEBench) (Garrido et al., 2023b) 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 as 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 on 3DIEBench 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; 2) 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 a batch size of 512 on 128×128 resolution images. We use the standard training, validation and test splits, made publicly available by the authors (Garrido et al., 2023b) .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, we manually construct positive pairs by applying compositions of these handcrafted augmentations. We consider three transformation groups: crop, blur and color. Similar to 3DIEBench, we consider combinations of two groups for each training run. We use the standard training, validation and test splits.Among the invariant self-supervised approached, we compare our approach to VICReg (Bardes et al., 2022) and and SimCLR (Chen et al., 2020a) . For each method, comparisons are drawn using their originally proposed architectures.For the equivariant baselines, we consider EquiMOD (Devillers and Lefort, 2023), SIE (Garrido et al., 2023b) and SEN (Park et al., 2022) . Similar to Garrido et al. (2023b) , 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 N + . Amongst these, we report the best performing variant in our results. For N + , we conduct analysis in two distinct settings: 1) a 'no context' or c = 0 invariant condition, and 2) a fully contextualized setting with a context length of 126.To ensure a fair comparison across different algorithms for each dataset, we use a standardized neural network backbone. Precisely, for our encoder, we use a ResNet-18 backbone pre-trained on ImageNet. For CONTEXTSSL, output features from the encoder are transformed into the context sequence, which is then processed by the decoderonly Transformer (Vaswani et al., 2017) from the GPT-2 Transformer family (Radford et al., 2019) . Our model configuration includes 3 layers, 4 attention heads, and a 2048dimensional 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 (x i , a i ) is masked out, context length L corresponds to 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 -5 and weight decay 1e-3. For baseline self-supervised approaches, in their original architecture, we use a learning rate of 1e -3 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. (2023b) , we report hyper-parameters and architectures specific to each method:• SimCLR (Chen et al., 2020a) We train using a 2048-2048-2048 dimensional multi-layered perceptron (MLP) based projection head with a temperature of 0.5.• VICReg (Bardes et al., 2022) We train using a 2048-2048-2048 MLP for the projection head and use weight of 10 for both the invariance loss and variance loss and 1 for covariance loss.• SEN (Park et al., 2022) Similar to other approaches we use a projection head of dimension 2048-2048-2048 and temperature 0.1.• EquiMod (Devillers and Lefort, 2023) We use the standatd projection head of dimensions 1024-1024-128 and use equal weighing of the invariance and the equivariance loss.• SIE (Garrido et al., 2023b) We use two 1024-1024-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 weightIn-Context Symmetries: Self-Supervised Learning through Contextual World Models of 10 for both the invariance loss and variance loss, 1 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.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 k (H@k) to 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 is calculated as the average reciprocal rank of the target embedding within these nearest neighbors. Hit rate-k (H@k) assigns a score of 1 if the target embedding is within the k-nearest neighbors of the predicted embedding and 0 otherwise. Similar to Garrido et al. (2023b) , 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.C.1. Quantitative Assessment of Adaptation to Task-Specific SymmetriesIn this section, we present additional results on the quantitative assessment of model performance on 3DIEBench, including the evaluation of learned representations on equivariant tasks (rotation and color prediction) to predict individual latent values. In contrast, the results in Table 1 focus on predicting relative latent values between pairs of image embeddings as inputs.As shown in Table 4 , 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 is much worse compared to CONTEXTSSL. Besides, aligning them with different targeted symmetry groups requires retraining the entire model. In contrast, CONTEXTSSL em-ploys a 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 Table 1 , when the context corresponds to pairs of data with transformations sampled from the rotation (color) group, the model adaptively learns to be invariant to color (rotation) while retaining equivariance to rotation (color).Results in Table 1 are the average value over three random seeds. We provide the standard deviation for rotation and color prediction of CONTEXTSSL in Table 5 and Table 6 .Similar to ??, we provide the performance of CONTEXTSSL on MRR and H@k compared to baseline methods with trained equivariance to rotation. While ?? uses the validation set data as the retrieval library, Table 7 provides the results using the training set data. CONTEXTSSL outperforms the baseline models, and its performance on all the metrics consistently improves with increasing context length, showing adaptation to rotation-specific features.In this section, we provide additional results for the role of context mask and auxiliary predictor.In addition to Figure 3 , 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 at a masking probability of 90%.Results in Figure 3 and Figure 7 are the average value over three random seeds. We provide the standard deviation for rotation and color prediction of CONTEXTSSL in Table 8 and Table 9 .We provide the complete results corresponding to Figure 4 in Table 10 to demonstrate that the auxiliary predictor is crucial for the model to achieve equivariance. In its absence, while the model retains its performance on the invariant classification task, it fails to learn equivariance, performs similarly to the invariant models, and cannot effectively adapt to different contexts. Unlike 3DIEBench where meta-latents for each data are available, we manually construct positives by applying aug-mentations like crop and blur on CIFAR10. The results for the combinations of crop and blur are reported in Table 2 . We additionally provide the results for the combinations of crop and color in Table 12 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 equivariance 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 We analyze the difference between the performance on representation and the performance on predictor embedding for both the invariance (classification) task and equivariance 17 and Table 18 . CON-TEXTSSL 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 is low.Apart from adaptively learning equivariance to a subset of transformation groups and invariance to the rest as shown in Table 1 , we extend CONTEXTSSL to operate within environments characterized by a 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. If the environment corresponds to enforcing equivariance, we construct our context in the same way as before i.e. pairs of positives transformed using augmentations sampled from the transformation group. However, if the environment corresponds to enforcing invariance, we maximize alignment between positives transformed by augmentation sampled from the transformation group without conditioning on that augmentation. Take rotation in 3DIEBench as an example. As shown in Table 19 , similar to our results in two transformation setting (rotation and color) in Table 1 , CONTEXTSSL effectively adapts to enforce invariance and equivarance to rotation depending on the context. Results for predicting individual latents are provided in Table 20 . Proceedings of the 1 st Workshop on In-Context Learning at the 41 st International Conference on Machine Learning, Vienna, Austria. 2024. Copyright 2024 by the author(s).Here, the concept of equivariance is used in 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 a group, and exact equivariance is not well-defined.In Table1, both the CONTEXTSSL models are the same and the performance is reported depending on whether the context corresponds to rotation or color augmentation group. 


Decision: 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

Reviews: review: paper_summary: paper_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.
strength_weakness: strength_weakness: 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?

2. 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.

3. 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.

References:

[1] https://arxiv.org/abs/2103.00020

[2] https://arxiv.org/abs/2304.07193 

## Comments:

* 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.
* Clipped sentence: “We further test this at For all our equivariant baselines on 3DIEBench”.
* The text “Contextual Self-Supervised Learning” in Figure 1 should be rotated by 180 degrees.
questions: questions: What is the path towards making this method discover or adapt to naturally occurring symmetries?
limitations: 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.
paper_summary: paper_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.
strength_weakness: strength_weakness: ⁃	How 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.

⁃	Does 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.

⁃	From 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. 

Minor: 

⁃	Figure 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.

⁃	Context length should be headed in the tables.
questions: questions: ⁃	I’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. 

⁃	See weaknesses
limitations: limitations: The limitations are appropriately addressed
paper_summary: paper_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.
strength_weakness: strength_weakness: **[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.

**[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.

**[W3]** It seems that ContextSSL can be trained on a single augmentation type, while other Equivariant-based approaches benefit from multiple augmentations.
questions: 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.

[1] Chen, Ting, et al. "A simple framework for contrastive learning of visual representations." International conference on machine learning. PMLR, 2020.
limitations: limitations: They addressed the limitations.

Comments: title: Reply to Reviewer iSU1 (Continued)
Reply to Reviewer iSU1
Gentle reminder to respond to our rebuttal
Gentle reminder to respond to our rebuttal
Thanks to Reviewer y5eJ
Response to Rebuttal
comment: We sincerely appreciate the reviewer’s continued engagement with our work and their thoughtful review. Their feedback has been very valuable in refining our paper. We also thank them for their decision to raise their score.
Dear reviewer WVFD,

Understanding that you may be busy, and with the author-reviewer discussion period coming to a close, we would like to take this last opportunity to summarize the major updates we made during the rebuttal to address your concerns.

1. In response to your concern about the applicability of ContextSSL beyond synthetic augmentations, we demonstrated that ContextSSL extends to naturally occurring symmetries and sensitive features in fairness and physiological datasets such as MIMIC III [1] and UCI Adult [2].
2. Based on your concerns regarding Table 1, we replaced it with a version featuring clearer annotations and captions, highlighting key strengths of our approach, ContextSSL.
3. We clarified the definition of equivariance and its connection to H(A|Z).
4. We provided additional clarification around Table 2 and added other limitations and future directions of our work. 
 
*We hope these revisions address your concerns and would be happy to answer any further questions. If you find our revisions satisfactory, we hope that you would kindly consider re-evaluating our work.*

Best,

Authors
Thank you for your response. I have no additional questions. I will carefully consider these issues and make a final rating.
> For the W1, it would be better to clarify the evaluation setup in the final manuscript. 

We thank the reviewer for their suggestion. For better clarity, we have already elaborated this in our rebuttal by adding new tables and figures with captions clearly stating the evaluation configuration for all baselines and ContextSSL (as shown in the attached rebuttal document). We will make sure to add this clarification to our revised manuscript too. 

---


>  Similarly, for Table 4 of the attached pdf, the performance of SimCLR is also underestimated to justify the benefit of understanding the context compared to the SSLs.

We would like to highlight that our context merely consists of **unlabeled data** of the form $(x, a, x^+)$ and, thus, does **not** provide any information about the downstream label ($y$). These unlabeled data pairs are generated with data augmentations of training data and thus do not provide extra supervision to the model. Therefore, it is indeed a fair comparison with SimCLR/VICReg since both representations are obtained **with only unlabeled data**. We apologize for this confusion and will definitely make it clearer throughout the paper.

---


> I agree with the author that SSLs can benefit from understanding context. but, I still couldn't find the benefit of in-context learning from the experimental justifications. SSLs can utilize linear or simple MLP headers to adapt the new tasks in a few-shot manner (like few-context), and this ability to adapt to the tasks may be better with recent state-of-the-art SSLs, like Dino (v1 or v2) or MoCo v3 too.

In our evaluation, we trained a linear head for color prediction and an MLP head for rotation prediction on SimCLR features, following [1]. As shown in Table S1, even with fine-tuning on the same few-shot samples, SimCLR falls short in performance on equivariant tasks of color and rotation prediction compared to ContextSSL.

As discussed in the paper, ContextSSL fundamentally differs from previous SSL paradigms (SimCLR, DINO, MoCo); Previous SSL methods learn features with **a given set of augmentations**, which requires the model to ignore certain augmentations, like color and rotation. Instead, ContextSSL does not drop any information during pretraining but models the feature symmetry to be dependent on the context. The experiments above clearly demonstrate this — ContextSSL significantly outperforms SimCLR in preserving color and rotation information. We believe that this comparison provides strong evidence to demonstrate that our method has clear gains over existing SSL paradigms.

---


*We thank the reviewer again for their constructive feedback, which has helped us clarify and improve our work. If the reviewer finds our revisions satisfactory,  we hope that they would kindly consider re-evaluating our work.*



---
We thank the reviewer for their response. We are happy to address their remaining concerns below and hope that this revision serves as a basis for a positive review.

---


> However, I think the zero-shot performance of SimCLR or VICReg underestimates the value of their representations: They can perform with the context to train a linear classifier or regressor (e.g., ContextSSL also uses the few labeled contexts too, which is not indeed zero-shot, but train-free - If I'm wrong, please correct the statement and please clarify about the unlabeled contexts in the final manuscript.) to predict the rotation/color.

The key difference between ContextSSL and SimCLR/VICReg is that ContextSSL leverages knowledge from context to adapt to the “inductive bias” of downstream tasks (e.g., preferring features with rotation invariance). This is the exact purpose of our study— to train a model that can dynamically adapt to a range of downstream tasks according to the feature invariance as required. In contrast, methods like SimCLR produce features with a “fixed inductive bias”, making them less adaptable to downstream tasks. 

**New comparison and setup.** To highlight this difference, we additionally finetune the SimCLR model with 128 examples (same examples used as context for ContextSSL) for rotation prediction and/or color prediction. This matches the total amount of information available to ContextSSL and SimCLR.  The results are shown in the table below. *SimCLR-ft(128) rot+color* is the SimCLR model finetuned with both rotation and color prediction loss. *SimCLR-ft(128) rot* is finetuned with rotation prediction only, and *SimCLR-ft(128) color* is finetuned with color prediction only. As with all results in our paper and consistent with [1], rotation prediction $R^2$ is calculated by training an MLP over frozen features, while color prediction $R^2$ and classification accuracy are based on linear probing over frozen features. 

**Analysis.** From the table, we can see that **ContextSSL significantly outperforms fine-tuned SimCLR in both rotation and color prediction by a large margin,** indicating clear gains of ContextSSL over previous SSL methods with fixed inductive bias.


**Table S1**
| $\mathcal{G}$ | Method                    | Rotation Prediction ($R^2$) | Color Prediction ($R^2$) | Classification (top-1) |
|---------------|---------------------------|-----------------------------|--------------------------|------------------------|
|               | SimCLR                    | 0.506                       | 0.148                    | 85.3                   |
|               |                           | *Higher is better*            | *Higher is better*         |                        |
| rot+color     | SimCLR-ft(128) rot+color  | 0.546                       | 0.319                    |  82.1                      |
|               |                           | *Higher is better*            | *Lower is better*          |                        |
| rot           | SimCLR-ft(128) rot        | 0.560                            |    0.222                      |  83.7                      |
| rot           | ContextSSL, rot. context  | 0.744                       | 0.023                    | 80.4                   |
|               |                           | *Lower is better*             | *Higher is better*         |                        |
| color         | SimCLR-ft(128) color      | 0.490                            | 0.411                         | 76.9                       |
| color         | ContextSSL, color context | 0.344                       | 0.986                    | 80.4                   |


[1] Self-supervised learning of split invariant equivariant representations. ICML 2023.

---
I'm sorry for the late response to the authors' answer.

For the W1, it would be better to clarify the evaluation setup in the final manuscript. However, I think the zero-shot performance of SimCLR or VICReg underestimates the value of their representations: They can perform with the context to train a linear classifier or regressor (e.g., ContextSSL also uses the few labeled contexts too, which is not indeed zero-shot, but train-free - If I'm wrong, please correct the statement and please clarify about the unlabeled contexts in the final manuscript.) to predict the rotation/color. Similarly, for Table 4 of the attached pdf, the performance of SimCLR is also underestimated to justify the benefit of understanding the context compared to the SSLs.

Overall, I agree with the author that SSLs can benefit from understanding context. but, I still couldn't find the benefit of in-context learning from the experimental justifications. SSLs can utilize linear or simple MLP headers to adapt the new tasks in a few-shot manner (like few-context), and this ability to adapt to the tasks may be better with recent state-of-the-art SSLs, like Dino (v1 or v2) or MoCo v3 too.
Dear reviewer iSU1,

As the discussion period is drawing to a close, we wanted to kindly request your feedback on our rebuttal. In our response, we have carefully tried to address all your concerns and also included additional experiments to further demonstrate the strengths of our work. We would greatly appreciate it if you could provide your feedback at your earliest convenience.

Thank you for your time.

Best regards,

Authors
Dear reviewer WVFD,

As the discussion period is drawing to a close, we wanted to kindly request your feedback on our rebuttal. In our response, we have carefully tried to address all your concerns and also included additional experiments to further demonstrate the strengths of our work. We would greatly appreciate it if you could provide your feedback at your earliest convenience.

Thank you for your time.

Best regards,

Authors
Thank you for the prompt response and recognition of our new experiments and discussions. We will incorporate all additional experiments and clarifications into our revised manuscript. The important concerns raised by you have been very valuable in enhancing the clarity of our work. We fully concur with the suggestion to include more vision benchmarks and are currently testing our approach on them. We will ensure that they are included in the revised manuscript.
Many thanks for the detailed response that addresses many of the weaknesses identified and questions raised. I appreciate the addition of naturally occurring symmetries, and the rebuttal results demonstrate the capability of the method to adapt to such settings. It would also be beneficial to see more visual benchmarks given this had been the main focus of the paper, however, I understand that time restrictions do not permit this. However, for future revisions comparisons on benchmark datasets to compare against methods such as EquiMod, E-SSL, CARE, would improve the findings. Additionally, the clarification on the context length is a useful addition. 

I emphasise that all clarifications made during this rebuttal should be made in any revised manuscript to improve clarity of the work.

Given my already positive review, I for now maintain my score.