SPARTAN: A Sparse Transformer World Model Attending to What Matters
Authors: Anson Lei
Bernhard Schölkopf
Ingmar Posner
Venue: NeurIPS 2025 poster ()
Keywords: World Models
Attention
Local Causal Model
Primary area: deep_learning
TL;DR: We develop a sparsity-regularised transformer world model which learns attention patterns that accurately captures the interactions between entities, resulting in improved interpretability, robustness and adaptation.

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.

SPARTAN: A Sparse Transformer World Model Attending to What Matters
Anson Lei Applied AI Lab University of Oxford, UK ansongrobots. ox. ac.uk
Bernhard Schölkopf MPI for Intelligent Systems Tübingen, Germany
bs@tue. .aa.ee..pp.... de
Ingmar Posner Applied AI Lab University of Oxford, UK ingmardrobots. ox. ac. uk
Abstract
Capturing the interactions between entities in 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. this end, we present the SPARse TrANsformer World model (SPARTAN), a Transformer-biee based world model that learns context-dependent interaction structures between entities in a scene. By applying sparsity regularisation the attention patterns between object- factored tokens, SPARTAN learns sparse, context-dependent interaction graphs that accurately predict future object states. We further extend model to adapt to sparse interventions with unknown targets in the dynamics of the environment. This results in highly interpretable world model that efficiently adapt changes. Empirically, we evaluate SPARTAN against the current state-of-the-as object-centric world models in observation- based environments and demonstrate that our model can learn local causal graphs that accurately reflect the underlying interactions between objects, achieving significantly improved few-shot adaptation to dynamics changes, as well as robustness against distractors.
Introduction
World Models [1]] have emerged in recent years as a promising paradigm to enable a wide range of downstream tasks such video prediction (50 [3), physical reasoning 771 and model-based RL [2] Recent advances have employed the transformer architecture
(48) to develop world models capable of longer horizons in increasingly complex settings 30 38
performing accurate predictions Owing to the flexibility of attention mechanisms, the transformer architecture can be combined with object-centric representations (28) accurately capture interactions between objects, achieving state-of-thed prediction performance [53] However, the ability to generalise and adapt to changes in the environment that transformer world models can be sensitive to changes in distractors [39], further highlighting the need for models to capture robust and generalisable interactions. As such, recent works have
data-efficient manner remains a significant challenge. Evidence suggests
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

Global Causal Graph
Time cependent Graph
SPARTAN
641
Hard adtention Sparsty Yogg.
In the context of modelling physical interactions, global causal graph is often uninfor- mative and close to fully-commected. A time-dependent local causal graph better captures the sparse nature of interactions between entities. We present SPARTAN, a Transformer-based world model that discovers local causal structure using hard attention with sparsity regularisation.
Figure 1:
argued that explicitly representing interaction structures between entities, i.e. when and how entities influence each other, is critical in improving generalisation 335 and adaptation efficiency 25
To this end, the intersection between causality and machine learning 40 offers a promising frame- work for building structured models that capture interactions in principled manner. Through this lens, the forward dynamics can be factorised into independent causal mechanisms, where prediction for each entity is made based solely other entities that have causal influence it. Here, the Sparse Mechanism Shift hypothesis [41 posits that natural changes in the data distribution can be explained by sparse changes in these causal mechanisms. This suggests that world model reflecting the causal structure of the world can adapt efficiently to changes by updating only a small part of the model while most of the model remains invariant. This causal view of world modelling has motivated prior works which explored the benefits of structured world models [25 14]. These approaches aim learn a graph that encodes how entities in an environment influence each other. However, these methods learn fixed graph to fit an entire dataset, which is unsuitable for modelling realistic physical interactions, such as collisions between objects, that manifest as temporally sparse events. To remedy this, recent approaches
42 15 formulate the problem as learning deteedeeenddnt local causal graphs (Fig. [T]. Despite their success in simple state-based settings, scaling prior local causal discovery methods to image observations with complex dynamics remains significant bottleneck. This hinders the application of these methods to realistic scenarios where, for example, objects are represented via learned embeddings and the number of objects can vary from scene to scene.
Taking inspiration from the conceptual framework of local causal models, the goal of this work is to develop a transformer-based world model that faithfully reflects the interaction structure between objects while maintaining the flexibility and accuracy afforded by the transformer architecture. The key insight is that attention mechanisms serve as natural starting point for constructing state- dependent graphs. By applying sparsity regularisation, an inductive bias used in prior local causal discovery methods (15), on modified transformer architecture with hard attention, we develop SPARTAN, a Transformer- based world model with a learnable local causal graph. We evaluate our model on observation-baat based environments with physical interactions and real-life traffic motion prediction dataset [44) To the best of our knowledge, SPARTAN the first method that learns local causal structures beyond simple state-based settings. Compared to prior transformer world models, our results show that SPARTAN able to accurately identify causal edges, resulting in significantly improved robustness against distractors, few-shot adaptation, and interpretability.
2 Background
We begin this section by introducing our problem setting and providing brief outline of causal graphical models and causal discovery methods. These are then discussed relation to our motivation in section 2.3 where we note in particular the role of causality in world models. The goal of our work is to learn a causal model that predicts future states from observation. In order to model interventions, we train our model on a set of intervened environments with unknown intervention targets. Specifically, the training data of our model consists of observation sequences sampled from different environments (x" x ..... x" [1]. where x' is the observation at timestep t and I is the

---

environment index. Here, an intervention amounts to altering the dynamics such as changing the friction of an object. We assume access to representation model that maps observations to factored latent states. 8 = S; x x Sv, such that each S; is an embedding that represents the state of an entity such as an object in the scene. Examples of such representation models include object-centric models [28] [0] or disentangled causal representations [1 27 24 where each S, is an object slot or a causal variable respectively. Our aim is to learn a transition function p(s (s*) that captures the
local causal structure.
2.1 Causal Graphical Models
A causal graphical model (CGM) 33 consists of a set of random variables V = (Vi,..., .,. Vx J. a directed acyclic graph = (V, E) random variable V, PVV)). where Pa(V,) is the set of parents of V, Each edge in the graph (i,5) means that Vi is cause of Vj. The joint distribution over the variables can be factorised as
over the set of nodes V and conditional distribution for each
N
= II p(us Pa(V)) 1=0
One important feature of CGMs is that they support interventions which are local changes to the distribution that correspond external changes to the data generation process. Given the intervention and its targets - C V, the joint distribution is
p (Ui,,..,, = II p(u:|P((V))
II (((y))Pa(VV))). My ETT
(2)
where ('(.) is the intervened conditional distribution. In contrast to a probabilistic graphical model, given a set of interventions, CGM defines a family of distributions.
Recent works in causal discovery, which aim to learn causal graph structure from data, formulate the identification of causal structures amounts to pruning spurious edges and interventions and finding a model that can explain the data
mewssttyrregllarisee optimisation [55 6 Intuitively, this
with the sparsest causal graph. Building on these works, sparsity regularisation as inductive bias to learn local causal graphs.
2.2 Local Causal Models
In the context of modelling physical systems, since object interactions often temporally sparse, we further define local causal models [34] 42 [5]. Suppose we observe the values of a subset of variables, Dobs V, yobs obs For some neighbourhood L containing gbb the induced local causal graph g is constructed by removing all edges (i.j) from the global causal graph G where within the neighbourhood, i.e., the local graph is a subgraph of the global
V, IL (a,P), graph obtained by removing any "inactive" edges based on the observed state. Note that the structure of the local graph now dependent on the observed variables. Concretely, in the context of transition functions, the set of observed variables is the states of the current timestep s As simple example, consider two billiard balls. If it observed that they are far apart, then there is no local edge, since they cannot influence each other. However, when they are about to collide, their subsequent states will depend on each other, and there is now a local edge between them.
Analogously to standard causal discovery, our goal is to infer the sparsest graph at each timestep, such that the next observation can be explained. In the local case, the optimisation objective for local causal discovery is to minimise the expected number of edges and targets with respect to the data distribution, where the causal edges and intervention targets are stateddepnndent and are built dynamically at each timestep. [15] provides theoretical analysis of the conditions under which local causal structures can be identified using sparse regularisation.
2.3 Causality and World Models
Conceptually, we argue that world models should be structured to reflect the underlying sparse and local causal structure of the dynamies to perform efficient adaptation. illustrate our point, consider the setting of traffic behaviours: changes in traffic rules arise when travelling to countries where
3

---

vehicles drive on the opposite side of the road. Although the relative position of vehicles given the road boundaries may change, most traffic behaviours (e.g., stopping at a red light or lane keeping) would remain the same. A world model that can efficiently adapt should update only a small subset of learnt dynamics (i.e., the relative position of vehicles to the road boundary) while leaving other dynamics unchanged. Our proposition is motivated by the sparse mechanism shift hypothesis [31 41 which states that naturally occurring distribution shifts can be attributed to sparse interventions on causal mechanisms. Formally, this means that reasonable changes in the environment can be modelled by changing a small subset of the conditional distributions in Eq. | while the rest remain invariant. In the special case where interventions act only on variables that are not causally related to model predictions, e.g., removing irrelevant objects in a scene [39] models that reflect the correct causal structure would remain robust.
While learning global causal graph sufficient in some settings, in the context of dynamics models, such a graph is often close to fully connected since any entities that interact with each other at any point are connected in the graph regardless of how unlikely the interaction is. Returning to the traffic example, a global causal graph would connect every vehicle in the scene, as they can all influence each other when close together. However, events such as "vehicle A causes vehicle B to stop" can be more appropriately captured by local causal graphs with the edge A - B but not by edges to other vehicles that are irrelevant at the time. As such, we argue that local causal graphs are more suitable for models to fully exploit the sparse structure of the problem at hand. In the following section, present SPARTAN, a transformer-bast based instantiation of
local causal discovery.
3 Sparse Transformer World Models
Our goal is to develop a world model that learns local causal models as transition functions. We build the Transformer architecture 48
it achieves sate-offthe-art prediction performance in faject-factored world models [53). Moreover, its attention mechanism provides a natural way to control the flow of information between object tokens. While prior work [34 argues that soft attention is a strong enough inductive bias for learning local causal graphs and proposes a thresholding heuristic on the attention patterns, we posit that this is insufficient beyond simple state-based scenarios and that appropriate sparsity regularisation and hard attention play a crucial role in scaling this up to more complex, observation- based environments. To demonstrate with a simple example, consider a system with only one causal edge, of + 5 and many other irrelevant nodes. A soft-attention model can predict sft as long as the attention value for st is non-zero, i.c., information can flow from the i-th token to the j-th token. However, the model is free to have non-zero attention values over other tokens. In this case, there is no guarantee that the attention value on sf is the highest amongst all tokens, and therefore, applying a threshold on the attention values may 'catch' spurious edges or miss necessary edges. Instead, we need model that masks the information flow between the tokens and penalises connections between tokens. Any model that contains spurious connections would have a high loss due to the sparsity penalty, while models that do not contain the s - 8551 would suffer from low prediction accuracy. Under this scheme, the model that achieves the best training loss would be the one that contains only the correct causal edge. In the following section, we present SPARrse TrANsformer World models (SPARTAN), an instantiation of a masked transition model using a Transformer-based architecture with hard attention.
3.1 Learnable Sparse Connections
We start with the architecture of a single layer before extending the apporach to the multi-layer case. Given object representations sf., as input, we obtain the keys, queries, and values, (ki, Gs v;) via linear projection, as in the standard Transformer. Using these, we sample an adjacency matrix,
Aij Bern kg)),
where Bern..) is the Bernoulli distribution and a is the sigmoid function. The hidden features are then computed using the standard scaled dot- propuct attention with the adjacency matrix masks.
h Σ Aijerple Kj/ Vdk
8/ LMMMLPhh; + 8:),
E, expla k2 Vi.)
where at+1 is the prediction for the i-th object in the next timestep. Here, the adjacency matrix acts as mask that disallows information flows from tokens that are not in the parent set of the querying

---

token. Note that the adjacency matrix is dependent on the current state via the key-query pairs. The learnt adjacency matrix can therefore be interpreted as the local causal graph, i.e. A, means 53 E (stt)) The sampling step made differentiable via Gumbel softmax trick [17]. Typically, multiple Transformer layers are stacked sequentially for expressiveness. This presents problem for the learnt adjacency matrix, as information can flow between tokens a multi-hop manner across different layers: for example, token can influence token j via i k + I without the edge (i, j) being present in any of the adjacency matrices. SPARTAN mitigates this problem by tracking the number of paths from token i j. Concretely, for L layers, we compute the path matrix,
= (A-+1 + ()....A42 + TT) IT,
(5)
where A' is the adjaceney matrix at layer and I the identity matrix. The identity matrix is added due to residual connections. The path matrix A has the property that Aiş is the number of paths that lead from to i. In this case, sf is local causal parent of s/+1 iff Aij >= 1.
3.2 Interventions
We can extend our model to represent interventions in cases where the model is trained on a dataset with a set of environments with intervened dynamics During training, the model has access to the environment index for each transition but does not know which object is affected by the intervention. Conditioned on the current state &i:v and the environment index J, the model needs to identify a subset of IIN as intervention targets and change the predictions (cf. Eq. 2) To do this, we keep a learnable codebook of intervention tokens, TiKK Rdxk where d is the dimension of the state tokens and K is the number of interventions. We append the corresponding intervention token Ti to the object tokens and proceed as described in the previous subsection. Here, A1(N+1) 0 indicates that there is path from the intervention token to s!+1 and therefore s is not an intervention target, i.e., the prediction for 8 is not affected by the changes in this environment.
f+1
At test time, the model adapts to unknown interventions by observing sequences without the environ- ment index. We perform adaptation by finding intervention token Todapt E _ that best fits the observed data via gradient descent in the token space. Note that the adaptation intervention token need not be in the discrete set of intervention tokens from training and can therefore model unseen environments. In Sec. 4 we show that this approach can generalise to the composition of previously seen interventions and leads to efficient adaptation.
3.3 Training
Our aim is fit the data distribution with the sparsest possible model in terms of causal edges and interventions. The training objective is to minimise the expected loss with sparsity regularisation,
min Ea 1 ((5 Aa1
(6)
where are the model parameters, including the parameters for the Transformer and the intervention tokens, L is loss function such as MSE, and X, denotes the regularisation hyperparameter. Note that due to the intervention token being part of the input during training, |A] is the sum of the number of causal edges and the intervention targets. In practice, the model is sensitive to the choice of À since high À can lead to mode collapse. We alleviate this via Lagrangian relaxation, which schedules the regularisation weight. The details of this setup and other training information are in App. A
4 Experiments
Our experiments are designed to investigate the following guiding questions:
1. Does sparsity enable local causal discovery from observations?
2. Can the sparse model match the prediction accuracy of fully-connected models?
3. Does learning interventions lead to improved robustness and adaptation sample efficiency?
Note that this is an extension, rather than requirement, to our approach to capture nulti-envvronment training. In single environment cases, SPARTAN can be directly applied as described above.
5

---

Internettiral Pooi
0:3
0:3
0:3
0:3
0:3
0:3
STARDAN
Greund-ruth
Tanuforer
Figure 2: Example rollouts in the two simulated environments with the learnt local causal graph visualised. Blue and red arrows between icons indicate the learnt causal edges and intervention targets respectively. In the Interventional Pong example, the intervention is that the ball slows down in the middle. SPARTAN correctly identifies the same causal dependencies as the ground-truth (e.g. ball causes the paddles to follow). The Transformer baseline learns edges that do not correspond to the ground-truthe truth.. Similarly in CREATE, SPARTAN learns the correct causal edges (e.g. green ball bounces off the blue plank) while Transformer learns many spurious interventions.
Walkk Cerect Alarion Alantinn
Figure 3: Visualisation of the causal relationships learned by the models compared to human labeled data. The ego vehicle is blue, Transparent rectangles means that the vehicle is not a causal parent of the ego vehicle. SPARTAN learns similar attention pattern to human data (e.g. focus vehicles that are moving in the same direction) whereas Transformer learns many spurious edges.
Datasets We evaluate our model in three domains: Interventional Pong, CREATE, and Traffic. The Interventional Pong dataset (27], standard benchmark for causal representation learning, is based on the Pong game with interventions. The CREATE [66] environment is a 2D physics simulation that consists of interacting objects such ladders, cannons, and balls. This is similar to the PHYRE dataset 2 on which SlotFormer (53] is evaluated, but it has more interaction types, allowing for more interventions better showcase the capabilities of SPARTAN. For these two domains, object slot representations by providing ground-truth masks for objects and encoding each object separately using a VAE (2T]. In order evaluate our model on more realistic tasks, the Traffic domain uses the Waymo Open Dataset 44 which is collected in real life. The task is to predict the motion of the ego vehicle given the observed vehicle trajectories (i.e., past positions) and map layout lines. For evaluating the accuracy of the learnt causal graphs, we compare them against ground-trui truth causal graphs for the two simulated domains. For the Traffic domain, we compare against habelelabelled causal graphs [39] In the simulated domains, we also train the models using interventional data, as described in Sec. 3.2 Interventions are defined as changes to the simulated dynamics such as changing the strength of gravity. See App. B for details.
obtain

---

Interventional Pong
CREATE
1000
ADDA
10x
4003
Seans Envirormets
Unswer Envirommnnss
Seert Envisemenns
Unswen Environments
Figure Adaptation emrors on the two datasets. Each model has access 5 trajectories with unknown environment index. "Unseen environments" refers to interventions that are not in the training set. Red lines indicates the prediction error if the environment index provided, which serves as a lower bound. SPARTAN (blue) consistently achieves the lowest errors across environments.
Baselines In order to investigate whether SPARTAN can maintain the state-of-the-ar prediction accuracy achieved by Transformer-based models, compare method against Transformer baseline using a Transformer for learning dynamics. This can be seen as an instantiation of the SlotFormer [53] architecture with one-step context length, except that the slots are learnt using ground-truth object masks. For the Traffic domain, we use MTR 45], state-of-the-an transformer- based model designed for traffic data as the base architecture. In terms of causal discovery, existing methods [51 15
only operate on state observations where each state represented as a scalar and are therefore not applicable our setup, which observes object embeddings. We provide more detailed discussion on the applicability of these prior approaches and offer extra experimental comparisons in App. C Closest to our work CoDA KE which argues that soft attention in Transformers alone is sufficient for learning local graphs and proposes a thresholding heuristic, We compare against this approach and show that sparsity regularisation in SPARTAN is crucial in local causal discovery. To corroborate that local causal graphs are more suitable for modelling physical interactions, we further compare against the Global Graph baseline, which is based on the transition function proposed in VCD [25) and AdaRI. [14], which learns fixed global causal graph using masked MIPs. The details of the baseline models are provided in App. A
Table Rollout prediction error and average SHD between the learned graphs and the ground-truth causal graphs. Lower is better.
Intv. Pong
CREATE
Traffic
Method
Pred. Err.
SHD
Pred. Err.
SHD
Pred. Err.
SHD
SPARTAN (Ours)
8.60
1.51
246.6
1.17
0.52
6.84
Transformer
8.83
6.37
265.5
6.29
0.51
10.6
Global Graph
11.36
5.42
277.4
8.77
4.1 Graph Learning
The prediction accuracy of the models is shown in table Overall, SPARTAN is able to make predictions at a similar level of accuracy to and often outperforms the Transformer baseline. We investigate whether casual discovery via sparsity regularisation is required to identify correct local causal graphs the setting of dynamics modelling. While (34] suggests that extracting causal graphs by thresholding attention patterns is sufficient in simple environments, here demonstrate that in more complex environments, learning hard attention patterns via sparsity regularisation is necessary for accurate local causal discovery. Fig. visualises example rollouts with the local causal graphs Here, we observe that SPARTAN can reliably recover the relevant causal edges and intervention targets each timestep. In the Interventional Pong environment, our method identifies that the paddles are constantly influenced by the ball as they follow it. In contrast, the Transformer baseline
'When evaluating the Transformer baseline, we pick a threshold that achieves the best structural hamming distance. Note that this is possible in practice without access to ground-tron information.
7

---

gives erroneous edges such as "Score + Ball" (frame 1). Similarly, in the CREATE environment. SPARTAN accurately discovers sparse interactions between objects. In the Traffic domain, Fig. El shows that SPARTAN learns to attend to vehicles in neighbouring lanes, which is largely consistent with the human labels, whereas the Transformer baseline attends to many irrelevant vehicles. In this case, heuristics based on spatial proximity would also fail since there are vehicles that are far ahead that are relevant while vehicles that are close by but in the opposite lane, for example, should be ignored. This further highlights the importance of learning local causal graphs. Quantitatively, we evaluate the Structural Hammiing Distance, a commonly used metric in graph structure learning, between the learned graphs and the ground-truth. Table that SPARTAN achieves significantly lower distance compared to the baselines across all domains. We do not apply the Global Graph baseline to the Traffic domain since the fixed graph approach cannot be extended to include different number of objects.
4.2 Robustness and Adaptation
Table 2: Percentage change of error when non-causal entities are removed. Lower is better.
SPARTAN
Transformer
Local Attn.
Global Graph
Intervention Pong
24.5 + 4.4
1140.2 + 15
22.2 + 1.3
CREATE
28.6 + 1.7
138.2 + 5.6
73.6 + 1.3
Traffic
13.9 + 0.2
32.8 + 0.4
25.5 + 0.3
proposed evaluating the robustness of models by considering the change in the prediction error when non-causal entities are removed from the scene. This can be seen as a special case of interventions that target only variables that are not the causes of the predicted variable. Concretely, for each object in the scene, we remove any object that is not in the ground-truth parent set and compute the mean absolute value of the percentage change in the prediction error. Here, model that correctly captures the causal structure would be robust, whereas a model learning spurious correlations would have large fluctuations. Table showws that, compared to non-lagularised Transformer, SPARTAN is significantly more robust, suggesting that the Transformer overfits to spurious correlations between entities. In the Interventional Pong domain, the ground-truth causal graph is largely consistent across time steps. As such, the Global Graph baseline is sufficient in capturing the causal structure and hence remains robust. However, this degrades in the more complex CREATE domain, where the local causal graph is more detteddeend,nt,, while SPARTAN achieves the best robustness. In the Traffic domain, we also implement a local attention variant of the Transformer baseline, proposed in MTR [43. which applies a spatial K-nearest neighbour mask on the attention pattern between the tokens. Whilst this heuristic improves the model, SPARTAN remains significantly more robust to non-causal changes. A wide range of model architectures is evaluated in [39], with reported changes ranging from 25% 38%, which is consistent with our findings. Remarkably, the authors in (39 propose a data augmentation scheme using ground-truth information, resulting in change of 22%, which is still significantly worse than what SPARTAN achieves, highlighting the efficacy of our method.
39
To investigate adaptation efficiency, we train the models to adapt to a sample of five trajectories from an intervened environment. For SPARTAN, we perform adaptation as described in Sec. 3.2 We follow a similar procedure of optimising over intervention tokens for the Global Graph baseline. For the Transformer baseline, adaptation performed via gradient finetuning on the provided trajectories. We also adapt the models to previously unseen environments where the interventions performed are not included in the training set. In the Interventional Pong dataset, the new environments are obtained by combining interventions from the training set, e.g. intervening on the ball and the motion of the paddle at the same time. In CREATE, we change the scene composition as well as the dynamics, i.e. we change the number of objects in the scene. In this scenario, the Global Graph baseline cannot be applied. We present the results in Fig. 4 which shows that SPARTAN consistently outperforms the baselines in this shot adaptation setting.
8

---

5 Related Works
Our work is situated within the context of learning world models
particular, we build
on methods that learn dynamics models over object-centric representations 10 28 52]. In this space, various architectures have been proposed to model the interactions between object slots as message passing, including GNNs [22 45 RNN [47 18] and pair-wise interactions 49]. SlotFormer [53] achieves state-of-the-art results by applying a Transformer-based dynamics model on pre- trained object-slots 28 We build on this approach by extending the Transformer architecture with learnable sparse masks, resulting in improved robustness, adaptability, and interpretability without compromising prediction performance or the flexibility to model varying numbers of objects.
Our approach is motivated by recent advances in causal machine learning
40
19
In particular,
37 provides theoretical account of how identifying causal world model leads to generalisable agents. Concretely, our sparsity regularisation-bant based approach is motivated by optimisation- -based causal discovery methods [6] that use sparsity as an inductive bias for learning causal relationships. In this vein, VCD [25] and AdaRL [14] learn fixed causal graphs using sparsity regularisation in the context of world models. We extend these approaches by learning state-dependent local graphs that are better suited for modelling physical interactions. The notion local causal model for dynamics
is discussed in 35 34 11], which demonstrates the potential benefits of local causal models in terms of data- efficience,, robustness, and exploration. The CoDA family of works (3 35 I uses local causal graphs to enable the generation of counterfactual data for data augmentation. However, estimating a local causal model from data remains challenging. Several approaches to learning local causal structures have been proposed. CAT 42 uses conditional mutual information to estimate the local influence that an agent has on the environment. ELDEN [51] uses sparsity rrgulaiy-eeullrrised partial derivatives between states based on a learnt dynamics model to infer local causal graphs. FCDL 15 uses sparsity regularisation to learn a quantised codebook of possible causal connections. The key difference between these approaches and ours is that they require state observations, whereas our method can operate on object-centric image embeddings. Moreover, owing to the flexibility of transformers, SPARTAN is able to learn in more realistic scenarios where the number of objects or the scene composition changes between data samples. We provide a more detailed discussion on the applicability of existing methods in our setting and present additional experiments in App. C Closest to our work is [34] which uses attention patterns in transformer layers as causal edges. We empirically show that our model achieves significantly improved accuracy in terms of the discovered graph structure.
In a broader context, enforeing sparse connections between tokens in Transformers has also been explored in other settings, such as NLP 9 g 26 and computer vision (56). These methods require pre-defined masks based on domain knowledge, such as paragraph structures. SPARTAN differs from these approaches as it does not require any pre-defined masks and can be seen as a way to learn these masks from data. To this end, investigating the application of sparsity-regularise hard attention in these domains outside of world modelling offers a tantalising avenue for future research.
On a conceptual level, our work is also related to the notion of the Consciousness Prior 3 which advocates for an attention-like mechanism that selects a sparse bottleneck of 'active' entities in any given scene. This idea is realised in (54, showing that the sparsity induced by the Conseiousness Prior leads to improved generalisation in model-based planning tasks. Our approach differs from this idea in two principal ways: 1) (54) requires pre-specifying the number of active entities, whereas the sparse attention in our approach flexibly attends to a varying number of objects in a state-dependent manner. 2) In the Conseiousness Prior framework, the model sparse in the sense that it chooses a small subset of objects that is active. In contrast, we consider a complementary kind of sparsity. where we model the dynamics of all objects but use sparse attention to determine how each object depends on other objects.
6 Conclusion
We tackle the problem of adaptation in world models through the lens of local causal models. To this end, we propose SPARTAN, structured world model that jointly performs dynamics model learning and causal discovery. We show on image-based datasets that attention mechanisms alone are not sufficient for discovering causal relationships and develop novel sparsity regularisation scheme

---

that learns accurate causal graphs, resulting in significantly improved interpretability and few-shot adaptation capabilities without compromising prediction accuracy,
Limitations and future work. There are several limitations to our approach that serve as pointers for future investigations. 1) We empirically show that local causal graphs can be learnt from data. While the use of sparsity regularisation is grounded in prior works that have theoretical guarantees [5]. we do not make theoretical statements about the identifiability of local causal graphs in the setting where the scene composition can vary between each data sample. Future work should explore the conditions under which local graphs can be identified. 2) During adaptation. SPARTAN adapts by optimising over the learnt 'intervention space'. We have shown that this is sufficient to generalise to combinations of seen interventions or different numbers of objects. However, in more extreme cases where the test time environment contains completely new behaviours (e.g. ball teleporting), there might not be a corresponding intervention token in the space, while a finetuning approach would be able to converge to the right behaviour (given enough data). One avenue of exploration is to consider procedurally generating intervened environments, akin to domain randomisation [32 (1)) to cover the space of all meaningful interventions. 3) Our approach relies pre-disentangled object representations. Note that while we have used ground-truth segmentation masks for objects in our evaluations, pre-trained object-centric representations can also be utilised. We provide further discussion on this in App. D An interesting extension of our work is to investigate whether local sparsity can induce the emergence of disentangled causal representations 25 23 an encoder with the dynamics model.
by jointly training
Acknowledgments
This research supported by an EPSRC Programme Grant (EP/V000748/1). We would also like to thank the University of Oxford Advanced Research Computing (ARC) hutp:///dx.doi.or//11.. 1022////ooooo22555 and the SCAN facility in carrying out this work.
References
[1] Ezra Ameperosa, Jeremy A. Collins, Mrinal Jain, and Animesh Garg, Rocoda: Counterfactual data augmentation for data- robot learning from demonstrations, 2024. URL https
//arxiv. org/ www.//gg/////////////.////eeee// 2411 16959
[2] Anton Bakhtin, Laurens van der Maaten, Justin Johnson, Laura Gustafson, and Ross Girshick. Phyre: A new benchmark for physical reasoning. 1rXiv:1908. 056566, 2019.
[3] Yoshua Bengio. The consciousness prior. arXiv preprint 1rXiv:1709.0883 2017.
[4] Yoshua Bengio, Tristan Deleu, Nasim Rahaman, Nan Rosemary Ke, Sebastien Lachapelle, Olexa Bilaniuk, Anirudh Goyal, and Christopher Pal. A meta-transfer objective for learning to disentangle causal mechanisms. In International Conference on Learning Representations, www.:eee/ooooo
2020. URL URLphttss:/pppeereiiie htt.tt/ee.pppppeeeeeeeee///////r
[5] Dimitri P Bertsekas. Constrained optimization and Lagrange multiplier methods. Academic press, 2014.
[6] Philippe Brouillard, Sébastien Lachapelle, Alexandre Lacoste, Simon Lacoste-Julien, and Alexandre Drouin. Differentiable causal discovery from interventional data. Advances in Neural Information Processing Systems, 33:21865-21888 2020.
[7] Zhenfang Chen, Jiayuan Mao, Jiajun Wu, Kwan- Yee K Wong, Joshua B. Tenenbaum, and Chuang Gan. Grounding physical concepts of objects and events through dynamic visual reasoning. In International Conference on Learning Representations, 2021.
[8] Rewon Child, Scott Gray, Alec Radford, and Ilya Sutskever. Generating long sequences with sparse transformers. arXiv preprint arXiv:1904. 10509, 2019.
[9] Gonçalo Correia, Vlad Niculae, and André FT Martins. Adaptively sparse transformers. In Proceedings of the 2019 Conference Empiricul Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNIP--JCNIP)), pages 2019.
10

---

[10] Martin Engelcke, Oiwi Parker Jones, and Ingmar Posner. GENESIS- V2: Inferring Unordered Object Representations without Iterative Refinement. arXiv preprint arXiv:2104. 09958, 2021.
[11]
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. Ad- vances in neural information processing systems, 31, 2018.
[12] Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering atari with discrete world models. arXiv preprint arXiv:2010.02193. 2020.
[13] Anthony Hu, Lloyd Russell, Hudson Yeo, Zak Murez, George Fedoseev, Alex Kendall, Jamie Shotton, and Gianluca Corrado. Gaia- A generative world model for autonomous driving, 2023. wttpt:tttt/aaaa
org/ 2309 17000
[14] Biwei Huang, Fan Feng, Chaochao Lu, Sara Magliacane, and Kun Zhang, AdaRL: What, where, and how to adapt in transfer reinforcement learning, In International Conference on Learning Representations (ICLR), 2022.
[15] Inwoo Hwang, Yunhyeok Kwak, Suhyung Choi, Byoung- Tak Zhang, and Sanghack Lee. Fine- grained causal dynamics learning with quantization for improving robustmess reinforcement learning. In Proceedings of the 4/th International Conference on Machine Learning, 2024.
[16] Ayush Jain, Andrew Szot, and Joseph Lim. Generalization new actions in reinforcement learning. In Hal Daumé III and Aarti Singh, editors, Proceedings of the 37th International Conference on Machine Learning, volume 119 of Proceedings of Machine Learning Research, URL BtLLhttps:/// /tpp::tt////rrreeedngss mlr wwt.psesssvii/
pages jain20b. html
PMLR, 13-18 Jul 2020.
[17] E. Jang, S. Gu, and B. Poole. Categorical reparameterizatioo with gumbel-softmax. In 5th International Conference on Learning Representations, ICLR 2017, Toulon, France, April 24-26. 2017, Conference Track Proceedings, 2017. URL https::/ wtwprreviiwweneefoorreee www.:eetffouumm id rkE3y85ee
[18] Jindong Jiang, Sepehr Janghorbani, Gerard de Melo, and Sungjin Ahn. Scalor: Generative world models with scalable object representations. Proceedings of ICLR 2020. OpenReview.net, 2020. URL URL
wwt. pddp://dd
[19] Jean Kaddour, Aengus Lynch, Qi Liu, Matt Kusner, and Ricardo Silva. Causal machine learning: A survey and open problems, 2022.
[20] Diederik P. Kingma and Jimmy Ba. Adam: A method for stochastic optimization. Yoshua Bengio and Yann LeCun, editors, 3rd International Conference on Learning Representations, ICLR 2015, San Diego, CA, USA, May 7-9, 2015, Conference Track Proceedings, 2015. URL http:://are http:://axxvv.ro abs/ 1412.692988
[21] Diederik Kingma and Max Welling. Auto-encoding variational bayes, 2022.
[22] Thomas Kipf, Elise van der Pol, and Max Welling. Contrastive learning of structured world models. arXiv preprint arXiv: 1.119.111.27,.... 2019.
[23] Sébastien Lachapelle, Pau Rodriguez, Yash Sharma, Katie E Everett, Rémi Le Priol, Alexandre Lacoste, and Simon Lacoste-Julien. Disentanglement via mechanism sparsity regularization: A new principle for nonlinear ica. In Conference on Causal Learning and Reasoning, pages 428-4844 484. PMLR, 2022.
[24] F. Leeb, G. Lanzillotta, Y. Annadani, M. Besserve, S. Bauer, and B. Schölkopf. Structure by architecture: Structured representations without regularization. Eleventh International Conference Learning Representations, 2023. 0_1FCPaF48E
net/pdf?id
[25] Anson Lei, Bernhard Schölkopf, and Ingmar Posner. Variational causal dynamics: Discovering modular world models from interventions. Transactions on Machine Learning Research, 2023. ISSN 2835-8856. URL http:p/:/..pppaaeeeeeeeeee//r/ts

---

[26] Yaniv Leviathan, Matan Kalman, and Yossi Matias. Selective attention improves transformer,
2024. URL [ttL/tt:://axxi/ org/ herggabb/24410 02703
[27] Phillip Lippe, Sara Magliacane, Sindy Löwe, Yuki M Asano, Taco Cohen, and Stratis Gavves. CITRIS: Causal identifiability from temporal intervened sequences. In Kamalika Chaudhuri, Stefanie Jegelka, Le Song, Csaba Szepesvari, Gang Niu, and Sivan Sabato, editors, Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 2557713363.. PMLR, 17-23 Jul 2022. URL https
//proceedings.. mlr. wwwps/ wttpsssvvv2//1pppp//e html
[28] Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf. Object-centric learning with slot attention. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 1 Associates, Inc., 2020. URL /t/pttpps//prrccciiigg .nwuries...p... pddf
Curran fitaaperrfllleeppperr
2020/file/851 http:///5......0000000aeeaaaaeee
[29] Sindy Löwe, David Madras, Richard Zemel, and Max Welling. Amortized causal discovery: Learning to infer causal graphs from time-series data. Causal Learning and Reasoning (CLeaR), 2022.
[30] Vincent Micheli, Eloi Alonso, and François Fleuret. Transformers sample-efficient world models. In The Eleventh International Conference on Learning Representations, 2023. URL https:///penreevee..u wwwp://tee/ooooodddddddduueeeee
[31] G. Parascandolo, N. Kilbertus, M. Rojas-Carulla, and B. Schölkopf. Learning independent causal mechanisms. In Proceedings of the 35th International Conference on Machine Learning (ICML), volume 80 of Proceedings of Machine Learning Research, pages 4033-4041. PMLR,
2018. URL URL htttp://proceedings mlr press/v80/ pttp:/aadaaccaooaooa/a/aa ppress//00/paraaccaallllll
html
[32] Xue Bin Peng, Marcin Andrychowicz, Wojeiech Zaremba, and Pieter Abbeel. Sim-to-real transfer of robotic control with dynamics randomization. In 2018 IEEE international conference on robotics and automation (ICRA), pages 3803-3810. IEEE, 2018
[33] J. Peters, D. Janzing, and B. Schölkopf. Elements of Causal Inference Foundations and Learning Algorithms. MIT Press, Cambridge, MA, USA, 2017.
[34] Silviu Pitis, Elliot Creager, and Animesh Garg. Counterfactual data augmentation using locally factored dynamies. In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin, editors, Advances in Neural Information Processing Systems, volume 33, pages 3976-3990. Curran Associates, Inc., 2020, URL httptt:p:ppppcccedintee wwwp:/eeeeeeiinnnnnnnnieeeeeeeee pdf
[35] Silviu Pitis, Elliot Creager, Ajay Mandlekar, and Animesh Garg. MocoDA: Model-based counterfactual data augmentation. In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, Advances in Neural Information Processing Systems, 2022. URL nttps://openrevieer hww.ppaawevieeeeeteefoorrraaaaae wwwt://orrrniiisstttiioro...
[36] Danilo Jimenez Rezende and Fabio Viola. Taming vaes, 2018.
[37] Jonathan Richens and Tom Everitt. Robust agents learn causal world models. In The Twelfth International Conference on Learning Representations, 2024. URL httpss//opeerreiii
[38] Jan Robine, Mare Höftmann, Tobias Uelwer, and Stefan Harmeling. Transformer-based world models are happy with 100k interactions. In The Eleventh International Conference on Learning Representations, 2023. URL httpss//openreeiii.ce http:///warrroeeeeoo.eeeeoraaaaa aet:/forumPAee-TTBDDG000p
[39] Rebecca Roelofs, Liting Sun, Ben Caine, Khaled S. Refaat, Ben Sapp. Scott Ettinger, and Wei Chai. Causalagents: A robustness benehmark for motion forecasting using causal relationships. 2022. URL httph:ttt////cct ://axii/ org/ wrg/abs////e 035886.
12

---

[40] Bernhard Schölkopf, Francesco Locatello, Stefan Bauer, Nan Rosemary Ke, Nal Kalchbrenner, Anirudh Goyal, and Yoshua Bengio. Toward causal representation learning. Proceedings of the IEEE, 109(5)x612-634, 2021.
[41] Bernhard Schölkopf. Causality for machine learning. 2019. doi: 10. 10.114553007774.3301555
[42] Maximilian Seitzer, Bernhard Schölkopf. and Georg Martius. Causal influence detection for improving efficiency in reinforcement learning. Advances in Neural Information Processing Systems, 34:22905- 229188, 2021.
[43] Shaoshuai Shi, Li Jiang, Dengxin Dai, and Bernt Schiele. Motion transformer with global inten- tion localization and local movement refinement. Advances in Neural Information Processing Systems, 2022.
[44] Pei Sun, Henrik Kretzschmar, Xerxes Dotiwalla, Aurelien Chouard, Vijaysai Patnaík, Paul Tsui, James Guo, Yin Zhou, Yuning Chai, Benjamin Caine, Vijay Vasudevan, Wei Han, Jiquan Ngiam, Hang Zhao, Aleksei Timofeev, Scott Ettinger, Maxim Krivokon, Amy Gao, Aditya Joshi, Yu Zhang, Jonathon Shlens, Zhifeng Chen, and Dragomir Anguelov. Scalability in perception for autonomous driving: Waymo open dataset. Proceedings of the IEEF/CVF Conference Computer Vision and Patlern Recognition (CVPR), June 2020.
[45] Andrea Tacchetti, H. Francis Song, Pedro A. M. Mediano, Vinicius Zambaldi, János Kramár, Neil C. Rabinowitz, Thore Graepel, Matthew Botvinick, and Peter W. Battaglia. Relational Forward Models for Multi-Agent Learning, In International Conference on Learning Represen- tations, December 2018.
[46] Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojeiech Zaremba, and Pieter Abbeel. Domain randomization for transferring deep neural networks from simulation to the real world. In 2017 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS), pages 23-30, 2017. doi: 10.1 1091/RO/ http:/ROSS.00....... 2017.8202133
[47] Sjoerd van Steenkiste, Michael Chang, Klaus Greff, and Jürgen Schmidhuber. Relational neural expectation maximization: Unsupervised discovery of objects and their interactions. In International Conference on Learning Representations, 2018. URL htttt:tttttsstreeeeee sst/forum?id-ryyHioo
[48] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, f. ukasz Kaiser, and Illia Polosukhin. Attention all you need. In I. Guyon, U. Von Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, ed- itors, Advances in Neural Information Processing Systems, volume 30, Curran Associates, Inc., 2017. URL htttps::/proceeedtiiii wwwpeeeddddiiinnnneeeeeeeeeeeeee cc/ wacpc/pppe__eeeee wtt.:fflls//apee//1/////e/////// Paper. .dd
[49] Rishi Veerapaneni, John D Co-Reyes, Michael Chang, Michael Janner, Chelsea Finn, Jiajun Wu, Joshua Tenenbaum, and Sergey Levine. Entity abstraction in visual model-based reinforcement learning. In Conference on Robot Learning, pages 1439-1456. PMLR, 2020.
[50] Ruben Villegas, Arkanath Pathak, Harini Kannan, Dumitru Erhan, Quoc V Le, and Honglak Lee. High fidelity video prediction with large stochastic recurrent neural networks. Advances in Neural Information Processing Systems, 32, 2019.
[51] Zizhao Wang, Jiaheng Hu, Peter Stone, and Roberto Martín-Martín. Elden: Exploration via local dependencies. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine, editors, Advances in Neural Information Processing Systems, volume 36, pages 15456-15474, Curran Associates, Inc., 2023. URL UtL/httpt///prrceeding nwwrips. cctp:ccccppppeeeeee httc/ppperfflles/////ppp files/ 10t
Paper- Saapeei6655dd-Paea----reeeeeee
pdd
[52] Yizhe Wu, Oiwi Parker Jones, and Ingmar Posner. Obpose: Leveraging pose for object-centric scene inference and generation in 3d, 2023.
[53] Ziyi Wu, Nikita Dvomií, Klaus Greff, Thomas Kipf, and Animesh Garg. Slotformer: Unsuper- vised visual dynamics simulation with object-centric models. In The Eleventh International Conference on Learning Representations, 2023. URt.tttpss///ppeeeeeee
//tpp//evp/eeeeee/ee wwwprrrviiieeeeeefrfuueee net/forum?
id=IFbwW61OVLg
13

---

[54] Harry Zhao, Zhen Liu, Sitao Luan, Shuyuan Zhang, Doina Precup, and Yoshua Bengio. A sussciousness-isspiree planning agent for model- reinforcement learning. In A. Beygelz- imer, Y. Dauphin, P. Liang, and J. Wortman Vaughan, editors, Advances in Neural Information Processing Systems, 2021. U ntt/teruuniiieee hate/foruumdddjj11aaaHH000
[55] Xun Zheng, Bryon Aragam, Pradeep Ravikumar, and Eric P. Xing. DAGs with NO TEARS: Continuous Optimzation for Structure Learning. In Advances in Neural Information Processing Systems, 2018.
[56] Shihao Zhou, Duosheng Chen, Jinshan Pan, Jinglei Shi, and Jufeng Yang. Adapt or perish: Adaptive sparse transformer with attentive feature refinement for image restoration. In Pro- ceedings of the IEEF/CVF Conference on Computer Vision and Pattern Recognition, pages 2952-2963, 2024.
14

---

A Model Details
A.1 Hyperparameters
SPARTAN and the Transformer baseline are implemented as stacked transformer encoder layers with residual connections. The Global Graph baseline is implemented as an ensemble of MLPs with learnable adjacency matrix masks. This is similar [25]]]] [4] to the modification that our baseline works on object embeddings rather than scalar states. Each MLP predicts a separate object embedding. For each MLP, the input object tokens are masked according to a learnable adjacency matrix before being passed through the model. A sparse regularisation is applied to the adjacency matrix. The hyperparameters for SPARTAN and the baselines are shown in table and 4 For the experiments on the simulated datasets, all models are trained on single GPUs (mixture of Nvidia V100 and RTX 6000) and converge within 3 days.
For the Traffic domain, predictions need to be multimodal since each vehicle has multiple possible trajectories. This is common for the dataset we consider. To overcome this, we use MTR [43 as the base architecture. MTR has two stages: first, it uses self-attention between tokenised map lines and vehicle trajectories; second, it applies cross attention using multiple queries, each representing one possible mode of motion. The output for each query is then used as a Gaussian mixture model to predict multimodal motion patterns. We adapt SPARTAN to this architecture by swapping all attention layers with the sparsified version proposed in this paper. The original work also proposes a local attention masking scheme for each layer, which only allows each token to attend to the k-nearest
neighbour based on the position of each object. We refer to this as Local Attn in table 2 In the traffic domain, the models are trained in parallel on 4 GPUs due to the size of each scene (roughly 1000 tokens). Training takes less than one week for the baseline MTR model and under two weeks for SPARTAN.
Table 3: Hyperparameters for SPARTAN and the Transformer Baseline.
Hyperparameter
Interventional Pong
CREATE
Token Dimension
32
64
Embedding Dimension
512
512
transformer layers
3
MLP hidden dimension
512
1024
n. MLP layers per transformer layer
3
Ir
5e-5
5e-5
Optimiser
Adam 20
Adam
A.2 Lagrangian Relaxation
As discussed in Sec. 3.3, the model is sensitive to the choice of X. Since removing non-causal edges should not deteriorate the prediction accuracy, we formulate constrained optimisation problem where we minimise |A] under the constraint that L < r, where r is the target loss,
min
E [[]
s.t. B[MSE($H +11 s ST.
Table 4: Hyperparameters for the Global Graph Baseline.
Hyperparameter
Interventional Pong
CREATE
Token Dimension
32
64
MLP hidden dimension
1024
1024
MLP layers
5
5
5e-5
5e-5
Optimiser
Adam
Adam
15

---

We set the target as the loss achieved by a fully connected model. Applying Lagrangian relaxation we obtain the following min-max objective,
[[AA + (E[MSF(s*)) 8+1) -7
max min XSO
(8)
We solve this by taking gradient steps on and an alternating manner. To ensure À is positive, we perform updates in the form of À <- erp(MSE 7) * X, where is the step size. Intuitively, this has the effect of increasing when the error is higher than the target, i.e. weighting the error term when the error is high. This is analogous to the GECO [36] framework for tuning the KL loss VAE. In practice, we initialise to be high and set T to the error achieved by the fully connected model. We rewrite the training objective to (MSE r) + |Al/ so that the loss value more stable. This acts scheduling scheme where the optimisation focusses on learning the dynamics and switches to pruning redundant edges once the error is below the target. Fig. showws example training curves demonstrating the training dynamics of SPARTAN. At the start of the training, log(A) remains low as the model prediction error improves. This allows an increase in the number of active edges. As the prediction error becomes low enough, the sparsity penalty automatically increases, and the number of edges gradually decreases while maintaining prediction accuracy.
mnoving-averagee estimator for the MSE
term for stability. We also
Interventioaal Pung
15
26
25
65
36
300
35
80
2.
25
dleg
48ps
mage
Create
25
drps
Figure 5: Example training curves.
B Datasets
B.1 Interventional Pong
The Interventional Pong environment is based on the Pong game and was originally developed in
27 for investigating causal representation learning. In the original work, the set of interventions
2:4
Figure 6: Example observations in the two environments.
16

---

Table 5: Environment configurations for Intervention Pong.
Env ID
Description
0
No intervention
1
Ball follows curved trajectories
2
Left paddle moves away from ball instead of towards Right paddle mirrored along y-axis
Seen
3
4
Strong friction on the ball in the middle section of the field Changed bouncing dynamics when paddles hit the ball Turn on gravity on the left side of the field
5
6
7
Gravity the left side changed bouncing dynamics with paddles
8
Strong Friction in the middle changed bouncing on left side only Curved ball trajectory + left paddle moves away from ball Gravity on the left side right paddle moves away from ball
Unseen
9
10
Table 6: Environment configurations for CREATE
Env ID
Objects
Description
0
No intervention
Disabled gravity
Seen 2
(Red, Blue, Green) balls, wall, ladder cannon
Balls go down ladders instead of up Increased elasticity of walls
Disabled cannon
(Red, Blue, Green) balls, Wall, LLLdders
No interventions
Unseen 6
(Red, Blue, Green] balls, wall, LaLdde (Red, Blue, Green). Wallls,, 2*Cannons
Balls go down ladders instead of up Increased elasticity of walls
included randomly perturbing the positions of the balls and the paddles. We modify the interventions model dynamics changes, such adding gravity changing the way the ball bounces off the paddles, as these are more appropriate in the context of world model learning. A similar setup is also used in 14 in the context of few-shot adaptation. The original Interventional Pong dataset is under the BSD 3- 3-Clause Clear License.
The environment contains 4 objects: the left paddle, the right paddle, the ball, and the score. Here, local causal edges include ball + left, right paddles as the paddles follow the ball; paddles - ball when the ball collides with the paddles; ball - score when the ball crosses the left right boundary lines. Observations are provided as 32 x 32 x RGB images with masks for each object. To capture velocity information for the ball, the ball image shows the position of the ball over consecutive timesteps in different colours. Each object is separately encoded into 32-dimensional latent object-slots using VAE (21). The configurations of different intervened environments shown in 5
B.2 CREATE
The CREATE environment is built using the CREATE simulator [16] which consists of physical interactions between various objects. The simulator under the MIT License. In our experiments, the objects in the environment include (blue, red, green) balls, wall, a ladder, and cannon. The balls collide elastically. All objects are initialised at random positions with random orientations. The wall remains static for each sequence. The ladder object allows the balls to climb up. The cannon object projects the balls at a fixed velocity the direction points at. Here, local edges capture when the objects collide or interact with each other. Observations are given as 84 x 84 x images with masks for each object. To capture the velocities of the balls, trail of the past positions (up to two timesteps) is shown in the image for each ball. The intervened environments are shown in table
17

---

6 In the unseen environments, we test the models using different scene composition, i.c. different
numbers of objects.
B.3 Traffic
We use the Waymo Open Dataset 44 for our Traffic domain. The dataset consists of real world scenes, which include map lines (such as lane markings) and vehicle trajectories. The map lines are provided as polylines (i.e. each line is represented as set of positions). The history of the state of each vehicle is provided as a list of positions, headings and velocities. Up to one second of state is given for each vehicle, and the task is to predict 9 seconds into the future for the centre vehicle. The number of map lines and vehicle trajectories is different for each scene, with roughly 800 map lines and 50 vehicles per scene. This presents major challenge for prior approaches in causal world model learning. Since the dataset is not generated by simulation, we do not apply interventions to the data.
C Relation to existing approaches and extra baselines
The aim of this work is integrate the inductive bias of sparse interactions. common in causal discovery, into a transformer-based architecture. As such, we focus our evaluations on environments that commensurate in complexity and scale the settings models are evaluated. While there exist prior approaches for local causal learning, as discussed in the related works section, these methods cannot be readily scaled the environments in the experiments presented in this paper. This is due to two principal reasons: first, prior methods operate directly on state-based observations, such as x, y positions of objects, whereas in our environments, objects are represented as learnt embeddings (typically 64 dimensions) that are mapped from images; second, and more fundamentally, our environments, in particular CREATE and traffic domains, require the world model to operate different numbers of objects and different compositions of objects across samples, e.g. different numbers of vehicles across scenes. transformer-base accommodate these requirements by treating each object embedding as a separate token. Here, we discuss the particular architectural choices of prior works that hinder their application in these environments.
which prior transformer world
world model can
ELDEN [5T] explores the role of local causal models in the context of intrinsic motivations for RL. Here, the forward dynamics model is trained with sparsity regularisaion, proxied by using the L1 norm, on the partial derivatives between the states. principle, this technique can be extended beyond scalar states to the case in which each object is represented by higher dimensional embedding vectors. However, we implemented ELDEN for our experiments and found that, since the loss a function of many partial derivatives, the optimisation process becomes prohibitively memory-intensive and highly unstable.
FCDL [15 investigates the robustness improvements afforded by local causal models. This approach learns a codebook of possible local causal graphs and infers the correct structure conditioned on the state. Sparse regularisation is then applied to the set of local causal graphs represented by the codebook. The drawback of this approach is that the local graphs must have the same set of nodes across samples, meaning that all scenes must have exactly the same composition of objects.
CAI 42 considers related but different setting from our method, where the goal is to estimate the causal influence of an agent on the states of the objects in the scene. This approach relies on testing the conditional independence using conditional mutual information estimation between the agent's action and the objects based on a learned dynamics model. However, estimating conditional mutual information challenging the multivariate case.
CODA 34] uses threshold heuristics to determine local structures. We compare against this in our main experiments and show that attention alone is not able to reliably learn the correct dependency structure.
Beyond the local causal structure learning literature, the idea of learning context-dependent causal structures is also explored in Amortised Causal Discovery (ACD) 29 which uses a GNN backbone to infer causal structures conditioned on the history of states in a time-series setting. This approach can be adapted to the state-deperdent local causal structure case by conditioning the graph encoder on the current state rather than on an entire tine-series. In table 7
we compare the graph learning
18

---

Table Average SHD between the learned graphs and the ground truth. Lower better.
SPARTAN
ACD
Sparse ACD
Intervention Pong
1.50
10.27
2.67
Create
1.17
39.77
5.45
accuracy between SPARTAN and ACD. We note that ACD requires setting a prior level of expected sparsity. Here, we compare against naive choice of 50% sparsity (labelled ACD) as well as the sparse baseline, which uses the ground-truth sparsity as prior (labelled Sparse ACD). The results show that SPARTAN significantly outperforms this baseline, even if the true sparsity level, which is in practice not accessible, is specified. This further corroborates the efficacy of our model.
D Using pretrained object-centrie representations
Many existing works that aim to capture object interactions, including SlotFormer [53], use pre-trained object-centric representations [28] that handle the segmentation and tracking problems. Our approach can be readily applied to any such pretrained representations. In our experiments, we opted to use ground-truth masked object representations for the principal reason that they enable the evaluation of the learned graphs: quantitatively evaluating the SHD learned graphs requires mapping each latent embedding to the corresponding ground-truth object. Learned object embeddings make this mapping noisy and ambiguous due to imperfect segmentations. Since learning high-quality representations is not the main focus of this work, we use ground-truth masked object representations to ensure the fairness and validity of our evaluation. To demonstrate the generality of our results, have trained our model using a learned slot representations in the pong environment, Here, to measure the SHD between the learned graph and the ground- ground-truth causal graph, we use the visualised attention masks from Slot Attention map the latent object representations to ground-truth objects. In table 8 we present the results, which are consistent with the main findings in the paper: our model achieve the same level of prediction while capturing the underlying graph structure more accurately.
Table 8: The prediction error and SHD of the learned graphs in the Intervention Pong environment.
SPARTAN
Transformer
Prediction Error (+ SE)
7.21 + 0.69
6.90 + 0.83
SHD (= SE)
2.69 + 0.02
8.52 + 0.02
19

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our main claim is that improve current transformer world models with sparse regularisation. This is supported by our experiments.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
is fine to include aspirational goals as motivation as long as is clear that these goals not attained by the paper.
2. Limitations
Question: Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification:
We have included limitations discussion in the conclusion section.
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution is low images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness.
While the authors might fear that complete honesty about limitations might used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't 't acknowledged judgment and recognize that individual actions in favor of transparency play impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
the paper. The authors should use their best
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
20

---

Justification:
We do not include theoretical results.
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced the statement of any theorems.
The proofs can either appear the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
The experiment details are provided in the appendix. The code and model weights to reproduce our results are attached supplementary materials.
Justification:
Guidelines:
The answer NA means that the paper does include experiments.
If the paper includes experiments, No answer this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
21

---

Answer: [Yes]
Justification: The code for reproducing the main results are included in the supplementary materials. We plan to open source our code upon acceptance.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.cm ec/
for more details.
While encourage the release of code and data, understand that this might not be possible, so "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurlPS code and data submission guidelines (https
//nips. //nips.cc/pil wwcppbblc/guuiicuccccccccoecoopp
http://oooSuuuisss.ooeee Stt.SSbbbiiiiooooiiiii.oo for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: training details explained in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Error bars and standard error are reported where applicable.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
22

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that are out range (e.g. negative еπхо rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification:
The compute resources are disclosed in the appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make it into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics
Answer: [Yes]
Justification: NA
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special circumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., there a special consid- eration due to laws or regulations in their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper presents work whose goal is to advance the field of Machine Learning. There are many potential societal consequences of work, none which we feel must be specifically highlighted here.
Guidelines:
The answer NA means that there no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
23

---

The conference expects that many papers will be foundational research and not tied particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of MI..
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: NA
Guidelines:
The answer NA means that the paper poses such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Lieenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: We have, to the best of our ability, cited the data we used in this work.
Guidelines:
The answer NA means that the paper does existing assets.
The authors should cite the original paper that produced the code package or dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (c.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification:
We do release new assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: NA
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: NA
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
25

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: The core method development does not involve LLMs.
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs any important, original, non-standard components.
Please refer to our LLM policy for what should or should not be described.
26


Reviews:
[{"review_id": "yaRMN7v3s3", "reviewer": "Reviewer_AuXs", "summary": "The 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.", "questions": "Questions 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)?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 2, "originality": 2, "strengths_and_weaknesses": "Strengths\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.", "was_revised": true, "final_justification": "The authors have addressed most of my concerns during the rebuttal. My main remaining concerns are that the main technical contribution is the path-counting regularization, and the method depends on existing approaches like SAVi to operate on raw visual observations.", "rebuttal": "We thank reviewer AuXs for their thorough review. The reviewer raised several concerns which we address below.\n\n**Requiring object masks**: We would like to point out that many existing works, including SlotFormer [1], use pre-trained object-centric representations such as SAVi [2] which handles the segmentation and tracking problems. Our approach can be readily applied to any such pretrained representations. In our experiments, we opted for using ground-truth masked object representations for the principal reason that it enables evaluation of the learned graphs: quantitatively evaluating the SHD of learned graphs requires mapping each latent embedding to the corresponding ground-truth object. Learned object embeddings can make this mapping noisy and ambiguous due to imperfect segmentations. Since learning high-quality representation is not the main focus of this work, we use ground-truth masked object representations to ensure the fairness and validity of our evaluation. Nevertheless, we understand that more evaluations can strengthen our claims. To demonstrate the generality of our results, we have trained our model on an unsupervised learning slot representation for the pong environment. The results are consistent with the main results in the paper: our model can achieve the same level of prediction while capturing the underlying graph structure more accurately (see below).\n| | SPARTAN | Transformer |\n|-|-|-|\n| Prediction error ($\\pm$ SE) | 7.21 $\\pm$ 0.69 | 6.90 $\\pm$ 0.83 |\n| SHD  ($\\pm$ SE) | **2.69** $\\pm$ 0.02 | 8.52 $\\pm$ 0.02 |\n\n**Details missing**: The total number of training episodes are 10k, 30k and 10k for the pong, CREATE and waymo datasets respectively. The prediction error metric measures the L2 distance between the prediction and the embedded ground-truth in the latent space. We will update the paper to include all experimental details. We will also open source the training code and the data-generation scripts to ensure reproducibility of our results.\n\n**Limited technical contribution**: We highlight our main technical contribution: we propose a method to perform sparsity regularisation on transformers via a novel path-counting regulariser as well as masked hard attention. While sparsity regularisation has been proposed in prior works on local causal discovery, existing methods remain difficult to scale beyond simple environments (see, for example, our discussion in App. C). Here, our approach of using a path-counting matrix is a key enabler for applying such sparsity penalty on the transformer architecture which is ubiquitous across many domains. As such, our method solves a crucial scalability problem in the field of local structure learning. Importantly, we note that our method is architecturally similar to the standard transformer **by design**, so that it can be readily used as a slot-in replacement for transformers.\n\nThe reviewer also raised three questions regarding the practical aspects of our method. \n1. In the waymo dataset, this is done by preprocessing the data to ensure no new vehicles enter the scene during the prediction horizon. Vehicles that no longer exist after a timestep are assigned a special token. \n2. The model can operate on unsupervised learning based slot representations that are noisy and imperfect. Please see the weakness section for a more detailed discussion and extra results.\n3. The transition model does not explicitly deal with tracking and assumes consistent indexing. In practice, object segmentation models such as SAVi [2] can perform tracking jointly with object discovery. \n\nIn general, we point out that these problems are open challenges in object-centric world modeling that are orthogonal to our main contribution, namely, improving robustness, interpretability and adaptability via learning local structures. We note that our method can be used as a slot-in replacement for the standard transformer which is highly versatile. As such, we expect that our core contributions can be readily applied to any future or existing methods that can handle these problems. \n\n[1] Wu et al. 2024, SlotFormer: Unsupervised Visual Dynamics Simulation with Object-Centric Models.\n\n[2] Kipf et al. 2022, Conditional Object-Centric Learning from Video"}, {"review_id": "YXqUZQZBGv", "reviewer": "Reviewer_hr2c", "summary": "The 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.", "questions": "1. 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).", "limitations": "The paper has discussed the limitations in Section 6.", "rating": 5, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 2, "originality": 1, "strengths_and_weaknesses": "### 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.", "was_revised": true, "final_justification": "I am not an expert in this domain, but the authors have addressed my concerns, and I have no major concerns with this work.", "rebuttal": "We thank reviewer hr2c for their thoughtful comments and suggestions. Below we address the concerns and questions raised.\n\n**Weakness 1 - Pseudo code**: We update the paper with the relevant pseudocode for training and adaptation for better clarity.\n\n**Weakness 2 - Theoretical analysis**: In terms of theory, our work is grounded in the theoretical results developed in local causal discovery. The FCDL framework [1], for example, establishes theoretical results showing that sparsity regularisation will converge to the correct local causal graphs. These can be seen as theoretical justification of our approach, which provides a scalable method to perform sparsity regularisation on transformers. What these theories do not cover, however, is the case where the variables can change across scenes (as in the Waymo dataset). Establishing a rigorous theoretical framework for this case would require significant extension of current theory. This is beyond the scope of our current submission. We will update the text to add a more thorough discussion on how our work links to prior theoretical results. \n\nWe also highlight that one of our main technical contributions is the path-counting algorithm over adjacency matrices which efficiently tracks interactions between tokens across layers. To the best of our knowledge, ours is the first method that successfully applies sparse regularisation to transformers in the context of discovering local structures. This solves a significant scalability challenge in the field of local causal discovery, where prior approaches can only work on simple environments. (See app C for a more detailed discussion). In this light, we sincerely invite the reviewer to reconsider their assessment of the originality of our work.\n\n**Question 1 - Waymo dataset**: In the waymo open dataset, there are two kinds of observations: for vehicles, the observations are x,y positions and heading angles for one second, sampled at 10Hz; for the lane markings, the observations are x y points that form a line. Both of these kinds of observations are essentially lines. Following the MTR architecture [2], these lines are encoded via a learnable polyline encoder which serves as the ‘tokeniser’. We will expand our description of the environment in App. B.3. to better explain the setting.\n\n**Question 2 - Effect of Env index**: We clarify that the environment index is not a prerequisite for training the model. We note that the Waymo experiment, for example, does not use labelled environments. The intervention token is only designed for the case where the model has access to multi-environment data where the behaviour of objects can be drastically different. This is inspired by the notion of soft interventions from causality which tells us which object is behaving differently, thereby providing an extra level of interpretability and adaptability. \n\n**Question 3 - Tuning Lambda**: The lambda parameter controls the strength of the sparsity regularisation. A value that is too high can result in a model that makes poor predictions as it does not use any connection. In practice, we use a lagrangian optimisation scheme (detailed in App. A.2) which serves as a schedule that automatically tunes the value of lambda to ensure that the model performance does not degrade due to sparsity. We provide a more detailed explanation of the formulation as well as example training dynamics in App. A. 2.\n\n**Question 4 - Number of layers**: The number of layers, much like in the standard transformer architecture, controls the expressivity of the model. Tuning this hyperparameter is a balance of how much data and compute is available. In our experiments, we have picked the number of layers to match what we use for the standard transformer baseline to ensure the comparison is fair. We note that the adjacency matrix regularisation term counts the number of paths between tokens and effectively counts the edges at every layer. We highlight that, empirically, our method works for a wide range of number of layers from three layers (in Pong) up to 12 layers (in Waymo).\n\n[1] Hwang et al., 2024, Fine-Grained Causal Dynamics Learning with Quantization for Improving Robustness in Reinforcement Learning\n\n[2] Shi et al., 2023, Motion Transformer with Global Intention Localization and Local Movement Refinement"}, {"review_id": "Qo1DQMWbhR", "reviewer": "Reviewer_7Ngu", "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.", "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.\n\n\nThe idea on simple local models can be traced back to [3], which should be discussed properly.\n\n[1] The Consciousness Prior\n[2] A Consciousness-Inspired Planning Agent for Model-Based Reinforcement Learning\n[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?", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "I find the approach to be simple and frankly quite amazed about the empirical performance of the Bernoulli binary masks.\n\nI commend the authors for the use of performance metrics other than task returns. These make the experiments much more convincing.", "was_revised": true, "final_justification": "The authors did not significantly improve the submission, despite having clarified one of my misunderstandings", "rebuttal": "We thank reviewer 7Ngu for their comments and are delighted that the reviewer finds our experiments convincing. The reviewer raised two main questions, one concerning the connections to prior works, the other concerning the use of the environment index. Below we address these concerns. \n\n**Related works**: We agree that the consciousness prior is similar to our framework on a conceptual level. However, we note the key difference between our work and this line of work: we consider different kinds of sparsity. In the consciousness prior framework, the model is sparse in the sense of picking out a small subset of objects that are ‘active’ while leaving the other objects stationary. In contrast, in our work we model the dynamics of all objects, but use sparse attention to determine how each object depends on other objects. We thank the reviewer for the suggestion and will add a more detailed discussion in the updated text. We will also add a discussion on simple local models to add context to the current work.\n \n**Environment index**: We clarify that the environment index for the intervention token is only designed for cases where the model is trained on multi-environment datasets, but is not required for training our model. In the waymo dataset, for example, we do not have access to the environment index and therefore omit it. The use of the environment embedding (intervention token) is primarily for interpreting which object has changed its behaviour in each environment. This is akin to identifying the intervention targets in the causal discovery literature."}, {"review_id": "Ikfxa4bCwz", "reviewer": "Reviewer_gSAB", "summary": "This 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.", "questions": "Please refer to the weaknesses section above.", "limitations": "NA", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "**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.", "was_revised": true, "final_justification": "The rebuttal has addressed most of my concerns, so I will increase my score to 4.", "rebuttal": "We thank reviewer gSAG for their comments. The reviewer has raised three main questions which we address below. \n\n**Extending to physics**: We first highlight that our model is able to capture non-contact effects, as evidenced in the Waymo dataset experiments, where interactions are mostly long-range. To answer the question more generally, we clarify that the proposed method is not just a binary attention, but rather a binary gate applied to the normal softmax attention. In the limit where the entities are fully connected, our architecture is equivalent to the standard transformer. Hence, we expect our model to predict as well as any other transformer-based architecture, albeit not learning sparse structures (because there is no sparse structure to be learned). A more nuanced point is that even when entities have non-zero influence on each other, there is still a tradeoff between interpretability and accuracy, and it is often still insightful to model emergent structures at the cost of negligible accuracy loss. In the gravity example, it is natural to think of a solar system as an independent system even though strictly speaking other solar systems have non-zero influences on its planets. In fact, all physical entities affect each other via gravity and other long-range forces, yet it is still useful to think of objects as separate entities that only interact when they make contact. Our lagrangian optimisation scheme (see App. A.2) formulates this problem as ‘finding the sparsest model that achieves a pre-determined error level’, which offers a flexible and practical algorithm for balancing this trade-off between accuracy and sparse structures.\n\n**Weighting paths equally**: The key desiderata for the sparsity measure is that it needs to track the existence of information flow from one token to another. In this perspective, we need to measure whether two tokens are connected at all and we do not necessarily care about whether the connections are multi-hop or direct. This is the main motivation for our path-counting matrix: we want to drive as many entries to zero as possible. It could be true that weighting different types of paths differently could offer practical benefits in terms of the optimisation process. However, for the experiments considered in this paper, the path-counting matrix was sufficient to induce good graph structures. We leave the investigation of whether other more involved regularisation schemes are useful to future work.\n\n**Requiring object masks**: Discovering and segmenting objects from images remains an active field of research that is orthogonal to our contribution. Since our approach can be readily applied to any pretrained representations such as slot-attention object embeddings, we expect the capability of our method to benefit from advances in this field. We clarify that, in our experiments, we opted for using ground-truth masked object representations primarily because it enables evaluation of the learned graphs: quantitatively evaluating the SHD of learned graphs requires mapping each latent embedding to the corresponding ground-truth object. To demonstrate the generality of our results, we have trained our model on an unsupervised learning based slot representation for the pong environment. Here, we pick a permutation that maps latent slots to ground-truth objects that results in the least SHD for reporting. The results are consistent with the main results in the paper: our model can achieve the same level of prediction while capturing the underlying graph structure more accurately (see below).\n\n| | SPARTAN | Transformer |\n|-|-|-|\n| Prediction error ($\\pm$ SE) | 7.21 $\\pm$ 0.69 | 6.90 $\\pm$ 0.83 |\n| SHD  ($\\pm$ SE) | **2.69** $\\pm$ 0.02 | 8.52 $\\pm$ 0.02 |\n\nWe also thank the reviewer for the other minor suggestions and will fix these in an updated version of the paper."}]

Author rebuttal:

Meta-review: This paper presents a transformer-based world modeling approach that explicitly enforces sparse structure in the causal interactions between elements in visual scenes. The reviewers acknowledged several strengths of this work, including clear and well-written presentation and intuitive examples that effectively justify the importance of the addressed problem.

The reviewers initially raised a few concerns but the authors adequately addressed them. Some reviewers also questioned whether the technical contribution was limited given that the approach builds upon existing developed components. However, the integration and application of these components to enforce sparse causal structure in world modeling represents a meaningful advance to me.
The work demonstrates both novelty in its approach to sparse world modeling and methodological rigour in its execution. Therefore, I recommend acceptance.
Decision: Accept (poster)