Title: SPARTAN: A Sparse Transformer World Model Attending to What Matters
Venue: NeurIPS 2025
Accepted: yes
Keywords: World Models

Attention

Local Causal Model
Abstract: Capturing the interactions between entities in a structured way plays a central role in world models that flexibly adapt to changes in the environment. Recent works motivate the benefits of models that explicitly represent the structure of interactions and formulate the problem as discovering local causal structures. In this work, we demonstrate that reliably capturing these relationships in complex settings remains challenging. To remedy this shortcoming, we postulate that sparsity is a critical ingredient for the discovery of such local structures. To this end we present the SPARse TrANsformer World model (SPARTAN), a Transformer-based world model that learns context-dependent interaction structures between entities in a scene. By applying sparsity regularisation on the attention patterns between object-factored tokens, SPARTAN learns sparse, context-dependent interaction graphs that accurately predict future object states. We further extend our model to adapt to sparse interventions with unknown targets on the dynamics of the environment. This results in a highly interpretable world model that can efficiently adapt to changes. Empirically, we evaluate SPARTAN against the current state-of-the-art in object-centric world models on observation-based environments and demonstrate that our model can learn local causal graphs that accurately reflects the underlying interactions between objects and achieve significantly improved few-shot adaptation to dynamics changes as well as robustness against distractors.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces SPARTAN, a Transformer-based world model designed to learn the underlying dynamics of multi-object systems from observations. The central argument is that world models should capture the sparse and local nature of physical interactions to achieve better generalization and adaptation. To this end, SPARTAN modifies the standard Transformer attention mechanism by replacing it with a learnable, sparse \"hard attention\" mask. This mask, representing a state-dependent local causal graph, is learned by applying a sparsity regularizer on the connections between object-centric tokens. The authors evaluate SPARTAN on simulated physics environments and a real-world traffic dataset, showing that it learns more accurate causal graphs, which in turn leads to superior robustness against distractors and more efficient adaptation compared to standard Transformer baselines.\n\nSTRENGTHS AND WEAKNESSES:\n**Strengths**\n\n- The paper is built on the clear and intuitive hypothesis that physical interactions are best modeled as sparse, local causal events rather than through dense, fully-connected graphs. The paper is also very clearly written, with effective visualizations, and is easy to follow.\n\n- The experimental results are thorough and convincing. SPARTAN not only matches or exceeds the prediction accuracy of dense Transformer baselines but also demonstrates a significantly better ability to identify the correct causal graph structure (lower SHD).\n\n-  The core contribution—a learnable, sparse hard-attention mechanism within a Transformer—is a simple and direct implementation of the paper's central hypothesis. \n\n**Weaknesses**\n\n- The binary attention mechanism seems well suited to the contact-based interactions dominant in the chosen datasets. How do the authors see this approach extending to physical systems governed by continuous fields of influence, like gravity, where interactions are continuous and never truly zero? \n\n- The matrix $\\bar{A}$ treats all paths equally in the sparsity calculation. This seems to conflate direct influence with indirect, multi-hop influence. Have the authors considered alternative regularization schemes that might be more sensitive to the structure of these paths? For example, one could penalize direct connections in the final layers more heavily than longer, multi-hop paths that are resolved in earlier layers.\n\n- The entire framework relies on the availability of high-quality, pre-disentangled object representations. In the simulated domains, this is achieved using ground-truth object masks to train a VAE. This is a significant assumption, as the performance of SPARTAN in a real-world setting would be fundamentally capped by the performance of an upstream object discovery and segmentation model, which is itself a challenging open problem.\n\n- (Minor) Images are not referred to in the text.\n\n- (Minor) Figure 2 is hard to parse. It would be nice to have a detailed subsection explaining it in text or in the caption.\n\nQUESTIONS:\nPlease refer to the weaknesses section above.\n\nLIMITATIONS:\nNA", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-15"}

comments: SUMMARY:
SPARTAN is a transformer-based method for learning sparse interaction patterns among entities, with a role of being a world model for RL agents.

STRENGTHS AND WEAKNESSES:
I find the approach to be simple and frankly quite amazed about the empirical performance of the Bernoulli binary masks.

I commend the authors for the use of performance metrics other than task returns. These make the experiments much more convincing.

QUESTIONS:
The ideas of this work (regarding sparsity, generalization, etc) are closely related to [1] and its follow-up [2], which are early works by Bengio's group, where sparse local interactions among entities are highlighted and then captured with transformer-based models. This paper would strongly benefit from acknowleging the commonalities and the differences from that line of work.


The idea on simple local models can be traced back to [3], which should be discussed properly.

[1] The Consciousness Prior
[2] A Consciousness-Inspired Planning Agent for Model-Based Reinforcement Learning
[3] Simple Local Models for Complex Dynamical Systems

LIMITATIONS:
The requirement for knowing the environment index I is not good for generalization. Can this be relaxed?
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 2
DATE: 2025-06-16

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a method to learn structured world models to recover the underlying casual structure of the environment. For this, the latent embeddings can be factored into independent causal mechanisms where with the data distribution shifts the changes to the causal structure is sparse. The paper presents a method to extend frameworks for state-based inputs to image observations and with different numbers of objects across scenes. The method, called SPARTAN, uses a multi-layer transformer where the adjacency matrix is sampled by sampling from Bernoulli distribution (to induce hard attention). The experiments on a few environments show that SPARTAN learns sparse graphs and is robust to change in dynamics.\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths\n1. The experiments support the claims proposed in the paper of using sparsity to accurately recover causal structure of the next state.  \n2. Overall the paper is well written and easy to follow.\n\n### Weaknesses\n1. A pseudocode of the training and inference steps would help better understand the algorithm. Also, this would highlight certain differences like the environment index is passed during training whereas the adaption token is obtained at test-time.\n2. The proposed idea is pretty simple and uses a Bernoulli distribution to induce sparsity and hard attention. There is no theoretical justification for why the model should work. Although the papers mentions this in limitations but it would be good to have some analysis.\n\nQUESTIONS:\n1. For the Waymo open dataset (Traffic), what does the observation include? Does it include the position of all objects or is the observation pixel-based? \n2. What is the effect of environment index I during training (as discussed in Sec 3.2)? Is there any ablation on training without that information?\n3. What is the value of $\\lambda_s$ defined in Eq 6? How was this hyperparameter tuned and what is the effect of different values of this coefficient? \n4. What is the effect of number of layers L? This looks like an important parameter as it also affects the adjacency matrix term for loss in Eq 6. \n\nTypos\n- Line 160– ‘the’ is mentioned two times.\n- The legend in Figure 4 shows STRAW (for blue).\n\nLIMITATIONS:\nThe paper has discussed the limitations in Section 6.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 2, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 1, "DATE": "2025-06-29"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes a Transformer-based world model that learns time-dependent local causal structures among entities from object-centric visual observations. Building on the idea that attention mechanisms can capture causal dependencies, the authors introduce hard attention combined with sparsity regularization. The resulting model, SPARTAN, is capable of adapting to interventions in the environment, even when the targets of those interventions are unknown. The method is evaluated across three domains, including Interventional Pong, CREATE, and the Waymo Open Dataset, and shows improved performance over baselines in terms of predicting causal structure, robustness to non-causal changes, and few-shot adaptation to novel dynamics.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths\n1. The paper is clearly written and well motivated.\n2. The examples in Sections 2.2 and 2.3 are intuitive and help clarify the rationale behind the modeling approach.\n3. Using the Waymo traffic dataset is a strong choice, as it reflects realistic, complex dynamics not easily captured in simulation environments.\n4. The visualizations in Figure 2 effectively illustrate the discovered local causal structure and demonstrate close alignment with ground-truth graphs.\n5. The adaptation mechanism, achieved through gradient-based optimization of intervention tokens, is interesting and useful for generalizing to new interventions with limited data.\n\nWeaknesses\n1. The assumption of access to ground-truth object segmentation for training is quite strong. Prior work, such as SlotFormer, learns from raw visual input without relying on ground-truth masks or object factorization. \n2. Details on the experimental setup are insufficient. For example, the number of training episodes and the exact meaning of the “Pred. Err.” metric are not thoroughly described.\n3. The approach primarily builds on existing components, including Gumbel-softmax sampling for hard attention and sparsity regularization. While the combination is effective, the technical contribution may be limited.\n\nQUESTIONS:\nQuestions for the Authors\n1. How does the method handle changing numbers of objects across time, such as vehicles entering or leaving the scene?\n2. Can the model deal with noisy or imperfect object segmentations, including missing or extra entities?\n3. Does the method require consistent object indexing over time (i.e., object tracking)?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-07-03"}