Joint Embedding Predictive Architectures Focus on Slow Features

Paper: arXiv:2211.10831 Code: vladisai/JEPA_SSL_NeurIPS_2022 Code reference: main @ 03eb1730 (2022-10-27)

1. Motivation (研究动机)

这篇短论文讨论的是一个非常具体但对 world model 很关键的问题：如果我们不用像 Dreamer / PlaNet / RSSM 那样重建像素，而是用 Joint Embedding Predictive Architecture (JEPA) 在 latent space 里预测下一步表示，模型到底会优先学什么？作者的结论不是“JEPA 总是更语义”这种乐观叙事，而是指出一种反直觉失败模式：在 reward-free offline setting 中，VICReg-JEPA 和 SimCLR-JEPA 很容易把容量用在时间上最稳定、最容易预测的慢变化特征上；如果这些 slow features 恰好是静态背景噪声，模型可以把动态主体完全忽略。

当前 pixel-based world model 的常见路线是 generative/reconstruction objective：模型从观测 $o_{t}$ 和动作 $a_{t}$ 学 latent dynamics，再用 decoder 重建 $o_{t + 1}$ 或整段未来图像。这类目标的优点是监督信号密集，但缺点也直接来自像素级重建：目标函数并不知道哪些像素与控制、规划、奖励有关。若背景、纹理、照明、噪声占据大量像素，reconstruction loss 会推动模型花容量解释这些视觉细节，出现 object-vanishing，即小而关键的可控物体在 latent representation 中被大面积背景压过。JEPA 被提出为一种 reconstruction-free 替代方案：不在 input space 复原像素，而是在 embedding space 预测未来/目标表示，希望 latent prediction 自动过滤掉不可预测或无关的细节。

本文要解决的不是 JEPA 的工程扩展问题，而是 JEPA 的归纳偏置问题：当只给一批离线的 observation-action sequences、没有 reward、没有 downstream task label、也没有 image augmentation 时，JEPA objective 是否真的会保留“对未来动态有用”的状态信息？更窄地说，若图像里同时存在一个 action-controlled moving dot 和若干 background distractors，编码器 $g_{ϕ}$ 与 forward model $f_{θ}$ 是否能保留 dot location，使 frozen representation 之后能被 linear probe 读出？这个 probe 目标在训练 world model 时不可见，因此它测试的是预训练表示自身的信息含量，而不是 end-to-end supervised fitting。

这个问题值得研究的原因在于它击中了 JEPA/world-model 作为通用预训练方法的假设边界。现实场景中有大量慢变化、可预测、但任务无关的因素：dash-cam 里的云层、道路纹理、摄像头污渍，机器人视频里的桌面背景、固定家具、光照模式。如果 JEPA 的 latent prediction objective 天然偏向这些 slow features，那么单纯从像素序列做 reward-free pretraining 可能得到一个“预测得很好但不会控制”的 latent model。反过来，如果能明确这个失败条件，后续设计就可以加入 hierarchy、object-centric bottleneck、difference/optical-flow input、bisimulation constraint 或 task-aware auxiliary loss，避免把最容易预测误认为最重要。

2. Idea (核心思想)

核心洞察可以压缩成一句话：JEPA 的预测目标会偏好时间上稳定且容易对齐的特征，而“稳定”不等于“有用”；当最慢的特征来自固定背景噪声时，VICReg/SimCLR 的防 collapse 机制仍然允许一个零预测误差、非坍缩但任务无关的解。论文的新意在于把这个偏置用一个极简 moving-dot 环境和一个可写出 loss 为零的理论构造展示出来，而不是只报告某个复杂 benchmark 的失败。

方法层面的创新是把 JEPA、reconstruction、IDM、supervised/random/center baselines 放在同一个 offline sequence setup 中，用同一套 frozen representation probing 协议比较它们保留 dot location 的能力。JEPA 部分分别用 VICReg 的 variance/covariance regularization 和 SimCLR 的 InfoNCE objective 防止 trivial collapse；实验再系统改变 distractor 的类型：uniform vs. CIFAR-10 structured noise，以及时间属性：每帧变化 vs. episode 内固定。

它与既有 world model/reconstruction 方法的根本差异在于优化目标的“可预测对象”不同。Reconstruction 要解释像素，因此只要 dot 仍然影响像素误差，它通常会保留 dot；但它也会被大量无关像素牵制。JEPA 不重建像素，理论上应更容易忽略不可预测噪声；本文证明它确实能忽略每帧变化的噪声，但它也会过度拥抱固定噪声。IDM 则走另一个极端：只关注能解释动作的因素，因此 action-controlled dot 表现很好，但 stationary dot 这种有信息却与动作无关的因素会被忽略。

3. Method (方法)

3.1 Overall framework: offline latent dynamics + probing

Figure 1 解读：左图是 VICReg-based JEPA，encoder 把每个 observation 编成 $s_{t}$ ，forward model 根据 $s_{t}$ 和动作 $a_{t}$ 预测 $\tilde{s}_{t + 1}$ ；预测表示与真实下一帧 encoder output 对齐，同时对表示施加 variance/covariance regularization。右图是 SimCLR-based JEPA，同样预测下一步 embedding，但用 InfoNCE 把同一时间步的 prediction/encoding 作为 positive pair，把 batch 内其它样本作为 negatives。两者都不解码像素，因此图中没有 pixel reconstruction branch；这正是论文想测试的 reconstruction-free 优势与风险。

论文形式化为一个 MDP $M = (O, A, P, R)$ 。 $O$ 是 observation space，在实验中是 $28 \times 28$ 图像； $A$ 是 action space； $P = Pr (o_{t} ∣ o_{t - 1}, a_{t - 1})$ 是 transition probability； $R : O \times A \to R$ 是 reward function。但训练时没有直接访问 MDP，也没有 reward，只给定离线序列 ${(o_{t}, a_{t})}$ 。目标是学习 encoder 与 forward model： $g_{ϕ} (o_{t}) = s_{t}, g_{ϕ} : O \to R^{D}$ $f_{θ} (s_{t}, a_{t}) = \tilde{s}_{t + 1}, f_{θ} : R^{D} \times A \to R^{D} .$ 给定初始 observation $o_{1}$ 和动作序列 $a_{1}, \dots, a_{T}$ ，模型先得到 $\tilde{s}_{1} = g_{ϕ} (o_{1})$ ，然后自回归展开： $\tilde{s}_{t} = f_{θ} (\tilde{s}_{t - 1}, a_{t - 1}), t = 2, \dots, T + 1.$ 测试时冻结 encoder/predictor，只训练一个线性 probe $q (\tilde{s}_{t}) : R^{D} \to R^{Q}$ 去恢复状态属性；本文的 $Q = 2$ ，即 dot 的二维位置。这个 protocol 很重要：如果 probe 失败，说明 latent sequence 中没有足够线性可读的位置状态，而不是下游头太弱或训练没 end-to-end 对齐。

直觉上，JEPA 的吸引力来自“只预测 latent 中可预测的结构”。对于每帧重新采样的随机/structured noise，背景在时间上没有稳定对应关系，forward model 很难从 $s_{t}, a_{t}$ 推出下一帧背景，因此 JEPA 有动机忽略它。危险也来自同一逻辑：若背景在 episode 内固定，它就是最容易预测的未来信息；编码器只复制固定噪声，predictor 学 identity，就能把 latent prediction objective 做得很好，同时 completely miss action-controlled dot。

3.2 VICReg-JEPA objective

VICReg adaptation 把每个时间步的 encoder output 或 projector output 作为 regularized representation，同时用 prediction loss 对齐 forward model output 和真实下一步 embedding。论文 appendix 给出的损失为： $L_{VICReg} = α L_{prediction} + β L_{variance} + L_{covariance} .$ 预测项是跨时间的 MSE： $L_{prediction} = \frac{1}{NT} \sum_{t = 1}^{T} \sum_{i = 1}^{N} ∥ f_{θ} (S_{t, i}, A_{t, i}) - g_{ϕ} (O_{t + 1, i}) ∥_{2}^{2} = \frac{1}{NT} \sum_{t = 2}^{T + 1} \sum_{i = 1}^{N} \tilde{S}_{t, i} - S_{t, i}_{2}^{2} .$ variance 项要求每个维度的 batch std 不低于阈值 $γ$ ，避免所有样本 collapse 到同一点： $L_{variance} = \frac{1}{( T + 1 ) D} \sum_{t = 1}^{T + 1} \sum_{j = 1}^{D} max (0, γ - Var (S_{t, :, j}) + ϵ) .$ covariance 项惩罚不同维度之间的 off-diagonal covariance，让表示维度不要冗余： $L_{covariance} = \frac{1}{( T + 1 ) ( N - 1 )} \sum_{t = 1}^{T + 1} \sum_{i = 1}^{D} \sum_{j = i + 1}^{D} (S_{t} S_{t}^{⊤})_{i, j} .$ released code 中对应实现位于 vicreg.py。VICRegPredMultistep.forward 先把所有帧 flatten 后过 backbone，再用 models.RNNPredictor.predict_sequence 从第一帧 encoding 和 action sequence 自回归预测未来 encoding。若 repr_loss_after_projector=true，MSE 对齐 projector space；若为 false，则直接对齐 encoder space。实际 best configs 对 VICReg 多数使用 projector: id、predictor: rnn、arch: menet5、embedding_size: 512、skip_first_step_vc: true，这说明实验最终采用的是较直接的 latent MSE + VC regularization，而不是固定的 paper-default projector 设置。

def vicreg_jepa_step(backbone, projector, predictor, states, actions, cfg):
    # states: [T + 1, batch, C, H, W], actions: [T, batch, 2]
    enc = backbone(states.flatten(0, 1)).view(states.shape[0], states.shape[1], -1)
    proj = projector(enc.flatten(0, 1)).view(enc.shape[0], enc.shape[1], -1)
 
    current = enc[0]
    pred_encs = predictor.predict_sequence(enc=current, actions=actions, h=None)
 
    repr_loss = 0.0
    pred_projs = []
    for t, pred in enumerate(pred_encs):
        pred_proj = projector(pred)
        pred_projs.append(pred_proj)
        target = proj[t + 1] if cfg.repr_loss_after_projector else enc[t + 1]
        source = pred_proj if cfg.repr_loss_after_projector else pred
        repr_loss = repr_loss + F.mse_loss(source, target)
    repr_loss = repr_loss / len(pred_encs)
 
    vc_targets = proj[1:] if cfg.skip_first_step_vc else proj
    std_loss, cov_loss = 0.0, 0.0
    for z in vc_targets:
        z = z - z.mean(dim=0)
        std = torch.sqrt(z.var(dim=0) + 1e-4)
        std_loss = std_loss + F.relu(1.0 - std).mean()
        cov = (z.T @ z) / (z.shape[0] - 1)
        cov_loss = cov_loss + off_diagonal(cov).pow(2).sum() / z.shape[1]
 
    return cfg.sim_coeff * repr_loss + cfg.std_coeff * std_loss + cfg.cov_coeff * cov_loss

3.3 SimCLR-JEPA objective

SimCLR adaptation 保持同一个 encoder/predictor 框架，但把 predicted embedding 和 target embedding 当成两种 view，使用 InfoNCE。论文中定义： $expsim (u, v) = exp (\frac{u ^{⊤} v}{τ ∥ u ∥∥ v ∥}) .$ 对于同一时间步的 $S_{t}$ 和 $\tilde{S}_{t}$ ，positive pair 是同一个 batch index 的 target/prediction，其它 index 是 negative。总体损失对两个方向求平均： $L_{SimCLR} = \frac{1}{2 T} \sum_{t = 2}^{T + 1} InfoNCE (S_{t}, \tilde{S}_{t}) + InfoNCE (\tilde{S}_{t}, S_{t}) .$ code 中 simclr.py:SimCLR.info_nce_simclr 的实现更接近标准 SimCLR：先 torch.cat([pred, encs], dim=0) 形成 $2 N$ 个 features，用相同 sample id 构造 label matrix，去掉主对角线后把 paired view 放进 positives，其余放入 negatives，再做 cross entropy。SimCLR.forward 对每个未来步 states_proj[i+1] 与 pred_proj 累加 info_nce_simclr，最后除以 $T - 1$ 。如果 normalize_z=true 才显式 normalize；实际 best config 示例中 normalize_z: false，温度 loss_temp 会随噪声设置调参，例如 fixed structured $α = 1$ 的 SimCLR config 使用 loss_temp: 0.25，changing structured $α = 1$ 使用 loss_temp: 16.0。

def simclr_jepa_step(backbone, projector, predictor, states, actions, cfg):
    enc = backbone(states.flatten(0, 1)).view(states.shape[0], states.shape[1], -1)
    proj = projector(enc.flatten(0, 1)).view(enc.shape[0], enc.shape[1], -1)
    pred_encs = predictor.predict_sequence(enc=enc[0], actions=actions, h=None)
 
    loss = 0.0
    for t, pred_enc in enumerate(pred_encs):
        pred_proj = projector(pred_enc)
        target_proj = proj[t + 1]
        features = torch.cat([pred_proj, target_proj], dim=0)
        if cfg.normalize_z:
            features = F.normalize(features, dim=-1)
        sim = features @ features.T / cfg.loss_temp
        labels = torch.arange(pred_proj.shape[0], device=pred_proj.device).repeat(2)
        positive_mask = labels[:, None].eq(labels[None, :])
        positive_mask.fill_diagonal_(False)
        logits, targets = build_simclr_logits(sim, positive_mask)
        loss = loss + F.cross_entropy(logits, targets)
    return loss / len(pred_encs)

3.4 Spurious correlation / slow-feature failure construction

论文最关键的理论解释是一个“看似正常但任务无关”的零损失解。假设输入中有 episode 内固定的高斯 distractor noise。编码器可以直接复制这个噪声： $g_{ϕ} (o_{t}) = s \sim N (0, σ^{2} I), s \in R^{D} .$ 因为噪声在同一个 episode 内不变， $g_{ϕ} (o_{t}) = g_{ϕ} (o_{t + 1}) = s$ 。如果 forward model 收敛到 identity， $f_{θ} (s, a) = s$ ，则 prediction loss 为零： $L_{prediction} = \frac{1}{TN} \sum_{t = 1}^{T} \sum_{i = 1}^{N} ∥ S_{t, i} - S_{t + 1, i} ∥_{2}^{2} = 0.$ 如果 $σ$ 足够大，variance loss 也为零： $Var (s_{t}) = \frac{1}{N - 1} \sum_{i = 1}^{N} (s_{i} - \overset{s}{ˉ})^{2} = σ^{2},$ $L_{variance} = \frac{1}{( T + 1 ) D} \sum_{t = 1}^{T + 1} \sum_{j = 1}^{D} max (0, γ - Var (S_{t, :, j}) + ϵ) = 0.$ 由于不同 episode 的噪声变量彼此独立，off-diagonal covariance 也可为零： $L_{covariance} = \frac{1}{( T + 1 ) ( N - 1 )} \sum_{t = 1}^{T + 1} \sum_{i = 1}^{D} \sum_{j = i + 1}^{D} (S_{t} S_{t}^{⊤})_{i, j} = 0.$ 这说明 VICReg 的 anti-collapse 只保证 batch 内表示有方差、维度不冗余，并不保证表示与 action-controlled state 有关。SimCLR 同样受影响：若 positive pairs 完全对齐，并且 encoder output 在单位球面上足够 uniform，InfoNCE 在 infinite negatives 极限下被最小化；固定噪声复制解满足“同一 episode 前后帧 aligned”和“跨 episode noise diverse”这两个条件。因此 JEPA 不是简单忽略所有背景，而是偏向最容易被未来预测的背景。

3.5 Dataset generation and probing protocol

Figure 4 解读：数据生成先采样初始位置 $c_{1}$ 与动作序列 $a_{1}, \dots, a_{T}$ ，再递推得到 $c_{1}, \dots, c_{T + 1}$ ，最后把连续坐标渲染成带 Gaussian blur 的小图像。这个设计让 ground-truth dot location 完全可知，但预训练不使用它；只有 probe 阶段用它检查 latent 是否保留状态。

单点环境的 square side length 为 1，dot 坐标 $c_{t} = (c_{t}^{x}, c_{t}^{y})$ ，每一维在 $[0, 1]$ 内。动作 $a_{t} = (a_{t}^{x}, a_{t}^{y})$ 是位置增量，范数受最大步长 $D$ 限制；论文实验中 $D = 0.14$ 。状态更新为 clipping 后的连续运动： $c_{t + 1} = max (0, min (c_{t} + a_{t}, 1)) .$ 动作方向不是独立 uniform，而是先采样一个 episode-level 方向 $ω \sim Uniform (0, 2 π)$ ，再用 Von Mises distribution 采样方向 $u_{t} \sim VonMises (ω, 1.3)$ ，最后乘以 $d_{t} \sim Uniform (0, D)$ 得到 $a_{t} = d_{t} u_{t}$ 。这样 dot 在一个 episode 内有连续运动趋势，不会因为完全随机方向而在期望上停在原地。渲染时对 dot 图像做 $σ = 0.05$ 的 Gaussian blur，分辨率为 $28 \times 28$ 。

Figure 2 解读：左侧展示 noise brightness coefficient $α$ 从 0 到 3 时，背景相对 dot 的强度变化；右侧展示 structured noise 的 changing/fixed 两种时间属性。Changing noise 每帧重新采样，因此对 forward prediction 是不可预测背景；fixed noise 在一个 episode 内保持不变，但跨 episode 重新采样，因此既能提供 batch diversity，又能被 identity predictor 轻松对齐。

released code 的 data/single.py:ContinuousMotionDataset 与论文描述一致：generate_actions 使用 torch.distributions.VonMises 采样方向，并用 self.max_step 缩放；generate_static_overlay 把同一张 random/CIFAR overlay repeat 到所有 time steps；generate_rnd_overlay 则每个 time step stack 新图。data/multiple.py:create_three_datasets 复用三个 ContinuousMotionDataset，第三个用 max_step=0 生成 stationary dot。

def generate_single_dot_episode(batch_size, n_steps=17, max_step=0.14, alpha=1.0, mode="fixed_structured"):
    direction_center = torch.rand(batch_size) * 2 * torch.pi
    von_mises = torch.distributions.VonMises(direction_center, torch.ones_like(direction_center) * 1.3)
    angles = von_mises.sample((n_steps - 1,)).T
    step_size = torch.rand(batch_size, n_steps - 1, 1) * max_step
    actions = torch.stack([torch.cos(angles), torch.sin(angles)], dim=-1) * step_size
 
    loc = sample_initial_location(batch_size)
    locations = [loc]
    for t in range(n_steps - 1):
        loc = torch.clamp(loc + actions[:, t], min=0.0, max=1.0)
        locations.append(loc)
    dot_images = gaussian_render(torch.stack(locations, dim=1), image_size=28, sigma=0.05)
 
    if mode.startswith("fixed"):
        overlay = sample_noise(batch_size, structured="structured" in mode).unsqueeze(1).repeat(1, n_steps, 1, 1, 1)
    else:
        overlay = torch.stack([sample_noise(batch_size, structured="structured" in mode) for _ in range(n_steps)], dim=1)
    return dot_images + alpha * overlay, torch.stack(locations, dim=1), actions

probe 的目标是训练一个线性层 $q$ ，从自回归预测出的 latent sequence 恢复 dot location： $L_{probing} = \frac{1}{TN} \sum_{t = 1}^{T} \sum_{i = 1}^{N} q (\tilde{S}_{t, i}) - C_{t, i}_{2}^{2} .$ code 中 probing.py:probe_pred_position 冻结 backbone/predictor（除非 full_finetune=true），只优化 models.Prober；evaluation 中用 dataset.unnormalize_mse 还原到原始坐标尺度，再报告 RMSE。这一点让 Figure 3 和 Table 1 的 RMSE 可以解释为位置恢复误差，而不是 normalized latent loss。

3.6 Baselines and released-code gaps

Reconstruction baseline 在论文正文中描述为 decoder MSE： $L_{reconstruction} = \frac{1}{T} \sum_{t = 2}^{T + 1} ∥ d_{ξ} (\tilde{s}_{t}) - o_{t} ∥_{2}^{2} .$ IDM baseline 加入 inverse dynamics head $h_{ξ} (s_{t}, s_{t + 1}) = \tilde{a}_{t}$ ： $L = L_{prediction} + L_{IDM}, L_{IDM} = \frac{1}{TN} \sum_{t = 1}^{T} \sum_{i = 1}^{N} ∥ h_{ξ} (S_{t, i}, S_{t + 1, i}) - A_{t, i} ∥.$ 论文公式与 released code 实现差异：论文 appendix A.3.4 写明 reconstruction implementation 不使用 latent variable，只是对称 decoder 做 MSE；但 rssm.py:RSSMPredMultistep.forward 在 main@03eb1730 中实际构造 Normal(prior_mu, prior_var)、Posterior、kl_divergence(z, post_z)，并把 KL 与 reconstruction MSE 一起放进 loss。

因此本文笔记把 paper 中的 baseline 称为 Reconstruction/RSSM-family，并在解释 code 时按 released code 的 stochastic prior/posterior 路径处理。另一个实现差异是数据量：论文 §3 写 pretraining 用 1,000,000 sequences、prober 用 300,000 sequences、evaluation 用 10,000 sequences；

released reproduce_configs/*.best.yaml 中 val_dataset_size: 10000、epochs: 100 与 train.py:init_dataset 的 generator 长度组合确实给预训练约 $100 \times 10, 000 = 1, 000, 000$ generated samples，但 probing.py:ProbingConfig.epochs=3 对同一个 10,000-size generator 只对应约 30,000 probe samples per validation call，未直接体现 300,000 这个数字。

def rssm_reconstruction_step(backbone, predictor, posterior, decoder, states, actions, cfg):
    enc = backbone(states.flatten(0, 1)).view(states.shape[0], states.shape[1], -1)
    belief = enc[0]
    prior_mu = predictor.prior_mu_net(belief)
    prior_var = F.softplus(predictor.prior_var_net(belief))
    z = torch.distributions.Normal(prior_mu, prior_var).sample()
    loss = F.mse_loss(decoder(z, belief), states[0])
 
    for t in range(states.shape[0] - 1):
        belief, prior_mu, prior_var = predictor(z, actions[t], belief)
        prior = torch.distributions.Normal(prior_mu, F.softplus(prior_var) + cfg.min_var)
        z_prior = prior.sample()
        post_mu, post_var = posterior(enc[t + 1], belief)
        post = torch.distributions.Normal(post_mu, F.softplus(post_var) + cfg.min_var)
        kl = torch.distributions.kl_divergence(prior, post).mean().clamp_min(3.0)
        recon = decoder(z_prior, belief)
        loss = loss + kl + F.mse_loss(recon, states[t + 1])
        z = post.sample()
    return loss / (states.shape[0] - 1)

def idm_step(backbone, predictor, action_head, states, actions, cfg):
    enc = backbone(states.flatten(0, 1)).view(states.shape[0], states.shape[1], -1)
    pred_encs = predictor.predict_sequence(enc=enc[0], actions=actions, h=None)
    prediction_loss, action_loss = 0.0, 0.0
    for t, pred in enumerate(pred_encs):
        prediction_loss = prediction_loss + F.mse_loss(pred, enc[t + 1])
        pred_action = action_head(enc[t], enc[t + 1])
        action_loss = action_loss + F.mse_loss(pred_action, actions[t])
    return prediction_loss / len(pred_encs) + cfg.action_coeff * action_loss / len(pred_encs)

Code reference: main @ 03eb1730 (2022-10-27) — pseudocode and mapping based on this commit

Paper Concept	Source File	Key Class/Function
CLI/config parsing and training loop	`train.py`, `configs.py`	`TrainConfig`, `Trainer.init_dataset`, `Trainer.train`, `Trainer.validate`
Single moving-dot dataset with fixed/changing noise	`data/single.py`	`ContinuousMotionDataset.generate_actions`, `generate_static_overlay`, `generate_rnd_overlay`, `generate_multistep_sample`
Three-dot diagnostic dataset	`data/multiple.py`	`create_three_datasets`, `MultiDotDataset.__getitem__`
VICReg-JEPA	`vicreg.py`, `models.py`	`VICRegConfig`, `VICRegPredMultistep.forward`, `get_cov_std_loss`, `RNNPredictor.predict_sequence`
SimCLR-JEPA	`simclr.py`, `models.py`	`SimCLR.info_nce_simclr`, `SimCLR.forward`, `Projector`
Reconstruction/RSSM-family baseline	`rssm.py`, `models.py`	`RSSMPredMultistep.forward`, `RSSMPredictor`, `Posterior`, `MeNet5Decoder`
Linear probing	`probing.py`, `models.py`	`probe_pred_position`, `probe_enc_position`, `Prober`, `location_losses`
Best hyperparameter configs	`reproduce_configs/*.best.yaml`	per-noise/per-method YAML files such as `sweep_fixed_structured.(1).vicreg.best.yaml`

4. Experimental Setup (实验设置)

4.1 Datasets and scale

主实验是 single moving dot dataset。每个 sequence 长度为 17 frames，对应 16 actions；图像分辨率 $28 \times 28$ ，dot 是 Gaussian-blurred point，位置坐标在 unit square 中。训练世界模型时论文报告使用 1,000,000 generated sequences；训练 linear probe 时论文报告使用 300,000 sequences；evaluation 使用 10,000 sequences。背景 distractor 有两类：uniform random noise，以及从 CIFAR-10 裁剪/overlay 得到的 structured noise。每类 noise 又分为 changing 与 fixed：changing 每帧重采样，fixed 在 episode 内不变但跨 sequence 重采样。亮度系数 $α$ 控制 noise peak 相对 dot peak 的强度，Figure 2 展示了 $α = 0, 1, 2, 3$ 的视觉效果。

补充实验是 3-dots dataset：图像有三个 channel，红色 channel 是 action-controlled dot，绿色 channel 是 uncontrollable random moving dot，蓝色 channel 是 stationary dot。目标不是只恢复 agent dot，而是学习能捕获 action-controlled 与 stationary dot、忽略 random dot 的 representation。这个 dataset 用来说明 IDM 的局限：它只会保留能预测 action 的因素，对 stationary 但可观测的信息不敏感。

4.2 Compared baselines and metrics

比较方法包括 VICReg-JEPA、SimCLR-JEPA、Reconstruction/RSSM-family、Inverse Dynamics Modeling (IDM)、Supervised、Random、Center。Supervised 是 end-to-end 传播 probe error 到 encoder/predictor 的 lower-bound error；Random 固定随机 encoder/predictor，只训练 probe；Center 永远预测 dot 在中心位置，是上界式的 naive baseline。指标是 dot location RMSE，Figure 3/5 中 dots 表示 17 time steps 平均 RMSE，shaded area 是 3 random seeds 的 standard deviation；Table 1 报告 3-dots dataset 上 Average / Action / Random / Stationary 四个 RMSE。

4.3 Training configs and hardware

实验口径：best configs 而非默认值

训练细节不能只看 code default，因为每个 noise type/level 和每个 method 都单独调参。released repo 中 reproduce_configs/ 提供 144 个 *.best.yaml。所有这些 best configs 的 epochs 都是 100；n_steps 与 val_n_steps 都是 17；val_dataset_size 是 10,000；dataset_normalize: true。由于 train.py:init_dataset 使用 val_dataset_size // dataset_batch_size 作为 generator 长度，100 epochs 对应约 1,000,000 generated pretraining samples。

VICReg config 范围

实际 config 范围来自 reproduce_configs/*.best.yaml：VICReg configs 共 28 个，dataset_batch_size 在 32/64/128/256/512 中变化，embedding_size=512，base_lr 范围为 4.452906512659173e-05 到 3.783005893585235；

SimCLR config 范围

SimCLR configs 共 56 个，batch size 为 64/128/256/512，embedding_size 在 2/32/64/128/256/512/1024 中变化，base_lr 范围为 0.00025175077582298053 到 3.892307395784348；RSSM configs 共 32 个，batch size 为 64/128/256/512，embedding_size=512，learning_rate 范围为 7.289764181769937e-05 到 0.003673679181439055；

IDM-like config 范围

IDM-like configs 在 code 中表现为 model_type: VICReg 且 action_coeff>0、sim_coeff=std_coeff=cov_coeff=0，共 28 个，base_lr 范围为 0.12527133826086445 到 9.561615845868983。

具体 config 样例

几个具体 config 例子：reproduce_configs/sweep_changing_structured.(1).vicreg.best.yaml 使用 dataset_noise: 1.0、dataset_structured_noise: true、dataset_batch_size: 256、vicreg.projector: id、vicreg.predictor: rnn、vicreg.arch: menet5、vicreg.embedding_size: 512、vicreg.base_lr: 0.84878357903525、sim_coeff: 4.692709866259136、std_coeff: 1.332571312091042、cov_coeff: 1.0。

sweep_fixed_structured.(1).vicreg.best.yaml 改为 dataset_static_noise: 1.0、dataset_batch_size: 64、base_lr: 0.18902825366389228、sim_coeff: 0.10538301983832848、std_coeff: 0.9930085150161816。

sweep_changing_structured.(1).simclr.best.yaml 使用 simclr.projector: id、simclr.predictor: rnn、embedding_size: 512、base_lr: 0.3039359882655922、loss_temp: 16.0；而 sweep_fixed_structured.(1).simclr.best.yaml 使用 dataset_batch_size: 128、base_lr: 0.008377498284938412、loss_temp: 0.25。

sweep_changing_uniform.(1).rssm.best.yaml 使用 rssm.arch: menet5、rssm.embedding_size: 512、rssm.learning_rate: 0.0003205588547535079、rssm.min_var: 0.05。

模型结构与硬件环境

模型结构在 appendix 与 code 中基本一致：encoder 是 3 层 convolution + ReLU + BatchNorm，末尾 average pooling 和 linear layer；released best configs 多用 models.MeNet5，其 conv channels 为 16*width_factor -> 32*width_factor -> 32*width_factor，再接 nn.Linear(9*32*width_factor, output_dim)。predictor 是单层 GRU，action input dimension 为 2，hidden/embedding 通常为 512。论文报告所有实验在 AMD MI50 或 Nvidia RTX 8000 GPUs 上运行；每个 noise type/level 运行 100 个随机 hyperparameter settings，选最佳后再跑 3 seeds；单个 experiment 少于 1 小时 GPU time。

5. Experimental Results (实验结果)

5.1 Main single-dot results

Figure 3 解读：这是每个 method、noise type、noise level 都单独调参后的主结果。没有 distractor 时所有方法都能恢复 dot location；当 noise 每帧变化时，VICReg-JEPA 与 SimCLR-JEPA 能保持较低 RMSE，说明 latent prediction 确实能过滤不可预测背景；但当 noise 在 episode 内固定时，两种 JEPA 都迅速退化，且 structured/fixed 和 uniform/fixed 都出现类似失败。Reconstruction 在 $α \leq 1.5$ 的各设置中仍较稳健，IDM 在单点任务中几乎总是很好，因为 dot movement 正好是 action 解释对象。

主图没有在 paper text 中提供每个曲线点的数值表，因此不在笔记中臆造具体坐标；可确切陈述的是 caption 指定的统计方式：dots 是 17 steps mean RMSE，shaded area 是 3 random seeds standard deviation。论文正文给出的定性结论是：all methods 在无噪声时表现好；Reconstruction 在 $α \leq 1.5$ 表现好；JEPA-based methods 在 fixed structured/uniform noise 下失败；所有方法在 changing noise 下表现好；JEPA methods 在 changing noise 上不太需要重新调参，而 reconstruction 若不调参会明显变差。

5.2 Untuned hyperparameter robustness

Figure 5 解读：这里固定使用 no-noise setting 选出的 hyperparameters，不再针对每种 noise level/type 单独调参。对比 Figure 3 可以看到，VICReg/SimCLR 在 fixed noise 下依旧失败，并且在更低 noise level 就可能退化；但在 changing noise 下，JEPA 仍相对稳健。Reconstruction 的相反问题更明显：不调参时它在 changing uniform noise 上明显恶化，说明 pixel reconstruction 对背景强度与分布更敏感。

这个结果支持论文的核心区分：JEPA 的优势不是“所有噪声都好”，而是对不可预测的、每帧变化的 distractor 更自然地忽略；它的缺陷也不是普通过拟合，而是 objective 对 episode-level stable features 的偏好。调参可以改变退化发生的 noise level，但不能消除 fixed slow-feature 的基本风险。

5.3 Three-dot diagnostic results

Figure 6 解读：红点是 action-controlled dot，绿色点随机运动但动作未知，蓝点 stationary。这个环境把“可控动态信息”和“慢但有用的信息”分开：好的通用 representation 应该保留红点和蓝点，忽略不可预测绿点；只看 inverse dynamics 的 representation 会偏向红点，只看 slow feature 的 JEPA 会偏向蓝点。

Table 1 的精确 RMSE 如下，所有数值均为 17 steps 上的 RMSE，3 seeds 给出 standard deviation：

Method	Average	Action	Random	Stationary
VICReg	0.229±0.031	0.277±0.041	0.273±0.044	0.066±0.026
SimCLR	0.158±0.001	0.193±0.001	0.193±0.002	0.025±0.001
IDM	0.234±0.001	0.035±0.000	0.298±0.002	0.272±0.000
Supervised	0.104±0.000	0.010±0.001	0.180±0.001	0.005±0.000
Reconstruction	0.107±0.000	0.021±0.001	0.182±0.000	0.026±0.002
Random	0.260±0.001	0.235±0.002	0.278±0.002	0.265±0.004
Center	0.299±0.000	0.304±0.000	0.304±0.000	0.289±0.000

这个表给出比主图更清楚的分解。VICReg 和 SimCLR 对 stationary dot 的 RMSE 很低，分别为 0.066±0.026 和 0.025±0.001，但 action dot 分别高到 0.277±0.041 和 0.193±0.001，说明它们把表示容量投向最慢的蓝点。IDM 的 action dot RMSE 是 0.035±0.000，非常好，但 stationary dot 是 0.272±0.000，几乎忽略了与动作无关的慢特征。Reconstruction 同时保留 action dot (0.021±0.001) 和 stationary dot (0.026±0.002)，Average 0.107±0.000 接近 Supervised 0.104±0.000；不过 random dot 因为动作未知，本来就难以预测，Supervised 与 Reconstruction 也分别只有 0.180±0.001 和 0.182±0.000。

5.4 Limitations and takeaways

作者明确承认两类限制。第一，实验是 toy dataset，只有 moving dot 和人工构造的 distractor；虽然它精准暴露 objective bias，但不能直接证明复杂视频、大模型、真实 control setting 中失败程度相同。第二，JEPA objectives 只测试了 VICReg 和 SimCLR，未覆盖 BYOL、Barlow Twins、HJEPA 或后续 V-JEPA/I-JEPA 风格的 masking/predictor choices。因此这篇 paper 的结论应理解为“JEPA family 存在 slow-feature risk 的反例与机制”，而不是对所有 joint-embedding 方法的最终否定。

总体结论是：JEPA 的 reconstruction-free latent prediction 确实能在 changing distractors 下比 pixel reconstruction 更自然地忽略不可预测噪声；但同一个 objective 会偏好固定、缓慢、易预测的背景特征。对 world model 来说，预测性本身不是充分条件，因为最可预测的因素可能与 control/reward 无关。若要把 JEPA 用于复杂 agents，必须额外约束 representation：例如引入 hierarchy 来同时表示 fast/slow features，加入 object/action relevance，使用 image differences or optical flow 降低静态背景吸引力，或采用能够区分 controllable state 与 nuisance slow variables 的 objective。

Paper Notes

探索

Joint Embedding Predictive Architectures Focus on Slow Features

Joint Embedding Predictive Architectures Focus on Slow Features

1. Motivation (研究动机)

2. Idea (核心思想)

3. Method (方法)

3.1 Overall framework: offline latent dynamics + probing

3.2 VICReg-JEPA objective

3.3 SimCLR-JEPA objective

3.4 Spurious correlation / slow-feature failure construction

3.5 Dataset generation and probing protocol

3.6 Baselines and released-code gaps

4. Experimental Setup (实验设置)

4.1 Datasets and scale

4.2 Compared baselines and metrics

4.3 Training configs and hardware

实验口径：best configs 而非默认值

VICReg config 范围

SimCLR config 范围

IDM-like config 范围

具体 config 样例

模型结构与硬件环境

5. Experimental Results (实验结果)

5.1 Main single-dot results

5.2 Untuned hyperparameter robustness

5.3 Three-dot diagnostic results

5.4 Limitations and takeaways

目录