Transformers and slot encoding for sample efficient physical world modelling

## Abstract

Abstract World modelling, i.e. building a representation of the rules that govern the world so as to predict its evolution, is an essential ability for any agent interacting with the physical world.
Recent applications of the Transformer architecture to the problem of world modelling from video input show notable improvements in sample efficiency. However, existing approaches tend to work only at the image level thus disregarding that the environment is composed of objects interacting with each other.
In this paper, we propose an architecture combining Transformers for world modelling with the slot-attention paradigm, an approach for learning representations of objects appearing in a scene.
We describe the resulting neural architecture and report experimental results showing an improvement over the existing solutions in terms of sample efficiency and a reduction of the variation of the performance over the training examples.
The code for our architecture and experiments is available at https://github.com/torchipeppo/transformers-and-slot-encoding-for-wm

## 1 Introduction

World modelling is the ability of an artificial agent to build an internal representation of the world in which it operates.
This representation is employed by the agent to forecast the evolution of the world.
The problem of building a world model spans many branches of artificial intelligence, such as planning and reinforcement learning (Micheli et al., ; Paster et al., ), physics modelling and reasoning (Ding et al., ), and robotics (Wu et al., ).
An accurate representation of the world allows building simulations that, in turn, enable practitioners to gather additional data and test the performance of an artificial agent without interacting with their environment.
This is convenient because interacting with an agent’s environment can be time-consuming, risky due to possible failures of physical components, and sometimes even impossible due to the potential unavailability of the environment (e.g. experimenting with the exploration of Mars by a rover).

Recent applications of the Transformer architecture (Vaswani et al., ) to the task of world modelling from video input suggest that this family of architectures is not only it is capable of capturing the dynamics of the environment,
but it is also capable of learning with high sample efficiency (Micheli et al., ; Robine et al., ).
However, existing approaches typically
operate directly at the image level, with little regard for the objects contained within it.
Understanding how objects interact with each other and with the environment is of paramount importance, as it endows agents with an intuitive theory of object motion (McCloskey, ).
In neural architecture, this problem is addressed in a separate line of research (Locatello et al., ; Kipf et al., ; Wu et al., ), that focuses on learning object-based representations that allow the objects in the scene and their interactions to be modelled explicitly.
We hypothesise that Transformers may benefit from object-based representations to learn more accurate models of the world.

#### Problem statement

In this paper, we focus on physical world modelling through the analysis and prediction of synthetic videos. Learning a model of basic physical laws such as gravity and collision is very important for any agent working in a real environment to understand better the world’s evolution and the consequences of the agent’s actions.

Machine learning research on modelling intuitive physics aims to replicate the innate understanding of physical concepts that humans display since their first months of age (McCloskey et al., ; Baillargeon, ). Specifically, we approach the physical interaction output prediction problem, as defined in a recent survey (Duan et al., ). In this context, the agent is shown a video, composed of a sequence of frames $x_{1},...,x_{T}$ depicting several objects interacting in a world governed by physical laws such as gravity and collision. The agent is then asked to predict the final outcome of the situation, which requires estimating how the objects will behave after what is shown in the input video (Yi et al., ).

#### Contribution

We design an entirely transformer-based architecture for world modelling, inspired by the principles of representation learning with slot encoding (Kipf et al., ; Singh et al., ). The evaluation is based on how well the learned model predicts the outcome of the situations it gets shown. We show that this allows us to reap the benefits from both approaches (i.e. slot encoding and transformers), while also noticeably improving the stability of the training process.

#### Paper outline

We structure the rest of the paper as follows. section 2 discusses some previous works on the topics of world modelling and representation learning. section 3 describes our architecture and how it was trained. section 4 describes the evaluation experiments we performed. section 5 offers some final remarks on this work.

## 2 Related work

### 2.1 World modelling

The world modelling problem has received tremendous attention in the last few years. In reinforcement learning, being able to simulate the environment dynamics is especially useful, because it enables the agent to act and learn in its own simulated world without paying the cost of interacting with the "real" environment. The Dreamer algorithm, in its various versions (Hafner et al., , , ), has been relatively influential on the topic, with Wu et al. () being an application in a robotic domain, where real-world interaction has the unfortunate potential of breaking usually costly equipment, in addition to time costs.
Additional approaches in reinforcement learning include solutions based on causal discovery and reasoning (Yu et al., ), which aim at learning a causal model of the environment to better understand the interactions between the agent and the world, and even provide an explanation for the actions taken by the agent.
Transformer-based approaches are also studied, due to their generally good performance in different tasks and the sample efficiency they provide in this specific problem (Micheli et al., ; Robine et al., ).

In the case of agents acting in a real physical environment, learning a model of the basic physical laws of the world is essential to act effectively in the environment and understand the consequences of each move.
We refer to this problem as intuitive physics modelling.
For this reason, several solutions have been studied for this problem with approaches ranging from deep learning (Qi et al., ), to violation of expectation (Piloto et al., ), to causal reasoning (Li et al., ).

With this work, we aim to improve the general level of performance and stability of Transformer-based approaches by implementing an unsupervised representation learning module, specifically one based on the principles of slot encoding.

### 2.2 Object-oriented representation learning

In this work, we integrate concepts from the line of research on learning object-oriented representations from images and videos (Locatello et al., ; Zoran et al., ; Jia et al., ). In particular, slot encoders for video (Kipf et al., ; Singh et al., ) learn a representation that tracks the prominent objects in a video frame by frame. The structure of our architecture is based on that of slot encoders, but we try to streamline it by focusing exclusively on using Transformer modules, which allows us to convert the image to just one intermediate form, i.e. a sequence of tokens to be elaborated by the Transformers, while Singh et al. () requires two separate elaborations of the image: one to produce convolutional features and one to produce a sequence of tokens.

The idea of leveraging slot encoding mechanisms for world modelling has also been explored in Wu et al. (). However, while that work uses a single transformer as the dynamics modelling module, we experiment with the idea of keeping representation correction and dynamics advancement separate, where each step is learned by a different, smaller neural model.

This is not the only existing approach: an additional, earlier line of work focuses on using generative models with object-centric features for images (Eslami et al., ; Engelcke et al., ) and video (Kosiorek et al., ) to distinguish the objects present in a scene and improve image/video generation with the learned object awareness. Jiang et al. () applies this paradigm to world modelling. Other approaches include spatial attention (Lin et al., ) and latent space factorization (Kabra et al., ).

## 3 Method

Figure: Figure 1: Architecture diagram for the Future-Predicting Transformer Triplet for world modeling.
Refer to caption: extracted/5630022/architecture-diagram.drawio.png

We introduce a multi-stage architecture, called "Future-Predicting Transformer Triplet for world modeling" (FPTT), which aims to model the behaviour of objects in a set of videos so as to predict their evolution.

We frame the world modelling problem as a sequence learning one.
We use transformers as the fundamental building block of our architecture, due to their proven performance in world modelling (Micheli et al., ) as well as other tasks that can be reduced to modelling and manipulating a sequence of tokens (Vaswani et al., ; Esser et al., ).In particular, we leverage the recent work by Micheli et al. (), which showcases a sample-efficient application of transformers to world modelling.

We also borrow some ideas from the slot-attention mechanism (Kipf et al., ; Singh et al., ) which learns a compact representation for objects appearing in a video.

### 3.1 Notation

We indicate with $x_{t}$ the $t$-th frame of a video and with $z_{t}$ a sequence of tokens corresponding to $x_{t}$. $\Lambda_{t}(x)$ is the internal representation
of the input video $\mathbf{x}$ up to time $t$, i.e. given frames $x_{1},...,x_{t-1}$, while $\Lambda_{t}^{*}(x)$ is the "corrected" representation which also includes information from frame $x_{t}$.
The initial representation $\Lambda_{1}(x)$
is randomly determined for initialization purposes.

### 3.2 Architectural overview

figure 1 shows the high-level components of
FPTT.
Further details on the implementation, e.g. the hyperparameters of the architecture, can be found in Appendix A.

The architecture takes as input a sequence of $T$ frames of a video, i.e. $x_{t}$ with $t=1,...,T$.
The frames are processed sequentially, so that $\Lambda_{t+1}(x)$ is determined by combining the previous representation $\Lambda_{t}(x)$ with the new frame $x_{t}$.

As in previous works (Micheli et al., ), each frame $x_{t}$ is transformed into a corresponding sequence of tokens $z_{t}$ by a discrete Vector Quantized Variational Autoencoder (VQVAE)
(van den Oord et al., ; Esser et al., ), as transformers need to work on sequences of tokens.
Further details on the VQVAE can be found in section 3.3.

After this preliminary step, the sequence of tokens $z_{t}$ is processed by the core components of the architecture which are meant to predict the next representation $\Lambda_{t+1}(x)$ based on the current one $\Lambda_{t}(x)$ and $z_{t}$. These components, both based on the transformer architecture,
have the same high-level purpose as their counterparts in slot attention for video (Kipf et al., ; Singh et al., ) architectures:

- •
The corrector transformer (see section 3.4) which compares the previous (internal) representation $\Lambda_{t}(x)$ with the tokenized representation of the current frame $z_{t}$
in order to consistently align the internal representation with the actual evolution of the video;
- •
The predictor transformer (see section 3.5) which predicts the evolution of the world state and produces the representation of the next time step $\Lambda_{t+1}(x)$ on the basis of the result of the corrector, i.e., $\Lambda_{t}^{*}(x)$. $\Lambda_{t+1}(x)$ is then passed to the corrector for the next stage.

The result of the prediction at stage $t$, i.e. $\Lambda_{t+1}(x)$, is also passed to the decoder transformer (see section 3.6).
This component transforms the predicted internal representation $\Lambda_{t+1}(x)$ into a sequence of tokens $\hat{z}_{t+1}$.
Finally, the loss is calculated by comparing $\hat{z}_{t+1}$ with $z_{t+1}$, i.e. the sequence of tokens obtained from the input frame.
All the above steps (correction, prediction, decoding and loss calculation) are computed for each input frame except the last one, i.e. for $t=1,...,T-1$. The last frame is not processed at training time because it would require the existence of a frame $x_{T+1}$ to calculate the loss against, which is impossible to provide since the video only goes up to frame $x_{T}$ by definition.

### 3.3 Vector Quantized Variational Autoencoder for tokenization

The Vector Quantized Variational Autoencoder (VQVAE) transforms video frames into a format that subsequent transformers can process. This format is a sequence of $L$ tokens, with each token represented by a vector of the space $\mathcal{V}=\left\{v_{1},v_{2},...,v_{N}\right\}\subset\mathbb{R}^{d}$, where $d$ is defined as a hyperparameter of the architecture (see appendix A).

The VQVAE alternates residual convolutional layers, attention blocks, and convolutional downsampling layers to convert an image(^1^11We omit the $t$ subscript in this paragraph for ease of notation, since the VQVAE processes images as single entities, not as parts of a video.) $x\in\mathbb{R}^{W\times H\times 3}$ (W and H are the width and height of the image, respectively)
to a latent-space representation $z_{l}(x)=\left(z_{l,1}(x),z_{l,2}(x),...,z_{l,L}(x)\right)\in\mathbb{R}^{L
\times d}$. Then each latent vector is quantized into a token simply by picking the closest embedding vector in $\mathcal{V}$, that is to say, $z(x)\in\mathbb{R}^{L\times d}$ is such that $z_{i}(x)=argmin_{v\in\mathcal{V}}\left(||z_{l,i}(x)-v||_{2}\right)$ for each $i=1,...,L$.

A decoder network with a symmetrical structure to the encoder (not shown in figure 1) is used to reconvert a token sequence $z$ back into an image $\hat{x}(z)$ for the purposes of training the whole autoencoder pair.

### 3.4 Corrector transformer

The purpose of the corrector transformer is to avoid drifting, i.e. making the internal representation stick with the evolution of the video. This is achieved by updating the estimated representation $\Lambda_{t}(x)$ with the corresponding frame $z_{t}$ thus producing a corrected representation $\Lambda^{*}_{t}(x)$.
It is implemented by a transformer that produces the corrected representation ($\Lambda^{*}_{t}(x)$) by performing an unmasked cross-attention of the two inputs ($\Lambda_{t}(x)$ and $z_{t}$).

It is worth noting that this structure fits neither the transformer encoder nor the transformer decoder descriptions as traditionally defined in Vaswani et al. (), since we perform cross-attention (like a decoder) without including a causal mask (like an encoder).
This allows us to compare the two input sequences as a whole, without arbitrarily limiting the context.
We also note that the purpose of this transformer is not to perform autoregressive generation, so a non-causal flow of information causes no harm.

### 3.5 Predictor transformer

The predictor transformer performs self-attention on the representation $\Lambda^{*}_{t}(x)$ to estimate its advancement to the next time step $\Lambda_{t+1}(x)$.

The predictor and corrector transformers can be seen as two halves of one model, dedicated to predicting the next internal representation on the basis of the current representation and the current frame of the video being processed. For this reason, each of them has individually fewer layers compared to the decoder (see section 3.6).

### 3.6 Decoder transformer

The decoder transformer
converts a representation $\Lambda_{t+1}(x)$ into a sequence of tokens $\hat{z}_{t+1}$.
The loss is computed by comparing $\hat{z}_{t+1}$ with $z_{t+1}$, i.e. the sequence of tokens obtained from the input frame.

As it can be seen in figure 1, we position this stage between the predict step that generates $\Lambda_{t+1}(x)$ and the subsequent correct step.
Thus, the loss computes the error on the prediction of the frame $z_{t+1}$.
It is worth noticing that this solution diverges from the original slot-attention architectures (Kipf et al., ; Singh et al., ), which calculate the loss on the output of the corrector.
The calculation of the loss on the predicted representation has the effect of directing the model’s attention towards the accurate prediction of future events, rather than towards the representation of the current, thus emphasizing the world modelling objective.

### 3.7 Training

The VQVAE is trained in isolation with respect to the whole architecture.
To enhance the stability of the training process, we maintain a fixed configuration of the VQVAE parameters throughout the training of the other components.
Following Micheli et al. (), the loss is a combination of a mean absolute error and a perceptual loss (Johnson et al., ) on the reconstruction, as well as a commitment loss on the embeddings (van den Oord et al., ).

As for the corrector, predictor and decoder transformers, they are trained together in an end-to-end fashion, with the objective to minimize a cross-entropy loss on the (tokenized versions of the) predicted frames $\hat{z}_{2},...,\hat{z}_{T}$ with respect to the ground-truth ones $z_{2},...,z_{T}$.

Both parts of the architecture are trained in a self-supervised way, with unlabeled videos from a suitable dataset. See section 4.3 for further details on the dataset.

## 4 Experiments

The ability to model the world of the proposed architecture (FPTT) is assessed through a physical reasoning task (see Section 4.1) that requires the ability to predict how a set of objects moves in a given environment.
Specifically, we experiment with the PHYRE dataset which provides a benchmark containing a set of simple classical mechanics puzzles in a 2D physical environment (Bakhtin et al., ).
The performance of FPTT is compared against two baselines.

### 4.1 Physical reasoning task

We adhere to the definition of a physical reasoning task as outlined in the PHYRE benchmark (Bakhtin et al., ).
The task is set in a two-dimensional world that simulates simple deterministic Newtonian physics with a constant downward gravitational force and a small amount of friction. This world contains non-deformable objects, distinguished by colour, that can be static (i.e. they remain in a fixed position) or dynamic (i.e. they move if they collide with another object and are influenced by the force of gravity).

A task consists of an initial world state and a goal (see figure 2).
The initial world state is a pre-defined configuration of objects.
The goal for all tasks is the following: at the end of the simulation the green object must touch the blue object.
If the goal is achieved, the task succeeds (as in the PHYRE terminology).

Given video (as a sequence of frames), the objective of the world model is to build an internal representation that can be used to predict if the depicted task will succeed or fail.
The ability to predict
that
represents an auxiliary classification problem.
This allows us to indirectly assess the performance of the world models.
It is worth noticing that the same evaluation protocol is used in related work (Wu et al., ).

Figure: (a) Initial state
Refer to caption: extracted/5630022/phyre-example-frame-1.png

### 4.2 Baselines

The performance of the proposed architecture is evaluated in comparison to the following baselines:

- •
STEVE (Singh et al., ), a slot encoding architecture that is also based on the correction-prediction pattern.
The input frame is encoded using a convolutional neural network (CNN) which feeds a recurrent neural network (RNN) acting as corrector.
The result is then passed to the predictor, a single-layer transformer, which is then translated into a token sequence (i.e. $\hat{z}_{t}$) by a transformer decoder.
- •
A single decoder transformer, intended to replicate the approach to world modelling from visual data by Micheli et al. ().
It is worth noticing that Micheli et al. () work in a reinforcement learning setting, where the state of the world can also be influenced by the actions of an agent.
Therefore, we adapted this architecture to experiment with a pre-rendered video dataset where no agents interact with the environment.
Specifically, in Micheli et al. () the transformer takes as extra inputs and outputs, the action taken by the agent and the reward given by the environment, which we removed from the architecture to fit it in the physical reasoning problem.
As a result, in our experiments, a decoder transformer predicts the (tokenized version of the) next frame $z_{t+1}$ directly from the previous one $z_{t}$, without using an internal representation.
In the following, we refer to this architecture as decoder-only.

### 4.3 Experimental setup

Figure: Figure 3: Diagram of the experimental setup, showing how the classifier is positioned with respect to the world modelling architecture. Note: in the case of the decoder-only baseline, replace $\Lambda_{T}$ with $z_{T}$.
Refer to caption: extracted/5630022/classification-diagram.drawio.png

We experiment on a dataset of synthetic videos presented in Qi et al. (). This dataset was generated by rendering simulations from the PHYRE benchmark for physical reasoning (Bakhtin et al., ).
Specifically, we focus exclusively on videos from B-tier tasks, and within-template evaluation.
figure 3 shows the experimental setup.
The world model takes a video (as a sequence of frames) from the PHYRE task and builds a representation.
This representation is then passed to a classifier which predicts the result of the task (i.e. success or failure).
As for the classifier, we use a BERT-like encoder architecture (Devlin et al., ), trained in a supervised manner.

As for FPTT and STEVE, we proceed as follows.
Each video in the dataset represents a task that is labeled as either "success" or "failure".
The world model is given the first $N$ frames of a video whose total length is $T$ frames, with $N<T$. The remaining $T-N$ frames are kept hidden from the model.
In order to obtain the representation of the whole video, including both the given section and an estimation for the following hidden one, the $N$ given frames are processed as usual, updating the representation in the correct step and advancing it to the next timestep in the predict step.
Afterwards, the remaining $T-N$ steps are projected by simply repeating the predict step, skipping the correction for the hidden frames (see figure 4). In the experiments, $N$ is set to 5, while $T$ varies depending on each video, ranging from 7 to 18, with many videos being 12-15 frames long.

We follow a similar approach for the decoder-only architecture, accounting for the lack of an internal representation in this case.
The world model sees the first N frames and the remaining are generated autoregressively by the transformer.
The final frame, i.e. the sequence of tokens $z_{T}$, is passed to the classifier instead of a representation.

All the experimenting architectures employ the same pre-trained VQVAE for transforming each input frame into a sequence of tokens (see section 3.3).

Overall, the dataset contains 50K videos: 47.5K of them (95%) are used for training purposes, and 2.5K for evaluation (5%).
We consider the following classification metrics: accuracy, precision, recall, and F1 score.
The experiments were run on a server with an NVIDIA A100 graphics card, with 40GB memory.

Each experiment on FPTT and STEVE took about an hour to complete, while experiments on the decoder-only baseline took about double the time. We attribute this to the fact that, without an internal representation, the transformer needs to deal with the longer $z_{t}$ sequences, and the time and memory requirements for the self-attention operation scale quadratically with sequence length.

Figure: Figure 4: Illustration of the process described in section 4.3 for FPTT and STEVE. Notation has been simplified with respect to figure 1. C represents the corrector transformer, P stands for the predictor one.
Refer to caption: extracted/5630022/differenza-corrpred-vs-monopred.drawio.png

### 4.4 Results

We report on the result of our experiments in figure 5.
Each experiment was repeated 5 times for statistical significance.

It can be observed that FPTT and STEVE are comparable in terms of performance, with the decoder-only exhibiting a lower performance.
Upon closer analysis, especially of the F1 score (5(b)) and recall (5(d)), we observe that FPTT tends to be more stable than STEVE as it shows a narrower error margin.
We also point out that FPTT achieves peak performance earlier than STEVE.
These considerations allow us to conclude that FPTT is a more stable and sample-efficient world model than the baselines.

The only exception to the above is the precision metric (5(c)), where STEVE achieves 1 after seeing very few examples.
However, its performance on recall is significantly more erratic and lower on average, resulting in an overall lower F1 score.

In addition to qualitative analysis, we support the sample efficiency claim with quantitative evidence.
Based on Gu et al. (), we set a performance threshold at 0.95 on the F1 score and measure the number of training steps required in each experimental run to reach this threshold for the first time.
As a consistency condition, we require the threshold to be exceeded for 4 consecutive training epochs (each epoch has 500 training steps, followed by an evaluation phase.)
We report the result in table 1, noting that our world model shows a $35\%$ improvement with respect to the STEVE baseline with about half the error margin.

**Table 1: Quantitative measure of sample efficiency by reporting the number of training steps required for the F1 score to be above 0.95 consistently for 4 epochs in a row. We report the mean result for each architecture, and the standard error of the mean as a measure of confidence. *Note: For the decoder-only baseline, only one out of the 5 runs cleared the threshold.**
|  | Mean $\pm$ Standard error |
| --- | --- |
| FPTT (Ours) | $\mathbf{5500\pm 758.3}$ |
| STEVE | $8500\pm 1483.2$ |
| Decoder only | $29000$ *(see caption) |

Figure: (a) Accuracy
Refer to caption: extracted/5630022/EVAL-accuracy.png

## 5 Discussion

In this section, we discuss the limitations of our approach, outline possible future research directions and draw our conclusions from this study.

### 5.1 Limitations

Despite the observed performance improvements, the representation remains opaque and lacks interpretability.
Our preliminary attempts at replicating the object segmentation displayed by slot-attention architectures (Kipf et al., ; Singh et al., ) have not yet yielded positive results.
However, this may be overturned by more systematic experimentation in the future.

We must also note that, despite the effort to keep the number of layers of the components low, the current version of the FPTT architecture is memory-demanding, with a requirement of about 22 GB of GPU RAM at training time, including the tokenizer and the transformer triplet and excluding the classifier.

Finally, we acknowledge that we only experimented on a simplistic synthetic dataset for the initial tests of this new model.
Although we claim that the presented experiments demonstrate the benefit of the proposed architecture, we do plan to extend the experiments to more complex video datasets such as MOVi-E (Greff et al., ) and Physion (Bear et al., ), which will test its performance in more realistic scenes as well as its generalization capabilities.

### 5.2 Conclusion

We propose a new architecture, the Future-Predicting Transformer Triplet for world modeling (FPTT), which leverages the power of transformers for sequence learning to model the behaviour of objects in a set of videos and predict the evolution of the environment.

We experimentally show that our architecture outperforms transformer-based world models (Micheli et al., ), and improves on slot-attention methods (Kipf et al., ; Singh et al., ) in terms of sample efficiency and stability during the training process.

In the future, we intend to conduct further experiments with the architecture in more interactive environments, in which objects can be moved by agents.
Moreover, we would like to study applications to causal discovery problems (Yu et al., ), where learning a compact representation that can be interpreted causally might help in understanding complex scenarios.

## Acknowledgments and Disclosure of Funding

This work has been funded by the Italian National PhD Program on Artificial Intelligence run by Sapienza University of Rome in collaboration with the Italian National Council for Research, and by the Italian PNRR MUR project PE0000013-FAIR.

## Appendix A Hyperparameters and configuration

This whole work was implemented in Pytorch (Paszke et al., ); the code is open source and published at the following URL:
[https://github.com/torchipeppo/transformers-and-slot-encoding-for-wm](https://github.com/torchipeppo/transformers-and-slot-encoding-for-wm)

### A.1 VQVAE for tokenization

See table 2 below. The implementation is based on Esser et al. () and Micheli et al. ().

**Table 2: Hyperparameters for the VQVAE, both encoder and decoder**
| Hyperparameter | Value |
| --- | --- |
| Video resolution (pixels) | $64\times 64$ |
| Number of tokens per frame | $64$ |
| Channels in convolution | $64$ |
| Number of residual conv. layers | $10$ |
| Number of self-attention layers | $3$ |

### A.2 Transformers and slot encoding

See table 3 below. The implementation is based on nanoGPT (Karpathy, ). The transformers involved, i.e. the corrector-predictor-decoder triplet and the task success classifier, use the same values for the hyperparameters unless otherwise specified.

**Table 3: Hyperparameters for the transformer triplet**
| Hyperparameter | Value |  |
| --- | --- | --- |
| Vocabulary size | $50304$ |  |
| Number of tokens per frame | $64$, as above |  |
| Token embedding dimension | $768$ |  |
| Number of layers | (corrector) | $2$ |
|  | (predictor) | $2$ |
|  | (decoder) | $6$ |
|  | (task classifier) | $2$ |
| Number of attention heads | $12$ |  |
| Number of slots | $4$ |  |
| Given video frames $(N)$ | (task classifier) | $5$ |

### A.3 Training process

See table 4, 5, and 6.

**Table 4: General configuration for the training process**
| Configuration |  | Value |
| --- | --- | --- |
| Epochs | $100$ |  |
| Batch size $(BS)$ | $50$ |  |
| Batches per epoch $(BPE)$ | $10$ |  |
| Training steps per epoch | $BS\times BPE=500$ |  |
| Data samples | Training | 47500 |
|  | Evaluation | 2500 |

**Table 5: Optimizer for the VQVAE**
| Hyperparameter | Value |  |
| --- | --- | --- |
| Type | Adam | (Kingma and Ba, ) |
| Leaning rate | $10^{-4}$ |  |

**Table 6: Optimizer for each transformer**
| Hyperparameter | Value |  |
| --- | --- | --- |
| Type | AdamW | (Loshchilov and Hutter, ) |
| Leaning rate | $6\cdot 10^{-4}$ |  |
| Weight decay | $0.1$ |  |
| $\left(\beta_{1},\beta_{2}\right)$ | $\left(0.9,0.95\right)$ |  |

### A.4 Classification experiments duration

See table 7. We remind that the experiments were run on a server with an NVIDIA A100 graphics card, with 40GB memory.

**Table 7: Training time of the classifier for each underlying world modeling architecture.**
| Architecture | Average time (hours) |
| --- | --- |
| FPTT | 1.077 |
| STEVE | 0.8604 |
| Decoder-only | 1.999 |

## Appendix B Existing asset attribution

The following implementations have been referenced during this work:

- •
nanoGPT (Karpathy, ) for the Transformer implementation, licensed under the MIT license;
- •
Esser et al. () for the VQVAE implementation, released under the MIT license;
- •
Micheli et al. () for further details about the Transformer and VQVAE implementations, as well as the "decoder only" baseline, licensed under the GPL;
- •
Singh et al. () for the STEVE baseline, licensed under the MIT license.

The PHYRE video dataset has been generated by Qi et al. () from the PHYRE simulator (Bakhtin et al., ). It was downloaded following the instructions on the author’s GitHub repository: [https://github.com/HaozhiQi/RPIN/blob/master/docs/PHYRE.md##11-download-our-dataset](https://github.com/HaozhiQi/RPIN/blob/master/docs/PHYRE.md##11-download-our-dataset)
Correspondence with the author has confirmed that the dataset is released under the same license as PHYRE itself, i.e. the Apache license.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: My main question is the use of "slot attention" - as far as I can understand there is actually no slot attention in this model, am I right? it seems that the corrector just uses cross-attention and not slot attention? (the difference would be the soft-max axis).
soundness: 2
strengths: Originality:
The model presented is a very mild variation on previously published works - using VQ-VAE encodings is nice (though probably requires a bit more analysis) and the general recurrent setup is appealing.

Quality:
The proposed model variants (pre and default) are interesting and probably a good step towards analyzing the model's behaviour.

Clarity:
The paper is nicely structured and well written.

Significance:
The context of the work is important, but see below for criticism.
summary: This paper present a slotted recurrent network model which uses transformers as the main backbone for "world modeling". In this context the resulting model is an "object centric" learning model which cross attends into VQ-VAE encoded input frame, updates the current state and then predicts the next state using a transformer. The model is trained for state prediction (with two variants, either next state prediction or current state prediction) and is demonstrated to mildly work better than a single external baseline (STEVE) and one ablation model (decoder only, where there is not explicit state representation, just prediction and decoding). The experiments are run on a physical reasoning task (PHYRE) and the output is a classification readout.
weaknesses: Unfortunately the paper suffers from several weaknesses.

Experimental validation:
The method is only validated on one task and even on that task results are not very convincing. The models perform very closely to one another and the claims for efficient learning with the model are not well supported.

Analysis:
In general I don't mind when results of a model are not competitive with baselines or ablations as long as there is good analysis of why that is the case, and how this can improve our understanding of the model or problem. Here, however, these are absent - there's very little analysis of what the model learns, how it does that and what determines its performance.

Presentation:
The experimental result figures are not to the level I would expect to see in an ICLR paper - raw training curves are fine if they tell a clear story. Here, however, they do not - there is very little signal there to observe. Export quality is also quite low and does not at the level I would expect.

Novelty:
While usually I don't think novelty is a determining factor for a paper, I feel here this is quite lacking and the proposed model is indeed quite close to existing literature (SAVI++, PARTS and more ). These are cited in the paper, so I have no complaints on that side, but given the generally weak results and analysis I think this hurts the paper.

---

confidence: 4
contribution: 1
initial_rating: 3
presentation: 1
questions: * __Slot encodings.__ How and where do the authors utilize slot encodings, and what specific benefits do they offer in this context?

* __Experimental design.__ Why do the authors limit their evaluation to a single dataset? Additionally, what is the rationale for selecting STEVE as the sole baseline, excluding other relevant related works?
soundness: 1
strengths: The authors address an important problem of physical world modeling using structured latent representations.
summary: The paper proposes the Future-Predicting Transformer Triplet (FPTT), an architecture aimed at modeling of physical world dynamics from video data. It employs Transformers to learn object-centric representations, enabling the model to predict physical interactions between objects more effectively. The architecture is tested on synthetic video dataset PHYRE. The authors also perform an ablation study to understand the contribution of different components.
weaknesses: * __Missing slot encodings and object-centricity.__ 

While the paper includes _slot encoding_ in its title, the approach itself appears to lack this feature. As I understand it, $\Lambda$ was intended to serve as slot encodings, but it is not even referred to as such. From the description, $\Lambda$ seems more like a standard intermediate representation in transformer layers rather than a distinct slot encoding.

In Appendix A2, the authors reference [1] for their transformer implementation, where they also mention using four slots. However, the referenced implementation does not include a parameter for the number of slots, leaving it unclear how slot encodings are actually integrated into the proposed approach, or if they are implemented at all.

Furthermore, the authors state that _"the representation remains opaque and lacks interpretability,"_ which raises questions about the motivation for using _slot encodings_ in the first place.


* __Experimental methodology.__ 

Firstly, the authors evaluate their model on a single, very simplistic dataset, while using a relatively large number of parameters. This limited evaluation setup may not provide sufficient empirical evidence to support their claims.

A more significant issue lies in their positioning among related works and choice of baselines. The authors overlook most recent related work (e.g., [2, 3, 4, 5]) and rely solely on STEVE as a baseline, aside from variations of their own approach.


* __Presentation.__ 

In addition to unclear explanations of their approach and its novelty, the authors fail to position it effectively within the existing literature, lacking a comparative analysis with prior work.

All the figures also present issues: they are unnecessarily large, some are in low resolution, and it is often unclear what the authors aim to demonstrate.

\
References: 


[1]: Andrej Karpathy. nanoGPT: The simplest, fastest repository for training/finetuning mediumsized GPTs (Generative Pretrained Transformers), 2023. URL https://github.com/karpathy/nanoGPT. \
[2]: Nakano, A., Suzuki, M. and Matsuo, Y., 2023. Interaction-based disentanglement of entities for object-centric world models. In The Eleventh International Conference on Learning Representations.\
[3]: Villar-Corrales, A., Wahdan, I. and Behnke, S., 2023, October. Object-centric video prediction via decoupling of object dynamics and interactions. In 2023 IEEE International Conference on Image Processing (ICIP) (pp. 570-574). IEEE.\
[4]: Wu, Z., Dvornik, N., Greff, K., Kipf, T. and Garg, A., 2022. Slotformer: Unsupervised visual dynamics simulation with object-centric models. arXiv preprint arXiv:2210.05861.\
[5]: Daniel, T. and Tamar, A., DDLP: Unsupervised Object-centric Video Prediction with Deep Dynamic Latent Particles. Transactions on Machine Learning Research.

---

confidence: 4
contribution: 2
initial_rating: 3
presentation: 3
questions: - How would the architecture compare against Dreamerv3 or Slotformer in world modeling?
- What happens if u try to make the decoder only model more efficient by reducing the number of tokens or dimensionality of the token?
- How would the paper compare against the baselines in benchmarks proposed in Steve or DreamerV3 or SlotFormer?
soundness: 3
strengths: - The paper focuses on world modeling which is an important problem.
- The writing and presentation is clean, which makes understanding the paper easy.
- The paper compares efficiency and accuracy, which helps understand the trade-offs.
summary: The paper addresses the challenge of creating sample-efficient models for physical world modeling, focusing on predicting object interactions in dynamic environments.

The authors propose an architecture that combines Transformers with slot encoding to improve sample efficiency and stability in world modeling. Unlike existing models that operate at the image level, this model incorporates object-based representations, enabling it to capture and predict interactions more accurately.

Their model, named Future-Predicting Transformer Triplet (FPTT), uses a corrector-predictor-decoder triplet of Transformers. The corrector aligns the internal state representation with the actual video evolution to prevent model drift, the predictor forecasts the next state based on the corrected representation, and the decoder converts this predicted state back into tokens for further training.

Experiments using the PHYRE dataset (a benchmark for physical reasoning) show that FPTT achieves greater sample efficiency and training stability compared to baseline models like STEVE. The model’s structured approach enables it to generalize well in physical environments simulated with basic Newtonian physics.


In summary, the paper presents a architecture that leverages the strengths of Transformers and slot encoding for efficient and stable world modeling, demonstrating improvements in tasks requiring understanding and predicting object dynamics in a physical environment​
weaknesses: - The contribution is not significant having an internal representation of the previous timesteps is common in world modeling architectures, for instance: Dreamer: https://arxiv.org/pdf/2301.04104, Slotformer: https://arxiv.org/pdf/2210.05861. It's unclear to me how this work is a better architecture than Dreamer or Slotformer or other recent works.
- The evaluations and baselines are weak, the paper only compares against STEVE. I dont think STEVE is a fair comparision as their objective was to get interpretable object representations and not necessarily the metric/benchmark paper uses for evaluation. Further the Decoder-Only model seems to perform as well as the proposed architecture on almost all tasks except efficiency.
- Lastly the work only compares on a single benchmark, which is not being used in the baseline works such as STEVE. I think a fair thing to do would be to compare on benchmarks shown in prior baselines, so we assume they are tuned well.

---

confidence: 5
contribution: 1
initial_rating: 3
presentation: 1
questions: - It appears that slot attention (or inverted attention) is absent, with only cross attention being mentioned. Is this an oversight in the explanation, or is it indeed absent? If it's truly absent, how can we be confident that it captures object-level dynamic understanding?
soundness: 2
strengths: - The authors' approach of capturing slot-like internal representations from VQ tokens, instead of CNN embeddings, was intriguing and showed promising results.
- They propose a novel evaluation protocol for testing world model architecture, utilizing the shared benchmark with different protocols.
summary: This paper proposes a world-modeling architecture that captures object-level interactions in the scene, instead of the scene itself. The architecture consists of three transformer-based models: a corrector, a predictor, and a decoder, alongside a VQ-VAE tokenizer for image encoding.

To evaluate the proposed model’s performance as a world model, the authors provide a physical reasoning task using the PHYRE benchmark, demonstrating that their model outperforms STEVE, the baseline, in terms of prediction accuracy and sample efficiency.
weaknesses: - **Lack of Novelty and Justification**: The main idea and direction of the paper have been explored in several existing works (OCVT[1], SlotFormer[2]). Although the authors are likely aware of this, they fail to convincingly demonstrate why their approach is unique and necessary for the proposed direction, compared to previous works.
- **Architecture and Design Choices**: The proposed architecture appears to be a combination of SAVi[3] and STEVE[4] architectures, but with a predictive loss function instead of a reconstructive loss function. While this variant may be promising, the paper lacks sufficient details to justify the design choices, such as thorough ablation studies. Furthermore, it is unclear whether the proposed model can outperform existing works, as it is not comprehensively compared.
- **Lack of Clarity in Architecture and Experiment Description**: The architecture section of the paper lacks clarity and detail, particularly in the description of the core components: corrector transformer, predictor transformer, and decoder transformer. Although the author provides a high-level overview of these architectural concepts, the explanation is insufficient given the emphasis on this part as the paper's core contribution. To thoroughly understand and investigate the proposed architecture, a detailed formulation of these components is necessary, including their implementation details and mathematical representations.
    
    Furthermore, the experiment section lacks sufficient details about the metrics used in the evaluation. To ensure transparency and reproducibility, it is essential to provide a clear explanation of each metric, including how the metric is calculated and what it represents.
    
- **Limited Evaluation of Proposed Architecture:** The author only provides a single task to evaluate the proposed architecture, which is insufficient to demonstrate its generality and versatility. To thoroughly assess the world modeling ability of the proposed architecture, it is essential to evaluate it on a diverse range of tasks that require the model to infer and understand the relationships between objects and scenes. Additionally, to facilitate a fair comparison with existing works, the authors may consider including several generation tasks (e.g. OBJ3D[1], CLEVR[5], Physion[6]), as has been done in prior research.
    
    To further demonstrate the effectiveness of the proposed model, the author could compare it with a broader range of baselines, such as SlotFormer, OCVT, SAVi, and other relevant models mentioned in the paper. Although these models are typically used for generation tasks, their predicted representations can be evaluated using the same protocol as the proposed model. Additionally, comparing with image-based world models would illustrate the advantages of object-level world models over their image-based counterparts. This approach would provide a more comprehensive understanding of the proposed model's performance and allow for a more accurate assessment of its strengths and limitations relative to existing approaches.
    
- **Ablation Results Raise Questions about Proposed Model:** The ablation results indicate that the ‘decoder-only’ model performs comparably to the proposed models. This suggests that VQ-tokenization and predictive loss might be sufficient to drive performance without explicitly enforcing object-level representations. This outcome seems misaligned with the paper's main theme, which emphasizes the importance of object-level representations. Consequently, this misalignment raises questions about the necessity and effectiveness of the proposed model's architecture.

[1] Wu, Yi-Fu, Jaesik Yoon, and Sungjin Ahn. "Generative video transformer: Can objects be the words?." International Conference on Machine Learning. PMLR, 2021.

[2] Wu, Ziyi, et al. "Slotformer: Unsupervised visual dynamics simulation with object-centric models." arXiv preprint arXiv:2210.05861 (2022).

[3] Kipf, Thomas, et al. "Conditional object-centric learning from video." arXiv preprint arXiv:2111.12594 (2021).

[4] Singh, Gautam, Yi-Fu Wu, and Sungjin Ahn. "Simple unsupervised object-centric learning for complex and naturalistic videos." Advances in Neural Information Processing Systems 35 (2022): 18181-18196.

[5] Johnson, Justin, et al. "Clevr: A diagnostic dataset for compositional language and elementary visual reasoning." Proceedings of the IEEE conference on computer vision and pattern recognition. 2017.

[6] Bear, Daniel M., et al. "Physion: Evaluating physical prediction from vision in humans and machines." arXiv preprint arXiv:2106.08261 (2021).

Meta-review: This paper introduces a framework for world modelling using tokenised (object-centric) representations. A key part of the method is the combination of a transformer architecture with a slot-attention-like mechanism. Experiments are run on the PHYRE simulated physical reasoning benchmark.

The reviewers acknowledged that this paper addresses an important problem, namely whether structured representations can benefit physical world models.

All reviewers raised significant concerns about clarity and quality of presentation, validity and strength of experimental evaluation, and novelty of the method, resulting in a clear reject decision.

Decision: Reject