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

arXiv: https://arxiv.org/abs/2405.18193 · HTML: https://mlsanity.com/paper/2405.18193/

## Key facts

- Cited by 5 papers in the map (data through 14 Oct 2026).
- Placed in Latent predictive world models, one of 11 directions in World models.
- Builds on 41 of its 51 references that are in the map.
- First posted on arXiv 28 May 2024.

## Placed in

- World models / [Latent predictive world models](https://mlsanity.com/wm/latent-predictive-world-models/)

## Builds on (most cited first, 30 of 41)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [A Simple Framework for Contrastive Learning of Visual Representations](https://mlsanity.com/paper/2002.05709/) (2020)
- [Momentum Contrast for Unsupervised Visual Representation Learning](https://mlsanity.com/paper/1911.05722/) (2019)
- [Masked Autoencoders Are Scalable Vision Learners](https://mlsanity.com/paper/2111.06377/) (2021)
- [Emerging Properties in Self-Supervised Vision Transformers](https://mlsanity.com/paper/2104.14294/) (2021)
- [Bootstrap your own latent: A new approach to self-supervised Learning](https://mlsanity.com/paper/2006.07733/) (2020)
- [Prefix-Tuning: Optimizing Continuous Prompts for Generation](https://mlsanity.com/paper/2101.00190/) (2021)
- [Exploring Simple Siamese Representation Learning](https://mlsanity.com/paper/2011.10566/) (2020)
- [BEiT: BERT Pre-Training of Image Transformers](https://mlsanity.com/paper/2106.08254/) (2021)
- [Barlow Twins: Self-Supervised Learning via Redundancy Reduction](https://mlsanity.com/paper/2103.03230/) (2021)
- [Dream to Control: Learning Behaviors by Latent Imagination](https://mlsanity.com/paper/1912.01603/) (2019)
- [Unsupervised Representation Learning by Predicting Image Rotations](https://mlsanity.com/paper/1803.07728/) (2018)
- [World Models](https://mlsanity.com/paper/1803.10122/) (2018)
- [Big Self-Supervised Models are Strong Semi-Supervised Learners](https://mlsanity.com/paper/2006.10029/) (2020)
- [Mastering Diverse Domains through World Models](https://mlsanity.com/paper/2301.04104/) (2023)
- [SimMIM: A Simple Framework for Masked Image Modeling](https://mlsanity.com/paper/2111.09886/) (2021)
- [VICReg: Variance-Invariance-Covariance Regularization for Self-Supervised Learning](https://mlsanity.com/paper/2105.04906/) (2021)
- [Self-Supervised Learning from Images with a Joint-Embedding Predictive Architecture](https://mlsanity.com/paper/2301.08243/) (2023)
- [iBOT: Image BERT Pre-Training with Online Tokenizer](https://mlsanity.com/paper/2111.07832/) (2021)
- [Self-Supervised Learning of Pretext-Invariant Representations](https://mlsanity.com/paper/1912.01991/) (2019)
- [Learning Representations by Maximizing Mutual Information Across Views](https://mlsanity.com/paper/1906.00910/) (2019)
- [GAIA-1: A Generative World Model for Autonomous Driving](https://mlsanity.com/paper/2309.17080/) (2023)
- [Learning Interactive Real-World Simulators](https://mlsanity.com/paper/2310.06114/) (2023)
- [Understanding Dimensional Collapse in Contrastive Self-supervised Learning](https://mlsanity.com/paper/2110.09348/) (2021)
- [Planning to Explore via Self-Supervised World Models](https://mlsanity.com/paper/2005.05960/) (2020)
- [Images Speak in Images: A Generalist Painter for In-Context Visual Learning](https://mlsanity.com/paper/2212.02499/) (2022)
- [Learning Representations for Automatic Colorization](https://mlsanity.com/paper/1603.06668/) (2016)
- [Visual Prompting via Image Inpainting](https://mlsanity.com/paper/2209.00647/) (2022)
- [What Should Not Be Contrastive in Contrastive Learning](https://mlsanity.com/paper/2008.05659/) (2020)

## Built on by (most cited first, 5 of 5)

- [WorldPrediction: A Benchmark for High-level World Modeling and Long-horizon Procedural Planning](https://mlsanity.com/paper/2506.04363/) (2025)
- [seq-JEPA: Autoregressive Predictive Learning of Invariant-Equivariant World Models](https://mlsanity.com/paper/2505.03176/) (2025)
- [Context and Diversity Matter: The Emergence of In-Context Learning in World Models](https://mlsanity.com/paper/2509.22353/) (2025)
- [Any-Subgroup Equivariant Networks via Symmetry Breaking](https://mlsanity.com/paper/2603.19486/) (2026)
- [Self-Supervised Learning from Structural Invariance](https://mlsanity.com/paper/2602.02381/) (2026)

## Abstract (arXiv)

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.