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

TL;DR

seq-JEPA introduces a novel world modeling framework that learns separate invariant and equivariant representations from sequences, improving flexibility and performance on diverse downstream tasks without trade-offs.

Contribution

It proposes architectural biases in joint-embedding predictive models to simultaneously learn invariant and equivariant representations from sequences, avoiding dual predictors or loss terms.

Findings

Strong performance on both invariance- and equivariance-demanding tasks

Effective in sequence-based tasks like path integration and eye movement prediction

Outperforms existing SSL methods in flexibility and downstream adaptation

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.