PaperlayerView exact revision on arXiv

Source-linked AI summary

QQWorld: Quantile-Quantile Matching for World Model Regularization

Zhoushun Yu, Xiaoyu Hu, Xiangyu Xu

arXiv:2607.28415v1cs.LGcs.AIcs.CVcs.MMcs.RO

TL;DR

LeWM’s EP regularizer can leave heavy-tailed latent deviations insufficiently controlled because its corrective gradients vanish in the tails. QQWorld replaces EP with rank-matched Gaussian quantile alignment and improves latent Gaussianity and downstream planning performance.

  • Problem

    EP’s corrective force rapidly vanishes for tail samples, leaving heavy-tailed latent deviations insufficiently controlled despite Gaussian regularization.

  • Method

    QQWorld replaces EP with quantile-quantile matching that aligns ordered projected latents with corresponding Gaussian quantiles and provides stronger tail correction.

  • Results

    QQWorld improves latent Gaussianity and downstream planning performance across the reported experiments.

  • Takeaways & Limitations

    Effective world-model regularization requires training objectives with informative, well-behaved gradients throughout optimization, not merely distributional discrepancy measures.

  • Takeaways & Limitations

    Cross-batch QQ has a bias-variance trade-off: excessively long queues can introduce staleness bias and reduce planning performance.

Abstract

from arXiv · show

Latent world models enable efficient planning by predicting future states in a compact representation space, but their performance depends critically on the quality of the learned latent distribution. LeWorldModel (LeWM) regularizes its latents toward an isotropic Gaussian using the Epps-Pulley (EP) objective. We show that the corrective gradients of EP rapidly vanish for isolated tail samples, leaving heavy-tailed deviations insufficiently controlled. To address this limitation, we propose QQWorld, which replaces EP with a quantile-quantile matching objective that directly aligns projected latent samples with rank-matched Gaussian quantiles, thereby maintaining effective corrective gradients in the tails. We further develop cross-batch QQ, which enlarges the effective ranking pool using detached samples from previous batches, and characterize its bias-variance trade-off. Across four control environments, QQWorld effectively improves the average planning success rate of LeWM, while consistently yielding better Gaussian alignment and thinner latent tails.

1 Introduction

Latent world models regularize learned latents toward an isotropic Gaussian, but LeWM’s EP regularizer provides weak correction for far-tail samples, leaving heavy-tailed deviations insufficiently controlled. QQWorld replaces EP with rank-matched Gaussian quantiles, preserving effective tail gradients and improving Gaussian alignment and downstream planning performance.

  • Motivation: LeWM uses the Epps–Pulley test as a distribution-level regularizer to match learned latent marginals toward an isotropic Gaussian.This extends beyond matching only low-order moments.
  • Problem: LeWM latents exhibit pronounced heavy tails, which can push learned dynamics into poorly represented regions and amplify multi-step rollout errors.Heavy tails also increase mismatch with the target Gaussian prior.
  • Problem: EP’s corrective force rapidly vanishes for tail samples beyond the kernel’s interaction scale, leaving heavy-tailed deviations insufficiently controlled.The EP penalty can be reinterpreted as a unit-bandwidth kernel discrepancy from N(0, 1).
  • Method: QQWorld replaces EP with quantile–quantile matching that aligns ordered latent values with corresponding Gaussian quantiles and maintains effective corrective gradients in both tails.Its rank-matched transport signal grows linearly with quantile deviation and is a drop-in LeWM replacement without additional hyperparameters.
  • Results: QQWorld reduces mean QQ RMSE from 0.157 to 0.121 compared with LeWM, indicating closer agreement with standard Gaussian quantiles.Figure 1 describes QQ RMSE across one-dimensional latent projections.
  • Results: QQWorld consistently improves downstream planning performance and introduces a memory-efficient cross-batch ranking strategy with a characterized bias–variance trade-off.The cross-batch strategy enlarges the effective ranking pool without significant GPU memory overhead.

2 Related Work

Related work spans pixel-space and latent-space world models, representation-collapse prevention, and classical distribution-matching methods. QQWorld instead targets the distribution-matching objective by replacing EP with rank-aligned Gaussian quantile matching for latent world-model regularization.

  • World Models: World models predict future dynamics to evaluate candidate actions, while pixel-space prediction makes planning rollouts expensive by modeling high-dimensional observations and irrelevant visual details.JEPA offers a reconstruction-free alternative by predicting future representations directly in latent space.
  • Latent Regularization: Representation collapse maps different observations to identical or weakly varying embeddings; variance control helps, but moment-based regularization does not determine the full latent distribution.LeWM addresses distributional regularization by introducing the EP test objective.
  • QQWorld: EP corrective gradients rapidly vanish for tail samples, whereas QQWorld matches projected samples to rank-aligned Gaussian quantiles to address this limitation.QQ-based objectives had rarely been explored for latent world-model regularization before this work.
  • Distribution Matching: Classical distribution-comparison tools include KS, EP, and MMD, with EP equivalently interpretable as MMD using a Gaussian kernel.Wasserstein methods provide another discrepancy family, while sliced Wasserstein extends the comparison to high dimensions through one-dimensional projections.
  • QQWorld: Concurrent sliced Wasserstein work targets self-supervised JEPA representation learning rather than world-model planning and requires an additional variance term for scale control.QQWorld directly replaces the EP regularizer in an end-to-end latent world model using a single quantile-matching objective.

3 Method

QQWorld regularizes latent representations by matching Gaussian quantiles along random projections, addressing EP’s vanishing tail gradients. Cross-Batch QQ expands the ranking pool with detached historical features, trading lower rank-estimation variance for possible staleness bias.

  • Random-projection regularization: Normality is enforced along random one-dimensional projections, whose finite-direction average approximates a sliced discrepancy targeting the joint latent distribution.The method projects batch latents onto directions uniformly sampled from the unit sphere and penalizes departures from N(0, 1).
  • EP limitation: EP’s Gaussian-kernel interaction scale causes its restoring force for extreme projected latents to decay toward zero, allowing heavy tails to persist.For an outlier h with other samples in a bounded bulk, the pairwise term becomes asymptotically negligible and the restoring gradient decays superexponentially after its maximum.
  • QQ matching: QQWorld replaces EP with direct matching between sorted projected samples and rank-matched Gaussian quantiles.The QQ loss can also be interpreted as a quadrature approximation to the squared 2-Wasserstein distance, rather than a limited-bandwidth kernel discrepancy.
  • QQ matching: Each projected latent receives a restoring signal of magnitude 2|xn −qρ(n)|, which strengthens with quantile discrepancy and alleviates heavy tails.Sorting does not create practical optimization difficulty because tied configurations are locally repelling and samples assigned different ranks separate toward distinct Gaussian quantiles.
  • Distributional control: QQ provides stronger distributional control than EP: driving empirical order statistics toward Gaussian quantiles drives EP discrepancy to zero, but small EP loss can coexist with arbitrarily large tail deviations.The asymptotically negligible Gaussian-kernel penalty on distant observations is the source of this one-way control.
  • Cross-Batch QQ: Cross-Batch QQ increases the effective ranking set from N to M using detached historical features without increasing retained computation graphs, reducing variance but introducing possible staleness bias.Increasing the queue length K reduces variance from order N −1 to order {N(K + 1)}−1, while encoder changes can make historical distributions differ from the current one.

4 Experiments

Across four control environments, QQWorld improves LeWM’s planning success while producing better Gaussian alignment and thinner latent tails. Cross-batch QQ retains much of the benefit with lower memory, though excessively long queues introduce staleness bias.

  • Experimental setup: QQWorld replaces LeWM’s EP regularizer with QQ matching, using the same datasets and evaluation protocol and a QQ regularization weight of 3.5.Experiments cover Two-Room, PushT, Reacher, and OGBench-Cube; all methods use six random seeds.
  • Planning performance: 85.08% average planning success improves the LeWM baseline by 5.33 percentage points, with QQWorld outperforming LeWM in all four environments.The reported values are means across six random seeds.
  • Gaussian alignment: 15.8% lower KS and 31.4% lower EP statistics show that QQWorld achieves consistently better Gaussian alignment than LeWM.QQWorld also achieves lower QQ RMSE, despite LeWM being trained directly with the EP objective.
  • Tail behavior: QQWorld progressively suppresses heavy radial tails toward the standard Gaussian reference, whereas LeWM retains substantially heavier tails throughout training.The tail calibration error εtail measures deviation from the Gaussian tail, whose reference probability is 0.10.
  • Task-relevant information: QQWorld consistently matches or outperforms LeWM when probing agent location, block location, and block angle from frozen latents.Both linear predictors and lightweight MLPs are evaluated on PushT, with results averaged over six random seeds.
  • Cross-batch QQ: 83.50% success with N = 32 and K + 1 = 3 exceeds LeWM’s 79.75%, while reducing batch size by 4× and GPU memory usage by 73%.This configuration uses detached samples from previous batches to enlarge the ranking pool while backpropagating only through the current batch.
  • Cross-batch QQ: Increasing the queue from three to four batches lowers success from 83.50% to 81.75%, consistent with variance reduction from moderate queues and staleness bias from excessive ones.Adding historical batches also fails to improve performance for N = 64 or N = 128.

5 Conclusion

QQWorld replaces the EP regularizer by matching projected latent samples directly to Gaussian quantiles, strengthening correction for tail deviations. Cross-batch QQ improves quantile estimation under limited GPU memory, while experiments show gains in latent Gaussianity and downstream planning performance.

  • 5 Conclusion: QQWorld replaces EP with direct matching between projected latent samples and Gaussian quantiles, providing stronger correction for tail deviations.The method is presented as a simple replacement for the EP regularizer in latent world models.
  • 5 Conclusion: Cross-batch QQ improves quantile estimation when GPU memory is limited.It does so by addressing quantile estimation under limited GPU memory.
  • 5 Conclusion: Extensive experiments show that QQWorld improves latent Gaussianity and downstream planning performance.The conclusion reports improvements in both representation quality and planning outcomes.
  • 5 Conclusion: The results indicate that a statistic effective for measuring distributional discrepancy is not necessarily an effective training objective for world models.This broader conclusion motivates distinguishing discrepancy measurement from training-objective effectiveness.
Loading 2607.28415v1…