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

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.

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

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

Rating

5

Confidence

3

Soundness

3

Presentation

3

Contribution

3

Limitations

They addressed the limitations.

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.

Strengths

⁃ The paper is generally well presented and written, describing a clear and rationale problem statement supported by appropriate examples. ⁃ The 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. ⁃ Empirical 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. ⁃ Extensive details to support replication are provided.

Weaknesses

⁃ 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

⁃ 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

Rating

7

Confidence

3

Soundness

3

Presentation

3

Contribution

4

Limitations

The limitations are appropriately addressed

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.

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

What is the path towards making this method discover or adapt to naturally occurring symmetries?

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

3

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.

Response to Rebuttal

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.

Gentle reminder to respond to our rebuttal

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