JEPA 能不能真正理解观测背后的生成机制?
传统的 JEPA 只能保证学到的隐表征不坍缩,但却完全没有办法对应到真实生成观测数据的那些隐变量。
来自 UIUC、冷泉港实验室、布朗大学、马普所的研究者给出了肯定的回答。他们提出的 DSReg(Dependency-Sparsity Regularization)可以直接加在已经训练好的 JEPA 上,让学到的每一维都对应一个真实的世界变量,比如场景里的太阳、月亮、温度和风。整个过程不需要重建观测,不需要解码器,也不需要任何标签。
【资料图】
研究团队同时为 DSReg 建立了完整的理论框架,由此得到首个完全可识别的 JEPA。DSReg 可以应用在预训练表征上,直接复用现有的 LeJEPA checkpoint,达到和联合训练相似的效果。
预测得准,不等于理解得对
JEPA 不去重建像素,而是在表征空间里预测未来。它因此省掉了昂贵的解码器,也成了构建世界模型的热门路线。但少了重建,也就少了一个把表征和真实世界对齐的 “锚”。
在一定假设下,LeJEPA 学到的表征和真实世界状态之间只差一个未知的线性变换。这个变换让模型完全无法理解真实的生成机制,因为表征的每一维往往是好几个未知物理因素的混合。
这些未知的混合有时候对预测并没有影响,但却完全阻碍了模型对环境的理解。比如机器人要把杯子往左推,同时让旁边的刀保持不动。如果某个学到的变量同时混着杯子和刀的位置,调这个变量就不再只是移动杯子,而是也可能把刀一起带动,给用户带来危险。很多时候我们不仅需要让机器看上去做得对,同时也要保证它们的认知没有偏差,不然永远会有隐藏的风险。
DSReg:给预训练表征加一个旋转
DSReg 的思路很直接。LeJEPA 学到的表征经过白化后,和真实变量之间只差一个旋转,而旋转不改变 JEPA 的训练损失。所以只要找到正确的那个旋转,就能把变量拆开,同时不损失任何预测能力。
但是如何找到正确的旋转?DSReg 将目标放在隐变量和观测变量之间的结构上。每个真实变量只影响一部分观测,混合变量则会影响它所有成分涉及的观测。因此,让观测依赖的变量最少的那个旋转,就是把变量拆开的旋转。具体分三步:
整个过程不训练解码器,也不优化重建误差,可以直接用在任意 LeJEPA checkpoint 上。局部 Jacobian 也不需要显式存下来,整个框架在一张 48GB 显卡上就可以做到 8192 维。
没有两片雪花是一样的
当然,不做任何假设,是不可能从观测里还原隐藏世界的。DSReg 需要生成过程满足 Structural Diversity:任意两个变量影响的观测集合不能完全相同,就像没有两片雪花是一样的。
这个条件很宽松。不同隐变量影响的观测变量可以重叠,甚至可以一个完全包含另一个,可以说是所有的结构化假设里面最弱的一个。
理论证明,在 Structural Diversity 下,DSReg 能把所有隐变量都恢复出来,表征里的每一维都和一个真实隐变量一一对应,最多只差顺序和正负号。文章同时还证明了可辨识性对于观测误差的鲁棒性,并且用 Lean 对理论结果做了形式化验证。
让世界模型理解真实世界
在合成数据上,LeJEPA 的表征始终是混合的,加上 DSReg 后,恢复效果在各个维度上都接近上限。
在从像素训练的编码器上,DSReg 达到了使用标签拟合的监督上界,而 PCA、Varimax、FastICA 等旋转方法几乎没有效果。在外部渲染的 3DShapes 上,DSReg 同样从混合提升到接近上限;在更难的 dSprites 上,也超过了 β-VAE 和 β-TCVAE。
DSReg 让表征对应真实的隐变量,从而对各个下游任务也都有显著的提升。在视觉编辑、稀疏控制、rollout 预测和意外检测共六个下游任务上,DSReg 都优于 LeJEPA,并且原有的预测性能保持不变。
DSReg 也具有非常高的可扩展性,只需要单卡就可以处理 8192 个隐变量,并且表现并不会随着维度升高而降低
让机器 “知行合一”
JEPA 学到了不坍缩的隐表征,DSReg 则让这个表征里的每一维都对应一个具体的世界隐变量。建在这样的世界模型上,机器人就可以精确理解观测场景背后的每一个变量。
因此,DSReg 让机器在和环境交互时,能够真正做到 "知行合一"。这对机器人尤其重要:LLM 如果没有真正理解环境,可能只是说错一句话;机器人如果没有真正理解,就可能带来无法预料的灾难性后果。
下一代智能体要想在各种新环境里都可靠地工作,不能只靠把在训练环境里学到的知识迁移过去。它们必须理解观测背后的生成机制,才能做到真正意义上的泛化。DSReg 让 JEPA 有了这种能力,这或许正是走向这个目标的一个起点。
作者介绍
郑雨嘉,伊利诺伊大学厄巴纳 - 香槟分校(UIUC)助理教授。博士毕业于卡内基梅隆大学。研究方向为因果 AI,包括因果表征学习、可信机器学习、世界模型等,主要围绕构建能理解因果关系的下一代智能体。个人主页:yjzheng.com