JEPA — Visual Experiments¶

Goal: Explore what changes when prediction happens in a learned representation space rather than directly in observation space.

What we'll explore¶

#	Experiment	Question answered
1	Predictable abstraction	How can a target encoder remove unpredictable detail?
2	Conditioning	Why must the predictor know which target is requested?
3	Multiple futures	How can a latent variable represent multimodal outcomes?
4	Invisible collapse	Why can prediction loss look perfect while representations fail?

---

Linked theory: research/06-jepa.md

In [1]:

import numpy as np
import matplotlib.pyplot as plt

np.random.seed(42)
plt.rcParams.update({
    'figure.dpi': 110,
    'font.size': 11,
    'axes.titlesize': 12,
    'axes.labelsize': 11,
    'figure.facecolor': 'white',
    'axes.facecolor': '#f8f9fa',
    'axes.grid': True,
    'grid.alpha': 0.35,
})

def effective_rank(features, eps=1e-12):
    centered = features - features.mean(axis=0, keepdims=True)
    eigenvalues = np.linalg.eigvalsh(centered.T @ centered / max(len(features) - 1, 1))
    eigenvalues = np.clip(eigenvalues, 0, None)
    probabilities = eigenvalues / max(eigenvalues.sum(), eps)
    entropy = -np.sum(probabilities * np.log(probabilities + eps))
    return np.exp(entropy)

import numpy as np import matplotlib.pyplot as plt np.random.seed(42) plt.rcParams.update({ 'figure.dpi': 110, 'font.size': 11, 'axes.titlesize': 12, 'axes.labelsize': 11, 'figure.facecolor': 'white', 'axes.facecolor': '#f8f9fa', 'axes.grid': True, 'grid.alpha': 0.35, }) def effective_rank(features, eps=1e-12): centered = features - features.mean(axis=0, keepdims=True) eigenvalues = np.linalg.eigvalsh(centered.T @ centered / max(len(features) - 1, 1)) eigenvalues = np.clip(eigenvalues, 0, None) probabilities = eigenvalues / max(eigenvalues.sum(), eps) entropy = -np.sum(probabilities * np.log(probabilities + eps)) return np.exp(entropy)

---

Experiment 1: A learned target can discard unpredictable detail¶

Question: Why can latent prediction be easier than reconstructing the full observation?

Each target contains a semantic signal $u$ that is available from context and independent texture noise $n$. The best deterministic pixel predictor recovers $u$ but cannot know the particular realization of $n$. A target encoder that maps $y=u+n$ to $s_y=u$ removes that irreducible error.

In [2]:

np.random.seed(1)
positions_exp1 = np.linspace(0, 2 * np.pi, 160)
semantic_exp1 = np.sin(positions_exp1) + 0.35 * np.sin(2 * positions_exp1 + 0.4)
noise_levels_exp1 = np.linspace(0, 1.2, 25)
trials_exp1 = 500
pixel_mse_exp1 = []
latent_mse_exp1 = []

for noise_level in noise_levels_exp1:
    noise = np.random.normal(0, noise_level, size=(trials_exp1, len(positions_exp1)))
    observations = semantic_exp1[None, :] + noise
    pixel_predictions = np.broadcast_to(semantic_exp1, observations.shape)
    latent_targets = observations.mean(axis=0, keepdims=True)
    latent_predictions = semantic_exp1[None, :]
    pixel_mse_exp1.append(np.mean((pixel_predictions - observations) ** 2))
    latent_mse_exp1.append(np.mean((latent_predictions - latent_targets) ** 2))

example_noise_exp1 = np.random.normal(0, 0.75, size=(4, len(positions_exp1)))
example_observations_exp1 = semantic_exp1[None, :] + example_noise_exp1

np.random.seed(1) positions_exp1 = np.linspace(0, 2 * np.pi, 160) semantic_exp1 = np.sin(positions_exp1) + 0.35 * np.sin(2 * positions_exp1 + 0.4) noise_levels_exp1 = np.linspace(0, 1.2, 25) trials_exp1 = 500 pixel_mse_exp1 = [] latent_mse_exp1 = [] for noise_level in noise_levels_exp1: noise = np.random.normal(0, noise_level, size=(trials_exp1, len(positions_exp1))) observations = semantic_exp1[None, :] + noise pixel_predictions = np.broadcast_to(semantic_exp1, observations.shape) latent_targets = observations.mean(axis=0, keepdims=True) latent_predictions = semantic_exp1[None, :] pixel_mse_exp1.append(np.mean((pixel_predictions - observations) ** 2)) latent_mse_exp1.append(np.mean((latent_predictions - latent_targets) ** 2)) example_noise_exp1 = np.random.normal(0, 0.75, size=(4, len(positions_exp1))) example_observations_exp1 = semantic_exp1[None, :] + example_noise_exp1

In [3]:

fig, axes = plt.subplots(1, 2, figsize=(13.5, 5.2))
fig.suptitle('Representation targets can quotient out unpredictable observation detail', fontsize=14, fontweight='bold')

ax = axes[0]
for index, observation in enumerate(example_observations_exp1):
    ax.plot(positions_exp1, observation, alpha=0.35, linewidth=1.2, label='Possible observations' if index == 0 else None)
ax.plot(positions_exp1, semantic_exp1, color='#212529', linewidth=2.8, label='Predictable semantic signal')
ax.set_xlabel('Observation coordinate')
ax.set_ylabel('Signal value')
ax.set_title('Same semantics, different texture noise')
ax.legend()

ax = axes[1]
ax.plot(noise_levels_exp1, pixel_mse_exp1, 'o-', color='#e76f51', linewidth=2.2, label='Observation-space MSE')
ax.plot(noise_levels_exp1, latent_mse_exp1, 's-', color='#457b9d', linewidth=2.2, label='Abstract-target MSE')
ax.set_xlabel('Unpredictable noise standard deviation')
ax.set_ylabel('Prediction mean squared error')
ax.set_title('Pixel error grows with irrelevant uncertainty')
ax.legend()

plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 2, figsize=(13.5, 5.2)) fig.suptitle('Representation targets can quotient out unpredictable observation detail', fontsize=14, fontweight='bold') ax = axes[0] for index, observation in enumerate(example_observations_exp1): ax.plot(positions_exp1, observation, alpha=0.35, linewidth=1.2, label='Possible observations' if index == 0 else None) ax.plot(positions_exp1, semantic_exp1, color='#212529', linewidth=2.8, label='Predictable semantic signal') ax.set_xlabel('Observation coordinate') ax.set_ylabel('Signal value') ax.set_title('Same semantics, different texture noise') ax.legend() ax = axes[1] ax.plot(noise_levels_exp1, pixel_mse_exp1, 'o-', color='#e76f51', linewidth=2.2, label='Observation-space MSE') ax.plot(noise_levels_exp1, latent_mse_exp1, 's-', color='#457b9d', linewidth=2.2, label='Abstract-target MSE') ax.set_xlabel('Unpredictable noise standard deviation') ax.set_ylabel('Prediction mean squared error') ax.set_title('Pixel error grows with irrelevant uncertainty') ax.legend() plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: Many observations share the same smooth semantic signal but differ in fine noise. Observation-space error rises quadratically with noise, while the abstract target stays predictable.
2. Why it looks this way: A deterministic predictor cannot recover independent target noise. The target encoder defines an equivalence class in which those realizations map to the same representation.
3. Key takeaway: JEPA can spend capacity on shared structure, provided the target encoder actually removes nuisance detail rather than useful state.

---

Experiment 2: Conditioning specifies the prediction query¶

Question: Why does a predictor need position, time, action, or another conditioning variable?

The same context is used to predict several target locations on a structured signal. Without a position token, one predictor output must serve every location and converges to the average target.

In [4]:

np.random.seed(2)
locations_exp2 = np.linspace(0, 1, 120)
target_signal_exp2 = (
    0.8 * np.sin(2 * np.pi * locations_exp2)
    + 0.35 * np.cos(6 * np.pi * locations_exp2)
    + 0.25 * locations_exp2
)
unconditioned_prediction_exp2 = np.full_like(target_signal_exp2, target_signal_exp2.mean())
basis_exp2 = np.column_stack([
    np.ones_like(locations_exp2),
    locations_exp2,
    np.sin(2 * np.pi * locations_exp2),
    np.cos(2 * np.pi * locations_exp2),
    np.sin(6 * np.pi * locations_exp2),
    np.cos(6 * np.pi * locations_exp2),
])
coefficients_exp2 = np.linalg.lstsq(basis_exp2, target_signal_exp2, rcond=None)[0]
conditioned_prediction_exp2 = basis_exp2 @ coefficients_exp2
squared_error_unconditioned_exp2 = (unconditioned_prediction_exp2 - target_signal_exp2) ** 2
squared_error_conditioned_exp2 = (conditioned_prediction_exp2 - target_signal_exp2) ** 2

np.random.seed(2) locations_exp2 = np.linspace(0, 1, 120) target_signal_exp2 = ( 0.8 * np.sin(2 * np.pi * locations_exp2) + 0.35 * np.cos(6 * np.pi * locations_exp2) + 0.25 * locations_exp2 ) unconditioned_prediction_exp2 = np.full_like(target_signal_exp2, target_signal_exp2.mean()) basis_exp2 = np.column_stack([ np.ones_like(locations_exp2), locations_exp2, np.sin(2 * np.pi * locations_exp2), np.cos(2 * np.pi * locations_exp2), np.sin(6 * np.pi * locations_exp2), np.cos(6 * np.pi * locations_exp2), ]) coefficients_exp2 = np.linalg.lstsq(basis_exp2, target_signal_exp2, rcond=None)[0] conditioned_prediction_exp2 = basis_exp2 @ coefficients_exp2 squared_error_unconditioned_exp2 = (unconditioned_prediction_exp2 - target_signal_exp2) ** 2 squared_error_conditioned_exp2 = (conditioned_prediction_exp2 - target_signal_exp2) ** 2

In [5]:

fig, axes = plt.subplots(1, 2, figsize=(13.5, 5.2))
fig.suptitle('A predictor must know which compatible target is being requested', fontsize=14, fontweight='bold')

ax = axes[0]
ax.plot(locations_exp2, target_signal_exp2, color='#212529', linewidth=2.6, label='Target representation')
ax.plot(locations_exp2, unconditioned_prediction_exp2, '--', color='#e76f51', linewidth=2.2, label='No position condition')
ax.plot(locations_exp2, conditioned_prediction_exp2, color='#457b9d', linewidth=2.2, label='Position-conditioned')
ax.set_xlabel('Requested target position z')
ax.set_ylabel('Target representation value')
ax.set_title('One context, many target locations')
ax.legend()

ax = axes[1]
ax.plot(locations_exp2, squared_error_unconditioned_exp2, color='#e76f51', linewidth=2.2, label='No condition')
ax.plot(locations_exp2, squared_error_conditioned_exp2, color='#457b9d', linewidth=2.2, label='With position')
ax.set_xlabel('Requested target position z')
ax.set_ylabel('Squared prediction error')
ax.set_title('Missing query information creates averaging error')
ax.legend()

plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 2, figsize=(13.5, 5.2)) fig.suptitle('A predictor must know which compatible target is being requested', fontsize=14, fontweight='bold') ax = axes[0] ax.plot(locations_exp2, target_signal_exp2, color='#212529', linewidth=2.6, label='Target representation') ax.plot(locations_exp2, unconditioned_prediction_exp2, '--', color='#e76f51', linewidth=2.2, label='No position condition') ax.plot(locations_exp2, conditioned_prediction_exp2, color='#457b9d', linewidth=2.2, label='Position-conditioned') ax.set_xlabel('Requested target position z') ax.set_ylabel('Target representation value') ax.set_title('One context, many target locations') ax.legend() ax = axes[1] ax.plot(locations_exp2, squared_error_unconditioned_exp2, color='#e76f51', linewidth=2.2, label='No condition') ax.plot(locations_exp2, squared_error_conditioned_exp2, color='#457b9d', linewidth=2.2, label='With position') ax.set_xlabel('Requested target position z') ax.set_ylabel('Squared prediction error') ax.set_title('Missing query information creates averaging error') ax.legend() plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: The unconditioned predictor returns one horizontal average for all target locations. A position-conditioned predictor recovers the location-specific target.
2. Why it looks this way: Context alone does not identify which compatible $y$ is requested. The conditioning variable turns an ambiguous mapping into a well-posed query.
3. Key takeaway: Mask positions, time offsets, and actions are not decoration; they define the prediction being asked of the world model.

---

Experiment 3: A latent variable can separate multiple futures¶

Question: What does a deterministic squared-error predictor do when two future states are equally plausible?

The context is identical for every sample, while the target moves either left or right. We compare the conditional mean with a predictor that receives a one-bit latent mode.

In [6]:

np.random.seed(3)
samples_exp3 = 600
modes_exp3 = np.random.choice([-1.0, 1.0], size=samples_exp3)
targets_exp3 = np.column_stack([
    modes_exp3 + np.random.normal(0, 0.10, size=samples_exp3),
    0.65 + np.random.normal(0, 0.08, size=samples_exp3),
])
deterministic_prediction_exp3 = targets_exp3.mean(axis=0)
conditioned_predictions_exp3 = np.column_stack([modes_exp3, np.full(samples_exp3, 0.65)])
mse_deterministic_exp3 = np.mean((targets_exp3 - deterministic_prediction_exp3) ** 2)
mse_conditioned_exp3 = np.mean((targets_exp3 - conditioned_predictions_exp3) ** 2)
candidate_modes_exp3 = np.array([[-1.0, 0.65], [1.0, 0.65]])

np.random.seed(3) samples_exp3 = 600 modes_exp3 = np.random.choice([-1.0, 1.0], size=samples_exp3) targets_exp3 = np.column_stack([ modes_exp3 + np.random.normal(0, 0.10, size=samples_exp3), 0.65 + np.random.normal(0, 0.08, size=samples_exp3), ]) deterministic_prediction_exp3 = targets_exp3.mean(axis=0) conditioned_predictions_exp3 = np.column_stack([modes_exp3, np.full(samples_exp3, 0.65)]) mse_deterministic_exp3 = np.mean((targets_exp3 - deterministic_prediction_exp3) ** 2) mse_conditioned_exp3 = np.mean((targets_exp3 - conditioned_predictions_exp3) ** 2) candidate_modes_exp3 = np.array([[-1.0, 0.65], [1.0, 0.65]])

In [7]:

fig, axes = plt.subplots(1, 2, figsize=(13.5, 5.2))
fig.suptitle('Latent conditioning represents multiple plausible target modes', fontsize=14, fontweight='bold')

ax = axes[0]
ax.scatter(targets_exp3[:, 0], targets_exp3[:, 1], s=14, alpha=0.25, color='#6c757d', label='Observed target latents')
ax.scatter(*deterministic_prediction_exp3, s=180, marker='X', color='#e76f51', label='Deterministic mean')
ax.scatter(candidate_modes_exp3[:, 0], candidate_modes_exp3[:, 1], s=150, marker='*', color='#2a9d8f', label='z-conditioned modes')
ax.set_xlabel('Target latent dimension 1')
ax.set_ylabel('Target latent dimension 2')
ax.set_title('The conditional mean lies between valid futures')
ax.legend()

ax = axes[1]
ax.bar(['Deterministic\nmean', 'Conditioned\nby z'], [mse_deterministic_exp3, mse_conditioned_exp3],
       color=['#e76f51', '#2a9d8f'])
ax.set_xlabel('Predictor type')
ax.set_ylabel('Mean squared prediction error')
ax.set_title('One bit of mode information resolves ambiguity')

plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 2, figsize=(13.5, 5.2)) fig.suptitle('Latent conditioning represents multiple plausible target modes', fontsize=14, fontweight='bold') ax = axes[0] ax.scatter(targets_exp3[:, 0], targets_exp3[:, 1], s=14, alpha=0.25, color='#6c757d', label='Observed target latents') ax.scatter(*deterministic_prediction_exp3, s=180, marker='X', color='#e76f51', label='Deterministic mean') ax.scatter(candidate_modes_exp3[:, 0], candidate_modes_exp3[:, 1], s=150, marker='*', color='#2a9d8f', label='z-conditioned modes') ax.set_xlabel('Target latent dimension 1') ax.set_ylabel('Target latent dimension 2') ax.set_title('The conditional mean lies between valid futures') ax.legend() ax = axes[1] ax.bar(['Deterministic\nmean', 'Conditioned\nby z'], [mse_deterministic_exp3, mse_conditioned_exp3], color=['#e76f51', '#2a9d8f']) ax.set_xlabel('Predictor type') ax.set_ylabel('Mean squared prediction error') ax.set_title('One bit of mode information resolves ambiguity') plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: The deterministic MSE optimum sits between the left and right clusters, where almost no valid target occurs. Conditioning on a binary $z$ recovers both modes.
2. Why it looks this way: Squared error estimates a conditional mean. A latent mode changes the predictor from one-valued to a family of compatible predictions.
3. Key takeaway: Multimodal world models need controlled latent capacity; otherwise they average futures or let $z$ bypass context entirely.

---

Experiment 4: Perfect prediction does not imply a useful representation¶

Question: Can context and target encoders collapse together while latent prediction loss remains zero?

We construct perfectly predicted embeddings and progressively remove either all variance or most feature dimensions. The predictor tracks the target exactly at every point.

In [8]:

np.random.seed(4)
samples_exp4 = 800
dimensions_exp4 = 12
base_features_exp4 = np.random.normal(size=(samples_exp4, dimensions_exp4))
alphas_exp4 = np.geomspace(1.0, 1e-4, 40)
uniform_variance_exp4 = []
uniform_prediction_loss_exp4 = []
dimensional_rank_exp4 = []
dimensional_prediction_loss_exp4 = []

for alpha in alphas_exp4:
    uniformly_collapsing = alpha * base_features_exp4
    uniform_variance_exp4.append(np.var(uniformly_collapsing, axis=0).mean())
    uniform_prediction_loss_exp4.append(np.mean((uniformly_collapsing - uniformly_collapsing) ** 2))

    dimension_scales = alpha ** np.linspace(0, 3, dimensions_exp4)
    dimensionally_collapsing = base_features_exp4 * dimension_scales
    dimensional_rank_exp4.append(effective_rank(dimensionally_collapsing))
    dimensional_prediction_loss_exp4.append(np.mean((dimensionally_collapsing - dimensionally_collapsing) ** 2))

uniform_variance_exp4 = np.asarray(uniform_variance_exp4)
dimensional_rank_exp4 = np.asarray(dimensional_rank_exp4)

np.random.seed(4) samples_exp4 = 800 dimensions_exp4 = 12 base_features_exp4 = np.random.normal(size=(samples_exp4, dimensions_exp4)) alphas_exp4 = np.geomspace(1.0, 1e-4, 40) uniform_variance_exp4 = [] uniform_prediction_loss_exp4 = [] dimensional_rank_exp4 = [] dimensional_prediction_loss_exp4 = [] for alpha in alphas_exp4: uniformly_collapsing = alpha * base_features_exp4 uniform_variance_exp4.append(np.var(uniformly_collapsing, axis=0).mean()) uniform_prediction_loss_exp4.append(np.mean((uniformly_collapsing - uniformly_collapsing) ** 2)) dimension_scales = alpha ** np.linspace(0, 3, dimensions_exp4) dimensionally_collapsing = base_features_exp4 * dimension_scales dimensional_rank_exp4.append(effective_rank(dimensionally_collapsing)) dimensional_prediction_loss_exp4.append(np.mean((dimensionally_collapsing - dimensionally_collapsing) ** 2)) uniform_variance_exp4 = np.asarray(uniform_variance_exp4) dimensional_rank_exp4 = np.asarray(dimensional_rank_exp4)

In [9]:

fig, axes = plt.subplots(1, 2, figsize=(13.5, 5.2))
fig.suptitle('Prediction agreement alone cannot diagnose representation collapse', fontsize=14, fontweight='bold')

ax = axes[0]
ax.plot(alphas_exp4, uniform_variance_exp4, color='#457b9d', linewidth=2.4, label='Mean feature variance')
ax.plot(alphas_exp4, uniform_prediction_loss_exp4, '--', color='#e76f51', linewidth=2.2, label='Prediction loss')
ax.set_xscale('log')
ax.set_yscale('symlog', linthresh=1e-12)
ax.invert_xaxis()
ax.set_xlabel('Representation scale alpha (collapse toward the right)')
ax.set_ylabel('Variance or loss')
ax.set_title('Complete collapse: variance vanishes, loss stays zero')
ax.legend()

ax = axes[1]
ax.plot(alphas_exp4, dimensional_rank_exp4, color='#2a9d8f', linewidth=2.4, label='Effective rank')
ax.plot(alphas_exp4, dimensional_prediction_loss_exp4, '--', color='#e76f51', linewidth=2.2, label='Prediction loss')
ax.set_xscale('log')
ax.invert_xaxis()
ax.set_xlabel('Dimension scale alpha (collapse toward the right)')
ax.set_ylabel('Effective rank or loss')
ax.set_title('Dimensional collapse is also invisible to agreement')
ax.legend()

plt.tight_layout()
plt.show()

fig, axes = plt.subplots(1, 2, figsize=(13.5, 5.2)) fig.suptitle('Prediction agreement alone cannot diagnose representation collapse', fontsize=14, fontweight='bold') ax = axes[0] ax.plot(alphas_exp4, uniform_variance_exp4, color='#457b9d', linewidth=2.4, label='Mean feature variance') ax.plot(alphas_exp4, uniform_prediction_loss_exp4, '--', color='#e76f51', linewidth=2.2, label='Prediction loss') ax.set_xscale('log') ax.set_yscale('symlog', linthresh=1e-12) ax.invert_xaxis() ax.set_xlabel('Representation scale alpha (collapse toward the right)') ax.set_ylabel('Variance or loss') ax.set_title('Complete collapse: variance vanishes, loss stays zero') ax.legend() ax = axes[1] ax.plot(alphas_exp4, dimensional_rank_exp4, color='#2a9d8f', linewidth=2.4, label='Effective rank') ax.plot(alphas_exp4, dimensional_prediction_loss_exp4, '--', color='#e76f51', linewidth=2.2, label='Prediction loss') ax.set_xscale('log') ax.invert_xaxis() ax.set_xlabel('Dimension scale alpha (collapse toward the right)') ax.set_ylabel('Effective rank or loss') ax.set_title('Dimensional collapse is also invisible to agreement') ax.legend() plt.tight_layout() plt.show()

What you see¶

1. What the plot shows: Feature variance can approach zero and effective rank can fall toward one while the prediction loss remains exactly zero.
2. Why it looks this way: If context and target branches collapse consistently, the predictor can match a useless target perfectly. Agreement measures compatibility, not information content.
3. Key takeaway: Every JEPA implementation needs independent health metrics such as variance, covariance spectrum, effective rank, probes, or an explicit anti-collapse objective.

---

Summary / Key Takeaways¶

Experiment	Key insight
1. Predictable abstraction	Target encoders can remove irreducible observation noise.
2. Conditioning	Position, time, and action specify which target relation to predict.
3. Multiple futures	Latent modes avoid averaging incompatible outcomes.
4. Invisible collapse	Zero prediction loss does not prove that embeddings retain information.

Connection to the broader theory¶

JEPA is an architectural contract: encode context, encode a target, and predict the target representation under a condition. I-JEPA and V-JEPA make different masking and data choices inside that contract.