Learning and Leveraging World Models in Visual Representation Learning

Paper: arXiv:2403.00504

1. Motivation (研究动机)

这篇论文要回答的问题不是“再做一个 ImageNet self-supervised pretraining 方法”，而是：视觉表征学习里已经存在的 encoder-predictor 结构，到底能不能被解释为、并进一步训练成一个可复用的 world model。在 RL 中，world model 的典型角色是预测一个 action 之后环境状态如何变化，并在下游规划或控制中继续使用；但在视觉 SSL 中，MAE/I-JEPA/BYOL/SimSiam 等方法虽然都有某种“预测变换后输入/表征”的模块，训练结束后通常只保留 encoder，predictor/decoder 被丢掉。作者认为这很浪费：如果 predictor 真的学会了在 latent space 里执行“变换”，它就应该不仅能帮助预训练，还能作为下游任务的可微、可调适模块。

Figure 2 解读：这张图给出本文的概念定位。Generative World Models 在像素空间预测/反演变换，Joint Embedding 方法通常没有显式可用的 world model，只让不同 view 的 embedding 接近；JEPA World Models 则在 latent space 里训练 predictor，使它根据 action-like transformation 参数和 mask 位置预测目标 embedding。IWM 位于 JEPA World Models 这一格：它仍然用 latent prediction 避免像素重建的低级细节负担，但把 predictor 明确当成可复用世界模型，而不是一次性 pretraining head。

现有方法的瓶颈有三层。第一，经典 augmentation-invariant SSL 的目标是让不同增强 view 的表示相同，容易得到“可线性读出”的抽象语义，但这种 invariant representation 往往主动丢掉颜色、纹理、局部细节等信息；这对线性分类很好，却未必适合需要 richer representation 的任务。第二，Masked Image Modeling / I-JEPA 通过 mask prediction 学习，但原始任务主要是补局部缺失，动作空间几乎只包含“哪些 patch 被遮住”，不足以检验模型是否理解更一般的 photometric transformation。第三，很多视觉 SSL 的 predictor/decoder 没有被设计成下游可调适对象：即便它在训练中学到了一些结构，也没有明确机制让它在分类、分割、多任务设置中继续发挥作用。

本文的具体目标是提出 Image World Model (IWM)：在 JEPA 框架内，把图像增强视为 action，将 source view 到 target view 的变换参数 $a_{x\to y}$ 显式喂给 predictor，让 predictor 在 latent space 中预测目标表征。这个目标值得研究，因为一旦成立，它会把视觉 SSL 从“只学 encoder”推进到“同时学 encoder 与可复用 world model”：encoder 提供状态表征，predictor 学会如何根据 action 改变状态，并可通过 predictor finetuning 低成本迁移到分类、分割、多任务预测。

更深一层的动机是控制表征的抽象层级。作者发现 world model 的能力强弱与表征是 invariant 还是 equivariant 直接相关：弱 predictor 不能反演增强，只能迫使 encoder 丢弃不可预测信息，得到更抽象、更适合线性 probe 的表示；强 predictor 能在 latent space 中执行增强，使 encoder 保留更多输入信息，虽然线性读出较难，但在 predictor finetuning 或复杂 head 下峰值性能更好。这使 IWM 不只是一个新模型，而是一种调节 representation abstraction 的工具。

2. Idea (核心思想)

核心 insight：把图像增强从“制造两个应当相同的 view”重新看成一种 action，并训练一个 transformation-conditioned JEPA predictor 去预测 action 后的 latent state；这个 predictor 如果足够强，就不应在预训练后被丢弃，而应作为下游任务可 finetune 的 world model。换言之，IWM 的新意不在于多加一个 augmentation，而在于把 predictor 的条件输入、容量、预测难度系统化，使它真的学会在表征空间里执行变换。

与 I-JEPA 的关键差异是 conditioning 与任务范围：I-JEPA 主要预测 masked target patch 的 latent representation，mask 位置是核心几何条件；IWM 仍继承 latent inpainting，但额外预测全局 photometric transformations（color jitter、blur、grayscale、solarization 等）对表征的影响，并把从 source 到 target 的逆变换参数 $a_{x\to y}$ 输入 predictor。与 BYOL/SimSiam/MoCo 这类 invariant SSL 相比，IWM 不要求增强 view 全部压成同一个表征；它允许 predictor 负责解释增强差异，从而让 encoder 在需要时保留更多信息。

这篇文章的创新可以压缩成三个判断。第一，没有 transformation conditioning 时，world model 会退化成传统 invariant SSL，MRR 接近 0；加入 sequence/feature conditioning 后，MRR 可以达到 0.82/0.79。第二，预测任务必须足够难：只做弱 color jitter 不足以迫使 predictor 学世界模型，加入 destructive augmentations 或 stronger jitter 后 MRR 大幅上升。第三，容量必须匹配任务难度：${\text{IWM}}_{18,384}^{\text{Equi}}$ 在默认增强下 MRR 为 0.79，而 ${\text{IWM}}_{12,384}$ 在同设置下只有 0.09。

3. Method (方法)

3.1 总体框架：把增强看成 action 的 JEPA world model

IWM 从一张原图 $I$ 生成 source view $x$ 和 target view $y$。target 经过随机 crop、horizontal flip 和 color jitter，但不施加强破坏性增强；source 从 target 继续施加额外 color jitter、grayscale、Gaussian blur、solarization，并按 I-JEPA 风格随机遮掉若干 patch。这样 $x$ 是信息被进一步扰动和遮挡的上下文，$y$ 是 predictor 要在 latent space 中恢复/预测的目标。

Figure 1 解读：这张架构图展示 IWM 的完整训练路径。source view 被 mask 后进入 online encoder $f_{\theta }$，target view 进入 EMA teacher encoder $f_{\theta }^{\text{EMA}}$；predictor $p_{\varphi }$ 同时接收 source embeddings、待预测位置的 mask tokens，以及从 source 到 target 的 transformation 参数 $a_{x\to y}$，最终输出 masked target positions 的 latent predictions。关键点是 target representation 来自 EMA encoder，用来避免 collapse；predictor 的输入不只是“哪里缺失”，还包括“做了什么 photometric action”。

设 $z_x=f_{\theta }(x)$，$z_y=f_{\theta }^{\text{EMA}}(y)$。source mask 为 $M_x$，待预测位置为 $M_x^C$，每个待预测位置对应一个包含几何位置的 mask token $m_a$。predictor 的目标是 ${\hat{z}}_y=p_{\varphi }\left(f_{\theta }(x),a_{x\to y},m_a\right),$ 训练损失是 masked target positions 上的 squared $L_2$ distance： $L(x,y)={\sum }_{i\in M_x^C}{\Vert p_{\varphi }{\left(f_{\theta }(x),a_{x\to y},m_a\right)}_i-f_{\theta }^{\text{EMA}}(y{)}_i\Vert }_2^2.$ 直觉上，encoder 不再被迫把所有增强后的 view 都压成同一个点。如果 predictor 能根据 $a_{x\to y}$ 恢复 target latent，那么颜色、模糊、灰度化等信息可以在 representation 中保留下来；如果 predictor 太弱或没有 action 条件，encoder 才会倾向于删掉这些难以预测的信息，形成 invariant representation。因此 IWM 的“world model 能力”和“representation abstraction level”不是两个独立指标，而是一枚硬币的两面。

3.2 Source / target / action 的构造

论文中的 target $y$ 从原图经过随机水平翻转、crop 与 color jitter 得到。source $x$ 在 target 的基础上进一步增强：额外 color jitter、grayscale、blur、solarization，并随机 mask patch。action $a_{x\to y}$ 不是外部环境动作，而是从 source 还原到 target 的 transformation parameters：包括 color jitter 差异，以及 destructive augmentation 是否被应用。作者有意不在 target 上施加 destructive augmentations，因为 target 应尽量保留可预测信息；destructive augmentations 只放在 source 侧，使任务变成“从受损 view 预测较完整 target latent”。

附录给出默认增强细节：source/target 先各自随机 crop，scale 在 $[0.3,1]$，horizontal flip 概率 0.5。target color jitter 概率 0.8，brightness/contrast/hue/saturation 最大强度分别为 0.4/0.4/0.1/0.2。source 使用同样 color jitter，并额外以 0.2 概率施加 Gaussian blur（半径 0.1 到 2）、solarization、grayscale。mask $M_x$ 是 4 个矩形 mask 的 union，每个面积在图像的 0.15 到 0.2 之间，aspect ratio 在 0.75 到 1.5 之间。

3.3 Predictor conditioning：Sequence vs Feature

IWM 的第一个关键组件是让 predictor 真正知道 action。作者比较两种方式：Sequence conditioning 把 transformation tokens 加入 predictor 输入序列，但为了让 transformer 能区分不同 token，需要用 token-specific linear layers 破坏普通 permutation equivariance；Feature conditioning 则把 transformation scalars 与 mask token 在 feature dimension 上拼接，再经 1x1/fully-connected mixing 网络映射回 predictor 维度。论文实际采用 Feature conditioning，因为它下游表现更好。

Conditioning	None	Sequence	Feature
MRR	0.00	0.82	0.79

这个表说明，没有 conditioning 时，predictor 几乎完全不能在 latent space 中执行变换，MRR 为 0.00；而 sequence 和 feature 两种 conditioning 都能让 world model 学会对应关系。Feature 的 MRR 比 Sequence 略低（0.79 vs 0.82），但由于下游性能更优，成为默认设置。

3.4 预测难度与 predictor 容量

作者将 world model 质量用 MRR 衡量：生成 256 个 augmented target images 的表示库，让 predictor 从 clean image representation 加 transformation 参数去预测某个 target 表示，再看真正 target 在 nearest-neighbor rank 中排第几。若 $r_n$ 是第 $n$ 个样本的目标 rank，则可理解为 $\text{MRR}=\frac{1}{N}{\sum }_{n=1}^N\frac{1}{r_n}.$ MRR 接近 1 表示 predictor 能准确“执行 action”；接近 0 表示它只学到了 invariant 或无效映射。

Predictor / setting	I-JEPA (12,384)	IWM (12,384)	IWM (18,384)
Jitter	0.00	0.11	0.25
+ Destructive	0.00	0.09	0.79
+ Strong Jitter	0.00	0.81	0.85

这个 ablation 显示，难度和容量必须同时匹配。I-JEPA 在三种设置下 MRR 都是 0.00，因为它没有显式学习 photometric action；IWM 12-layer 在默认 destructive setting 下只有 0.09，但遇到 strong jitter 时可达 0.81；IWM 18-layer 在默认 destructive setting 下达到 0.79，在 strong jitter 下达到 0.85。作者还报告：12-layer predictor 在 jitter equivariance 上 5 次只有 1 次成功，18-layer 有 4 次成功。这说明等变 world model 不是自动出现的，它需要足够难的动作空间和足够大的 predictor。

Figure 1 解读：这张可视化把 predictor 的 latent prediction 放到 256 个 augmented images 的表示库里做 nearest-neighbor retrieval。若 predictor 学会了 action，检索到的图像应接近施加对应变换后的 target。图中显示 IWM 能较好执行颜色/模糊/恢复等 transformation，但在反转 grayscale 时会有误差，因为灰度化本身丢失颜色信息，不是完全可逆动作。

3.5 Leveraging：Predictor finetuning

视觉 world model 的预训练任务（colorize、undo blur、inpaint）与下游分类/分割不同，因此 predictor 不能直接输出分类结果。作者的做法是冻结 encoder，把 predictor 接在 EMA teacher encoder 之上，并让它在下游任务中重新训练；分类时在 predictor output 上加 attentive head，分割时接 UperNet 等 segmentation head。predictor finetuning 的核心假设是：如果 predictor 已经学会理解 latent space 和 action-conditioned transitions，那么它比随机初始化的大 head 更容易适配下游任务。

Predictor finetuning task	ImageNet Top-1	Gap
Default	82.9	-
+ Teacher	83.2	+0.3
+ Null latents	83.3	+0.1
+ Pred only one token	82.8	-0.5

默认下游 predictor finetuning 会使用 teacher network，预测 whole target image，并使用 null transformation parameters。表中最重要的结论不是某一个 trick，而是任务定义相当鲁棒：teacher 稍有帮助，null latents 更灵活且略好，只预测一个 token 约掉 0.5 点。也就是说，下游阶段不必复刻预训练的每个 photometric action，关键是让 pretrained predictor 继续处理 frozen encoder 的 latent space。

Figure 3 解读：这张图比较 finetuned parameters 与 ImageNet Top-1。结论是，当可训练参数数量相近时，IWM 的 predictor finetuning 比 encoder finetuning 更有效；论文文字中给出对比：predictor finetuning 相比 MAE encoder finetuning 约高 1 point，相比 IWM encoder finetuning 约高 1.5 points。它说明 IWM 的 world model 不是只在预训练时有用，而能作为一个参数效率更高的下游适配器。

3.6 Multitask predictor tuning

作者进一步把 predictor finetuning 扩展到多任务：为每个任务添加 learned task token，并接 task-specific head/loss；训练 batch 在任务间均匀采样，所有 loss 相加后更新共享 predictor 与任务 head。这类似 LLM 的 instruction tuning / prefix tuning：同一个 pretrained predictor 被不同 task token 指示执行不同判别任务。

Figure 7 解读：多任务图中，每个 batch 被均匀分配到多个任务，输入 predictor 前加入 task token；predictor 共享，但每个任务有自己的 head 和 loss。这样做的意义是摊销 predictor 的推理/存储成本：训练四个 single-task predictors 会得到四个模型，而 multitask predictor 只保留一个共享世界模型适配器。

多任务实验使用 ${\text{IWM}}_{18,384}^{\text{Equi}}$，任务包括 ImageNet、iNaturalist18、SUN397、Places205。结果为：ImageNet 79.6 vs single-task 80.8（-1.2），iNat18 72.0 vs 72.4（-0.4），SUN397 78.2 vs 75.6（+2.6），Places205 64.1 vs 64.8（-0.7），平均 73.5 vs 73.4（+0.1）。这说明一个共享 predictor 能近似替代多个单任务 predictor，且 SUN397 甚至明显提升。

3.7 表征抽象层级：Invariant 与 Equivariant 的可控谱系

IWM 的一个重要解释是：world model 能力控制 representation abstraction。若 predictor 近似 identity 或无力执行动作，encoder 只能保留 source 和 target 共同的信息，即更高层语义；这类似 contrastive methods 的 invariant representation。若 predictor 有足够容量执行动作，encoder 可以保留更多低/中层信息，让 predictor 负责变换，得到 equivariant representation。

Figure 4 解读：横轴可理解为 world model equivariance 程度，纵向展示不同评估协议的表现关系。线性评估偏好更 invariant、更抽象的 representation；predictor finetuning 偏好更 equivariant、更丰富的 representation；attentive probing 与 equivariance 的相关性较弱。图的核心含义是“好 representation”不是单一标量，不同下游协议读取的是不同层级的信息。

Figure 5 解读：这张图把 contrastive、MIM、IWM 放到 abstraction spectrum 上。Contrastive 位于高抽象端，信息容易线性读出但峰值受限；MIM 位于信息丰富端，线性 probe 较弱但复杂 finetuning 强；IWM 通过调节 world model 能力在两者之间移动，${\text{IWM}}^{\text{Inv}}$ 更像 contrastive，${\text{IWM}}^{\text{Equi}}$ 更像 MIM/可复用世界模型。

附录还用 latent averaging 检验另一种 invariance：给定随机增强集合 $A$，用 predictor 对多个增强做边际化，形成 $z_x^{\text{Inv}}=\frac{1}{N}{\sum }_{i=1}^N{p}_{\varphi }\left(f_{\theta }(x),A_i,m_{A_i}\right).$ 这个 representation 的最近邻确实接近 clean image，但表格显示 $N=1,8,16,32,128$ 时 ImageNet linear accuracy 分别为 64.5、64.3、64.6、64.6、64.4，基本没有提升。作者据此强调：contrastive invariance 有效，不只是因为“不同增强表示相同”，还因为 encoder 通过删除增强相关信息形成更语义化的抽象；保留全部信息后再做平均，不一定带来同样收益。

Figure 9 解读：这张图展示用 256 个 latent-space augmentations 边际化得到的 invariant representation。最近邻首先是原始未增强图像，其次是轻微变换图像，说明 predictor averaging 可以产生“clean-like” representation；但后续线性评估没有提升，说明 clean-like 不等于 task-useful abstraction。

Figure 8a–8c 解读：三联图分别对应 ${\text{IWM}}_{12,384}^{\text{Inv}}$、${\text{IWM}}_{18,384}^{\text{Equi}}$ 与 I-JEPA。每张图对同一批图像各做 16 次 augmentation，然后计算 embedding similarity matrix；黄色块表示同一原图的不同增强高度相似。Invariant IWM 的对角块最集中，说明增强被压到相近表示；Equivariant IWM 与 I-JEPA 的块内变化更大，说明 augmentation 信息仍留在 representation 里。

3.8 论文级伪代码与代码搜索状态

代码搜索未找到开源实现。已检查 arXiv 页面、GitHub/web 搜索（title、arXiv ID、Image World Models/IWM/JEPA、作者相关关键词）、GitHub code search、CatalyzeX/Papers with Code 入口；未找到官方或可信 released implementation。因此下面伪代码是根据论文公式与 appendix protocol 写出的 paper-level PyTorch-style sketch，不是 released code 对齐版本；代码版本锚点省略，论文公式与 released code 实现差异 不适用。

import torch
import torch.nn.functional as F
 
 
def make_iwm_views(image, aug, mask_sampler):
    # Paper-level sketch: no released implementation found.
    base = aug.random_resized_crop(image, scale=(0.3, 1.0))
    base = aug.horizontal_flip(base, p=0.5)
 
    target = aug.color_jitter(base, p=0.8,
                              brightness=0.4, contrast=0.4,
                              saturation=0.2, hue=0.1)
 
    source = aug.color_jitter(base, p=0.8,
                              brightness=0.4, contrast=0.4,
                              saturation=0.2, hue=0.1)
    source = aug.gaussian_blur(source, p=0.2, radius=(0.1, 2.0))
    source = aug.solarize(source, p=0.2)
    source = aug.grayscale(source, p=0.2)
 
    mask = mask_sampler(num_blocks=4, area=(0.15, 0.20), aspect=(0.75, 1.5))
    action = aug.inverse_transform_params(source, target)
    return source, target, mask, action

 
def iwm_pretrain_step(encoder, ema_encoder, predictor, batch, optimizer, aug, mask_sampler):
    sources, targets, masks, actions = [], [], [], []
    for image in batch:
        x, y, mask, a_xy = make_iwm_views(image, aug, mask_sampler)
        sources.append(x); targets.append(y); masks.append(mask); actions.append(a_xy)
 
    z_x = encoder(torch.stack(sources))
    with torch.no_grad():
        z_y = ema_encoder(torch.stack(targets))
 
    # mask tokens contain target positions; feature conditioning mixes action params.
    mask_tokens = predictor.make_conditioned_mask_tokens(masks, actions)
    pred = predictor(z_x, actions, mask_tokens)
    loss = F.mse_loss(pred[masks.complement()], z_y[masks.complement()])
 
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()
    update_ema(ema_encoder, encoder)
    return loss

 
def predictor_finetune_step(frozen_teacher, predictor, attentive_head, batch, labels, optimizer):
    with torch.no_grad():
        z = frozen_teacher(batch)
    null_actions = torch.zeros(len(batch), predictor.action_dim, device=batch.device)
    full_target_tokens = predictor.all_patch_tokens(batch.shape[0], device=batch.device)
    h = predictor(z, null_actions, full_target_tokens)
    logits = attentive_head(h)
    loss = F.cross_entropy(logits, labels)
    optimizer.zero_grad(set_to_none=True)
    loss.backward()
    optimizer.step()
    return loss

 
def multitask_predictor_step(frozen_teacher, predictor, task_heads, task_batches, optimizer):
    total_loss = 0.0
    for task_name, (images, labels) in task_batches.items():
        with torch.no_grad():
            z = frozen_teacher(images)
        task_token = predictor.task_token(task_name).expand(images.size(0), -1)
        h = predictor(z, action=None, mask_tokens=None, task_token=task_token)
        logits = task_heads[task_name](h)
        total_loss = total_loss + F.cross_entropy(logits, labels)
    optimizer.zero_grad(set_to_none=True)
    total_loss.backward()
    optimizer.step()
    return total_loss

 
def compute_mrr(predictor, encoder, clean_images, augmented_bank, actions):
    ranks = []
    z_clean = encoder(clean_images)
    z_bank = encoder(augmented_bank)  # shape: [B, 256, tokens, dim]
    for i in range(clean_images.size(0)):
        pred = predictor(z_clean[i:i+1], actions[i], mask_tokens=None).flatten()
        bank = z_bank[i].flatten(start_dim=1)
        distances = torch.cdist(pred[None, :], bank)[0]
        rank = torch.argsort(distances).tolist().index(0) + 1
        ranks.append(1.0 / rank)
    return torch.tensor(ranks).mean()

Paper concept	Released source file/class	Status
IWM pretraining data pipeline	无	代码搜索未找到开源实现；依据论文 Method 与 Appendix 写出 paper-level sketch
Feature conditioning predictor	无	无 released code；训练细节来自论文 appendix
Predictor finetuning / multitask tuning	无	无 released code；数值来自论文 tables
Paper-vs-code discrepancy	N/A	没有 released code，无法比较实现差异

4. Experimental Setup (实验设置)

数据集与规模

论文核心预训练与主评估在 ImageNet-1k 上完成。作者没有在正文逐一列出所有数据集样本数；下面规模按这些 benchmark 的标准设置标注，具体训练 epoch/协议按论文报告：ImageNet-1k 约 1.28M train / 50k val、1000 类；ADE20K semantic segmentation 标准设置约 20,210 train / 2,000 val、150 类；iNaturalist18 约 437k train / 24k val、8142 类；SUN397 约 108k images、397 场景类；Places205 约 2.45M train images、205 类。论文在 OOD attentive probing 中使用 iNat18、SUN397、Places205，在分割中使用 ADE20K。

Baselines

主要 baseline 包括 MAE、I-JEPA、MoCoV3、DINO、iBOT，以及 IWM 的不同变体：${\text{IWM}}_{12,384}^{\text{Inv}}$、${\text{IWM}}_{12,384}^{\text{Equi}}$、${\text{IWM}}_{18,384}^{\text{Equi}}$；大模型附录还比较 ViT-L/16 下的 I-JEPA、${\text{IWM}}_{18,384}^{\text{Inv}}$、${\text{IWM}}_{36,512}^{\text{Equi}}$。

Evaluation metrics

• MRR：world model 质量指标。对每个样本生成 256 个 augmented targets 的 representation bank，让 predictor 预测目标 latent，按 nearest-neighbor rank 计算 reciprocal rank 后平均；越接近 1，说明 predictor 越能执行 transformation。
• ImageNet Top-1：分类 Top-1 accuracy，分别在 linear probing、attentive probing、encoder finetuning、predictor finetuning、end-to-end finetuning 下报告。
• ADE20K mIoU：semantic segmentation validation mean IoU，使用 MMSegmentation + UperNet。
• OOD attentive probing accuracy：在 iNat18、SUN397、Places205 等数据集上用 attentive probing 评估 representation transfer。

Training config

论文没有报告 GPU 型号/数量或总 wall-clock；因此 hardware 记为“论文未详细说明”。预训练配置为 ViT-B/16 encoder，在 ImageNet 上训练 300 epochs；optimizer 为 AdamW，learning rate $1\times 10^{-3}$，${\beta }_1=0.9$、${\beta }_2=0.999$；40 epochs linear warmup 后 cosine annealing；scheduler 的 iteration-per-epoch scale 为 1.25；weight decay cosine schedule 从 0.04 到 0.4。默认 predictor 维度通过符号 ${\text{IWM}}_{X,Y}^Z$ 表示，例如 ${\text{IWM}}_{18,384}^{\text{Equi}}$ 是 18 layers、384 dim。

Linear probing：ImageNet 90 epochs，random crop scale $[0.08,1]$，horizontal flip 0.5；sequence average pooling 后接 linear layer；batch size 16,384，LARS optimizer，learning rate 6.4，10 epochs warmup，cosine schedule，weight decay 0，momentum 0.9。Attentive probing：90 epochs，crop scale $[0.3,1]$，source augmentations（不含 mask），batch size 1024，AdamW learning rate $1\times 10^{-4}$，weight decay 0.01。

Encoder finetuning：ImageNet 100 epochs，RandAugment rand-m9-mstd0.5-inc1，CutMix/MixUp，random erasing 0.25，mixup $\alpha =0.8$，cutmix $\alpha =1$，label smoothing 0.1；AdamW learning rate $2\times 10^{-3}$，5 epochs warmup，weight decay 0.005，batch size 1024，drop path 0.2，layer-wise LR decay 0.65。

Predictor finetuning：100 epochs，batch size 1024，AdamW LR $1\times 10^{-3}$，5 epochs warmup，weight decay 0.1，drop path 0.2；若 predictor 已预训练，其 LR 除以 10。End-to-end finetuning：encoder LR 也除以 10，组合 predictor+encoder 的 layer decay 为 0.9，LR $2\times 10^{-3}$，其余同 predictor finetuning。

Multitask predictor tuning：等价 50 ImageNet epochs，每个任务 batch size 512。ADE20K segmentation：UperNet，160k iterations，学习率在 $1e$-5、$2e$-5、$3e$-5 中择优，weight decay 0.01，linear LR decay。

5. Experimental Results (实验结果)

5.1 World model 质量：conditioning、难度、容量缺一不可

最直接的结论来自 MRR。Conditioning ablation 中，无 conditioning 的 MRR 为 0.00，Sequence 为 0.82，Feature 为 0.79；说明 predictor 必须知道 action，否则只会退回传统 invariant SSL。Architecture/augmentation ablation 中，I-JEPA 三项 MRR 都为 0.00；IWM 12-layer 对默认 destructive setting 只有 0.09，而 IWM 18-layer 达到 0.79；strong jitter 下 12-layer 和 18-layer 分别达到 0.81/0.85。这证明“更难的 transformation”不是噪声，而是逼迫 predictor 学 transition 的训练信号；但如果 predictor 容量不足，难度会变成无法解决的任务。

Figure 10 解读：附录随机检索样例显示，给定 blue bird 图像和 256 个 augmented views，IWM 的 latent prediction 大多能检索到与目标增强一致的图像。失败主要集中在 grayscale inversion 等不可逆或信息已丢失的操作上。

Figure 11 解读：penguin 样例强化了同一结论：IWM 不是仅记住语义类别，而是在 latent space 中保留并操纵颜色、模糊、solarization 等 augmentation 相关信息。检索结果越接近目标增强，说明 predictor 的 action-conditioned transition 越可靠。

Figure 12 解读：setter 图像中的多行 retrieval 展示 predictor 对不同强度 photometric transformations 的泛化。由于 source 侧有 destructive augmentations，而 target 侧保留较多信息，IWM 的输出体现了“根据动作恢复目标 latent”的能力。

Figure 13 解读：firefox 样例展示了颜色与局部纹理变化较丰富时的检索效果。图中多数预测仍能找到正确或相近增强视图，说明 MRR 高不是单个简单样例造成的，而在不同图像类别上都有定性支撑。

Figure 14 解读：detailed bird retrieval 进一步展开某个样本的多个 transformation 结果，能看到预测 latent 与 ground-truth augmented latent 的近邻关系。它有助于判断失败模式：若增强不可逆或过强，nearest neighbor 会偏离目标。

Figure 15 解读：detailed firefox retrieval 对比了更多具体增强组合，说明 world model 学到的是组合式 photometric action，而不仅是单一颜色校正。它也提醒：latent-space world model 的输出仍由 representation bank 间接可视化，不能等价于像素级完美重建。

Figure 16 解读：detailed penguin retrieval 展示了复杂背景和主体颜色变化下的预测。多数情况下检索结果保留语义与增强一致性，但 grayscale inversion 仍是弱点，符合论文对 destructive augmentation 不可逆性的解释。

5.2 下游分类：equivariant predictor finetuning 才真正利用 world model

ImageNet finetuning 主表如下：

Method	Epochs	Encoder	Tuned predictor (random)	Tuned predictor (pretrained)	End-to-end
MAE	300	82.7	82.4	82.7 (+0.3)	82.3
MAE	1600	83.6	83.0	83.1 (+0.1)	83.3
I-JEPA	300	83.0	79.1	80.0 (+0.9)	82.0
${\text{IWM}}_{12,384}^{\text{Inv}}$	300	83.3	80.5	81.3 (+0.8)	82.7
${\text{IWM}}_{18,384}^{\text{Equi}}$	300	82.9	81.5	83.3 (+1.8)	84.4

关键不是 IWM 每一列都赢，而是 pattern 很清楚：invariant IWM 的 encoder finetuning 更强（83.3），equivariant IWM 的 pretrained predictor finetuning 才释放 world model 价值（83.3，比随机 predictor 高 1.8），end-to-end 更达到 84.4。MAE 的 pretrained decoder/predictor 相比随机只带来 0.1 到 0.3 点，说明它的 world model 对分类几乎不可复用；I-JEPA 和 invariant IWM 的 predictor gain 不到 1 点；只有 equivariant IWM 的 predictor 学到了下游可利用的信息。

Peak performance 表进一步压缩了结论：DINO 1600 epochs frozen encoder 82.0 / any protocol 82.8；MoCoV3 300 epochs 76.4 / 83.2；iBOT 1600 epochs 83.0 / 84.0；MAE 1600 epochs 83.1 / 83.6；I-JEPA 300 epochs 80.0 / 82.0；${\text{IWM}}_{12,384}^{\text{Inv}}$ 81.3 / 83.3；${\text{IWM}}_{18,384}^{\text{Equi}}$ 83.3 / 84.4。也就是说，在允许使用 predictor 的协议下，equivariant IWM 给出最高峰值。

5.3 分割与大模型：趋势跨任务成立

ADE20K segmentation mIoU：

Method	Encoder	Predictor	End-to-end
I-JEPA	44.2	45.4	45.1
${\text{IWM}}_{12,384}^{\text{Inv}}$	45.6	45.7	46.5
${\text{IWM}}_{18,384}^{\text{Equi}}$	44.2	46.8	47.0

分割结果与分类一致：只看 encoder，invariant IWM 更强；但 finetune predictor 后，equivariant IWM 达到 46.8，end-to-end 达到 47.0，明显超过 I-JEPA 和 invariant IWM 的 encoder-only 设置。说明 IWM world model 的可复用性不局限于 ImageNet 分类。

大模型 ViT-L/16 附录中，I-JEPA encoder/predictor 为 84.1/79.9，${\text{IWM}}_{18,384}^{\text{Inv}}$ 为 84.3/81.5，${\text{IWM}}_{36,512}^{\text{Equi}}$ 为 83.7/85.0/end-to-end 85.4。作者还指出，ViT-B end-to-end 84.4 与 ViT-L encoder finetuning 84.3 相比，用 121M vs 307M 参数达到相近甚至略好结果，支持“复用 world model 比单纯扩大 encoder 更高效”的论点。

5.4 Linear / attentive / OOD：IWM 控制 representation abstraction

Method	Effective Epochs	Linear	Attentive
MoCoV3	300	76.3	76.4
MAE	300	60.2	73.5
MAE	1600	68.0	76.0
I-JEPA	300	70.0	75.0
${\text{IWM}}_{12,384}^{\text{Inv}}$	300	74.5	77.0
${\text{IWM}}_{18,384}^{\text{Equi}}$	300	67.5	75.1

这个表解释了为什么不能只用 linear probe 判断 representation。${\text{IWM}}^{\text{Inv}}$ linear 74.5，接近 MoCoV3 的 76.3，且 attentive 77.0 最好；${\text{IWM}}^{\text{Equi}}$ linear 只有 67.5，却在 predictor finetuning 和 OOD 上更有潜力。作者因此提出：linear probe 衡量的是“信息是否易于被简单线性层抽取”，而不是 representation 是否包含最多下游可用信息。

OOD attentive probing 表：MAE 在 ImageNet/iNat18/SUN397/Places205 为 73.5/50.1/70.2/60.3；I-JEPA 为 75.0/50.4/69.2/58.3；${\text{IWM}}_{12,384}^{\text{Inv}}$ 为 77.0/51.6/71.0/59.4；${\text{IWM}}_{18,384}^{\text{Equi}}$ 为 75.1/54.2/71.7/60.5。Equivariant IWM 虽然 ImageNet attentive 不如 invariant，但在 iNat18、SUN397、Places205 均为最高，分别比 invariant 高 2.6、0.7、1.1 点，说明 richer representation 对跨域任务更有利。

5.5 多任务、数据增强与局限

多任务结果显示，共享 predictor 的平均性能 73.5 与四个 single-task predictors 的 73.4 基本相同，代价是个别任务略降、SUN397 明显升高。这支持作者把 predictor 看作“可 instruction-tune 的视觉 world model”：同一个 latent transition module 可通过 task token 与 task head 适配多个任务。

数据增强附录揭示两个重要 caveat。第一，12-layer predictor 的 MRR 受增强强度强烈影响：默认弱 destructive setting 只有 0.09，strong jitter 可到 0.81，但 linear 从 74.5 降到 69.3；18-layer predictor 在多个 color setting 下 MRR 更稳定（0.79/0.85/0.85/0.83），但过强增强会让 linear 大幅下降到 34.3。第二，target/source augmentation asymmetry 很重要：作者报告 symmetric augmentations 对 invariant world model 会造成 ImageNet attentive 约 2 points、linear 约 1.5 points 的下降；equivariant predictor 对这个选择更稳。

作者明确或隐含的局限包括：一，论文没有 released code，复现需要自行实现 action parameterization、feature conditioning、masking 与 EMA schedule；二，硬件与训练成本没有详细披露；三，某些 destructive augmentations 不可逆，grayscale inversion 的定性错误说明 IWM 学到的是可预测 latent transition，而不是恢复已丢失像素信息；四，equivariant representation 在线性评估上较弱，实际使用必须配合 predictor finetuning、attentive head 或更复杂 protocol；五，world model 能力依赖 predictor capacity 和 augmentation recipe，过弱不能学 transition，过强/过难会损害抽象线性可读性。

总体结论：IWM 证明了在视觉表示学习中“学习 world model”可以不止服务于预训练损失。只要 predictor 被显式 action-conditioned、预测任务足够难且容量足够，它能成为 latent-space transformation model；这个 world model 既能提升下游 predictor finetuning/segmentation/multitask 效率，也能解释 invariant 与 equivariant representation 的 trade-off。对于后续工作，最值得延伸的是把 action 从 photometric augmentations 扩展到更接近物理/交互的 transformation，并公开实现以验证训练 recipe 的稳定性。