TensorX

TensorX · Paper Blog

JEPA-Anything: Learning Predictive Models across Different Worlds

Taoyong Cui, Zhongyao Wang, Xinyue Xu, Weiyang Liu, Zhaochen Yu, Yuying Zhang, Qiang Gao, Mengyue Yang, Wanli Ouyang, Pheng Ann Heng, Yingcheng Wu, Zhenfei Yin, Ling Yang (PhAI Labs, The Chinese University of Hong Kong, Fudan University, City University of Hong Kong, University of Bristol, Stanford University, University of Oxford, Princeton University) · 2026-09-18

World ModelsJEPARepresentation LearningScientific AIDynamical Systems

JEPA-Anything提出了正交预测因子分解机制,打破了以往单体表征预测的容量分配瓶颈。该框架横跨视觉、单细胞生物学、临床轨迹、偏微分方程物理场和分子动力学七大异构世界,构建起统一可逆的潜在世界模型。

💡 核心创新点

打破单体目标预测的容量倾轧假设

传统的联合嵌入预测架构(JEPA)将环境未来或未观测状态编码为一个单体嵌入向量(monolithic target embedding)。这种设计默认整个世界状态可以用单一梯度流联合更新。但在跨尺度、多实体的复杂物理或生物系统中,高方差或易预测的动态成分往往占据优化梯度主导地位,导致弱信号、长程或精细结构被严重忽略。

JEPA-Anything 打破了这一传统假设,提出了正交预测因子分解(Orthogonal Predictive Factorization,OPF)。它不再将世界状态视作不可分割的整体,而是通过一组正交基底投影矩阵,将高维目标表征投影到互不重叠的正交子空间中。每个子空间由专用的预测分支负责建模,从而在不依赖预设语义标签的前提下,自发涌现出互补的预测维度。

结构可逆合成与内在自防塌陷机制

与传统的特征解耦(Disentanglement)或对比学习方法不同,JEPA-Anything 构建了一套具备严格数学保障的双向闭环预测接口:

  • 因子正交化与伪逆合成:利用摩尔-彭若斯广义逆(Moore-Penrose pseudoinverse),各个独立预测因子的输出能够无损合成为完整的潜在世界状态,保障自回归多步推演(Rollout)的数值稳定度。
  • 因子与编码器双重活性保障:引入防塌陷标准差惩罚,直接向在线编码器和各投影分支反向传播防塌陷梯度,无需生成式像素重构即可杜绝信息退化。

这一机制差异让模型不仅能做表征读出,更可以直接支持闭环控制规划、长时域自回归演化与因果干预分析。

JEPA-Anything 架构与单体 JEPA 对比及三组核心评测概览

🔍 背景与动机

现有世界模型的领域割裂困境

智能体预测未知、规划干预和交互演进的核心基石是世界模型(World Models)。然而长久以来,世界模型的研究呈现严重的领域割裂:计算机视觉依赖像素生成或遮蔽补全,机器人控制依赖动作条件的状态转移网络,分子动力学与天气预测依赖等变神经网络或偏微分方程(PDE)求解器,生物医学则依赖专用的图神经网络或单细胞基础模型。

现有基于潜空间自监督预测的方案在处理异构世界时普遍失效。如果直接使用单体潜状态预测,遇到诸如流体漩涡、分子振动模式或细胞扰动响应这类多频段、多实体信号时,网络往往只拟合主频分量,次级演化模式在几步自回归迭代后迅速衰减发散。

问题的直观比喻:多乐器交响乐录制

可以将构建全能世界模型比作单声道录音与多轨分轨录音的对比。传统的单体 JEPA 就像是用一个麦克风将整支交响乐团的演奏录制进单个音频轨道:低音大鼓的声音极其响亮,瞬间掩盖了弱音小提琴的长音旋律与长笛的颤音。随着音乐向前推演,混音完全失真。

🎯 核心结论:JEPA-Anything 相当于为乐团的每个声部安排了独立的正交拾音器(正交子空间)与专属调音台(专用预测头),在独立精准捕捉各声部演进后,再通过调音母带(伪逆矩阵)精准复原出整场交响乐。

📊 数据集与实验环境

覆盖七大领域的宏大评测矩阵

为了检验同一套预测核心是否能适应不同的物理规律与数据形态,作者构建了横跨七大截然不同领域的实验基准,涵盖 10 项动力学任务、超 1,000 种未来临床疾病预测、以及 4 种材料的 100 步分子推演:

  1. 视觉场景绑定:基于 MuJoCo 渲染生成的受控干预视觉场景,利用 DINOv3 ViT-S 与 SigLIP2 Base 骨干网络,使用留一单元格(leave-one-cell-out)与图像互斥划分进行 10 随机种子评估。
  2. 单细胞生物学:约 800,000 个真实人类肾脏细胞预训练,在 PBMC-10K 上测试微调与零样本细胞聚类(AvgBIO 指标),并在 Adamson 及 Norman 基因扰动数据集上验证绝对表达预测的 Pearson 相关系数。
  3. 临床健康队列:基于 UK Biobank 真实去标识多模态队列,涵盖时间序列事件、生化指标与人口统计学变量,采用固定病人划分预测 1,000 余种罕见与常见临床风险事件(PRAUC 指标)。
  4. 受控干预动力学:CITRIS Interventional Pong 平台,评估单一干预及未见组合干预下的 6 步自由推演(4 通道 MSE)。
  5. 物理场与自回归动力学:涵盖 CausalWorld、DeepMind Control Suite (DMC)、PDEBench(Burgers、浅水波方程、反应扩散方程)、WeatherBench2 气象场以及 APEBench(一维粘性 Burgers 方程与混沌 Kuramoto-Sivashinsky 方程)。
  6. 连续运动控制:Gymnasium 物理引擎中的 Hopper-v5、Walker2d-v5 与 HalfCheetah-v5 连续控制环境,结合交叉熵方法(CEM)规划器评估累积真实环境回报。
  7. 力场自由分子模拟:TrajCast 架构下的水分子、α-石英、对乙酰氨基酚分子及苯分子系统,进行 100 步等变无力自由自回归推演。

JEPA-Anything 场景全景图(Scenario Atlas)

硬件与通用超参数设置

所有实验在统一的基线对比下严格控制了参数量与浮点运算量(FLOPs):

  • 硬件与随机种子:全实验采用 5 个固定随机种子(如动力学基准采用 {11, 23, 37, 53, 71},连续控制采用 {42, 123, 456, 789, 2026})进行重复统计。
  • 潜变量与因子维度:潜空间维度依据各领域灵活配置,但满足乘积关系。动力学任务设潜维度 128、4 个因子、每因子 32 维;CITRIS 设潜维度 160、5 个因子、每因子 32 维;单细胞任务设潜维度 512、4 个因子、每因子 128 维;临床模型与视觉 SigLIP2 设潜维度 768、4 个因子、每因子 192 维。
  • 正则化权重配置:动力学与物理场任务通用设置正交正则项系数 0.10、因子活性阈值惩罚系数 0.05、编码器活性系数 0.02;连续控制任务则配置正交系数 0.02、因子活性系数 0.10、编码器活性系数 0.02。

🛠️ 核心机制与研究方法

统一的抽象接口与领域适配器

JEPA-Anything 将任何物理或数字世界的预测建模过程统一拆分为两层结构:领域特异性边界(Domain-Specific Boundary)与共享正交预测核心(Shared JEPA-Anything Core)。

领域特异性边界与不变核心的精确划分

针对任意领域的原始观测输入,首先通过专用适配器提取结构:

x→Aδ(H,S)→Vδ(HC,SC,T,ST)x \xrightarrow{\mathcal{A}_{\delta}} (H, S) \xrightarrow{\mathcal{V}_{\delta}} (H_C, S_C, T, S_T)

在公式 (1) 中,原始数据经由适配器映射为内容标记集合与结构描述符集合(如坐标、时间戳、晶格等变标号),再由视图采样器切分出已知上下文集合与待预测目标索引。

系统采用非对称的双编码器拓扑,目标侧参数通过指数移动平均(EMA)进行无梯度更新:

θˉ←mθˉ+(1−m)θ,0≤m<1\bar{\theta} \leftarrow m \bar{\theta} + (1 - m)\theta, \quad 0 \le m < 1

在公式 (2) 中,动量系数平滑目标表征演进,而在线编码器则通过反向传播实时优化。

模块一:目标空间的低秩正交投影分解

在获得停止梯度的目标潜表征后,系统使用分析基底矩阵将其拆解为若干正交因子:

z~t=sg(zt),zt(k)=Pk⊤z~t,k=1,…,K\tilde{z}_t = \mathrm{sg}(z_t), \quad z_t^{(k)} = P_k^\top \tilde{z}_t, \quad k = 1, \dots, K

在公式 (3) 中,原目标表征被截断反向传播梯度后,分别乘以投影矩阵的转置,投射到互补的正交子空间中。每个投影矩阵的秩等于因子维度,确保子空间无冗余重叠。

每个预测因子配置独立的专用预测网络,从共享的上下文表征映射到该因子子空间中:

z^t(k)=qk(zc,st)\hat{z}_t^{(k)} = q_k(z_c, s_t)

在公式 (4) 中,预测器根据当前的在线上下文状态以及目标时空描述符,直接在该因子的专属几何子空间中拟合其演化轨迹。

模块二:伪逆算子驱动的完全态无偏合成

当模型完成各个独立因子的推演后,需要构建可供下游决策或下一步自回归滚动的完整全局状态。系统通过投影矩阵块拼接的摩尔-彭若斯广义逆完成重建:

z^t=(P⊤)†u^t∈Rd,P=[P1,…,PK]\hat{z}_t = (P^\top)^\dagger \hat{u}_t \in \mathbb{R}^d, \quad P = [P_1, \dots, P_K]

在公式 (5) 中,拼接的因子预测向量经由矩阵伪逆变换映射回原始潜状态空间。当投影矩阵之间保持严格正交时,伪逆变换直接退化为简单的转置乘积求和。

为了确保合成后的状态在幅值和方向上均无失真,预测损失直接在因子层面采用均方误差回归:

Lpred=1K∣T∣r∑t∈T∑k=1K∥z^t(k)−zt(k)∥22\mathcal{L}_{\mathrm{pred}} = \frac{1}{K|T|r} \sum_{t \in T} \sum_{k=1}^K \|\hat{z}_t^{(k)} - z_t^{(k)}\|_2^2

在公式 (6) 中,网络在所有采样目标时间步和正交因子上计算预测误差并执行梯度反传,从而让因子的物理含义由可预测性自然驱动形成。

模块三:流形正交性约束与防塌陷方差底线

若无外加几何约束,多个投影矩阵可能自发塌陷并学习到相同的方向。为此,JEPA-Anything 引入了块内正交与块间解耦的联合损失:

Lorth=∑k=1K∥Pk⊤Pk−Ir∥F2+∑1≤i<j≤K∥Pi⊤Pj∥F2\mathcal{L}_{\mathrm{orth}} = \sum_{k=1}^K \|P_k^\top P_k - I_r\|_F^2 + \sum_{1 \le i < j \le K} \|P_i^\top P_j\|_F^2

在公式 (7) 中,第一项约束每个因子内部的投影基底相互正交归一,第二项惩罚不同因子基底之间的互相关重叠,迫使因子的几何空间完全直和分解。

为了防止自监督模型陷入常数表征的平凡解(Representation Collapse),框架进一步引入两级激活方差监控机制:

Lfac=1Kr∑k=1K∑j=1rmax⁡(0,γfac−σk,jfac),Lenc=1d∑j=1dmax⁡(0,γenc−σjenc)\mathcal{L}_{\mathrm{fac}} = \frac{1}{Kr} \sum_{k=1}^K \sum_{j=1}^r \max(0, \gamma_{\mathrm{fac}} - \sigma_{k,j}^{\mathrm{fac}}), \quad \mathcal{L}_{\mathrm{enc}} = \frac{1}{d} \sum_{j=1}^d \max(0, \gamma_{\mathrm{enc}} - \sigma_j^{\mathrm{enc}})

在公式 (8) 中,第一项对每个因子各个坐标在批次内的样本标准差施加底线惩罚,第二项直接向在线编码器注入最小方差梯度,确保整条通路的信息通量始终充沛。

最终的端到端联合训练损失函数为:

LOPF=Lpred+λorthLorth+λfacLfac+λencLenc,Ltrain(δ)=Lbase(δ)+LOPF\mathcal{L}_{\mathrm{OPF}} = \mathcal{L}_{\mathrm{pred}} + \lambda_{\mathrm{orth}}\mathcal{L}_{\mathrm{orth}} + \lambda_{\mathrm{fac}}\mathcal{L}_{\mathrm{fac}} + \lambda_{\mathrm{enc}}\mathcal{L}_{\mathrm{enc}}, \quad \mathcal{L}_{\mathrm{train}}^{(\delta)} = \mathcal{L}_{\mathrm{base}}^{(\delta)} + \mathcal{L}_{\mathrm{OPF}}

在公式 (9) 中,整个优化目标将领域特定的基础损失与包含预测、正交、因子活性及编码器活性的正交预测分解(OPF)损失完全解耦并叠加执行。

📈 结果与深度分析

动力学基准与长程推演性能飞跃

在 10 项跨越连续控制、像素网格和偏微分方程的动力学任务中,JEPA-Anything 全面超越了参数量对齐的标准单体 JEPA 基线:

10 项连续动力学基准中的均方误差降幅与闭环控制表现

在由图表可印证的具体任务中,JEPA-Anything 取得了显著的误差缩减:

  • 物理场演化:在 Burgers 方程中,单步 MSE 相对误差下降达 39.7%;在浅水波方程(Shallow water)中误差下降 39.3%;在反应扩散方程中下降 15.2%;在 WeatherBench2 气象动力学中下降 10.5%。
  • 连续控制与机械臂:CausalWorld 机器人状态预测误差下降 20.6%,其闭环控制环境回报从标准 JEPA 的 -0.087 显著提升至 -0.055。
  • 干预泛化实验:在 CITRIS Interventional Pong 上,单干预预测误差降低了 34.8%(MSE 从 0.009541 降至 0.006218),未见组合干预误差降低 12.9%,且在 6 步自由滚动预测中误差下降 8.6%。

在分子动力学场景中,针对水、石英、对乙酰氨基酚和苯 4 种系统,JEPA-Anything 在 100 步自回归推演后均取得了最低的末态 RMSD,其中石英的末端位移均方根误差从 2.089 Å 压降至 1.877 Å。

几何消融实验:正交约束的不可替代性

为了验证正交正则项是否真的发挥了作用,作者在 CITRIS 平台上将正交分解与参数量完全相同的无约束多头结构(Unconstrained multi-head)进行了严格机制审计:

CITRIS 平台机制审计:无约束多头与正交分解的几何性质对比

💡 关键技巧:机制审计清晰表明,若没有正交正则化,无约束多头的子空间重叠度高达 0.4550,条件数飙升至 438.52,且转置合成归一化均方误差高达 0.7886,导致状态合成严重偏离真实流形;而正交因子分解的条件数精确收敛为 1.00005,重叠度与合成误差分别压降到超低极值,保证了潜在合成的数值稳定性与可逆性。

在 Hopper 机器人运动控制的因子消融实验中,轮流屏蔽 F1 至 F4 各个因子均引发 20 步推演误差的剧烈上升(MSE 增加量介于 +2.067 到 +4.462 之间),证明模型学到的各个因子均处于高度活跃且功能解耦的状态。

科学发现与第一性原理验证

在生命科学与天体物理的外推分析中,JEPA-Anything 展现出从数据中挖掘本质规律的独特能力:

  1. 生物湿实验验证:利用因子坐标分析识别并推荐了 IL-18 联合 NT5E/CD73 免疫靶点阻断方案,随后在患者来源的肝细胞癌类器官(PDO)及肿瘤小鼠实验中得到了真实验证,展现出显著增强的抗肿瘤杀伤效果与免疫细胞激活率。
  2. 物理定律重现:在未给予任何天体物理先验标签的情况下,通过对轨道运动轨迹进行因子的傅里叶谱分析,模型提取出的潜空间主频与轨道半长轴之间自发复现了开普勒第三定律,拟合幂律斜率达到 -1.4991(理论值为 -1.5000),决定系数 R2=0.9999999R^2 = 0.9999999。

局限性与未来展望

尽管表现亮眼,论文也指出了当前框架存在的客观局限:

⚠️ 局限性:正交因子的数学正交并不等价于因果机制上的绝对解耦;若缺乏针对性的干预信号引导,因子的物理可解释性可能随初始化产生漂移。此外,在个别控制环境(如 Hopper 规划)中,因子分解未能稳定胜过单体模型,表明对于自由度极低、强耦合的系统,强行因子化可能会引入额外的规划瓶颈。

未来,研究团队计划将 JEPA-Anything 作为基础预测底座,深度整合至 ScienceIDE 与 ScienceBuddy 科学发现智能体框架中,使智能体能够在闭环虚拟与物理实验中主动生成假设、调度仿真工具并更新对真实世界的内在认知。跨领域世界模型的统一时代,正在拉开序幕。

北京市昌平区探索星信息技术及软件开发工作室

京ICP备2026059466号