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

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.

---

Reporting results for invariant baselines for context length 0 to 128

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

---

Lower performance of ContextSSL on linear probe compared to baselines

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

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

• 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).

• 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%.

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

Please refer to the detailed response about key observations from Table 1 in the consolidated review above.

---

ContextSSL can be trained on a single augmentation type, while other Equivariant-based approaches benefit from multiple augmentations.

Similar to other invariant and equivariant baselines, ContextSSL is indeed trained with multiple augmentations by using multiple augmentations for generating positive samples in the context. As 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).

Please refer to the detailed response about key observations from this table in the consolidated review above.

---

ContextSSL benefits in few shot classification setting

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

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

---

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

Extension to complex augmentation strategies

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.

---

Does the choice of context length require supervision at inference?

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

Furthermore, 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.

---

Does our work still hard coding symmetries to some extent?

Indeed, 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.

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

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.

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

---

Minor issues and typographical errors

We thank the reviewer for their attention to detail. We will make the following corrections in our revised manuscript:

• Positioning of Figure 1: We will add more discussion on Figure 1 to highlight our approach clearly.

• 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).

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.

H(A|Z) as a definition of equivariance

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

[1] Dangovski, Rumen, et al. "Equivariant contrastive learning." arXiv preprint arXiv:2111.00899 (2021).

[2] Lee, Hankook, et al. "Improving transferability of representations via augmentation-aware self-supervision." NeurIPS (2021): 17710-17722.

[3] Xie, Yuyang, et al. "What should be equivariant in self-supervised learning." CVPR. 2022.

[4] Scherr, Franz, Qinghai Guo, and Timoleon Moraitis. "Self-supervised learning through efference copies." NeurIPS (2022): 4543-4557.

Confusion regarding Table 1 and explanation about the benefits of ContextSSL

We 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).

Please refer to the detailed response about key observations from this table in the consolidated review above.

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

---

Confusion regarding Table 2

We provide more details regarding Table 2 here and will add the corresponding discussion in the paper.

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

[1] Garrido, Quentin, Laurent Najman, and Yann Lecun. "Self-supervised learning of split invariant equivariant representations." arXiv preprint arXiv:2302.10283 (2023).

---

Extension beyond synthetic augmentations and towards adapting to naturally occurring symmetries

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

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.

---

Additional Limitations

We would like to highlight some additional limitations of our work.

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

---

Minor errors

We thank the reviewer for their attention to detail. We will make these corrections in our revised manuscript.

---

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