PaperlayerView exact revision on arXiv

Source-linked AI summary

World Modeling with Probabilistic Structure Integration

Klemen Kotar, Wanhee Lee, Rahul Venkatesh, Honglin Chen, Daniel Bear, Jared Watrous, Simon Kim, Khai Loong Aw, Lilian Naing Chen, Stefan Stojanov, Kevin Feigelis, Imran Thobani, Alex Durango, Khaled Jedoui, Atlas Kazemian, Dan Yamins

arXiv:2509.09737v1cs.CVcs.AIcs.LG

TL;DR

World models remain difficult to control and query flexibly, particularly in visual domains where language-based descriptions are sparse. PSI uses a probabilistic predictor to extract and integrate intermediate structures, and reports state-of-the-art results across novel-view synthesis and object manipulation benchmarks.

  • Problem

    World models lack rich, flexible control and query interfaces, while extracting visual features often requires supervised probes or separate task-specific models.

    Source: §1 · p. 3, §6.1 · p. 31

  • Method

    PSI trains a probabilistic predictor, extracts intermediate structures through zero-shot causal prompts, and integrates them as tokens used for conditioning and prediction.

    Source: §1 · pp. 3–4, §3.1 · p. 13

  • Results

    Ψ achieves state-of-the-art performance on novel-view synthesis and object manipulation, while flow-based segmentation reaches 0.57 mIoU on SpelkeEntitySeg.

    Source: §4.2 · p. 22, §4.3 · p. 24

  • Takeaways & Limitations

    Integrated structures serve as control surfaces for using flow to move objects, depth to steer viewpoints, and segments to select and edit.

    Source: §1 · p. 4

    The structures reported in the paper were created using hand-designed causal probes, leaving automated identification of useful probes an open question.

    Source: §6.4 · p. 36

Abstract

from arXiv · show

We present Probabilistic Structure Integration (PSI), a system for learning richly controllable and flexibly promptable world models from data. PSI consists of a three-step cycle. The first step, Probabilistic prediction, involves building a probabilistic graphical model Psi of the data, in the form of a random-access autoregressive sequence model. Psi supports a complete set of learned conditional distributions describing the dependence of any variables in the data on any other set of variables. In step 2, Structure extraction, we show how to extract underlying low-dimensional properties in the data, corresponding to a diverse set of meaningful "intermediate structures", in a zero-shot fashion via causal inference on Psi. Step 3, Integration, completes the cycle by converting these structures into new token types that are then continually mixed back into the training diet as conditioning signals and prediction targets. Each such cycle augments the capabilities of Psi, both allowing it to model the underlying data better, and creating new control handles -- akin to an LLM-like universal prompting language. We train an instance of Psi on 1.4 trillion tokens of internet video data; we use it to perform a variety of useful video prediction and understanding inferences; we extract state-of-the-art optical flow, self-supervised depth and object segmentation; and we use these structures to support a full cycle of predictive improvements.

1 Introduction

World models are difficult to control and query flexibly, especially for visual and other non-language data. PSI cycles through probabilistic prediction, zero-shot structure extraction, and integration to expand a model’s controls while improving its predictions.

1 Introduction

  • Existing world models often rely on coarse action sets or short text prompts, while visual understanding commonly requires separate generative and discriminative models.These interfaces do not readily support local manipulations such as changing a viewpoint or adjusting a specific object.Source: §1 · p. 3
  • PSI is a three-step cycle that trains a probabilistic predictor, extracts intermediate structures, and integrates them back into the predictor.The interface remains the same as cycles enlarge its control surfaces and predictive fidelity.Source: §1 · p. 3
  • PSI’s predictor Ψ models conditional distributions for local variables given arbitrary subsets of other variables, enabling flexible visual queries and plausible outcome sampling.For video, these variables are tokenized spatiotemporal RGB patches; Ψ is formulated as a random-access autoregressive transformer trained on many random orderings.Source: §1 · p. 3
  • PSI extracts low-dimensional structures through causal comparisons between predictions under factual and hypothetical interventions.The paper identifies optical flow, depth, and segments as examples of such structures.Source: §1 · pp. 3–4
  • PSI integrates extracted structures as new token types used both as conditioning signals and prediction targets, creating control surfaces for editing scenes.The same model can use flow to move objects, depth to steer viewpoints, and segments to select and edit.Source: §1 · p. 4

2 Probabilistic Prediction: Richly Controllable World Models

This section presents Ψ as a probabilistic graphical model of world data, using flexible conditioning to support diverse inference pathways and controllable generation. It describes how LRAS and entropy estimates make the approach tractable for visual data and support uncertainty-aware generation at speeds suited to different application priorities.

2 Probabilistic Prediction: Richly Controllable World Models

  • LRAS makes learning Ψ’s probabilistic graphical model efficient and scalable through Local Random-Access Sequence modeling.The model treats pointer-content pairs as traversals through the graph, recasting learning as autoregressive sequence prediction.Source: §2 · p. 5, §2.2 · pp. 6, 8

2.1 Defining the Probabilistic Relationships Between the Variables of the World

  • For any unspecified pointer, Ψ predicts its value distribution conditioned on a partially specified datum, enabling arbitrary patch-wise conditioning and completion.In video, sparse patches from a subsequent frame reveal motion cues while the preceding full frame supplies appearance information.Source: §2.1 · pp. 5–6

2.2 Local Random Access Sequence Modeling

  • Validation loss improved consistently from 100M to 7B parameters across three orders of magnitude, without saturation at 7B.The authors interpret this scaling behavior as evidence that LRAS transfers language-model scaling properties to structured visual data.Source: Figure 2 · p. 7, §2.2 · p. 8
  • LRAS interleaves pointer and content tokens so data can be serialized in arbitrary order, with pointers predicted or supplied to guide generation.Its Hierarchical Local Quantizer preserves patch locality, so local interventions behave predictably.Source: §2.2 · pp. 7–8

2.3 Diverse Inference Pathways Confer Rich Controllability

  • Without task-specific training, prompts enable patch-conditioned completion, counterfactual motion edits, and camera-conditioned novel views; combining constraints further narrows predictions.Often fewer than 10% of target-frame patches suffice to make completion near-deterministic; camera-conditioned predictions remain multimodal where scene scale, intrinsics, or occluded regions are ambiguous.Source: §2.3 · pp. 9–11

2.4 Entropy and Uncertainty Management

  • Entropy maps expose uncertain scene regions, and revealing patches progressively reduces uncertainty, supporting adaptive generation and active perception.Sequential decoding maximizes quality but requires multiple forward passes, while parallel decoding is faster but assumes conditional independence; grouped decoding trades quality for speed.Source: §2.4 · pp. 12–13
  • Ψ manages uncertainty and generation speed to support robot planning that prioritizes speed and video generation that prioritizes sequential consistency.The passage contrasts real-time robot planning, where speed matters more than perfect quality, with high-fidelity video generation, where sequential consistency is paramount.Source: §2.4 · p. 13

3 Structure Extraction: Prompts as Causal Inference

This section extracts intermediate visual structures as causal properties exposed by prompts to Ψ, then turns them into token interfaces that can be mixed with RGB. Across flow, object segmentation, and depth, it uses hypothetical interventions and the model’s distributional predictions.

3.1 Intermediate Structures Are Zero-Shot Causal Prompts

  • Causal prompts expose visual structures by perturbing a factor of interest while holding most of the scene fixed and measuring Ψ’s response.This zero-shot approach is proposed as an alternative to training task-specific probes for visual properties.Source: §3.1 · pp. 13–14, §3 · p. 13

3.2 Optical Flow via Tracer Counterfactuals

  • For optical flow, Ψ compares factual and perturbed conditional predictions using KL divergence, and the displacement to the most-shifted patch gives the flow estimate.The method can produce diverse, individually coherent flow predictions rather than averaging incompatible motion hypotheses.Source: §3.2 · p. 14
  • Ψ achieves state-of-the-art optical-flow results on TAP-DAVIS and TAP-Kubric, surpassing task-specific and world-model methods.Table 1 covers TAP-Vid First tracking on DAVIS and Kubric, where Ψ outperforms two-frame baselines.Source: Table 1 · p. 15, §3.2 · p. 16

3.3 Object Segments via Motion Hypotheticals

  • Ψ’s sharp, multimodal hypothetical motions support segmentation of non-rigid objects and allow different samples to reflect ambiguous object boundaries.The paper contrasts these predictions with blurred mean predictions from CWM and notes that plausible interpretations can vary by sample.Source: §3.3 · p. 16
  • Object segments are extracted by shifting a copied patch, generating a hypothetical frame, and grouping pixels with coherent induced motion.The resulting motion field reveals object boundaries without task-specific supervision.Source: §3.3 · p. 16, Figure 8 · p. 17

3.4 Depth via Viewpoint Hypotheticals

  • For depth, camera-translation prompts create hypothetical views; optical-flow disparity is converted to unscaled depth without depth supervision.The method uses disparity d_j = |u_j| and depth z_j = b/d_j.Source: §3.4 · p. 18, Figure 9 · p. 18
  • Sharp, geometrically consistent views support precise parallax measurements, while averaging estimates across samples can isolate depth from variable object dynamics.The distributional predictions also represent multiple valid depth interpretations in ambiguous regions.Source: §3.4 · p. 19

3.5 Extracted Structures Define New Token Types

  • Flow, depth, and segments are treated as local token types that Ψ can read or predict alongside RGB, and can be interleaved with RGB during continued training.These structures are described respectively as linking appearance to future position, mediating camera motion and parallax, and capturing object membership and coherent motion.Source: §3.5 · p. 19

4 Integrating Structure: Ratcheting Prompting and Control

PSI integrates extracted structures as tokens in a unified autoregressive model, expanding its control and extraction capabilities while also improving video prediction. The reported applications include precise object and camera control, cleaner structure extraction, higher-order motion properties, and more dynamic frame predictions.

4.1 A Universal Mechanism for Structure Integration

  • PSI integrates structures by tokenizing them, interleaving their tokens with frame tokens, and continuing training, so one model supports varied prediction and conditioning tasks without architectural changes.The sequence order and causal mask determine which tokens condition which predictions; the model can continue training from a stable checkpoint without catastrophic forgetting.Source: §4.1 · pp. 20–21, Figure 10 · p. 21
  • The paper presents integration as a recursive cycle in which improved predictions support cleaner extractions, and each added structure enables richer prompts and further abstractions.Its examples progress from pixels to motion to affordances, while treating each integrated structure as both a representation and a basis for higher-order properties.Source: §4.1 · pp. 21–22

4.2 Exposing Richer Control Surfaces

  • Sparse flow vectors provide precise motion controls, and the model propagates them into coherent motion fields for rigid objects and material-appropriate deformations for non-rigid ones.Flow tokens also let users specify exact pixel displacements for camera motion, avoiding the range of interpretations required by raw camera-pose tokens.Source: §4.2 · p. 22, Figure 11 · p. 23
  • On 3DEditBench, PSI achieves 22.73 PSNR and 0.797 Edit Adherence; on WildRGB-D, flow-based camera control outperforms specialized novel-view-synthesis models.The authors report state-of-the-art performance on both tasks; the supplied Table 2 caption identifies the benchmark comparisons.Source: §4.2 · p. 22, Table 2 · p. 23

4.3 Improving Structure Extractions

  • Integrated flow improves structure extraction by enabling pure-motion prompts and direct flow prediction, reducing segmentation confusion between movement, appearance changes, and shadows.Camera-conditioned flow also encodes parallax, which the paper uses to obtain sharper depth boundaries.Source: §4.3 · p. 24, Figure 12 · p. 24
  • Flow-based segmentation achieves 0.57 mIoU on SpelkeEntitySeg, and PSI reports state-of-the-art self-supervised depth results, particularly on dynamic scenes.Flow prompts isolate motion from appearance changes, while directly predicting flow avoids reliance on an off-the-shelf flow extractor.Source: §4.3 · p. 24, Table 3 · p. 25

4.4 Bootstrapping Further New Tokens

  • Flow-model beliefs yield higher-order properties: probability of motion estimates which patches may move next, while expected displacement captures likely motion direction and magnitude.The paper describes these as statistical properties derived from predictions over possible flow vectors, including before motion occurs.Source: §4.4 · pp. 25–26, Figure 13 · p. 25

4.5 Upgrading the Underlying Predictor

  • With flow integration, PSI outperforms its RGB-only variant and the COSMOS baseline on video prediction across multiple datasets, while flow-conditioned generation captures motion missed by direct RGB prediction.The authors attribute this to predicting flow first and conditioning RGB generation on it; the supplied Table 4 caption specifies single-frame prediction 500 ms ahead on DAVIS and YouTube.Source: §4.5 · pp. 26–27, Figure 14 · p. 27, Table 4 · p. 27

5 Example Applications

The applications illustrate PSI’s use for physical video editing, stability reasoning in Visual Jenga, and anticipatory motion mapping for robotics. Across these tasks, the model uses localized edits or motion predictions to reason about downstream physical changes.

5 Example Applications

  • The applications are illustrative rather than exhaustive, demonstrating one prompting interface for factual rollouts, counterfactual analysis, and physically consistent manipulation.The paper frames these as applications to existing computer-vision and physical-reasoning problems.Source: §5 · p. 28

5.1 Physical Video Editing with PSI

  • In physical video editing, PSI rolls out futures that reflect contact, momentum transfer, and gravity; translating a ball in the context frame produces a counterfactual miss with the pins remaining upright.The same editable-context mechanism supports factual, hypothetical, and counterfactual generations.Source: §5.1 · pp. 28–29

5.2 Visual Jenga

  • For Visual Jenga, PSI combines asymmetric bidirectional motion-dependency tests with flow-based amodal completion to identify a minimally disruptive object and check whether the remaining scene stays stable.The paper reports that this physics-aware method significantly outperforms inpainting baselines.Source: §5.2 · p. 29

5.3 Robotic Motion Map

  • From a single static image, PSI estimates which patches are likely to move, offering robotic planning cues about manipulable objects and stable regions without interaction or motion observation.The paper reports generalization to novel object configurations without task-specific training and uncertainty estimates for manipulation.Source: §5.3 · p. 30

6 Discussion

The discussion situates PSI’s cycle of probabilistic prediction, causal structure extraction, and integration in relation to prompting, generative modeling, and causal inference, while outlining open questions about its scope and development. It also raises questions about integrating contrastive techniques and memory, and about possible relations to human cognition and neuroscience.

6 Discussion

  • PSI’s three-step cycle uses pointer-indexed local sequences for probabilistic prediction, causal inference for zero-shot structure extraction, and token intermixing to improve prediction and control.The paper presents the framework in natural video and describes it as a general approach to modeling complex natural data.Source: §6 · pp. 30–31

6.1 Beyond Representation Learning: From Probes to Prompting to Integration

  • PSI extends prompting beyond language by extracting nonlinguistic structures with counterfactual programs and integrating them as model conditioners and targets.The authors avoid text conditioning initially because it could provide a shortcut to future prediction; they envision adding it after visual intermediates capture scene dynamics.Source: §6.1 · pp. 31–32

6.2 Related Work

  • PSI’s distributional autoregression supports inferences across potential futures, while localized conditioning enables precise spatial control and effective hypotheticals.Unlike diffusion models, which model one sample at a time, PSI makes quantities such as entropy and motion maps accessible; LRAS combines random traversal, strict token locality, and pointer-based addressing.Source: §6.2 · pp. 32–33, Figure 15 · p. 32
  • PSI treats data patches as proxies for actions, learning to perform hypotheticals without explicit action data, and can incorporate real action data when available.Relative to CWM, PSI replaces regression with distributional prediction and adds sequence-mixing integration, at the cost of substantially more model parameters.Source: §6.2 · pp. 33–34

6.3 Predictions, Interventions, Hypotheticals & Counterfactuals

  • PSI defines interventions as predictions conditioned on synthesized inputs and hypotheticals as sequential samples, allowing a spectrum of interventions rather than a strict prediction–intervention divide.Generalization to unobserved conditioning is empirical; failures may require adding examples of relevant interventions to the training data.Source: §6.3 · p. 34
  • Flow extraction is a stronger causal inference than segmentation because it can use more ground truth as a counterfactual, enabling early models to bootstrap weaker extractions such as segmentation.PSI’s intervention mechanisms remain uncertain, and its tracer analysis can mistake a tracer for lens dirt; the authors note that real-world causal analysis also depends on learned empirical understanding.Source: §6.3 · pp. 34–35

6.4 Limitations and Key Open Problems

  • PSI has yet to show integration gains from global structures, and applying it beyond vision leaves open how to identify useful causal inferences in domains without established intermediates.The paper also identifies automated discovery of useful causal probes as a key open problem; a parameterized domain-specific language is one possible direction.Source: §6.4 · p. 36
  • Future work includes learning pointer sequences, distilling slow structure extraction, and determining whether PSI encoders support semantic category readouts.The paper suggests entropy information for pointer-sequence learning and sparse-to-dense prediction for distillation.Source: §6.4 · p. 36
  • Whether or to what extent contrastive techniques should be integrated with PSI remains an important open question.Source: §6.4 · p. 37
  • PSI can model sequences spanning at least a few seconds, but long-range generation and planning may require online-adaptive memory, whose best design remains uncertain.The authors suggest memory akin to hippocampal-cortical interaction as a possible direction.Source: §6.4 · p. 37

6.5 Does PSI have anything to do with human neuroscience or cognition?

  • Whether PSI relates to human cognition or neuroscience is unknown, and the authors present filling AI capability gaps through contravariance as a future-work strategy.They describe PSI as an exploration of novel algorithm spaces and say whether this strategy is sound remains to be seen.Source: §6.5 · p. 37
Loading 2509.09737v1…