BiJEPA: Bi-directional Joint Embedding Predictive Architecture for Symmetric Representation Learning

TL;DR

BiJEPA introduces a bi-directional, cycle-consistent SSL architecture with norm regularization, improving stability and representation quality across diverse data modalities.

Contribution

It proposes a novel bi-directional predictive framework with cycle consistency and regularization, enhancing stability and semantic capture in SSL.

Findings

Achieves stable convergence without collapse.

Captures semantic structure of chaotic systems.

Learns robust representations for generation and generalization.

Abstract

Self-Supervised Learning (SSL) has shifted from pixel-level reconstruction to latent space prediction, spearheaded by the Joint Embedding Predictive Architecture (JEPA). While effective, standard JEPA models typically rely on a uni-directional prediction mechanism (e.g. Context $\to$ Target), potentially neglecting the informative signal inherent in the inverse relationship, degrading its performance. In this work, we propose \textbf{BiJEPA}, a \textit{Bi-Directional Joint Embedding Predictive Architecture} that enforces cycle-consistent predictability between data segments. We address the inherent instability of symmetric prediction (representation explosion) by introducing a critical norm regularization mechanism on the representation vectors. We evaluate BiJEPA on three distinct modalities: synthetic periodic signals, chaotic Lorenz attractor trajectories, and high-dimensional image data (MNIST). Our results demonstrate that BiJEPA achieves stable convergence without collapse, captures the semantic structure of chaotic systems, and learns robust temporal and spatial representations capable of generation and generalisation, offering a more holistic approach to representation learning.