World-as-Graph: Relational World Modeling Through Latent Space Graphs

• JEPA predicts future representations directly in latent space, while object-centric models represent scenes as object-level latent states or slots.These approaches support predictive representation learning and separate modeling of object dynamics.Source: §1 · p. 1
• Existing object-centric predictors may leave object-object relations implicit, limiting relational and temporal structure modeling and object-centric dynamic memory.The paper identifies weak self-supervised signals for interaction learning and potential loss of historical cues for long-horizon prediction.Source: §1 · p. 2
• WAG improves performance by up to 30.03 percentage points, with 2.4× faster training and up to 16.4× lower peak GPU memory.The experiments cover visual reasoning and robotic manipulation tasks.Source: §1 · p. 2
• WAG combines time-varying latent graphs and relation-aware object masking with a temporal GNN transition encoder and memory-driven future predictor.The masker selects objects according to relational centrality or temporal neighborhood changes; the transition encoder uses neighboring relations and historical memory.Source: §1 · p. 2
• WAG introduces latent-space graphs as relational inductive bias for object-centric JEPA-style world models, explicitly modeling object interactions and temporal dynamics.The paper distinguishes graphs used for world modeling from world modeling on graph-structured data.Source: §1 · pp. 2–3

2.1 PRELIMINARY

• WAG represents historical object-centric latent states as time-indexed graphs and predicts future slot states over a defined prediction horizon.The model is defined by relation-aware structure induction and object-centric memory transition.Source: §2.1 · p. 3

2.2 WORLD-AS-GRAPH FRAMEWORK

• A frozen encoder extracts slots; WAG induces dynamic graphs and masks relationally informative nodes, then propagates memory to predict future representations autoregressively.The memory-driven predictor evolves object-level memory states after the observed history ends.Source: §2.2 · pp. 3–4, §2.1 · p. 3

2.2.1 RELATION-AWARE STRUCTURE INDUCTION

• A parameter-free cosine-similarity KNN builder links each slot to Top-K neighbors; masking targets relationally central nodes or nodes with changing neighborhoods.Centrality is measured by in-degree, while temporal dynamics masking tracks changes in neighbor sets over the observed history.Source: §2.2.1 · pp. 4–5

2.2.2 OBJECT-CENTRIC MEMORY TRANSITION

• A temporal GNN and memory GRU update per-object states from partially observed graphs and prior memory, preserving object history and neighboring relational context.The module addresses masked objects' missing states in later history frames and the absence of future observations during rollout.Source: §2.2.2 · p. 5

2.2.3 LEARNING OBJECTIVE

• Training combines masked-state recovery with future-state prediction using MSE, while inference omits masking and autoregressively rolls out from fully observed history.The two losses are weighted in the overall objective; inference repeats rollout until the target horizon is reached.Source: §2.2.3 · p. 6

3.1 EXPERIMENTAL SETUP

• Experiments span CLEVRER visual reasoning and PushT robotic planning, using frozen object-centric encoders and autoregressive object-trajectory rollouts for visual reasoning.CLEVRER evaluation reports VQA accuracy across four question types; PushT planning is evaluated by success rate.Source: §3 · p. 6, §3.1 · p. 7

3.2 PERFORMANCE OF WAG ON VISUAL REASONING AND ROBOTIC MANIPULATION

• 1.97 and 8.84 pp are WAGB’s average-accuracy gains over C-JEPA with VideoSAUR and SAVi, respectively.With SAVi, counterfactual and predictive accuracy gains are 15.26 and 14.74 pp; batch-shared masking also consistently outperforms instance-wise masking.Source: §3.2 · pp. 7–8, Table 1 · p. 7
• 90.70% success with 6 × 128 tokens puts WAG near DINO-WM’s 91.33% with 196 × 384 tokens and above object-centric baselines at the compact budget.At 6 × 128 tokens, WAG exceeds OC-JEPA by 14.70 pp and C-JEPA by 2.03 pp.Source: §3.2 · p. 8, Table 2 · p. 8

3.3 ABLATION STUDY OF WAG ON KEY COMPONENTS

• Full WAG reaches 91.37% average accuracy, 1.97 pp above C-JEPA; removing the memory predictor causes the largest ablation drop, especially on predictive questions.Predictive accuracy falls from 89.54% to 85.55% without the memory predictor; removing DyGraph also lowers performance, while relational masking and TGNN add consistent gains.Source: §3.3 · p. 8, Table 3 · p. 8
• WAG’s KNN latent dynamic graphs outperform IID and subset uniform graphs across CLEVRER question types, indicating gains beyond graph connectivity alone.The results point to the value of constructing informative latent-space relations.Source: §3.3 · p. 8, Figure 3 · p. 8
• Under SAVi, relational centrality improves average per-question accuracy by 8.84 pp and leads the reported metrics; adding dynamic graph information also improves random masking.Temporal-dynamics masking also consistently improves over random masking without dynamic graph information.Source: §3.3 · pp. 8–9, Table 4 · p. 8

3.4 IN-DEPTH ANALYSIS OF WAG

• 4.83 versus 6.66 × 10^-3 mean latent MSE, and 8.9 versus 12.1 peak error, favor WAG over C-JEPA during a two-step collision chain.The sequence has A strike B, then B strike C; WAG’s latent graph changes as physical relations evolve, especially around the second collision.Source: §3.4 · p. 9, Figure 4 · p. 9
• 143 seconds per epoch versus C-JEPA’s 341.4 makes WAG approximately 2.4× faster, with peak GPU memory reduced from 6712 MiB to 410 MiB.Performance remains relatively stable across parameter combinations; K=3 generally performs strongly, with VideoSAUR reaching 91.37% average accuracy at K=3, M=2.Source: §3.4 · p. 9, Figure 5 · p. 9

4 CONCLUSION

• WAG combines relation-aware structure induction with object-centric memory transition so relational information and historical object states jointly support future prediction.The model explicitly represents object relations and their temporal evolution in latent space.Source: §4 · pp. 9–10
• Experiments on visual reasoning and robotic manipulation show consistent improvements, with ablations and physical-interaction analysis further supporting dynamic-graph representations.Source: §4 · p. 10

AI USE STATEMENT

• The authors report using generative AI for manuscript editing, figure presentation, and software development and debugging, with authors reviewing and revising AI-assisted text.Source: AI USE STATEMENT · p. 10

ETHICS STATEMENT; REPRODUCIBILITY STATEMENT

• The paper reports no human subjects, personal or sensitive data, or other ethical concerns requiring specific consideration, and places additional reproduction details in the Appendix.Source: ETHICS STATEMENT · p. 10, REPRODUCIBILITY STATEMENT · p. 10

A APPENDIX

• The appendix collects additional related work, theoretical analysis, the training algorithm, and experimental details and comprehensive results.The supplied overview does not specify those results.Source: A APPENDIX · p. 14

Joint-Embedding Predictive Architectures

• JEPA predicts future representations in latent space; Causal-JEPA masks latent states but selects masked objects randomly.The cited discussion describes JEPA as learning predictive representations from observations and forecasting future world states from context and actions.Source: §A.1 · p. 14

Graph World Models

• Unlike graph world models that use graphs as the modeled state space, WAG uses evolving latent graphs to guide relational masking and object-level memory transition.The comparison distinguishes graph-structured environment states from graphs used as relational inductive biases for object-centric predictive learning.Source: §A.1 · p. 14

Dynamic Graph Representation Learning

• Dynamic graph methods model changing temporal relations; WAG’s temporal GNN component is related to this line, but its goal is not general-purpose dynamic graph representation learning.Examples include TGAT’s temporal attention over neighbors and TGN’s recurrent node memory.Source: §A.1 · p. 14

A.2 THEORETICAL ANALYSIS

• During rollout, memory errors can change neighbor selection and relational message passing, affecting later predictions; the analysis asks when KNN neighborhoods remain stable and how errors accumulate.The reference memory trajectory assumes node norms bounded below by r.Source: §A.2 · p. 15
• The analysis bounds normalized-node and cosine-similarity changes, yielding the sufficient condition e_t < rγ_t/8 for preserving neighbor sets.The bound uses the neighborhood margin γ_t and a lower node-norm bound r.Source: §A.2 · p. 15
• If the margin condition holds throughout rollout, neighborhood separation prevents perturbation-induced rewiring, while transition sensitivity controls error accumulation; for L < 1, the bound is uniform over the horizon.These are sufficient analytical conditions, not guarantees provided by recurrent memory alone.Source: §A.2 · p. 16

A.3 ALGORITHMS

• Algorithm 1 summarizes the training procedure for WAG.The supplied passage does not give the algorithm’s steps.Source: §A.3 · p. 16

• PushT planning uses CEM and success requires matching the goal within 20 pixels and π/9 radians; visual-reasoning comparisons use OC-JEPA and C-JEPA.The robotic-manipulation baselines include DINO-WM, DINO-WM-Reg., OC-DINO-WM, OC-JEPA, and C-JEPA.Source: §A.4 · p. 16
• VideoSAUR is used for CLEVRER and PushT, SAVi is additionally evaluated on CLEVRER, and PushT planning appends proprioception and action tokens to four object slots.Evaluation selects the checkpoint with lowest validation loss, and reported WAGI/WAGB CLEVRER configurations vary by encoder and λfutu.Source: §A.4 · pp. 16–17
• With temporal-dynamics masking, WAGB improves counterfactual and explanatory per-question accuracy by 2.35 and 1.57 points, respectively, with a 0.05-point descriptive decrease.Source: Table A4 · p. 17, §A.4 · p. 18
• WAGB also slightly improves random masking, while relational-centrality results are mixed: WAGB leads overall by 0.06 points, but WAGI gains 1.16 explanatory points.Source: §A.4 · pp. 17–18, Table A4 · p. 17
• Increasing graph neighborhood K from 2 to 4 lowers PushT success by about 15 points, while increasing MLP width from 2048 to 4096 modestly lowers performance.The reported success rates fall from 90.7 to 76.7 with the smaller MLP and from 88.7 to 73.3 with the larger MLP.Source: Figure A1 · p. 18, §A.4 · p. 18
• WAG’s predictor has O(TND^2) complexity in the stated setting, with 16× fewer FLOPs per sample and 9–16× lower peak memory than C-JEPA and OC-JEPA.The setting is D=128, N=7, and K≤6; Figure A3 presents computational-cost, memory, and running-time analyses.Source: §A.4 · pp. 18–19, Figure A3 · p. 19