← 返回 2026-08-03

QQWorld:用于世界模型正则化的分位数-分位数匹配 QQWorld: Quantile-Quantile Matching for World Model Regularization

Zhoushun Yu, Xiaoyu Hu, Xiangyu Xu 📅 2026-07-30 👍 30 2026-08-08 18:30
QQ匹配 Wasserstein距离 世界模型 分布匹配 强化学习规划 潜在正则化

用QQ匹配替代EP正则化,矫正潜在空间重尾偏差,提升规划性能。

前置知识

潜在世界模型与JEPA规划

潜在世界模型在紧凑的潜在表示空间中预测未来状态并据此规划,而非在像素空间重建。联合嵌入预测架构(JEPA)是其代表范式:直接在 latent 空间预测未来表示,避免高维像素重建的低层细节开销。LeWM(LeWorldModel)等近期工作在该框架下用 CEM(交叉熵方法)做目标条件规划。

QQWorld 是对 LeWM 的即插即用改进,必须先理解潜在世界模型为什么需要把 latents 正则化到各向同性高斯(这是 JEPA 规划能工作的前提),才能理解为什么正则化目标的选择会直接影响规划成功率。

Epps-Pulley (EP) 检验与正态性正则化

EP 检验是经典的一维正态性检验,基于经验特征函数 $\hat{\psi}(t)=?rac{1}{N}\sum_{n=1}^N e^{itx_n}$ 与标准正态特征函数的差异构造统计量 $L_{EP}(X)=N\int_{-\infty}^{\infty}|\hat{\psi}(t)-\psi(t)|^2 w(t)dt$,权重 $w(t)=e^{-t^2/2}$。LeWM 将其作为可微惩罚项约束 latent 分布逼近 $N(0,1)$。

本文的全部动机都建立在分析 EP 检验的梯度几何上:作者证明 EP 等价于单位带宽高斯核下的 MMD,其梯度在尾部超指数衰减。如果不熟悉 EP 与特征函数检验,就无法理解 LeWM 为何失败、QQ 为何有效。

最大均值差异 (MMD) 与高斯核带宽

MMD 是衡量两分布差异的核方法,定义为再生核希尔伯特空间中均值的距离。对高斯核 $k(x,y)=e^{-(x-y)^2/2}$(单位带宽),MMD 只对落在彼此交互尺度内的样本产生显著的成对力。带宽越小,能感知的样本间距离范围越窄。

作者用 MMD 等价视角($L_{EP}\propto ext{MMD}^2(\{x_n\},N(0,1))$)揭示 EP 失效的本质:单位带宽高斯核对远在尾部的样本几乎'看不见'。这是理解 Proposition 1(恢复力消失)的关键。

分位数-分位数 (QQ) 匹配与 Wasserstein 距离

QQ 匹配把一维样本排序后的顺序统计量 $\hat{x}_1\le\cdots\le\hat{x}_N$ 与对应的高斯分位数 $q_n=\Phi^{-1}(?rac{n-0.5}{N})$ 逐点对齐,损失 $L_{QQ}=\sum_{n=1}^N(\hat{x}_n-q_n)^2$。在一维下,平方 2-Wasserstein 距离正是分位数函数的平方 $L_2$ 距离,故 $L_{QQ}?pprox N\cdot W_2^2(\{x_n\},N(0,1))$。

QQ 损失是本文方法的核心。理解它作为 Wasserstein 距离求积近似的几何含义,才能理解为何其梯度 $?rac{\partial L_{QQ}}{\partial x_n}=2(x_n-q_{ ho(n)})$ 会随尾部偏差线性增长、不衰减。

Cramér-Wold 定理与随机投影切片

Cramér-Wold 定理指出:$\mathbb{R}^d$ 上的分布由其所有一维投影唯一确定。因此可通过在 $S$ 个随机单位方向 $u_s\sim ext{Uniform}(S^{d-1})$ 上施加一维正态性约束,蒙特卡洛逼近对整个 $d$ 维联合潜在分布的正则化。

QQWorld 不是直接在高维做分布匹配,而是切片后在多个一维投影上分别做 QQ 匹配,再平均。理解 Cramér-Wold 定理是理解该方法为何能约束整个潜在分布(而非仅边际分布)的理论基础。

研究动机

潜在世界模型通过在紧凑表示空间中预测未来状态来实现高效规划,其性能高度依赖学习到的潜在分布的质量。近期理论(Balestriero & LeCun, 2025)表明,将潜在边际分布正则化到各向同性高斯可以最小化下游预测风险。为此 LeWM(Maes et al., 2026)采用 Epps-Pulley(EP)检验作为可微的全分布正态性惩罚。然而作者观察到,即使模型被显式地惩罚非正态性,LeWM 学到的 latents 仍有显著的重尾现象(如图1右所示,尾部率高达 0.315,而高斯参考仅 0.10)。这种重尾对世界模型是有害的:极端潜在值会把学到的动力学推入表示不充分的区域,在多步 rollout 中放大误差,并加剧潜在分布与目标高斯先验的失配。作者深入分析发现这并非偶然——EP 检验在等价 MMD 视角下是单位带宽高斯核的核差异,其针对尾部样本 $h$ 的恢复力以 $\sqrt{\pi}h e^{-h^2/4}$ 形式衰减,在 $h=\sqrt{2}$ 处达峰后超指数衰减,导致一旦样本逃逸出核的交互尺度,EP 正则化就几乎'看不见'它。

本文的目标是本文的具体目标是设计一种潜在分布正则化目标,使其在尾部样本上仍能维持有效的纠正梯度,从而真正抑制重尾偏差,提升潜在空间的高斯性并最终提升下游规划成功率。该目标必须满足三个硬性要求:第一,作为 EP 的'即插即用'替换,不引入任何方法专属的新超参数(与 Sub-JEPA、SD-JEPA、SMWM 不同,后者都引入了需调参的新项);第二,对尾部偏差提供强且对称的纠正力,使重尾被有效压回主体;第三,在有限 GPU 显存下可高效训练,即使物理 batch size 较小也能获得准确的分位数估计。作者瞄准在四个控制环境上同时改进平均规划成功率与多项正态性指标(KS、EP、QQ RMSE、尾部率)。

与已有工作不同的是,现有改进 LeWM 的工作(Sub-JEPA 在多子空间做高斯正则、SD-JEPA 分解 progression/content 子空间、SMWM 引入逆动力学监督)都通过改变潜在空间结构或加辅助目标来改善表示学习,但几乎没有人重新审视'分布匹配目标本身在优化过程中的梯度几何'。本文的独特切入角度是提出一个深刻洞察:'一个对衡量分布差异有效的统计量,未必是用于训练世界模型的有效目标'——衡量能力(distinguishing distributions)与优化能力(providing informative gradients)是两回事。作者从梯度视角而非统计检验视角,把 EP 重新解释为单位带宽核 MMD,揭示其在尾部的梯度消失缺陷,进而从 2-Wasserstein 距离的几何视角构造秩匹配传输目标 QQ,并建立单向控制 $L_{EP}(X)\le C(L_{QQ}(X)+?rac{\log N}{N})$,从理论上证明 QQ 是比 EP 更强的分布控制形式。

核心方法

QQWorld 的整体思路很简洁:与其用本质上是有限带宽高斯核 MMD 的 EP 检验惩罚分布偏离,不如直接把投影后的一维潜在样本排序,与对应的标准高斯分位数逐点对齐。技术路线借鉴 Cramér-Wold 定理:在 $S$ 个随机单位方向上分别施加一维正态性约束再平均,蒙特卡洛逼近对整个 $d$ 维联合潜在分布的约束。每个切片损失取排序样本与高斯分位数的平方差 $L_{QQ}(X)=\sum_n(\hat{x}_n-q_n)^2$,$q_n=\Phi^{-1}((n-0.5)/N)$。直觉上每个样本被指派一个与其秩匹配的高斯目标并被直接拉向该位置。该构造严格视为平方 2-Wasserstein 距离的求积近似($L_{QQ}\approx NW_2^2$),继承最优传输的全局几何而非核方法的有限带宽局部几何。总正则项 $R=\frac{1}{S}\sum_s L_{QQ}$,权重固定 3.5,所有环境一致。

核心创新是把正则化目标从'有限带宽核差异度量'(EP/MMD)替换为'秩匹配量化传输目标'(QQ),本质区别在梯度行为。EP 对样本 $x_n$ 的梯度含 $\sqrt{\pi}x_n e^{-x_n^2/4}$ 项,对尾部样本 $h$ 渐近为 $\sqrt{\pi}h e^{-h^2/4}$、在 $h=\sqrt{2}$ 达峰后超指数衰减。QQ 的梯度则为 $\frac{\partial L_{QQ}}{\partial x_n}=2(x_n-q_{\rho(n)})$,直接指向秩匹配分位数、大小随偏差线性增长,尾部越偏纠正越强。作者还证明三点:(1) 排序的梯度不连续在实践中无害——并列边界是排斥的(损失对分离参数在并列处方向导数为 $-2(q_{k+1}-q_k)<0$),梯度下降主动增大样本分离,反具局部抗塌缩作用;(2) 单向控制 $L_{EP}\le C(L_{QQ}+\log N/N)$ 表明把 QQ 降为零会把 EP 也降为零,反之不然;(3) QQ 是 EP 的严格更强分布控制形式。

方法步骤详情

训练流程分五步。第一步潜在编码与投影:对批次潜在嵌入 $\{z_n\}\in\mathbb{R}^d$,从单位球面采样 $S$ 个随机方向 $u_s$,算一维投影 $x_n=\langle u,z_n\rangle$。第二步排序与分位数生成:对每个方向投影排序得顺序统计量 $\hat{x}_n$,算高斯分位数 $q_n=\Phi^{-1}((n-0.5)/N)$($\Phi^{-1}$ 为标准正态逆 CDF)。第三步 QQ 损失:$L_{QQ}=\sum_n(\hat{x}_n-q_n)^2$。第四步切片平均:$R=\frac{1}{S}\sum_s L_{QQ}$ 以固定权重 3.5 加入训练目标端到端反传,与 LeWM 唯一区别是把 $L_{EP}$ 换成 $L_{QQ}$。第五步可选 Cross-Batch QQ:维护先进先出队列保留前 $K$ 次迭代投影特征(已 detach),把当前 $N$ 个与历史 $KN$ 个合并为大小 $M=(K+1)N$ 的排序池,在 $M$ 个样本上算秩,梯度只通过当前 $N$ 个特征反传,不增显存即扩大有效排序集合。

技术新颖性

技术新颖性体现在四层。第一,从'优化几何'而非'统计检验'视角重审分布匹配目标,提出'统计有效≠优化有效'的洞察:不仅用 MMD 等价性($L_{EP}\propto\text{MMD}^2$)解释 EP 为何失败,更从梯度行为证明 QQ 为何成功,并以 Prop 1-3 作理论支撑。第二,正面回应'排序梯度不连续是否阻碍优化'的担忧,用并列排斥性证明说明不连续反带来局部抗塌缩益处,把缺点转为优点。第三,Cross-Batch QQ 把'排序池规模'与'反传批次大小'解耦,并用偏差-方差分解 $\text{MSE}(\hat{q})\approx[\frac{F_t(1-F_t)}{N(K+1)}+(\bar{F}_{t,K}-F_t)^2]/\phi^2$ 严格刻画'方差减少 vs 表示陈旧偏差'的权衡,把工程技巧提升为可分析设计。第四,与并发的 Wu et al. (2026)(切片 Wasserstein 做 JEPA 自监督、需额外方差项)相比,QQWorld 用单一分位数匹配目标端到端替换 EP,无新超参,更轻量且聚焦世界建模与规划。

实验结果

作者在 Two-Room、Reacher、PushT、OGBench-Cube 四环境六种子下系统验证。规划性能(Table 1):QQWorld 平均 85.08% 全场最高,比 LeWM(79.75%) 高 5.33 pp,四环境一致领先,亦优于 Sub-JEPA、SD-JEPA、SMWM 等变体。正态性(Table 2/Figure 1)最有说服力:QQWorld 把 KS 降 15.8%、EP 降 31.4%(119.909→82.294),且非 EP 训练却更优;QQ RMSE 与尾部率显著逼近高斯参考。Figure 2 显示训练 40k 步中 QQWorld 尾部概率稳步趋近高斯参考、LeWM 始终重尾,证明 QQ 提供持续有效的尾部纠正梯度。物理探测(Table 3):位置/角度线性与 MLP 探测与 LeWM 持平或略优,压平重尾未损害物理信息编码。Cross-Batch QQ(Table 4):$N=32$ 加历史批次把成功率从 65.42% 拉到 83.50%、显存几乎不变,缩 4× 批次、减 73% 显存仍超过 128 批次 LeWM;队列过长则因陈旧偏差反降。

Planning success rate (%) across four environments.
Table 1: Planning success rate (%) across four environments.
Normality comparison between LeWM and QQWorld.
Table 2: Normality comparison between LeWM and QQWorld.
Linear- and MLP-probing results for agent location, block location, and block angle.
Table 3: Linear- and MLP-probing results for agent location, block location, and block angle.
Effect of the batch size N and cross-batch queue length K+1 on planning success rate and GPU memory.
Table 4: Effect of the batch size N and cross-batch queue length K+1 on planning success rate and GPU memory.
Evolution of the tail over 40k training steps on Two-Room.
Figure 2: Evolution of the tail over 40k training steps on Two-Room.
Visualization of the learned latent representations on Two-Room.
Figure 3: Visualization of the learned latent representations on Two-Room.
查看结构化数据
任务指标本文基线提升
四环境平均目标条件规划成功率 Planning Success Rate (%),6 随机种子均值±标准差 85.08±(Two-Room 93.67、Reacher 85.33、PushT 91.00、OGBench-Cube 70.33) LeWM 79.75、Sub-JEPA 83.17、SD-JEPA 82.67、SMWM 83.00 比 LeWM 基线 +5.33 pp,四环境一致领先,全方法第一
潜在分布正态性(KS/EP 统计量,四环境平均) Kolmogorov-Smirnov 与 Epps-Pulley 统计量(越低越接近 $N(0,1)$) KS 0.032, EP 82.294 LeWM: KS 0.038, EP 119.909 KS 相对降低 15.8%,EP 相对降低 31.4%(且非 EP 训练却更优)
QQ RMSE 与径向尾部率(6144 随机投影,20000 latents) QQ RMSE 与 $P(\|z\|_2^2>q_{0.90}(\chi^2_{192}))$(越低越好,高斯参考 0.10) QQ RMSE 0.121, 尾部率 0.123 LeWM: QQ RMSE 0.157, 尾部率 0.315 QQ RMSE 降 22.9%,尾部率降 60.9%,接近高斯参考
Cross-Batch QQ 在 $N=32$ 下的内存-性能权衡 平均规划成功率 (%) 与 GPU 显存 (MB) $N=32,K+1=3$: 83.50%, 3435.0 MB LeWM (N=128): 79.75%; QQWorld 大批次 N=128: 85.08%, 12747.1 MB 以 83.50% 超过 LeWM 大批次,反传批次缩 4×,显存减约 73%
物理状态线性/MLP 探测(PushT) agent/block 位置与 block 角度的 MSE 与 Pearson r block 角度线性 MSE 0.172, r 0.910 LeWM: block 角度线性 MSE 0.176, r 0.908 三项探测任务均持平或略优,证明 QQ 正则化不损害物理信息编码

局限与改进

作者承认的局限主要有两点。其一,Cross-Batch QQ 的收益高度依赖于当前批次大小 $N$:当 $N$ 较小时(如 32)方差减少显著(成功率从 65.42% 跳到 83.50%),但当 $N$ 已足够大(如 64、128)再加历史批次几乎无收益甚至下降($N=128,K+1=2$ 反而降至 80.83%),表示陈旧偏差抵消了方差收益。其二,论文仅在四个相对小规模的离线控制环境上验证,且当前模型 128 批次显存(约 12.7 GB)已可承受,Cross-Batch QQ 的真正价值在于扩展到更大模型/更高分辨率观测/更长序列时才凸显,但本文未对此做大规模验证。从我自己的观察补充三点局限:(1) 缺少对更高维视觉控制任务(如 Atari、机器人操作)的验证,难以判断 QQ 在复杂动力学下的鲁棒性;(2) 理论分析(Prop 1-3)主要针对一维投影切片,多切片(有限 $S$)协同效应与 Cramér-Wold 蒙特卡洛误差的理论刻画较弱;(3) 排序操作在 GPU 上对大 $N$ 的计算效率与并列(ties)处理细节未充分讨论,可能影响极大规模训练的吞吐。

独立分析的弱点

弱点一:实验环境复杂度有限。四个任务状态空间低维,QQ 在高维视觉控制(Atari、机械臂、自动驾驶 latent world model)下的稳定性未检验——改进方向是扩展到 DMControl/BabyAI 等基准并考察与高分辨率观测的兼容性。弱点二:Cross-Batch QQ 的队列长度 $K$ 需手工选,论文未给自动准则——改进方向是设计基于表示陈旧度(相邻迭代 latent 分布 KL/MMD 监测)的自适应 $K$ 或依据偏差-方差分解式实时平衡。弱点三:QQ 损失依赖全批次排序,GPU 上对极大 $N$ 的排序开销与并列处理未充分讨论——改进方向是研究近似排序、近似最近分位数估计或分块匹配提升吞吐。弱点四:理论聚焦一维投影,多切片有限 $S$ 下的收敛速率与 Cramér-Wold 蒙特卡洛误差缺乏形式化保证——改进方向是建立多切片 QQ 非渐近收敛界与样本复杂度。弱点五:固定权重 3.5,未做权重敏感性分析——改进方向是补充权重消融与自适应调度。

未来方向

作者在结论中隐含的未来方向:将 QQ 正则化与 Cross-Batch QQ 扩展到更大规模世界模型、更高分辨率观测与更长序列,验证'解耦排序池规模与反传批次大小'这一设计在显存受限下的真正价值;并把'统计有效≠优化有效'这一思想推广到其它依赖分布匹配的目标。基于本成果可延伸的研究包括:(1) 把 QQ 匹配推广到非高斯先验或可学习目标分布(如混合高斯、流模型先验),探索秩匹配在任意目标下的传输性质;(2) 将 rank-matching 与其它表示学习目标(VICReg、对比学习、Barlow Twins)结合,研究其与坍缩防止机制的互补性;(3) 设计动态调整投影方向数 $S$ 与队列长度 $K$ 的自适应训练策略,利用方差-陈旧权衡实现每步最优配置;(4) 把 Cross-Batch 的 detach-queue 思想推广到其它需要大参考集合但显存受限的场景(如 MoCo 类对比学习、大 batch normalization 统计);(5) 探索 QQ 损失的二阶优化或排序感知的自适应优化器,进一步缓解排序梯度不连续带来的收敛抖动。

复现评估

复现评估总体较好。论文严格遵循 LeWM 的官方实现,在相同的数据预处理、模型训练与基于 CEM 的目标条件评估协议下进行,所有方法在相同的六个随机种子下评估,且对引入专属超参的基线(Sub-JEPA、SD-JEPA、SMWM)统一使用其最佳超参数,公平性有保障。作者声明 QQWorld 不引入新超参数、正则化权重固定为 3.5(所有环境一致),这大幅降低了调参成本。从 Table 4 的显存数据($N=32$ 约 3.4 GB、$N=128$ 约 12.7 GB)看,规模较小、单卡即可复现,算力门槛低。然而存在若干不确定因素:论文未明确说明代码或检查点是否开源;投影方向数 $S$ 的具体取值、Cross-Batch 队列的实现细节部分依赖 LeWM 默认设置未完全交代;Prop 1-3 的证明细节是否随正文/附录提供需核实。综合判断:若作者开源代码,复现难度低(单卡、固定超参、明确协议);若未开源需自行复现 LeWM 基线及其多个变体作为对照,复现难度中等。