Identifiable World Models from Pretrained Diffusion Representations

1 INTRODUCTION

• Generative fidelity alone does not establish that a diffusion model’s internal coordinates recover underlying state variables or causal interactions.The introduction frames this as an identifiability problem because multiple latent parameterizations can explain the same observations.Source: §1 · p. 1
• Under the stated assumptions, the aligned coordinates preserve dynamic-SCM factorization and parent structure, while lagged causal edges correspond to inverse transition-noise Jacobian sparsity.For additive-noise dynamics, the graph is estimated using sparsity of a learned forward-dynamics Jacobian.Source: §1 · p. 2
• The work asks whether frozen diffusion representations can acquire identifiable, causally meaningful coordinates without retraining the generative backbone.The aligned embedding is intended as a coordinate chart on the frozen representation, not a replacement generative representation.Source: §1 · pp. 1–2
• ConDA learns a lightweight auxiliary-aligned map from diffusion latents to coordinates that recover latent states up to permutation and componentwise invertible transformations.The framework separates the underlying state, diffusion representation, observed frame, and aligned embedding.Source: §1 · p. 2
• TCL- and GCL-based ConDA achieve strong, often near-perfect blockwise state recovery and recover lagged causal structure, including exact graph recovery in simulated falling-body dynamics.In controlled Microduck dynamics, the learned representations preserve responses to held-out control perturbations.Source: §1 · p. 3

World models and diffusion representations

• World models and diffusion representations support prediction or generation, but are not generally identifiable and may not recover physical states or causal dependencies.The paper therefore asks whether auxiliary-aligned embeddings of frozen pretrained diffusion representations can recover latent states and lagged dependencies.Source: §2 · p. 3

Contrastive learning and nonlinear ICA

• TCL uses temporal nonstationarity for nonlinear ICA, GCL uses auxiliary variables, and CEBRA aligns embeddings with behavioral, temporal, or neural variables.Auxiliary information can provide structure for identifiability, whereas nonlinear latent variables are not identifiable without additional assumptions (Hyvarinen & Morioka, 2016; 2017; Hyvarinen et al., 2019; Schneider et al., 2023; Hyvärinen & Pajunen, 1999; Locatello et al., 2019).Source: §2 · p. 3

Causal representation learning and causal discovery

• Causal discovery for dynamical systems uses temporal ordering and lagged dependencies, but many methods assume access to meaningful state variables.Here, lagged edges are characterized by sparsity in the Jacobian of the inverse transition-noise map.Source: §2 · p. 3

Positioning of this work

• This work studies identifiable, causally meaningful states from frozen diffusion representations, learning only a lightweight auxiliary-aligned map rather than jointly learning the encoder and temporal or causal model.It connects auxiliary-variable nonlinear ICA (Hyvarinen et al., 2019; Khemakhem et al., 2020), contrastive density-ratio estimation (Oord et al., 2018), and Jacobian-based causal discovery.Source: §2 · p. 4

Contrastive Diffusion Alignment

• ConDA learns an alignment map from frozen diffusion latents using temporal or auxiliary information, leaving the pretrained diffusion backbone fixed.Its contrastive variants include TCL, GCL, and CEBRA, with positive pairs selected by auxiliary-context similarity in the supervised formulation.Source: §3 · p. 5

Identifiability

• Under the stated TCL/GCL identification conditions, aligned coordinates recover latent states up to permutation and componentwise invertible transformations.The conditions include a smooth, injective diffusion representation and sufficient variation in auxiliary-context modulation coefficients.Source: §3 · p. 5

Interpretation and scope

• The coordinatewise-identification guarantee does not cover arbitrary contrastive training or the InfoNCE-based CEBRA variant, and approximate assumptions may yield mixtures or blocks.A generic pairwise similarity objective is also outside the stated coordinatewise-identification claim.Source: §3 · p. 5

Preservation of dynamic structure

• After coordinatewise identification and resolving permutation ambiguity, the learned dynamics preserve the original dynamic SCM's lagged parent sets and conditional factorization.This result assumes the latent dynamics factorize by parent sets and have mutually independent noises conditional on the auxiliary variable.Source: §3 · pp. 5–6

Lagged graph recovery

• Under smoothness, noise-map, differentiability, and faithfulness assumptions, the inverse transition-noise Jacobian's sparsity pattern characterizes the lagged causal graph.For additive-noise dynamics, the forward and inverse Jacobians have the same sparsity pattern, enabling forward-Jacobian magnitudes as empirical edge scores.Source: §3 · p. 6

Experimental Setup

• The evaluation covers CausalVerse physical-fall and robotics videos and Microduck bipedal locomotion, which includes externally imposed control commands.Microduck supports assessment of controlled dynamics and interventional prediction in addition to state and graph recovery.Source: §4 · p. 6
• ConDA uses CEBRA, GCL, and TCL objectives and is compared with IDOL, CaRiNG, TDRL, temporal SuaVE, and iVAE across five random seeds.Physical Fall and Microduck use 5,000 videos each, while Robotics uses 2,700; each dataset includes a held-out test split.Source: §4 · p. 6

Evaluation Pipeline

• The pipeline freezes a trained conditional latent diffusion model, aligns its representations, and evaluates state recovery and lag-1 graph recovery.Binary graph metrics use the reference parent count, while AUPRC measures threshold-free edge ranking; Microduck graphs use state- and control-dependent transition Jacobians.Source: §4 · pp. 6–7

Evaluation Metrics

• Identifiability is measured with coordinate-wise MCC and block-wise R2, while graph recovery uses F1, accuracy, SHD, AUPRC, and SID.Microduck intervention predictions compare matched rollouts under specified control commands with a zero-command baseline.Source: §4 · p. 7

Identifiability and Graph Recovery

• On the Robotics Study, ConDA-TCL/GCL reach MCC 0.679 and 0.683 with R2 = 1.000; iVAE performs slightly better on graph-recovery metrics.The TCL and GCL graph results remain competitive.Source: §5 · p. 7
• On Microduck, ConDA-GCL has the highest MCC (0.629), while ConDA-TCL leads in F1 (0.769), accuracy (0.888), SHD (0.111), and SID (0.076).ConDA-TCL's AUPRC is comparable to iVAE.Source: §5 · p. 7

Block-wise vs. Coordinate-wise Recovery

• ConDA-TCL/GCL show stronger block-wise R2 than coordinate-wise MCC, retaining state-reconstruction information even when coordinates do not correspond one-to-one with physical variables.Near-perfect block-wise recovery can support accurate lagged-dependency recovery without exact coordinate alignment in every setting.Source: §5 · pp. 7–8

Identifiability and Graph Recovery

• On Physical Fall, ConDA-TCL/GCL achieve R2 ≈1, F1/AUPRC = 1, and SHD/SID = 0; no method dominates every metric and dataset.Across the Robotics Study and Microduck, the ConDA variants remain competitive, with different methods leading particular metrics.Source: §5 · p. 7, Table 1 · p. 8

Dynamic Intervention Effect

• Under held-out Microduck controls, TCL recovers intervention direction and much of the longer-horizon response, but its immediate response is smoother and attenuated.For do(ux = −0.25), the response has the correct sign and temporal trend but attenuated magnitude; for do(uy = −0.17), it follows the simulator more closely and captures the dominant negative effect on vy.Source: §5 · p. 8

6 CONCLUSION AND LIMITATIONS

• TCL- and GCL-based alignment identify states up to permutation and componentwise invertible transformations, preserve dynamic-SCM structure, and support lagged graph recovery across physical and robotic systems.Microduck experiments also report preservation of intervention direction and temporal response structure under held-out controls.Source: §6 · pp. 8–9
• Figure 2 compares recovered lag-1 adjacency patterns across methods and systems; ConDA variants more closely match true-state sparsity, especially TCL and GCL.The reported intervention studies use controlled simulation and open-loop perturbations.Source: Figure 2 · p. 9, §6 · p. 9
• The theoretical guarantees are population-level and rely on informative auxiliary variables, separability, and regularity; intervention studies are limited to controlled simulation and open-loop perturbations.The paper identifies partially observed, higher-order, nonstationary, and real-world scientific systems as future extensions.Source: §6 · p. 9

A.1 IDENTIFIABILITY OF LATENT STATES FROM AUXILIARY-VARIABLE NONLINEAR ICA

• Conditional independence and exponential-family assumptions yield identifiability up to permutation and componentwise invertible transformations.The proof sketch attributes this result to auxiliary-variable nonlinear ICA.Source: §A.1 · p. 13

A.2 CONTRASTIVE LEARNING AS DENSITY RATIO ESTIMATION

• With an expressive score class and many negatives, the population-optimal supervised contrastive score is the positive-versus-negative pair log-likelihood ratio, up to a conditioning-only additive term.The proof frames the objective as classifying matched pairs drawn from ppos against mismatched pairs drawn from pneg.Source: §A.2 · p. 13

• A smooth inverse of g on its image preserves the state information available to the specified TCL or GCL procedure.Source: §A.3 · p. 13
• Under assumptions (2)–(3), the applicable identification result yields its stated equivalence class when applied to x_t.Source: §A.3 · p. 13
• When that class is coordinatewise after any required unmixing, the output satisfies Eq. (10).Source: §A.3 · p. 13

• Theorem 2 yields invertible scalar functions ϕ_i after resolving the permutation ambiguity.Source: §A.4 · p. 13
• The transformed mechanism depends on the same parent coordinates because componentwise invertible transformations do not mix coordinates.Source: §A.4 · p. 13
• Componentwise reparameterization preserves independence of the noise coordinates and the SCM’s parent sets.Source: §A.4 · p. 13

A.5 JACOBIAN-BASED GRAPH RECOVERY

• Under bijective noise mechanisms and structural faithfulness, lagged parent links correspond to non-canceling dependence of inferred noise on the relevant prior-state coordinates.For a non-parent coordinate, the inferred noise does not depend on it; faithfulness ensures parent dependence does not cancel almost everywhere, establishing the equivalence.Source: §A.5 · pp. 13–14

A.6 CAUSAL GRAPH FOR PHYSICAL FALL

• Physical Fall provides an analytically specified benchmark with three-dimensional position and velocity states and lagged position-to-position, velocity-to-velocity, and velocity-to-position dependencies.Gravity and time serve as auxiliary variables; learned states and estimated lagged adjacency are compared with the known system.Source: §A.6 · p. 14

A.7 PHYSICS-INFORMED REFERENCE LAG-1 DYNAMICS GRAPH FOR ROBOTICS STUDY

• The robotics reference graph uses a 25-dimensional state and specifies 25 self-lag edges plus seven joint-velocity-to-position cross-lag edges, totaling 32 edges.The graph is physics-informed; time, object position, and object velocity are auxiliary information for representation learning.Source: §A.7 · p. 14

A.8 CONTROLLED CAUSAL DYNAMICS FOR MICRODUCK

• Microduck enables intervention-focused evaluation by modeling the six-dimensional robot state as a controlled transition driven by three-dimensional commands.Forward and lateral velocity commands are expected to affect corresponding future body velocities, and yaw-rate commands future yaw rate; the remaining graph is derived from simulator dynamics and controlled perturbations rather than imposed arbitrarily.Source: §A.8 · p. 15

A.9 EVALUATION METRIC DEFINITIONS

• Evaluation uses component-level MCC, blockwise R2, and lagged-graph metrics, including thresholded accuracy and precision, AUPRC, and normalized SHD and SID.R2 = 1 indicates perfect blockwise predictability; SHD = 0 denotes exact structural recovery, and SID = 0 denotes no structural intervention errors.Source: §A.9 · pp. 15–16

A.10 DYNAMIC INTERVENTION EFFECT

• Microduck intervention effects are estimated from matched trajectories reconstructed to the same state, comparing positive and negative forward or lateral commands with zero-command baselines over 25 steps.The learned aligned representation is compared with both MuJoCo and a true-state dynamics baseline; intervention commands lie within the empirical training support.Source: §A.10 · p. 16

A.10.1 RESULTS

• The aligned TCL representation preserves the qualitative direction and temporal structure of Microduck responses, with stronger recovery for lateral than forward perturbations and attenuated forward effects.Forward-command responses had the correct sign but generally underestimated magnitude; lateral responses were more accurate. The true-state baseline was more accurate for forward interventions, while the representation model had lower RMSE than that baseline for the negative lateral intervention at H ≥5.Source: Figure 5 · p. 17, §A.10.1 · pp. 17–18

A.12 DETAILED RELATED WORK

• Prior world models prioritize prediction, planning, or control, while diffusion representations are optimized for generative fidelity; neither objective alone ensures identifiable or causally interpretable state coordinates.The paper studies whether auxiliary alignment of frozen diffusion representations can recover dynamical states and support structural causal analysis.Source: §A.12 · pp. 19–20

A.5 JACOBIAN-BASED GRAPH RECOVERY

• Auxiliary-aligned contrastive learning on frozen diffusion latents is connected to density-ratio estimation and auxiliary-variable nonlinear ICA to establish sufficient conditions for identifying latent dynamical states.The auxiliary information may include time, labels, experimental conditions, or context.Source: §A.12 · p. 20
• Componentwise identifiable dynamical coordinates preserve the underlying dynamic SCM’s factorization and lagged parent structure.This connects latent identifiability directly to preservation of causal structure in time-dependent systems.Source: §A.12 · p. 21