# Solving Reasoning Tasks with a Slot Transformer

Ryan Faulkner
Deepmind
<EMAIL_ADDRESS>
&Daniel Zoran
Deepmind
<EMAIL_ADDRESS>

###### Abstract

The ability to carve the world into useful abstractions in order to reason
about time and space is a crucial component of intelligence. In order to
successfully perceive and act effectively using senses we must parse and
compress large amounts of information for further downstream reasoning to take
place, allowing increasingly complex concepts to emerge. If there is any hope
to scale representation learning methods to work with real world scenes and
temporal dynamics then there must be a way to learn accurate, concise, and
composable abstractions across time. We present the Slot Transformer, an
architecture that leverages slot attention, transformers and iterative
variational inference on video scene data to infer such representations. We
evaluate the Slot Transformer on CLEVRER, Kinetics-600 and CATER datesets and
demonstrate that the approach allows us to develop robust modeling and
reasoning around complex behaviours as well as scores on these datasets that
compare favourably to existing baselines. Finally we evaluate the
effectiveness of key components of the architecture, the model’s
representational capacity and its ability to predict from incomplete input.

##  1 Introduction

Reasoning over time is an indispensable skill when navigating and interacting
with a complex environment. However, rationalizing about the world becomes an
intractable problem if we are incapable of compressing it into to a reduced
set of relevant abstractions. For this reason relational reasoning and
abstracting high-level concepts from complex scene data is a critical area of
machine learning research. Past approaches use relational reasoning and
composable scene representation and have met success on static datasets
(Santoro et al., 2017, 2018; Greff et al., 2020; Locatello et al., 2020;
Burgess et al., 2019) however, those gains must now be extended across time;
this introduces a large set of new complexities in the form of temporal
dynamics and scene physics. Some recent work has utilized transformers (Jaegle
et al., 2021; Vaswani et al., 2017; Ding et al., 2020) which provide a
potential path toward understanding the type of approach that can succeed at
solving this problem.

The goal of the approach presented in this paper involves learning to output
useful spatio-temporal representations of the input scene sequence which are
hypothesized to be critical to solving downstream tasks for scene
understanding and generalisation in domains containing complex visual scenes
and temporal dynamics. The target domains for the evaluation of this type of
model we have chosen involve scenes of synthetic and real world objects and
behaviours requiring complex scene understanding and abstract reasoning via
question and answer datasets. This work combines a number of existing ideas in
a novel way, namely slot attention (Locatello et al., 2020), transformers
(Vaswani et al., 2017) and iterative variational inference (Marino et al.,
2018); in order to better understand conceptual scene representation and how
it can be achieved and embedded in more complex systems. Several hypotheses
drive the direction for this work. First, information about scenes may be more
efficiently represented as independent components rather than in a monolithic
representation. Second, ingesting spatio-temporal input in such a way that
there is no bias to any step of the sequence is critical (see supplementary
material for analysis) when forming representations about space and time over
long sequences. Finally, processing information in an iterative fashion can
provide the means to recursively recombine information in a way that is useful
for reasoning tasks.

The main contributions of this work are: 1) present the Slot Transformer
architecture for spatio-temporal inference and reasoning and a framework for
learning how to encode useful representations, 2) evaluate this model against
downstream tasks that require video understanding and reasoning capabilities
in order to be solved, and finally 3) demonstrate the role the components of
the overall approach play in the problem solving capabilities induced during
training. We have chosen three tasks on which to evaluate the Slot
Transformer: CLEVRER (Yi et al., 2020), a video dataset where questions are
posed about objects in the scene, Kinetics-600 comprising of YouTube video
data for action classification (Carreira et al., 2018) and CATER (Girdhar and
Ramanan, 2019), object relational data which requires video understanding of
scene events to successfully solve. In each of these cases we evaluate the
Slot Transformer against current state-of-the-art approaches.

##  2 Related Work

We have drawn inspiration from models that induce scene understanding via
object discovery. In particular, the approaches taken by Locatello et al.
(2020); Greff et al. (2020); Burgess et al. (2019) all involve learning latent
representations that allow the scene to be parsed into distinct objects. Slot
attention (Locatello et al., 2020) provides a means to extract relevant scene
information into a set of latent vectors where each one queries scene pixels
(or ResNet encoded super pixels) and where the attention softmax is done
across the query axis rather than the channel axis, inducing slots to compete
to explain each pixel in the input. Another _slotted_ approach appears in
IODINE (Greff et al., 2020), which leverages iterative amortized inference
(Marino et al., 2018; Andrychowicz et al., 2016) to produce better posterior
latent estimates for the slots which can then be used to reconstruct image by
using a spatial Gaussian mixture using masks and means decoded from each slot
as mixing weights and component means respectively. We aim to extend this type
of approach by applying the same ideas to sequential input, in particular,
video input.

Self-attention and transformers (Vaswani et al., 2017; Parisotto et al., 2019)
are also central to our work and have played a critical role in the field in
recent years. Some of the latest applications have demonstrated the efficacy
of this mechanism applied to problem solving in domains that require a
capacity to reason as a success criteria (Santoro et al., 2018; Clark et al.,
2020; Russin et al., 2021). In particular, self-attention provides a means by
which to form a pairwise relationship among any two elements in a sequence,
however, this comes at a cost that scales quadratically in the sequence
length, so care must be taken when choosing how to apply this technique. One
recent approach involving video sequence data, the TimeSformer from Bertasius
et al. (2021), utilizes attention over image patches across time applied
successfully to sequence action classification (Goyal et al., 2017; Carreira
et al., 2018).

There has been a great deal of work on self-supervised video representation
learning methods (Wang et al., 2021; Qian et al., 2020; Feichtenhofer et al.,
2021; Gopalakrishnan et al., 2022) among the family of video representation
learning models. In particular in Qian et al. (2020) the authors have explored
using a contrastive loss and applied to Kinetics-600. We include some of these
results below in table 3 below. In Kipf et al. (2021) the authors present an
architecure similar to ours, albeit without generative losses, where they show
that their model achieves state-of-the-art performance on image object
segmentation in video sequences and optical flow field prediction. In Ding et
al. (2020) the authors use an object representation model (Burgess et al.
(2019)) as input to a transformer and achieve very strong results on CLEVRER
and CATER.

Neuro-symbolic logic based models (d’Avila Garcez and Lamb, 2020) have been
used to solve temporal reasoning problems and, in particular, to CLEVRER and
CATER. In section 4 we examine the performance of the Dynamic Context Learner
(DCL, Chen et al. (2021)) and Neuro-Symbolic Dynamic Reasoning (NS-DR, Yi et
al. (2020)) where NS-DR consists of neural networks for video parsing and
dynamics prediction and a symbolic logic program executor and DCL is composed
of Program Parser and Symbolic Executor while making use of extra labelled
data when applied to CLEVRER. Some components of these methods rely on
explicit modeling of reasoning mechanisms crafted for the problem domain or on
additional labelled annotations. In contrast, the intent of our approach is to
form an inductive bias around reasoning about sequences in general.

Recently Jaegle et al. (2021) have published their work on the Perceiver, a
model that makes use of attention asymmetry to ingest temporal multi-modal
input. In this work an attention bottleneck is introduced that reduces the
dimensionality of the positional axis through which model inputs pass. This
provides tractability for inputs of large temporal and spatial dimensions
where the quadratic scaling of transformers become otherwise prohibitive.
Notably the authors show that this approach succeeds on multimodal data,
achieving state-of-art scores on the AudioSet dataset (Gemmeke et al., 2017)
containing audio and video inputs. Ours is a similar approach in spirit to the
Perceiver in compressing representations via attention mechanisms however, we
focus on the ability to reason using encoded sequences and primarily reduce
over spatial axes.

##  3 Model

Figure 1: The general model architecture with the encoding, decoding and
iteration phases of the model. The encode phase computes an updated context
$c_{k+1}$, the decode phase produces a reconstruction
$\mathbf{x}^{{}^{\prime}}_{k}$. Finally, the iterate process repeats the
encode-decode steps compressing more information into the context at each
iteration. KL and likelihood losses are computed across all iterations.

We now present the Slot Transformer, a generative transformer model that
leverages iterative inference to produce improved latent estimates given an
input sequence. Our model can broadly be described in terms of three phases
(see Figure 1): an encode phase where a spatio-temporal input is compressed to
form a representation, then a decode phase where the representation is used to
reconstruct known data or predict unseen data; finally the iterate phase,
which encompasses the other two phases, whereby conditioning new
representations on those from the previous iteration in the sequence input
enable more useful representations to be inferred. This process flow is
similar in manner and inspired by the encode-process-decode paradigm (Hamrick
et al., 2018). We leverage the ideas of iterative attention used in slot
attention and IODINE (Locatello et al., 2020; Greff et al., 2020) based upon
iterative variational inference to model good, compact representations for
downstream tasks.

###  3.1 Encoder

The input sequence $\mathbf{X}\in\mathbb{R}^{T\times H\times W\times 3}$ is
encoded once from raw pixels with a residual network (He et al., 2015) applied
to each frame across the sequence: $\mathbf{e}=f_{\text{ResNet}}(\mathbf{X})$
where $\mathbf{e}\in\mathbb{R}^{T\times H^{\prime}\times W^{\prime}\times E}$
where $H^{\prime}<H$ and $W^{\prime}<W$. To the resulting image encodings
$\mathbf{e}$ we concatenate a spatial encoding basis comprised of Fourier
basis functions dependent upon the spatial coordinates of the pixels.

Next we define the context: a set of $T\times K$ vectors (or _slots_), each of
$C_{\text{context}}$ dimensions, that store contextual information about the
scene. We denote the context as $\mathbf{c}\in\mathbb{R}^{T\times K\times
C_{\text{context}}}$ where $T$ is the number of time steps in the input
sequence. We will perform iterative updates on the context over the entire
spatio-temporal volume with each forward pass through the model. At each
iteration, the encoder’s role is to infer an updated spatio-temporal context
via the context-transformer where the input is the combination of the context
from the previous iteration plus the output of the slot attention on the input
(more detail in 3.3).

We initialize the context by sampling a standard Gaussian of size $K\times
C_{\text{context}}$ then tiling this tensor over the same number of time steps
as the input 111It should be noted that while in this work the encoded
representations retain the full time dimension this is not strictly necessary
and this architecture may be modified to project to lower cardinality time
dimensions in a straightforward manner.. This yields the initial context
$\mathbf{c}_{0}$. Sampling the context this way we break symmetry across the
slots at each frame, and by tiling across time we encourage slots to have the
same role across adjacent time-steps. The slotted context provides an
inductive bias to learn structured state information of the scene dynamics
over time and space which the attention and context transform operations
facilitate.

####  3.1.1 Slot Attention

Given the initial context, $\mathbf{c}_{0}$, we next apply a shared MLP across
the slots to construct the queries, $\mathbf{q}_{k}$, one for each slot at
each time step. We use these queries to attend over the keys and values
decoded from the super pixels in $\mathbf{e}$ and their corresponding spatial
position encodings. The softmax step in this operation is applied across the
slot dimension such that slots "compete" on which pixels in the scene they
"explain" and losses that reward more efficient representations will induce
the model to learn representations that best explain the scene (Locatello et
al., 2020). One critical detail of this attention readout is that it occurs
batched across time steps, learning temporal relationships is handled by the
context transformer described below. Much of this detail is captured in figure
1.

The readout from the slot-attention step $\mathbf{a}_{k}$ for each slot at
each time step is combined with the context via GRU style gating (Hochreiter
and Schmidhuber, 1997; Cho et al., 2014) with the gating function
$\mathbf{c_{k}^{\prime}}=f_{\text{gate}}(\mathbf{c}_{k},\mathbf{a}_{k})$.

####  3.1.2 Context Transformer

Now that we have obtained the updated context $\mathbf{c^{\prime}}$ using the
input, it is passed through a transformer (Vaswani et al., 2017),
$\mathbf{c}_{k}=T_{\text{context}}(\mathbf{\mathbf{c^{\prime}_{k}}})$, where
positional encodings across the time sequence are applied but no masking. This
enables the model to form temporal connections across the sequence via the
context slots. At this step the same transformer weights are applied
separately to each slot across time meaning that elements of each slot only
communicates with other elements in that slot over the time axis.

The reason for this approach is two-fold, 1) this provides an inductive bias
for the slots to adopt specialized roles when explaining the input, 2) this
ensures that the complexity of the transform operation does not scale with the
number of slots. At this step we could choose to introduce a temporal
asymmetry as has shown promise in other work (Jaegle et al., 2021). This could
yield better scalability and representational power if done correctly and we
leave this as an avenue of future work.

One final note, the context transformer is only applied on the first iteration
of the encoder as this allows the model to scale more easily to greater
numbers of iterations, integrating information from the input through the slot
attention operation alone on subsequent iterations.

###  3.2 Decoder

Once the encoder has generated an updated context we may hang new losses off
of this representation by way of the decoder. From here there are many
possibilities for the way forward, for the scope of this this work we chose to
explore image reconstruction over all frames in the sequence. We hypothesize
that these will help induce the model to learn useful and composable
representations via the slots and to learn the temporal dynamics of the input
via the context transformer. We also decided to apply variational inference
(Doersch, 2021) as we posit that this should help the model to better
generalize by compressing the latent representations. Therefore, we use the
context to parameterize a Gaussian distribution (approximate posterior) by
linearly projecting the context into parameters
$\lambda_{k}\in\mathbb{R}^{2C_{\text{latent}}}$ and decode after sampling from
it: $\mu_{k},\log\sigma_{k}=\lambda_{k}=W_{\lambda}\mathbf{c}_{k}$ and
$\mathbf{z}_{k}\sim N(\mu_{k},\sigma_{k})$. While our choice of decoder is
certainly dependent upon this choice we discuss this in more detail in the
next section.

####  3.2.1 Spatial Broadcast Decoder & Masks

Since the slots contain information that is likely spatially entangled we want
to ensure that the encoder is left free to learn the spatio-temporal content
based representations among the slots that can then be used by the decoder to
rebuild the input. For this reason we use a spatial broadcast decoder (Watters
et al., 2019) to distill information from the slots with the key property that
the context channels are broadcast across spatial dimensions and decoded
without upsampling - in effect each slot is allowed to explain independent
parts of the image content. The spatial decoder is batch applied across the
batch, time and slot dimensions of the slotted latent to produce a one channel
mask, $\mathbf{m}_{k,t}$, and an RGB mean image, $\mathbf{r}_{k,t}$:
$\mathbf{m}_{k,t},\mathbf{r}_{k,t}=f_{\text{decoder}}(\mathbf{z}_{k,t})$.
These two elements are combined via weighted sum in a manner similar as is
done in Slot Attention:
$\mathbf{x^{\prime}}_{t}=\sum_{k=1}^{K}\mathbf{m}_{k,t}\mathbf{r}_{k,t}$.

###  3.3 Iterative Model & Losses

As mentioned in Section 3.2 variational inference is applied to the model by
projecting the context output from the encoder as Gaussian posterior parameter
estimates over a latent distribution from which we sample $z_{k}$.

We define a conditional prior, $\mathcal{P}(\mathbf{z}|\mathbf{x}_{1..p})$, by
feeding the first $p<T$ steps of the context to the context transformer
thereby limiting the prior to information within a subsequence of the input to
form the prior context. To predict the remaining $T-p$ steps of the latent
representation an auto-regressive model is used where the initial state is
formed via a self attention operation on the prior context concatenated with a
learnable token. After the self attention operation is executed the output
token is read and used as the initial state of the auto-regressive model. The
output of this model is combined with the prior context to form the full
definition of the latent prior parameters for the full sequence.

We make use of iterative inference (Marino et al., 2018) as applied by Greff
et al. (2020) where updates to the posterior parameters estimate,
$\lambda_{k}$, are computed dependent on the sampled latents $z_{k}$, the
input data $x$, the encoder architecture $f_{\phi}$, and auxiliary inputs
$a_{k}$:

$\lambda_{k+1}=\lambda_{k}+f_{\phi}(z_{k},x,a_{k})\\\ $ (1)

On each iteration we also compute the reconstruction loss in addition to the
Kullback-Liebler (KL) measure (Kullback and Leibler, 1951) of the posterior
parameters with respect to the prior (i.e the ELBO):

$\mathcal{L}_{\text{gen}}=\sum_{t=1}^{T}\sum_{k=1}^{K}{\mathcal{D}_{KL}(\mathcal{N}(\mu^{\text{post}}_{k,t},\sigma^{\text{post}^{2}}_{k,t})||\mathcal{N}(\mu^{\text{prior}}_{k,t},\sigma^{\text{prior}^{2}}_{k,t}))}-\log{L(\mathbf{x}_{t}|\mathbf{x}_{t}^{\prime})}$
(2)

We use a mixture distribution for the output likelihood
$\log{L(\mathbf{x}_{t}|\mathbf{x}_{t}^{\prime})}$, where the masks define a
categorical distribution of components over Gaussian data distributions.

###  3.4 Auxiliary Losses

In addition to the variational and supervised losses we define auxiliary
losses to help to induce better representations for scene understanding. For
this we use the information from the slot specialization and language
describing the scene when available.

####  3.4.1 Object Mask Prediction

This loss is formed via model predictions on the latent slot values for
selected target frames. This is similar to the self-supervision strategies
used in Ding et al. (2020). For this we take the samples from the final
iteration context, $\mathbf{z}_{K}$, then select $s<S$ slots at random masking
out the last $k$ steps, for $k$ a fixed parameter; this is then fed once
through the context transformer to compose $\mathbf{z}^{\prime}_{K}$. Next,
from the masked steps we randomly select a subset of step indices,
$\mathbf{T}_{\text{targets}}$, and compute the L2 norm over the difference
between the target predictions from $\mathbf{z}^{\prime}_{K}$ and the true
latents from $\mathbf{z}_{K}$:
$\mathcal{L}_{\text{object}}=\sum_{i}^{\mathbf{T}_{\text{targets}}}||\mathbf{z}_{K}-\mathbf{z}^{\prime}_{K}||$.
All of our top results were obtained when including this loss where
predictions are done across the last half of the input sequence for each of
$1$, $2$, and $3$ slots summed together.

####  3.4.2 Question Prediction

Given a Question-Answer dataset we hypothesize that we can learn better
representations if they are good enough to predict the original question when
conditioned on the answer. During training, given the answer and the model
latent $z_{K}$ we compose a new belief vector from these using a transformer
and a CLS token. This serves as the initial state of an auto-regressive model
where we use teacher forcing from the question tokens and compute the sequence
question word embeddings. The cross-entropy between the predicted question
embeddings and the true ones to form the question prediction loss:
$\mathcal{L}_{\text{question}}$.

###  3.5 Heads for Downstream Tasks

Figure 2: General form of the task head to compute answers given the output
latent, $\textbf{z}_{K}$ from the encoder on the final iteration, an encoded
question, the mask embeddings $\textbf{m}_{K}$ and a CLS token.

The context output from the final iteration, $c_{K}$, now forms an encoded
sequence that may be used as input for specialized tasks. The task head in
figure 2 depicts the general head architecture used for downstream tasks
defined in Section 4. The full input to the head consists of the question(s)
provided with the task (when available), the final context from the encoder
$\mathbf{c}_{K}$, the mask information from the last iteration of decoder
output $\textbf{m}_{K}$, and a classification token, CLS. As alluded to in
Section 3.6, we hypothesize at this point that $\mathbf{c}_{K}$ and
$\mathbf{m}_{K}$ should contain information about the objects in the scene and
the relations they share with eachother across the sequence. Note that the
time axis of $\mathbf{c}_{k}$ is now expanded to include all slots, so it is a
sequence of length $T\times K$. The token CLS is prepended to the input
sequence.

These elements are concatenated along the time axis, including absolute
position encodings and now form the input to a multi-layer gated transformer
(Parisotto et al., 2019). From the output the contents of the initial position
corresponding to the CLS token is then fed into an MLP (and softmaxed,
depending on the task) to produce logits which can then be compared to the
task labels via cross-entropy to produce a supervised loss
$\mathcal{L}_{\text{QA}}$.

The generative, auxiliary, and supervised losses are then combined with
scaling constants to yield the total loss:

$\mathcal{L}_{\text{total}}=\lambda_{\text{gen}}\mathcal{L}_{\text{gen}}+\lambda_{\text{question}}\mathcal{L}_{\text{question}}+\lambda_{\text{object}}\mathcal{L}_{\text{object}}+\lambda_{\text{QA}}\mathcal{L}_{\text{QA}}$
(3)

Refer to the appendix for the dataset specific settings for the loss scaling
constants.

###  3.6 Training

As discussed in section 3.3, we define our prior distribution model by feeding
only a proper subsequence of the input to the encoder while the posterior
"sees" the full input sequence. In practice the subsequence is chosen to be
all but a final fixed length segment of the sequence. Then sampling from the
posterior distribution via the re-parameterization trick (Kingma and Welling,
2019) we compute the log likelihood of the input sequence given the decoded
masks and outputs from the sample. Note that our final posterior estimates
also include at least one iteration of approximate inference (Marino et al.,
2018) and the full generative loss is composed of a sum over the iterations.
These loss terms are added to the supervised and auxiliary losses (eq. 3). The
KL term ensures that the conditional prior is capable of generating samples
that accurately model the unknown portion of the sequence (figure 7). Thus the
model must be able to encode information about the temporal dynamics of the
scene objects as well as the scene content in order to generalize well for
scene understanding.

Samples from the posterior on the final model iteration are used as input to
both the decoder and the task heads which also receive embeddings from the
decoded masks. For CLEVRER and CATER we use the LAMB optimizer (You et al.,
2019) with weight decay to compute update gradients while stochastic gradient
descent is applied to the Kinetics-600 dataset. The object prediction
auxiliary loss is used over all datasets, while question prediction is only
applicable to CLEVRER. We trained our models until they began to overfit which
typically occured by at most $10e6$ model updates. Detail on the experimental
parameter values used can be found in the supplementary material.

Figure 3: Input, reconstructions, and decoder masks $\mathbf{m}$ of an eight
slot model (only three shown for brevity). The slot masks partition the input
space into content based components which combine to produce accurate
reconstructions.

##  4 Experiments

We trained and evaluated our Slot Transformer on three video datasets and
tasks: 1) Kinetics-600: action classification on data from YouTube, CATER:
object localization, CLEVRER: question and answer. To solve these types of
tasks requires cognitive abilities such as object identification, object
permanence, and object relational reasoning across video sequences among
others.

###  4.1 CLEVRER

CLEVRER (Yi et al., 2020) is a visual question and answer video dataset that
requires reasoning about the objects contained in the scene to answer a given
question correctly. Much like its predecessor dataset, CLEVR (Johnson et al.,
2016), CLEVRER consists of scenes containing objects of various shapes and
colours that may move about a plain surface potentially resulting in
collisions. There are four types of questions posed: descriptive, explanatory,
predictive and counterfactual, each challenging different powers of reasoning.

We trained our model without any additional labelled data, simultaneously on
all four CLEVRER question types with separate MLPs for multiple choice and
descriptive question types. We measure test performance and compare to the
results from Chen et al. (2021) in Table 1. We computed results for
correctness on the full question, that is, choosing the correct word answer
for the descriptive task and selecting all the options correctly in the case
of multiple choice.

Table 1: CLEVRER per Question Accuracies (Chen et al. (2021))

Methods | Extra Labels | Descriptive | Predictive | Explanatory | Counterfactual
---|---|---|---|---|---
| Attr. | Prog. |  |  |  |
CNN+MLP |  |  | 48.4 | 18.3 | 13.2 | 9.0
CNN+LSTM |  |  | 51.8 | 17.5 | 31.6 | 14.7
Memory | No | No | 54.7 | 13.9 | 33.1 | 7.0
HCRN |  |  | 55.7 | 21.0 | 21.0 | 11.5
MAC (V) |  |  | 85.6 | 12.5 | 16.5 | 13.7
Slot Transformer (Ours) |  |  | 87.4 | 48.3 | 65.3 | 45.2
TVQA+ | Yes | No | 72.0 | 23.7 | 48.9 | 4.1
MAC (V+) |  |  | 86.4 | 22.3 | 42.9 | 25.1
IEP (V) |  |  | 52.8 | 14.5 | 9.7 | 3.8
TbD-net (V) | No | Yes | 79.5 | 3.8 | 6.5 | 4.4
DCL |  |  | 90.7 | 82.8 | 82.0 | 46.5
NS-DR |  |  | 88.1 | 79.6 | 68.7 | 42.4
NS-DR (NE) | Yes | Yes | 85.8 | 74.3 | 54.1 | 42.0
DCL-Oracle |  |  | 91.4 | 82.0 | 82.1 | 46.9

Figure 4: CLEVRER eight slot model trajectory chunk. These samples are drawn
from a sequence unseen during training and the quality of the reconstructions
remains qualitatively high. Objects and significant contiguous chunks of
frames are reconstructed accurately indicating that the model can infer the
missing scene information.

###  4.2 CATER

Figure 5: CATER ten slot model Top-1 training and evaluation accuracy over $1,2$ and $4$ model iterations. Models with fewer iterations tend to quickly overfit so while we see better performance on the training set the higher iteration models generalize significantly better and score higher on the unseen test set.  Table 2: CATER Results on Task 3: "Localization" Methods | Static Camera | Moving Camera
---|---|---
| Top-1 | Top-5 | L1 | Top-1 | Top-5 | L1
Random | 2.8 | 13.8 | 3.9 | - | - | -
R3D LSTM | 60.2 | 81.8 | 1.2 | 28.6 | 63.3 | 1.7
R3D + NL LSTM | 46.2 | 69.9 | 1.5 | 38.6 | 70.2 | 1.5
Aloe Ding et al. (2020) | 74.0 | 94.0 | 0.44 | 59.7 | 90.1 | 0.69
Slot Transformer (Ours) | 62.9 | 84.7 | 0.86 | 35.8 | 60.5 | 2.4

The Cater dateset (Girdhar and Ramanan (2019)) is composed of sequences of
rendered objects, also based off of the earlier CLEVR dataset. The dataset is
composed of three separate tasks: atomic action recognition, compositional
action recognition, snitch localization. In this paper we focus on the most
challenging of the three, "snitch localization". For this task, an object
called the snitch must be located on the final frame of the sequence on a
$6\times 6$ grid on the surface on which the objects move. The grid is not
actually visible in the image and the model must make it’s best guess. The
camera may be either moving or stationary.

In the case of CATER the final MLP on the task head outputs $36$ labels over
which we add a softmax layer (see Figure 2). Observing figure 5 it appears
that in order to achieve good generalisation performance it is critical to
apply multiple iterations. We compare our approach to existing baselines from
the dataset paper (Girdhar and Ramanan, 2019) and we see that our results are
competitive with existing state-of-the-art approaches. More details on the
baselines and experimental setup can be found in the supplementary material.

###  4.3 Kinetics-600

Kinetics-600 (Carreira et al., 2018) is an action recognition dataset based on
YouTube video. The model is presented with a frame sequence and asked to
classify the action in the scene. Some previous approaches have used scaled up
models or encoder pre-training setups (Kondratyuk et al., 2021; Arnab et al.,
2021) on larger YouTube or other video datasets prior to fine-tuning the
classification head. Our primary goal is to determine whether we can score in
the range of stronger models using a similar setup to what we have applied on
CLEVRER and CATER in the previous sections.

We compare our results with those reported from Carreira et al. (2018) on the
Inflated 3D ConvNet (Carreira and Zisserman, 2017), the TimeSformer (Bertasius
et al., 2021), and baselines from Qian et al. (2020). Our results are
approaching these baselines and it is noteworthy that the scores for these
architectures have been pre-trained whereas ours have not. We are intent on
exploring what effect such a training regime may have on the slot transformer
as future work.

Figure 6: Ablate Attention: Model accuracy when we ablate the slot attention,
and return ResNet codes in place of the attention readout $\mathbf{y}_{k}$.

##  5 Discussion

Scene and sequence representation. Core to this work is the hypothesis that
the combination of the context transformer and slot attention will 1) enable
modeling of pairwise relations over the scene frames and the transmission of
relevant information among the time-steps in the sequence and 2) provide a
means to extract necessary object information from the scene to form a basis
for scene understanding. The slot transformer is able to learn good quality
representations as evidenced by the model reconstructions and predictions
(figure 3). Figure 3 also suggests that the slot information partitions the
scene in a sensible manner as different objects and content, such as shadows,
are well defined in different slots as evidenced by the masks. The predictions
from the conditional prior in figure 7 (see appendix) demonstrate that the
generative model functions as a good regularizer over scene data capable of
encoding time sensitive scene data and using that information to infer missing
data, at times over a large number of contiguous frames.

Figure 7: Prior Sampling: Full 16 frame sampled rollouts from the test set
drawn from the auto-regressive prior at various points in training on CLEVRER.
As the model learns its ability to predict unseen portions of the sequence
given some information about the sequence improves markedly.

Ablating the Slot Attention. We ablate the attention by reducing to a single
slot where the attention readout now simply returns the encoded frames from
the ResNet. This is compared to a ten slot model where the slot size has been
normalized to ensure equal capacity for the comparison (ie. the "slot size" of
the one slot model is ten times as large) - the models are otherwise
equivalent. This ten slot model significantly outperforms the flat model by a
margin of $57.5\%$ to $40.3\%$ (See the supplementary material for figure).

Figure 6 shows that the slotted model clearly outperforms the flat model in
evaluation on CATER as generalisation performance plateaus early for the non-
attention model.

Effect of more Iterations. In the main text in section $5$ we noted that a ten
slot model was trained on CATER-Static over $1$, $2$, and $4$ iterations with
three seeds each. The results are depicted in figure 5 for $1$, $2$, and $4$
iteration models. When scoring the final model on the test set after training
for the associated models we see scores of $51.2\%\pm 3.4$, $42.3\%\pm 2.9$,
and $29.8\%\pm 2.1$ respectively. This demonstrates a boost in model
generalization when adding iterations to the model over training and
evaluation.

The Effect of the Iterative Model. We ran a ten slot model on CATER-Static for
$500K$ learner updates with batch size $32$ over $1$, $2$, and $4$ iterations
with three seeds each. After training with more iterations there is
substantially less overfitting for the higher iteration models (figure in
supplementary material). The $1$, $2$, and $4$ iteration models scored
$51.2\%$, $42.3\%$, and $29.8\%$ respectively. We see in figure 5 that the
iterative setup, in this case on CATER, helps generalisation performance on
reasoning tasks - similar results were observed on CLEVRER. Along with a boost
in generalization performance additional iterations also appear to help
stabilize performance likely due to the additional gradients and inference
estimates available during training.

Effect of the Generative Losses. We sought to analyze the contribution of the
generative losses to our model. We did this ablation on CLEVRER-Descriptive by
running the model in a deterministic mode where the output of the context
isn’t sampled but is instead passed on in a feed-forward manner. In our
ablation experiment we used a four iteration model averaged over three seeds
for each condition. Both models achieve near perfect QA accuracy on the
training set where on the test set the generative model scores $55.8\%$ while
the deterministic model scoring $43.8\%$. We hypothesize that generative
losses help to provide a means to regularize over the latent space, enabling
the model to learn better representations which are useful for reasoning
tasks.

Table 3: Kinetic-600 Results. No pre-training was carried out on the Slot Transformer in contrast to the baselines. We compare to models that benefit from pre-training: the TimeSformer (Bertasius et al., 2021), the Inflated 3D ConvNet (Carreira and Zisserman, 2017), and from Qian et al. (2020) Contrastive Video Representation Learning (CVRL), SimCLR, and ImageNet models where the top model flavour is reported. Methods | Top-1 | Top-5
---|---|---
TimeSformer (+pre-train) | 82.2 | 95.6
I3D (+pre-train) | 71.7 | 90.4
SimCLR (+pre-train) | 51.6 | -
ImageNet (+pre-train) | 54.7 | -
CVRL (+pre-train) | 72.9 | -
Slot Transformer (Ours) | 68.2 | 87.6

Effect of the Auxiliary Losses. When including both object prediction and
question prediction we achieved scores on the descriptive CLEVRER task of
$87.4\%$ compared with $45.4\%$ without any auxiliary losses, $53.7\%$ when
including only question prediction, and $63.2\%$ when including only object
prediction. The improved performance on the test set is congruent with our
hypothesis that these losses help to learn better latent representations for
generalisation on reasoning tasks.

Limitations. It is evident that our results while strong are below state-of-
the-art on some tasks, a point which we’d like to address. As stated earlier,
in all of our experiments we do not make use of any supplementary annotations,
labels and nor do we do any pre-training. Next, there are ways to better
optimize the model: 1) scaling to more iterations, 2) further tuning the
auxiliary losses, 3) optimizing better between generative, auxiliary and
supervised losses, 4) pre-training the generative model - these are all areas
of active experimentation along with applications to new domains. We are
encouraged by our success accurately modeling input data distributions, in
particular in sequence prediction on sequences not seen in training, as is
demonstrated in figures 3 and 7 in both the main text and the supplementary
material.

Ethical Considerations. While gaining an understanding and building systems to
solve reasoning tasks is exciting work there are ethical implications to
consider. We are categorically opposed to our work being used in the context
of warfare or in fraud however, we understand that the possibility exists that
a system which attempts to mimic reasoning could be used fraudulently. While
this is an inescapable risk when publishing work we hope that understanding
these systems will outweigh the potential negative impact, and that by gaining
deeper knowledge of it we can mitigate its misuse.

Conclusion. We have set out to demonstrate that we could learn to encode video
scenes into compressed representations that store relevant spatial and
temporal scene information and also that the biases induced by slot attention,
transformers, and iterative inference would yield suitable representations for
downstream processing in reasoning tasks with video input. Taking this
approach we have shown that we can achieve strong results on CATER and CLEVRER
without using labelled data or specialized components and competitive results
on Kinetics-600 with no model pre-training. We believe this work demonstrates
promising first steps in this direction and we are eager to scale this effort.

## References

  * Andrychowicz et al. (2016) Marcin Andrychowicz, Misha Denil, Sergio Gomez, Matthew W. Hoffman, David Pfau, Tom Schaul, Brendan Shillingford, and Nando de Freitas.  Learning to learn by gradient descent by gradient descent, 2016.
  * Arnab et al. (2021) Anurag Arnab, Mostafa Dehghani, Georg Heigold, Chen Sun, Mario Lučić, and Cordelia Schmid.  Vivit: A video vision transformer, 2021.
  * Bertasius et al. (2021) Gedas Bertasius, Heng Wang, and Lorenzo Torresani.  Is space-time attention all you need for video understanding?  _CoRR_ , abs/2102.05095, 2021.  URL https://arxiv.org/abs/2102.05095.
  * Burgess et al. (2019) Christopher P. Burgess, Loïc Matthey, Nicholas Watters, Rishabh Kabra, Irina Higgins, Matthew Botvinick, and Alexander Lerchner.  Monet: Unsupervised scene decomposition and representation.  _CoRR_ , abs/1901.11390, 2019.  URL http://arxiv.org/abs/1901.11390.
  * Carreira and Zisserman (2017) João Carreira and Andrew Zisserman.  Quo vadis, action recognition? A new model and the kinetics dataset.  _CoRR_ , abs/1705.07750, 2017.  URL http://arxiv.org/abs/1705.07750.
  * Carreira et al. (2018) João Carreira, Eric Noland, Andras Banki-Horvath, Chloe Hillier, and Andrew Zisserman.  A short note about kinetics-600.  _CoRR_ , abs/1808.01340, 2018.  URL http://arxiv.org/abs/1808.01340.
  * Chen et al. (2021) Zhenfang Chen, Jiayuan Mao, Jiajun Wu, Kwan-Yee Kenneth Wong, Joshua B. Tenenbaum, and Chuang Gan.  Grounding physical concepts of objects and events through dynamic visual reasoning, 2021.
  * Cho et al. (2014) Kyunghyun Cho, Bart van Merrienboer, Çaglar Gülçehre, Fethi Bougares, Holger Schwenk, and Yoshua Bengio.  Learning phrase representations using RNN encoder-decoder for statistical machine translation.  _CoRR_ , abs/1406.1078, 2014.  URL http://arxiv.org/abs/1406.1078.
  * Clark et al. (2020) Peter Clark, Oyvind Tafjord, and Kyle Richardson.  Transformers as soft reasoners over language, 2020.
  * d’Avila Garcez and Lamb (2020) Artur d’Avila Garcez and Luis C. Lamb.  Neurosymbolic ai: The 3rd wave, 2020.
  * Ding et al. (2020) David Ding, Felix Hill, Adam Santoro, Malcolm Reynolds, and Matt Botvinick.  Attention over learned object embeddings enables complex visual reasoning, 2020.  URL https://arxiv.org/abs/2012.08508.
  * Doersch (2021) Carl Doersch.  Tutorial on variational autoencoders, 2021.
  * Feichtenhofer et al. (2021) Christoph Feichtenhofer, Haoqi Fan, Bo Xiong, Ross Girshick, and Kaiming He.  A large-scale study on unsupervised spatiotemporal representation learning, 2021.  URL https://arxiv.org/abs/2104.14558.
  * Gemmeke et al. (2017) Jort F. Gemmeke, Daniel P. W. Ellis, Dylan Freedman, Aren Jansen, Wade Lawrence, R. Channing Moore, Manoj Plakal, and Marvin Ritter.  Audio set: An ontology and human-labeled dataset for audio events.  In _Proc. IEEE ICASSP 2017_ , New Orleans, LA, 2017.
  * Girdhar and Ramanan (2019) Rohit Girdhar and Deva Ramanan.  CATER: A diagnostic dataset for compositional actions and temporal reasoning.  _CoRR_ , abs/1910.04744, 2019.  URL http://arxiv.org/abs/1910.04744.
  * Gopalakrishnan et al. (2022) Anand Gopalakrishnan, Kazuki Irie, Jürgen Schmidhuber, and Sjoerd van Steenkiste.  Unsupervised learning of temporal abstractions with slot-based transformers, 2022.  URL https://arxiv.org/abs/2203.13573.
  * Goyal et al. (2017) Raghav Goyal, Samira Ebrahimi Kahou, Vincent Michalski, Joanna Materzyńska, Susanne Westphal, Heuna Kim, Valentin Haenel, Ingo Fruend, Peter Yianilos, Moritz Mueller-Freitag, Florian Hoppe, Christian Thurau, Ingo Bax, and Roland Memisevic.  The "something something" video database for learning and evaluating visual common sense, 2017.
  * Greff et al. (2020) Klaus Greff, Raphaël Lopez Kaufman, Rishabh Kabra, Nick Watters, Chris Burgess, Daniel Zoran, Loic Matthey, Matthew Botvinick, and Alexander Lerchner.  Multi-object representation learning with iterative variational inference, 2020.
  * Hamrick et al. (2018) Jessica B. Hamrick, Kelsey R. Allen, Victor Bapst, Tina Zhu, Kevin R. McKee, Joshua B. Tenenbaum, and Peter W. Battaglia.  Relational inductive bias for physical construction in humans and machines.  _CoRR_ , abs/1806.01203, 2018.  URL http://arxiv.org/abs/1806.01203.
  * He et al. (2015) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  _CoRR_ , abs/1512.03385, 2015.  URL http://arxiv.org/abs/1512.03385.
  * Hochreiter and Schmidhuber (1997) Sepp Hochreiter and Jürgen Schmidhuber.  Long Short-Term Memory.  _Neural Computation_ , 9(8):1735–1780, 11 1997\.  ISSN 0899-7667.  doi: 10.1162/neco.1997.9.8.1735.  URL https://doi.org/10.1162/neco.1997.9.8.1735.
  * Hudson and Manning (2018) Drew A. Hudson and Christopher D. Manning.  Compositional attention networks for machine reasoning, 2018.
  * Jaegle et al. (2021) Andrew Jaegle, Felix Gimeno, Andrew Brock, Andrew Zisserman, Oriol Vinyals, and João Carreira.  Perceiver: General perception with iterative attention.  _CoRR_ , abs/2103.03206, 2021.  URL https://arxiv.org/abs/2103.03206.
  * Johnson et al. (2016) Justin Johnson, Bharath Hariharan, Laurens van der Maaten, Li Fei-Fei, C. Lawrence Zitnick, and Ross Girshick.  Clevr: A diagnostic dataset for compositional language and elementary visual reasoning, 2016.
  * Kingma and Welling (2019) Diederik P. Kingma and Max Welling.  An introduction to variational autoencoders.  _CoRR_ , abs/1906.02691, 2019.  URL http://arxiv.org/abs/1906.02691.
  * Kipf et al. (2021) Thomas Kipf, Gamaleldin F. Elsayed, Aravindh Mahendran, Austin Stone, Sara Sabour, Georg Heigold, Rico Jonschkowski, Alexey Dosovitskiy, and Klaus Greff.  Conditional object-centric learning from video.  _CoRR_ , abs/2111.12594, 2021.  URL https://arxiv.org/abs/2111.12594.
  * Kondratyuk et al. (2021) Dan Kondratyuk, Liangzhe Yuan, Yandong Li, Li Zhang, Mingxing Tan, Matthew Brown, and Boqing Gong.  Movinets: Mobile video networks for efficient video recognition, 2021\.
  * Kullback and Leibler (1951) S. Kullback and R. A. Leibler.  On Information and Sufficiency.  _The Annals of Mathematical Statistics_ , 22(1):79 – 86, 1951.  doi: 10.1214/aoms/1177729694.  URL https://doi.org/10.1214/aoms/1177729694.
  * Le et al. (2020) Thao Minh Le, Vuong Le, Svetha Venkatesh, and Truyen Tran.  Hierarchical conditional relation networks for video question answering, 2020.
  * Lei et al. (2019) Jie Lei, Licheng Yu, Mohit Bansal, and Tamara L. Berg.  Tvqa: Localized, compositional video question answering, 2019.
  * Locatello et al. (2020) Francesco Locatello, Dirk Weissenborn, Thomas Unterthiner, Aravindh Mahendran, Georg Heigold, Jakob Uszkoreit, Alexey Dosovitskiy, and Thomas Kipf.  Object-centric learning with slot attention.  _CoRR_ , abs/2006.15055, 2020.  URL https://arxiv.org/abs/2006.15055.
  * Marino et al. (2018) Joseph Marino, Yisong Yue, and Stephan Mandt.  Iterative amortized inference.  _CoRR_ , abs/1807.09356, 2018.  URL http://arxiv.org/abs/1807.09356.
  * Parisotto et al. (2019) Emilio Parisotto, H. Francis Song, Jack W. Rae, Razvan Pascanu, Caglar Gulcehre, Siddhant M. Jayakumar, Max Jaderberg, Raphael Lopez Kaufman, Aidan Clark, Seb Noury, Matthew M. Botvinick, Nicolas Heess, and Raia Hadsell.  Stabilizing transformers for reinforcement learning, 2019.
  * Qian et al. (2020) Rui Qian, Tianjian Meng, Boqing Gong, Ming-Hsuan Yang, Huisheng Wang, Serge J. Belongie, and Yin Cui.  Spatiotemporal contrastive video representation learning.  _CoRR_ , abs/2008.03800, 2020.  URL https://arxiv.org/abs/2008.03800.
  * Russakovsky et al. (2015) Olga Russakovsky, Jia Deng, Hao Su, Jonathan Krause, Sanjeev Satheesh, Sean Ma, Zhiheng Huang, Andrej Karpathy, Aditya Khosla, Michael Bernstein, Alexander C. Berg, and Li Fei-Fei.  ImageNet Large Scale Visual Recognition Challenge.  _International Journal of Computer Vision (IJCV)_ , 115(3):211–252, 2015.  doi: 10.1007/s11263-015-0816-y.
  * Russin et al. (2021) Jacob Russin, Roland Fernandez, Hamid Palangi, Eric Rosen, Nebojsa Jojic, Paul Smolensky, and Jianfeng Gao.  Compositional processing emerges in neural networks solving math problems, 2021.
  * Santoro et al. (2017) Adam Santoro, David Raposo, David G. T. Barrett, Mateusz Malinowski, Razvan Pascanu, Peter Battaglia, and Timothy Lillicrap.  A simple neural network module for relational reasoning, 2017.
  * Santoro et al. (2018) Adam Santoro, Ryan Faulkner, David Raposo, Jack Rae, Mike Chrzanowski, Theophane Weber, Daan Wierstra, Oriol Vinyals, Razvan Pascanu, and Timothy Lillicrap.  Relational recurrent neural networks, 2018.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  _CoRR_ , abs/1706.03762, 2017.  URL http://arxiv.org/abs/1706.03762.
  * Wang et al. (2021) Jue Wang, Gedas Bertasius, Du Tran, and Lorenzo Torresani.  Long-short temporal contrastive learning of video transformers.  _CoRR_ , abs/2106.09212, 2021.  URL https://arxiv.org/abs/2106.09212.
  * Watters et al. (2019) Nicholas Watters, Loic Matthey, Christopher P. Burgess, and Alexander Lerchner.  Spatial broadcast decoder: A simple architecture for learning disentangled representations in vaes, 2019.
  * Yi et al. (2020) Kexin Yi, Chuang Gan, Yunzhu Li, Pushmeet Kohli, Jiajun Wu, Antonio Torralba, and Joshua B. Tenenbaum.  Clevrer: Collision events for video representation and reasoning, 2020\.
  * You et al. (2019) Yang You, Jing Li, Jonathan Hseu, Xiaodan Song, James Demmel, and Cho-Jui Hsieh.  Reducing BERT pre-training time from 3 days to 76 minutes.  _CoRR_ , abs/1904.00962, 2019.  URL http://arxiv.org/abs/1904.00962.
  * You et al. (2020) Yang You, Jing Li, Sashank Reddi, Jonathan Hseu, Sanjiv Kumar, Srinadh Bhojanapalli, Xiaodan Song, James Demmel, Kurt Keutzer, and Cho-Jui Hsieh.  Large batch optimization for deep learning: Training bert in 76 minutes, 2020.

##  Appendix A Model & Experimental Details

This section provides additional details and code for the model shown in
figures 1 and 2 of the main text.

###  A.1 Features Encoder

Given the input sequence, $\mathbf{x}\in\mathbb{R}^{T\times H\times W\times
3}$, we encode the input features using a $4$-Layer Residual Network (He et
al. [2015]): $\mathbf{e}=f_{\text{ResNet}}(\mathbf{x})$. The hyper-parameters
per layer of this encoder are: output channels $=(64,128,256,512)$, residual
blocks $=(2,2,2,2)$, and strides $=(2,2,2,1)$. We only do this once per input
and we append a spatial Fourier basis (Vaswani et al. [2017]) to the super-
pixels in $\mathbf{e}$ via $f_{\text{spatial-encoding}}$ (see algorithm 1).

###  A.2 Slot Attention

First we batch apply the query MLP:
$\mathbf{q_{k}}=f_{\text{MLP}}(\mathbf{c}_{k})$. We use a $2$-layer MLP of
sizes $(128,128)$. We apply queries independently across the sequence to their
corresponding embeddings according to $t$.

For our attention readouts we take the approach of Locatello et al. [2020].
Namely, we perform the softmax operation over the attention weights then take
the weighted mean of the weights before multiplying by the values:

$\displaystyle Q$ $\displaystyle=W_{Q}\mathbf{q}_{k}+\mathbf{b}_{Q}$
$\displaystyle K$ $\displaystyle=W_{K}\mathbf{e}+\mathbf{b}_{K}$
$\displaystyle V$ $\displaystyle=W_{V}\mathbf{e}+\mathbf{b}_{V}$
$\displaystyle W^{\text{attn}}$
$\displaystyle=\text{softmax}(\frac{KQ^{T}}{\sqrt{s}},\text{axis=slots})$
$\displaystyle W^{\text{attn}}$
$\displaystyle=\frac{W^{\text{attn}}_{i,k}}{\sum_{l=1}^{I}W^{\text{attn}}_{l,k}}$
$\displaystyle\mathbf{a}_{k}$
$\displaystyle=W_{\text{attn}}^{T}V\in\mathbb{R}^{T\times\text{slots}\times\text{s}}$

where $s$ is the slot size. This is denoted as $f_{\text{SlotAttention}}$ in
algorithm 1 with input $\mathbf{q}_{k}$ and $\mathbf{e}$.

###  A.3 Gating

After computing the attention readout we use gating in a fashion similar to a
Gated Recurrent Unit (Cho et al. [2014]). Given the context $\mathbf{c}_{k}$,
attention readout $\mathbf{a}_{k}$, and weights
$W_{z,c},W_{z,c},W_{z,a},W_{r,c},W_{r,a},W_{h,c},W_{h,a},\mathbf{b}_{z},\mathbf{b}_{r},\mathbf{b}_{h}$,
we compute the updated context, $\mathbf{c}^{{}^{\prime}}_{k}$ and do a final
LayerNorm:

$\displaystyle\mathbf{z}$
$\displaystyle=\sigma(W_{z,c}\mathbf{c}_{k}+W_{z,a}\mathbf{a}_{k}+\mathbf{b}_{z})$
$\displaystyle\mathbf{r}$
$\displaystyle=\sigma(W_{r,c}\mathbf{c}_{k}+W_{r,a}\mathbf{a}_{k}+\mathbf{b}_{r})$
$\displaystyle\mathbf{h}$
$\displaystyle=\tanh{(W_{h,c}\mathbf{c}_{k}+W_{h,a}\mathbf{a}_{k}+\mathbf{b}_{h})}$
$\displaystyle\mathbf{c}^{{}^{\prime}}_{k}$
$\displaystyle=(\mathbf{1}-\mathbf{z})\mathbf{c}_{k}+\mathbf{z}\mathbf{h}$
$\displaystyle\mathbf{c}^{{}^{\prime}}_{k}$
$\displaystyle=\text{LayerNorm}(\mathbf{c}^{{}^{\prime}}_{k})\in\mathbb{R}^{T\times\text{slots}\times\text{s}}$

Where $\sigma$ is the sigmoid function. This is denoted as $f_{\text{gate}}$
in algorithm 1 with input $\mathbf{q}_{k}$ and $\mathbf{e}$.

###  A.4 Context Transformer

When apply the context transformer to $\mathbf{c}^{{}^{\prime}}_{k}$ we use a
gated transformer (Vaswani et al. [2017], Parisotto et al. [2019]) with
absolute Fourier positional encodings concatenated to the input. The
transformer is applied to each set of slots separately over the time dimension
and the resulting sets are concatenated back together to form
$\mathbf{c}_{k}^{{}^{\prime\prime}}=T_{\text{context}}(\mathbf{\mathbf{c^{\prime}_{k}}})$
(see algorithm 1). The transformer hyper-parameters per task are shown in
appendix table 5. Note that during training we use a dropout rate of $0.1$ on
the attention weights Vaswani et al. [2017].

###  A.5 Latent Sampling, Spatial Broadcast Decoding

We first sample a latent from the posterior parameters derived from slotted
context output from the context transformer,
$\mathbf{c}_{k}^{{}^{\prime\prime}}$:

$\displaystyle\lambda_{k}$ $\displaystyle=W_{\lambda}\mathbf{c}_{k}$
$\displaystyle\mathbf{\mu}_{k},\log\mathbf{\sigma}_{k}$
$\displaystyle=\lambda_{k}$ $\displaystyle\mathbf{z}_{k}$ $\displaystyle\sim
N(\mathbf{\mu}_{k},\mathbf{\sigma}_{k})\in\mathbb{R}^{T\times\text{slots}\times\text{s}}$

To decode our latent we broadcast our latent to
$\mathbb{R}^{T\times\text{slots}\times H\times W\times\text{s}}$, concatenate
spatial Fourier encodings to all pixels, then batch apply a spatial decoder
over time and slots (Watters et al. [2019]). For the spatial broadcast decoder
we use $4$ conv layers with the following parameters: output channels
$=(64,64,64,64)$, kernel shapes $=(3,3,3,3)$, and strides $=(1,1,1,1)$.

###  A.6 Iterative Inference Update

In the paper, in equation 1, we note that in addition to the auxiliary ELBO
losses, $a_{k}$, we also compute an iterative inference update, an
approximation of the loss on the prior with respect estimated latent posterior
parameters, $\lambda_{k}$. This loss estimate is dependent on the latent
sample, $\mathbf{z}_{k}$, and the data $\mathbf{x}$. We have used the approach
from Marino et al. [2018] to compute an approximate update to our estimated
latent posterior parameters and to derive the context for the next iteration,
$\mathbf{c}_{k+1}$:

$\displaystyle\mathbf{\epsilon}_{x,k}$
$\displaystyle=\frac{\mathbf{x}-\mathbf{\mu}_{x}}{\mathbf{\sigma}_{x^{2}}}$
$\displaystyle\mathbf{\epsilon}_{z,k}$
$\displaystyle=\frac{\mathbf{z}-\mathbf{\mu}_{z}}{\mathbf{\sigma}_{z^{2}}}$
$\displaystyle\lambda_{k+1}$
$\displaystyle=\lambda_{k}+f_{\text{iterative}}(\mathbf{z}_{k},\mathbf{x}_{k},\mathbf{\epsilon}_{x,k},\mathbf{\epsilon}_{z,k})$
$\displaystyle\mathbf{c}_{k+1}$
$\displaystyle=\text{LayerNorm}(\text{MLP}(\lambda_{k+1}))\in\mathbb{R}^{T\times\text{slots}\times\text{s}}$

Where is $f_{\text{iterative}}$ is a Linear network of size $\text{num-
slots}\times\text{slot-size}(s)$. Finally, note that all parameters are shared
across iterations.

###  A.7 Downstream Task Heads

In figure 2 from the main text we show the general setup for the task head,
more details can be found in algorithm 2. To form the input for the QA
Transformer, $T_{\text{QA}}$, we project the slot values from the context, a
classification token, and the first and last frame embeddings from the ResNet
down to a common embedding size. We also conditionally generate word
embeddings for the input question as these only exist for CLEVRER tasks, these
are otherwise omitted. Once the $T_{\text{QA}}$ output has been computed, we
extract its first frame corresponding to the token output
$\in\mathbb{R}^{1\times\text{embed-size}}$, and feed that as input to the
final MLP, $\text{MLP}_{\text{QA}}$, to produce logits against which we can
compute a cross-entropy loss and accuracy. The hyper-parameters of the QA
Transformer and the final MLP to generate the label logits is shown in table
5. Note that for the transformer during training we use a dropout rate of
$0.1$ on the attention weights Vaswani et al. [2017].

###  A.8 Losses & Optimizer

For CATER and CLEVRER we used the Lamb optimizer (You et al. [2020]) with
learning rates and weight decay values reported in table 5. All runs use an
exponential decay rate for the first moment of past gradient of $0.9$, and for
the second moment $0.999$ with an eps constant of $1\text{e}^{-6}$. For
Kinetics-600 we used stochastic gradient descent with $8000$ burn-in steps to
linearly ascend to the maximum learning rate and with a cosine decay schedule
thereafter.

###  A.9 Hardware Resources

For our experiments we ran all learners on $4\times 4$ DragonDonut TPUs with
one full donut per learner. We ran evaluators concurrently on CPU, sampling a
batch from the evaluation set every 30 seconds with evaluation batch sizes of
$32$.

###  A.10 Datset, Training & Evaluation Details

Tables 4, 5, and 6 describe the various hyper-parameter and dataset statistics
relevant to our experiments.

Table 4: Dataset & task details.

Dataset | Sequence Length | Frame Size | # Training Set Size | # Evaluation Set Size | Label Size
---|---|---|---|---|---
CLEVRER-Descriptive | $64$ | $64\times 64$ | $10,000$ | $5,000$ | $42$
CLEVRER-Explanatory | $64$ | $64\times 64$ | $10,000$ | $5,000$ | $1$
CLEVRER-Predictive | $64$ | $64\times 64$ | $10,000$ | $5,000$ | $1$
CLEVRER-Counterfactual | $64$ | $64\times 64$ | $10,000$ | $5,000$ | $1$
CATER Localization Task | $80$ | $64\times 64$ | $3,850$ | $1,650$ | $36$
Kinetics-600 | $300$ | $256\times 256$ | $392,622$ | $68,617$ | $600$

Table 5: Hyper Parameters for the encode-decode-iterate model.

|  |  |  |  |  | Context Transformer |  |  |
---|---|---|---|---|---|---|---|---|---
Dataset | Learning Rate | Batch Size | Total Learner Steps | Weight Decay | Min. Mask Rate ($\rho$) | Layers | Heads | Embedding Size | MLP Layer Sizes | Slots | Slot Size | Iterations
CLEVRER | $1e^{-4}$ | $8$ | $500K$ | $1e^{-4}$ | $1e^{-1}$ | $8$ | $4$ | $128$ | $(64,128,64)$ | $8$ | $64$ | $4$
CATER | $1e^{-4}$ | $8$ | $500K$ | $1e^{-2}$ | $5e^{-1}$ | $8$ | $4$ | $256$ | $(64,256,64)$ | $10$ | $64$ | $4$
Kinetics-600 | $1e^{-2}$ | $1M$ | $16M$ | $1e^{-2}$ | $5e^{-1}$ | $16$ | $4$ | $256$ | $(128,256,128)$ | $8$ | $128$ | $4$

Table 6: Hyper Parameters for the task head.

|  | QA Transformer |
---|---|---|---
Dataset | Input Embedding Size | Layers | Heads | Attention Embedding Size | MLP Layer Sizes | Final MLP Layer Sizes
CLEVRER | $64$ | $8$ | $4$ | $128$ | $(128,128)$ | $(2048,2048,\text{label-size})$
CATER | $64$ | $8$ | $4$ | $256$ | $(256,256)$ | $(2048,2048,\text{label-size})$
Kinetics-600 | $128$ | $16$ | $4$ | $256$ | $(256,256)$ | $(2048,2048,2048,2048\text{label-size})$

##  Appendix B Code

We now include pseudo-code that matches our implementation using the
components described above. See algorithms 1 and 2.

Algorithm 1 Pseudo code for the Encode-Decode-Iterate phases.

1: Input: $\mathbf{x}\in\mathbb{R}^{T\times\text{H}\times\text{W}\times 3}$

2: $\mathbf{c}_{0}\sim N(0,1)\in\mathbb{R}^{T\times K\times
C_{\text{context}}}$

3: $\mathbf{e}\leftarrow f_{\text{ResNet}}(\mathbf{x})$

4: $\mathbf{e}\leftarrow\text{concat}([\mathbf{e},f_{\text{spatial-
encoding}}()],\text{axis=pixel})$

5: for k = 0 : K-1 do

6: $\mathbf{q_{k}}\leftarrow f_{\text{MLP}}(\mathbf{c}_{k})$ {Make Queries,
batched over time}

7: $\mathbf{a_{k}}\leftarrow
f_{\text{SlotAttention}}(\mathbf{q}_{k},\mathbf{e})$
{$f_{\text{SlotAttention}}$ batched over time}

8: $\mathbf{c_{k}^{\prime}}\leftarrow
f_{\text{gate}}(\mathbf{c}_{k},\mathbf{a}_{k})$

9: $\mathbf{c}_{k}^{{}^{\prime\prime}}\leftarrow
T_{\text{context}}(\mathbf{\mathbf{c^{\prime}_{k}}},\text{position\\_encode=True})$
{$T_{\text{context}}$ applied to each set of slots separately.}

10: $\lambda_{k}\leftarrow W_{\lambda}\mathbf{c}_{k}^{{}^{\prime\prime}}$
{Fetch posterior params and sample latent.}

11: $\mu_{k},\log\sigma_{k}\leftarrow\lambda_{k}$

12: $\mathbf{z}_{k}\sim N(\mu_{k},\sigma_{k})$

13: $\mathbf{m}_{k,t},\mathbf{r}_{k,t}\leftarrow
f_{\text{decoder}}(\mathbf{z}_{k,t})$ {Compute masks, assemble
reconstructions.}

14:
$\mathbf{x}_{t}^{\prime}\leftarrow\sum_{k=1}^{K}\mathbf{m}_{k,t}\mathbf{r}_{k,t}$

15:
$\lambda_{k+1}\leftarrow\lambda_{k}+f_{\text{iterative}}(\mathbf{z}_{k},\mathbf{x}_{k},\mathbf{\epsilon}_{x,k},\mathbf{\epsilon}_{z,k})$
{Iterative Inference Update}

16: $\mathbf{c}_{k+1}\leftarrow\text{LayerNorm}(\text{MLP}(\lambda_{k+1}))$
{New context}

17: end for

18:
$\mathcal{L}_{\text{ELBO}}=\sum_{t=1}^{T}\sum_{k=1}^{K}{\mathcal{D}_{KL}(\mathcal{N}(\mu_{k,t},\sigma_{k,t}^{2})||\mathcal{N}(0,1)}-\log{L(\mathbf{x}_{t}|\mathbf{x}_{t}^{\prime})}$

Algorithm 2 Downstream task heads.

1: Input 1: $\mathbf{c}_{K}\in\mathbb{R}^{T\times\text{slots}\times\text{s}}$

2: Input 2: $\mathbf{e}_{\text{first}}\in\mathbb{R}^{2\times
H_{\text{ResNet}}\times W_{\text{ResNet}}\times 3}$

3: Input 3: $\mathbf{e}_{\text{last}}\in\mathbb{R}^{2\times
H_{\text{ResNet}}\times W_{\text{ResNet}}\times 3}$

4: Input 4: $\mathbf{question}\in\mathbb{R}^{\text{num-words}\times 1}$

5: Input 5: $\mathbf{answer}\in\mathbb{R}^{\text{label-size}}$

6:
$\text{token}_{\text{in}}\leftarrow\mathbf{1}\in\mathbb{R}^{1\times\text{embed-
size}}$

7: $\text{token}\leftarrow\text{Bias}(\text{token}_{in})$ {Use a learnable
bias to shape the input token.}

8: if $\mathbf{question}$ then

9: $\mathbf{question}^{\text{embed}}\leftarrow
f_{\text{embed}}(\mathbf{question})\in\mathbb{R}^{\text{num-
words}\times\text{embed-size}}$

10: else

11: $\mathbf{question}^{\text{embed}}\leftarrow\emptyset$

12: end if

13:
$\mathbf{c_{\mathbf{K}}^{\text{embed}}}\leftarrow\text{Reshape}(\text{Linear}_{\text{context}}(\mathbf{c}_{K}))\in\mathbb{R}^{(T\times\text{num-
slots})\times\text{embed-size}}$

14:
$\mathbf{e}_{\text{first}}^{\text{embed}}\leftarrow\text{Linear}_{\text{frames}}(\text{Flatten}(\mathbf{e}_{\text{first}}))\in\mathbb{R}^{1\times\text{embed-
size}}$

15:
$\mathbf{e}_{\text{first}}^{\text{embed}}\leftarrow\text{Linear}_{\text{frames}}(\text{Flatten}(\mathbf{e}_{\text{last}}))\in\mathbb{R}^{1\times\text{embed-
size}}$

16:
$\mathbf{x}_{\text{QA}}\leftarrow\text{Concatenate}(\text{token},\mathbf{question}^{\text{embed}},\mathbf{c_{\mathbf{K}}^{\text{embed}}},\mathbf{e}_{\text{first}}^{\text{embed}},\mathbf{e}_{\text{last}}^{\text{embed}})$

17:
$\mathbf{y}_{\text{QA}}\leftarrow\text{T}_{\text{QA}}(\mathbf{x}_{\text{QA}},\text{position\\_encode=True})$
{Run the QA Transformer.}

18:
$\text{logits}\leftarrow\text{MLP}_{\text{QA}}(\mathbf{y}_{\text{QA}}[0,...])\in\mathbb{R}^{\text{label-
size}}$ {Final MLP to produce logits.}

19: $\mathcal{L}_{\text{QA}}=\text{CrossEntropy}(\text{logits},\text{labels})$

20: $\mathcal{L}_{\text{QA}}$, logits

##  Appendix C Additional Results & Analysis

We include performance curves below for train and evaluation on CLEVRER and
Kinetics-600 in figures 11, 12, and 13. We also include a figure illustrating
the full set of masks over sequence snippets for CLEVRER and both CATER
datasets in figure 10.

###  C.1 Ablating the Context Transformer

We carried out an additional analysis where we ablated the context transformer
as we posited that this component would be critical in modeling temporal
dynamics by allowing direct relations among any two frames in the sequence to
be formed. In order to measure its effect we replaced the Transformer with a
Deep LSTM (Hochreiter and Schmidhuber [1997]). The Deep LSTM consists of an
$4$-Layer model with hidden a state size of $256$, twice the size of the
attention in the case of the transformer model. The LSTM is rolled forward
separately for each slot in a similar fashion to as is done for the context
transformer and the slots are recombined to form the updated context. Both
were $2$ iteration models.

We ran the analysis on the CATER static camera dataset. In figures 8 and 9 we
see the models produce qualitatively different masks and reconstructions.
There is notably much more blur and missing elements such as shadows in the
LSTM model’s reconstructions and the slot masks for the LSTM model also tend
to be distributed among many objects in the scene. This may indicate the
representations learned in the LSTM model have been induced to model spurious
correlations which may be expected in the absence of direct pathways between
temporal events afforded by the transformer architecture. Finally we note the
performance of each model, we report scores for these models when the LSTM
model performance began to top out. For the LSTM model: top-5 accuracy
$60.7\%$, top-1 accuracy $36.9\%$, L1 distance $2.29$, and for the
Transformer: top-5 accuracy $71.9\%$, top-1 accuracy $39.6\%$, L1 distance
$1.77$. We believe that the advantage of the Transformer architecture would
only grow greater as we larger sequences and more complex scenes are used.

Figure 8: Masks comparison when ablating the context Transformer for an LSTM.
Figure 9: Reconstructions comparison when ablating the context Transformer for
an LSTM. Figure 10: All masks for $8$ slot models on CLEVRER & CATER tasks
with static and moving cameras. Figure 11: Clevrer Train & Evaluation Results
on an $8$ slot, $4$ iteration model on Descriptive and Explanatory tasks.
Figure 12: Clevrer Train & Evaluation Results on an $8$ slot, $4$ iteration
model on Predictive and Counterfactual tasks. Figure 13: Kinetics Top-5 &
Top-1 Accuracy and Train vs. Evaluation Losses.

##  Appendix D Baseline Details

In Section 4 of the main paper a number of baselines are compared against and
here we give a bit more detail on them. For CLEVRER more detail can be found
in Yi et al. [2020] and Chen et al. [2021], and in Girdhar and Ramanan [2019]
for CATER.

First a note on CLEVRER, some of the baselines take advantage of additional
ground truth data from video annotations and labels from the program executor.
These can include ground truth information about the image content such as
shapes, colours, positions, or motion data. As we are primarily concerned with
constructing a high quality sequence encoder that can infer these properties
we did not include them in our training. This notably puts our model at a
disadvantage when comparing to neuro-symbolic approaches however, we believed
it an important point to optimize for the cases where we may not be able to
expect such ground truth labeling exists. As a future proof of concept we are
considering experimenting using of this data.

Hierarchical Conditional Relational Networks (HCRN) Le et al. [2020] introduce
this model that imposes a hierarchy of CRN networks on clips that are built up
into the full video. CRNs take some number of input objects and a condition
(query) and generate an output where relations among the input objects may be
learned. In HCRNs the output of lower clip level CRNs is fed as input to the
higher video level CRNs to produce a final output that is decoded to an
answer.

Memory Attend Control Network (MAC V & MAC V+) Hudson and Manning [2018]
introduce the MAC network architecture adapted function on video input. The
network processes input via its input unit to produce encodings from a
question and video. Next, the MAC unit processes input by iteratively
attending to the output from the input unit. The control unit recurrently
attends to parts of the input, determining which parts to reason about, while
the memory stores stores the result of step $i$. Finally an output unit
combines the output from MAC cell iteration to produce a final output.

The slot transformer exceeds the MAC network on Predictive, Explanatory, and
Counterfactual tasks in CLEVRER both with and without labelled data and HCRN
on all CLEVRER tasks as well as CNN+MLP and CNN+LSTM stacks.

Localized, Compositional Video Question Answering (TVQA+) Lei et al. [2019]
introduce the TVQA dataset and a multi-stream neural network that can ingest
multimodal input for video QA. Each stream effectively is composed of an
embedding module (a CNN in the case of the video frames) followed by an LSTM.
Embedding streams include video, question, and answer data and are then fused
via elementwise products and concatenation before being softmaxed to produce
answer scores. This model also makes use of annotation data in CLEVRER.

I3D & R3D LSTM Inflated 3D ResNet models. Carreira and Zisserman [2017],
Girdhar and Ramanan [2019] use pooling and filtering layers increased in size,
or inflated, so as to retain additional capacity to model 3D spatio-temporal
input. They also pre-train on ImageNet. The model utilizes two streams: one
feed-forward on RGB frames, and another for optical flow data from the
sequence. For CATER, a ResNet is used instead of a CNN and obtains strong
results making use of optical flow and ImageNet pre-training (Russakovsky et
al. [2015]). Our model is competitive with theirs and in some cases exceeding
their scores on the training videos alone.

NS-DR & DCL We discuss neuro-symbolic logic models in the main paper in
section 2. In particular two strong baselines Neuro-Symbolic Dynamic Reasoning
(Yi et al. [2020]) and the Dynamic Context Learner (DCL/ DCL (Oracle)) (Chen
et al. [2021]). As discussed above we haven’t used ground truth data when
training our models so our scores are difficult to compare with the top
baseline however, despite this, the performance of the slot transformer
remains competitive.