seq-JEPA: Autoregressive Predictive Learning of Invariant-Equivariant World Models

TL;DR

seq-JEPA introduces a sequence-based world-modeling approach that reconciles invariance and equivariance in self-supervised learning by processing action-conditioned view sequences through a transformer aggregator, producing separate latent representations optimized for downstream classification and transformation prediction. It avoids explicit equivariance losses and dual predictors, instead achieving architectural disentanglement where per-view encodings become equivariant while the aggregated representation supports invariance. Empirically, seq-JEPA matches or surpasses state-of-the-art invariant and equivariant SSL on tasks requiring either property and demonstrates robust predictive learning across saccades and path integration, with favorable scaling properties as sequence length increases. The approach shows potential for efficient learning from partial observations, with promising transfer to larger datasets and multi-modal extensions, highlighting an architectural pathway to flexible, scalable world models.

Abstract

Joint-embedding self-supervised learning (SSL) commonly relies on transformations such as data augmentation and masking to learn visual representations, a task achieved by enforcing invariance or equivariance with respect to these transformations applied to two views of an image. This dominant two-view paradigm in SSL often limits the flexibility of learned representations for downstream adaptation by creating performance trade-offs between high-level invariance-demanding tasks such as image classification and more fine-grained equivariance-related tasks. In this work, we propose \emph{seq-JEPA}, a world modeling framework that introduces architectural inductive biases into joint-embedding predictive architectures to resolve this trade-off. Without relying on dual equivariance predictors or loss terms, seq-JEPA simultaneously learns two architecturally separate representations for equivariance- and invariance-demanding tasks. To do so, our model processes short sequences of different views (observations) of inputs. Each encoded view is concatenated with an embedding of the relative transformation (action) that produces the next observation in the sequence. These view-action pairs are passed through a transformer encoder that outputs an aggregate representation. A predictor head then conditions this aggregate representation on the upcoming action to predict the representation of the next observation. Empirically, seq-JEPA demonstrates strong performance on both equivariance- and invariance-demanding downstream tasks without sacrificing one for the other. Furthermore, it excels at tasks that inherently require aggregating a sequence of observations, such as path integration across actions and predictive learning across eye movements.

Figures (17)

• Figure 1: seq-JEPA is a self-supervised world model that leverages a sequence of action-observation pairs to learn architecturally distinct representations for downstream tasks requiring transformation invariance or equivariance.
• Figure 2: a. Transformations and observations used for training b. In predictive learning across saccades, image saliencies and inhibition of return help create more informative and less redundant patch sequences.
• Figure 3: Top-1 linear classification (invariance) vs. rotation prediction (equivariance) performance on 3DIEBench. seq-JEPA learns good representations for both tasks. Subscripts indicate training and inference sequence lengths.
• Figure 3: Evaluation of predictive learning across saccades on STL-10. Equivariance is measured with respect to fixation coordinates. Unless stated otherwise, $M_{tr} = M_{val} = 4$.
• Figure 4: Retrieval of nearest representations; given a query image, we extract the three nearest encoder representations in the validation set of 3DIEBench. The retrieved views of models with the highest quantitative rotation equivariance performance maintain the rotation of the query image across all retrieved views.