SimCLR-based JEPA: Predictive World Modeling

• The paper introduces SimCLR-JEPA, a reconstruction-free framework that aligns predicted and observed latent representations using a temporal InfoNCE loss.
• It employs a dual-module design with an encoder and a GRU-based predictor, eliminating the projection head to optimize instance-wise contrastive loss.
• Empirical evaluations reveal sub-pixel accuracy with changing distractors, but a collapse in representation when faced with static, persistent distractors.

SimCLR-based Joint Embedding Predictive Architecture (SimCLR-JEPA) is a reconstruction-free framework for self-supervised world model learning that leverages a temporal contrastive objective to align predicted and observed latent representations. In contrast to generative models trained with pixel-level reconstruction loss, SimCLR-JEPA avoids decoding and directly optimizes an instance-wise contrastive loss across temporally adjacent embeddings. Its primary distinguishing feature is its reliance on slow feature learning, which both enables robust performance under certain types of background variation and leads to a notable failure mode in the presence of temporally persistent distractors (Sobal et al., 2022).

1. Architectural Design

SimCLR-JEPA employs a dual-module structure:

• Encoder $g_\phi$: Implements three convolutional layers (channels 32→64→64; activation functions: ReLU and BatchNorm), followed by $2\times2$ average pooling and a linear projection to a $d=512$-dimensional latent space. This module converts each input image $o_t \in \mathbb{R}^{H\times W}$ into an instantaneous latent state $s_t = g_\phi(o_t) \in \mathbb{R}^d$.
• Predictor (Forward Model) $f_\theta$: Consists of a one-layer GRU with hidden size 512. For each time step, it consumes the previous latent and action $a_{t-1} \in \mathbb{R}^2$, generating a one-step latent prediction $\tilde{s_t} = f_\theta(\tilde{s}_{t-1}, a_{t-1})$. The prediction sequence initializes with $\tilde{s}_1 = s_1$. Notably, the architecture omits the standard SimCLR MLP projection head; $\tilde s_t$ serves as the second view in the contrastive framework.

2. Loss Function: SimCLR InfoNCE in JEPA

SimCLR-JEPA adapts the SimCLR InfoNCE loss to a predictive setting. At each time step $t = 2, ..., T+1$, two views are produced:

• $s_t = g_\phi(o_t)$ (“encoder view”)
• $\tilde s_t = f_\theta(s_{t-1}, a_{t-1})$ (“predictor view”)

The instancewise symmetric InfoNCE loss, applied with temperature $\tau > 0$, is defined as: $$\mathrm{expsim}(u, v) = \exp \left( \frac{u^\top v}{\tau \|u\| \|v\|} \right)$$

$$\mathrm{InfoNCE}(S_t, \tilde S_t) = -\frac{1}{N} \sum_{i=1}^N \log \frac{\mathrm{expsim}(s_{t,i}, \tilde s_{t,i})} {\sum_{k=1}^N \left[ \mathrm{expsim}(s_{t,i}, \tilde s_{t,k}) + \mathbf{1}_{k \neq i} \mathrm{expsim}(s_{t,i}, s_{t,k}) \right]}$$

The total symmetric objective, averaged over $T$ time steps, is: $$\mathcal{L}_\text{SimCLR} = \frac{1}{2T} \sum_{t=2}^{T+1} [\mathrm{InfoNCE}(S_t, \tilde S_t) + \mathrm{InfoNCE}(\tilde S_t, S_t)]$$ Positive pairs consist of time-aligned embedding and one-step predicted latent; negatives are all other embeddings from the same batch and time step.

3. Training Protocol and Environment

Training operates fully offline and without reward information. The setup incorporates:

• Dataset: One million pre-collected trajectories of the format $\{o_1, a_1, o_2, ..., a_T, o_{T+1}\}$, used exclusively for pretraining.
• Simulated Environment: Each episode is a 17-frame sequence in which a single Gaussian-blurred dot traverses a $[0,1]^2$ square. Actions are sampled as $$a_t \sim \text{Uniform}[0,D] \times \text{VonMises}(\mu, \kappa)$$ with $D=0.14$.
• Distractor Generation: Backgrounds consist of either “uniform” (pixel values $\sim \text{Uniform}[0,1]$) or “structured” (random CIFAR-10 images) noise. Distractors can be:
  • “Changing”: independently resampled each frame,
  • “Fixed”: held constant within an episode, resampled between episodes.
• A brightness coefficient $\alpha \in [0,1.5]$ regulates distractor visibility.
• No classical data augmentation is applied; all nondeterministic variation arises from distractor generation.

4. Empirical Performance and Evaluation

Representational fidelity is assessed via a linear “prober” trained on the (frozen) latent sequence $\tilde s_t$ to regress the true dot position $c_t \in \mathbb{R}^2$, reporting RMSE averaged across 17 frames.

Performance under changing background noise:

• SimCLR-JEPA and VICReg-JEPA achieve sub-pixel accuracy (RMSE $\approx 0.02–0.06$) as $\alpha$ increases up to $1.5$, without hyperparameter retuning.
• Pixel-level reconstruction baselines degrade (RMSE $\approx 0.2$) at high $\alpha$ if not tuned.

Performance under fixed background noise:

• SimCLR-JEPA and VICReg-JEPA experience representational collapse: RMSE increases to $\approx 0.17–0.28$ at $\alpha=1.5$, similar to random guessing.
• In contrast, reconstruction-based models remain robust (RMSE $\approx 0.07–0.11$ for $\alpha \leq 1.5$).

<table> <thead> <tr> <th>Condition</th> <th>SimCLR-JEPA RMSE</th> <th>Reconstruction RMSE</th> </tr> </thead> <tbody> <tr> <td>Fixed Uniform (α=1.0)</td> <td>≈0.19</td> <td>≈0.07</td> </tr> <tr> <td>Changing Uniform (α=1.0)</td> <td>≈0.03</td> <td>≈0.20</td> </tr> </tbody> </table>

This empirical evidence shows SimCLR-JEPA is highly robust against unpredictable, temporally varying distractors but fails catastrophically when distractors are temporally persistent.

5. Theoretical Analysis: Failure Mode under Fixed Distractors

SimCLR-JEPA’s contrastive predictive objective selects for “slow features” (features constant or slowly varying with time). With a fixed episode-wise distractor $Z$, which is constant within an episode but independent between episodes, the architecture admits a trivial solution: $$g_\phi(o_t) = s \sim \mathcal{N}(0, \sigma^2 I), \quad f_\theta(s, a) = s$$ With $s_t = s_{t+1} = s$ for all $t$, positive pairs align and the InfoNCE loss is minimized. The resulting embeddings are uniform on the hypersphere (in the large negative sample limit) and thus satisfy the contrastive uniformity condition (cf. Wang & Isola 2022, Thm 1). The construction achieves $\mathcal{L}_\text{SimCLR}\to 0$ without encoding any task-relevant rapid feature (the moving dot). A similar collapse argument applies to VICReg, as the prediction, variance, and covariance losses also vanish under this solution. This outcome demonstrates that SimCLR-JEPA cannot distinguish between static and dynamic content in the absence of additional constraints.

6. Practical Strengths, Weaknesses, and Research Implications

Strengths of SimCLR-JEPA:

• Reconstruction-free: Circumvents blurry pixel-MSE losses and can ignore highly unpredictable (changing) distractors without retuning.
• Simplicity: Employs only an encoder and a forward model—decoder is unnecessary.
• Hyperparameter robustness: The temperature parameter $\tau$ generalizes well across noise intensity, provided distractors are non-stationary.

Weaknesses:

• Susceptibility to static slow features: When any static or slow-varying background is present, the learned representation collapses onto those spurious distractors, neglecting dynamic, task-relevant features.
• Blindness to source dynamics: The architecture cannot inherently distinguish between dynamic and static sources; this limitation requires additional mechanisms, such as hierarchical time-scale modeling or penalties that discourage constant outputs.

Implications for world-model research:

Contrastive-predictive approaches such as SimCLR-JEPA show robustness to changing distractors and promise for pretraining in visually complex environments. However, prioritization of slow features results in catastrophic failure in the presence of persistent distractors, representing a fundamental Achilles’ heel. Mitigating this failure is likely to necessitate architectural inductive biases (e.g., hierarchical JEPA), enhanced objectives penalizing over-focus on static features, or alternative input regimes (e.g., frame-differences, optical flow) that circumvent the encoding of fixed backgrounds (Sobal et al., 2022).