Latent Variable Sequential Set Transformers for Joint Multi-Agent Motion Prediction
Authors: Roger Girgis
Florian Golemo
Felipe Codevilla
Martin Weiss
Jim Aldon D'Souza
Samira Ebrahimi Kahou
Felix Heide
Christopher Pal
Venue: ICLR 2022 Spotlight (2022)
Keywords: trajectory prediction
motion forecasting
transformers
latent variable models

Abstract: Robust multi-agent trajectory prediction is essential for the safe control of robotic systems. A major challenge is to efficiently learn a representation that approximates the true joint distribution of contextual, social, and temporal information to enable planning. We propose Latent Variable Sequential Set Transformers which are encoder-decoder architectures that generate scene-consistent multi-agent trajectories. We refer to these architectures as “AutoBots”. The encoder is a stack of interleaved temporal and social multi-head self-attention (MHSA) modules which alternately perform equivariant processing across the temporal and social dimensions. The decoder employs learnable seed parameters in combination with temporal and social MHSA modules allowing it to perform inference over the
entire future scene in a single forward pass efficiently. AutoBots can produce either the trajectory of one ego-agent or a distribution over the future trajectories for all agents in the scene. For the single-agent prediction case, our model achieves top results on the global nuScenes vehicle motion prediction leaderboard, and produces strong results on the Argoverse vehicle prediction challenge. In the multi-agent setting, we evaluate on the synthetic partition of TrajNet++ dataset to showcase the model’s socially-consistent predictions. We also demonstrate our model on general sequences of sets and provide illustrative experiments modelling the sequential structure of the multiple strokes that make up symbols in the Omniglot data. A distinguishing feature of AutoBots is that all models are trainable on a
single desktop GPU (1080 Ti) in under 48h.

Published as a conference paper at ICLR 2022
LATENT VARIABLE SEQUENTIAL SET TRANSFORMERS FOR JOINT MULTI-AGENT MOTION PREDICTION
Roger Girgis' 1. Florian Golemo 2.3 Felipe Codevilla Martin Weiss Jim Aldon D'Souza Samira F.. Kahou Felix Heide Christopher Pal 12,38,
PPolytechniqee Montréal, Mila, Quebcc AI Institute BlemeetAA/ / Service Now, 'Independent Robotics, 3 'Algolux, "MeGill University, "École de technologie supérieure, CCaadaa CTFAR AT Chair, PPrinceton University. Correspondence: roger.girgis gmail.com attps://Egole: github. wwt.:i///uttta /a//uuoooos/////oo
ABSTRACT
Robust multi-agent trajectory prediction is essential for the safe control of robotic systems. A major challenge is to efficiently learn a representation that approxi- mates the true joint distribution of contextual, social, and temporal information enable planning. We propose Latent Variable Sequential Set Transformers which are encoder-decoder architectures that generate scene-consistent multi-agent tra- jectories. We refer to these architectures as "AutoBots". The encoder is stack of interleaved temporal and social multi-head self-attention (MHSA) modules which alternately perform equivariant processing across the temporal and social dimensions. The decoder employs learnable seed parameters in combination with temporal and social MHSA modules allowing it to perform inference over the entire future scene in a single forward pass efficiently. AutoBots can produce either the trajectory of one ego-agent or distribution over the future trajectories for all agents in the scene. For the single-agent prediction case, our model achieves top results the global nuScenes vehicle motion prediction leaderboard, and produces strong results on the Argoverse vehicle prediction challenge. the multi-agent setting, evaluate on the synthetic partition of TrajNett++ dataset showcase the model's socially-comsisten predictions. We also demonstrate our model on general sequences of sets and provide illustrative experiments modelling the quential structure of the multiple strokes that make up symbols in the Omniglot data. A distinguishing feature of AutoBots is that all models are trainable on a single desktop GPU (1080 Ti) in under 48h.
INTRODUCTION
Many problems require processing complicated hierarchical compositions of sequences and sets. For example, multiple choice questions are commonly presented as query (e.g. natural language) and an unordered set of options (Kembhavi et al., 2017; Richardson et al., 2013). Machine learning models that perform well in this setting should be insensitive to the order in which the options are presented. Similarly, for motion prediction tasks (Biktairov et al., 2020; Sadeghian et al., 2018) where the inputs are agent trajectories evolving over time and outputs are future agent trajectories, models should be insensitive to agent ordering.
In this work, we focus on the generative modelling of sequences and sets and demonstrate this method on variety of different motion forecasting tasks. Suppose X = {(z)..... æx) _K.I....... ....., (II,..., Ix) ) 3 each with K E N elements. Allowing also that X evolves across some time horizon we denote this sequence of sets X (X1....,X8 Xr).
have set of M N sequences
This problem setting often requires the model to capture high-order interactions and diverse futures. Generally, solutions are built with auto-regressive sequence models which include steps of context. In our approach, to better model the multi-modal distribution over possible futures, we also allow to include discrete latent variable Z, to create sequential models of sets of the form
P( [e....,XX = Σ P(Z++1, ,I,,,..,,%
Z

---

Published as a conference paper at ICLR 2022
Figure Overview: Our model generates multi- ple possible futures for agents based on several input timesteps. The input trajectories are encoded into context tensor that captures their respective behav- ior and interaction with surrounding agents. While Transformer-basee architectures are often autoregres- sive, we decode all future steps at once. This is ac- complished via learnable seed parameters. Our de- coder transforms these, together with the map and the context into the possible future trajectories that we call "modes".
Contest breoo dioccion.
urknown
gren
Time
Icamed
Ausobut
Modee
pposutia
Sood Polanccerg
Time
Map
In this work, we propose this type of parameterized conditional mixture model over sequences of sets, and call this model a Latent Variable Sequential Set Transformer (affectionately referred as "AutoBot"). We evaluate AutoBot on nuScenes (Caesar al., 2020) and Argoverse (Chang et al., 2019), two autonomous driving trajectory prediction benchmarks, on the synthetic partition of the TrujNes+++ (Kothari et al., 2021) dataset, pedestrian trajectory forecasting benchmark, and Omniglor (Lake et al., 2015), a dataset of hand-drawn characters that use for predicting strokes. All of these provide empirical validation of our model's ability to perform robust multi-moda prediction. Specifically, we make the following contributions:
We propose novel approach for modelling sequences of set-structured continuous variables, and extend this latent-variable formulation to capture multi-modal distributions.
We provide theoretical analysis highlighting respect to different components of the input data structures.
model's permutation sensitivity with
We validate our method with strong empirical results on diverse tasks, including the large- scale muScenes and Argoverse datasets for autonomous driving (Caesar et al.. 2020; Chang al., 2019), the multi- multi-lingual Omniglot dataset of handwritten characters (Lake al.. 2015), and the synthetic partition of the + pedestrian trajectory prediction task (Sadeghian al., 2018).
BACKGROUND
We now review several components of the Transformer (Vaswani et al., 2017) and Set Transformer (Lee et al., 2019) architectures, mostly following the notation found in their manuscripts. For additional background, see their works and Appendix A.
Multi-Head Self Attention (MHSA) can be thought of as an information retrieval system, where query is executed against a key-value database, returning values where the key matches the query best. While (Vaswani al., 2017) defines MHSA three tensors, for convenience,
input single set-valued argument. Internally, MHSA then performs intra-set attention by casting the input set X to query, key and value matrices and adding residual connection, MHSA(X) - MMHSA X, X).
Multi-Head Attention Blocks (MAB) resemble the encoder proposed Vaswani et al. (2017), but lack the position encoding and dropout (Lee et al., 2019). Specifically, they consist of the MHSA operation described in Eq. (appendix) followed by a row-wise feed-forward neural network (rFFN), with residual connections and layer normalization (LN) (Ba al., 2016) after each block. Given an input set X containing n, elements of dimension d, and some conditioning variables C containing ne elements of dimension the MAB can be described by the forward computation.
MAB(X) LN(H + IFFN(H))
where H LN + MHSA XHSS((XX,X))
(I)
Multi-head Attention Block Decoders (MABD) were also introduced in Vaswani et al. (2017), and were used to produce decoded sequences. Given input matrices and representing sequences, the decoder block performs the following computations:
MABD(Y, X) LN(H + (FF((H))
= LN(H' + HH' X, where
and H' = LN(MHSA(Y))
Our model, described throughout the following section, makes extensive use of these functions.
2

---

Published as a conference paper at ICLR 2022
repeat
Inpul sequerce
Atsibiutee each
IFNN
MAB Time Encoding
MAB Social Encoding
Numbbr of arnenns net aggents
Hilden sloc
Nurbbrr ano sioo. layers Engoded (CQ) map intrrcctionn
por aenomtorr
6
Decnder contett
IFFN
MABD
MAB Finooling
Mall artertion
penams
Encoding
14580
Nul hoad anterton block (ecoclr
tonal anooding
Decoder
repeat L times
Figure 2: Architecture Overview. Our model takes input tensor of dimension K, M, A row-wisc food-forward network (rFFN) applied each row along the M plane transforming vectors dimension to dx. After adding positional encoding (PE) to the axis, the encoder passes the tensor through L repeated layers of multi-head attention blocks (MAB) that are applied to the time axis (time encoding) and the agent axis (social encoding) before outputting the context tensor. In the decoder, the encoded map and the learnable seed parameters tensor concatenated and passed through rFFN before being passed through L repeated layers of multi head attention block decoder (MABD) along followed by a MAB along the agent axis. This figure shows the process for predicting one possible future trajectory (" ("mode") "). When multiple modes are predicted, each mode has its own learnable seed parameters and the decoder is rolled out once for each mode.
axis (using context from the encoder)
LATENT VARIABLE SEQUENTIAL SET TRANSFORMERS
Latent Variable Sequential Set Transformers is a class of encoder-decoder architectures that process sequences of sets (see Fig 2). In this section, we describe the encoding and decoding procedures, the latent variable mechanism, and the training objective. Additionally, we provide proof of the permutation equivariance of our model in Appendix B, and implementation details in Appendix C as well as a discussion of the special case of AutoBot-Ego, which is similar to AutoBots but predicts future sequence for only one element in scene.
3.1 ENCODER: INPUT SEQUENCE REPRESENTATION
AutoBot takes as an input sequence of sets, X (X....., which in the motion prediction setting may be viewed as the state of a scene evolving over timesteps. Let us describe each set as having M elements (or agents) with K attributes (e.g. x/y position). To process the social and temporal information the encoder applies the following two transformations. Appendix B.2 provides proof of equivariance with respect to the agent dimension for this encoder.
First, the AutoBots encoder injects temporal information into the sequence of sets using a sinusoidal positional encoding function PE(.) (Vaswani al., 2017). For this step, we view the data collection of matrices, (Xu.... X that describe the temporal evolution of agents. The encoder processes temporal relationships between sets using MAB, performing the following operation each set where m E (1....,M}:
Sm MAB(FFFN(X..))
where (FFN(-) is embedding layer projecting each element in the set the size of the hidden state dg- Next, we process temporal slices of S, retrieving sets of agent states S, at some timestep 7 and is processed using S, (A,)(S,) These operations are repeated Lens times to produce context tensor C Ridx,M, summarizing the entire input scene, where is the number of timesteps in the input scene.
3.2 DECODER: MULTIMODAL SEQUENCE GENERATION
The goal of the decoder is to generate temporally and socially consistent predictions in the context of multi-moda data distributions. To generate N different predictions for the same input scene,
3

---

Published as a conference paper at ICLR 2022
the AutoBot decoder employs matrices of learnable seed parameters Q (dxx:T) where T is the prediction horizon, where i E {l,....c) c}. Intuitively, each matrix of learnable seed parameters corresponds to setting of the discrete latent variable in AutoBot. Each learnable matrix Q2 is then repeated across the agent dimension M times to produce the input tensor Q', (A, Additional contextual information (such as a rasterized image of the environment) is encoded using convolutional neural network to produce vector of features m, where i {1,..., c). In order to provide the contextual information to all future timesteps and to all set elements, we copy this vector along the M and dimensions, producing the tensor M, concatenated with M, along the dæ dimension, shown the left side of Figure 2 (bottom). This tensor is then processed using rFFN to produce the tensor H E I (dd7)
Each tensor Q', is
We begin decode by processing the time dimension, conditioning on the encoder's 's output C, and the encoded seed parameters and environmental information in H. The decoder processes each of m agents H separately with an MABD layer:
H MABD(Hmm Cm)
where H' is the output tensor that encodes the future time evolution of each element in the set independently. In order to ensure that the future scene socially consistent between set elements, we process each temporal slice of H', retrieving sets of agent states HH, at some future timestep T where each element h2 E H, is processed by a MAB layer. Essentially, the MAB layer here performs a per-timestep attention between all set elements, similar to the Set Transformer (Lee et al., 2019).
These two operations are repeated Laec times to produce final output tensor for mode i. The decoding process is repeated times with different learnable seed parameters Q and additional contextual information m,. The output of the decoder is a tensor O E (dd M,T,e) which can then be processed blement-wise using a neural network (-) to produce the desired output representation. In several of our experiments, we generate trajectories in (Ir, y) space, and such, & produces the parameters of a bivariate Gaussian distribution.
Note on the learnable seed parameters. One of the main contributions that make AutoBot's inference and training time faster than prior work is the use of the decoder seed parameters Q. These parameters serve dual purpose. Firstly, they account for the diversity in future prediction, where each matrix Q E (R( (where i (1,...cc c)) corresponds to setting of the discrete latent variable. Secondly, they contribute to the speed of AutoBot by allowing it to perform inference on the entire scene with single forward pass through the decoder without sequential sampling. Prior work (e.g., Tang and Salakhutdinov (2019)) uses auto-regressive sampling of the future scene with interlaced social attention, which allows for socially consistent multi-agent predictions. Other work (e.g., Messaoud et al. (2020)) employ single output MLP to generate the entire future of each agent independently with single forward pass. Using the attention mechanisms of transformers, AutoBot combines the advantages of both approaches by deterministically decoding the future scene starting from each seed parameter matrix. In Appendix C.5, we compare the inference speed of AutoBots with their autoregressive counterparts.
3.3 TRAINING OBJECTIVE
Given a dataset D = (x1,.... Xr) consisting of N sequences, each having T sets, our goal is to maximize the likelihood of the future trajectory of all elements in the set given the input sequence of sets, i.e., max p( (((+,1777 (xx)). In order to simplify notation, we henceforth refer to the ego-agent's 's future trajectory by y X:+1.7. As discussed above, AutoBot employs discrete latent variables, which allows us to compute the log-likelihood exactly. The gradient of our objective function can then be expressed as follows:
VoC(0) = Vo log Vo log
Σ Po()),
= Σ Pe (Z), X 1 V log (), X
As previously discussed in Tang and Salakhutdinov (2019), computing the posterior likelihood (o(2),, X is difficult in general and varies with 9. As discussed Bishop (2006), one can introduce distribution the latent variables, q (2), that is convenient to compute. With this
4

---

Published as a conference paper at ICLR 2022
distribution, now decompose the log- log-likelihood follows:
log (og = Σ log po (), X 1:t. (())
+ +Dxu((())llooZ))) X1:)),
(4)
where (.Il) is the Kullback-Leibler divergence between the approximating posterior and the actual posterior. A natural choice for this approximating distribution is g(Z) Pa.ia (Z X1a), where Ood corresponds to AutoBot's parameters before performing the parameter update. With this choice of (Z), the objective function can be re- re-written as
(()) = Q00, oota) + const.
Q Oota)
Σ Pa (1:)) log Po X1:e)
(5)
Σ 2
po ( (Z), X 12) log ()11,, log po (Z (Z[X::))
where the computation of pa (2 (1t) is described in Appendix C.2. Note that the posterior Po.ca (e.aa((]) (2 11x) can be computed exactly in the model we explore here since AutoBot operates with final objective function becomes maximize Q(0, Oota) and
discrete latent variables. Therefore, minimize D (pe..a ( X1:t) llpo (.)))
As mentioned, for experiments for each agent at every timestep. We found that to ensure that each mode's output sequence does not have high variance, we introduce mode entropy (ME) regularization term penalize large entropy values of the output distributions,
function & produce bivariate Gaussian distribution
T
LME = Ae max Σ (po 12, (1::,))
(6)
r=++1
As we can see, this entropy regularizer penalizes only the mode with the maximum entropy. Our results in Section D.2 show the importance of this term in learning to cover the modes of the data.
4 EXPERIMENTS
We now demonstrate the capability of AutoBots to model sequences of set-structured data across several domains. In subsection 4.1 we show strong performance on the competitive motion- prrdiinpreeittion dataset nuScenes (Caesar et al., 2020), a dataset recorded from a self-driving car, capturing other vehicles and the intersection geometry. In subsection 4.2, we demonstrate our model on the Argoverse autonomous driving challenge. In subsection 4.3, show results on the synthetic partition of the T (Sadeghian et al., 2018) 8) dataset which contains multi-agent scenes with high degree of interaction between agents. Finally in subsection 4.4, show that model can be applied to generate characters conditioning across strokes, and verify that AutoBot-Egoo is permutation-invare across an input set. See Appendix D.2, for an illustrative toy experiment that highlights our model's ability to capture discrete latent modes.
4.1 NUSCENES, A WORLD DRIVING DATASET
In this section, we present AutoBot-Ego 's results on the nuScenes dataset. The goal of this benchmark dataset is predict the ego-agent's 's future trajectory seconds at 2hz) given the past trajectory (up to 2 seconds at 2hz). Additionally, the benchmark provides access to all neighbouring agents' past trajectories and birds-eye-viee RGB image of the road network in the vicinity of the ego-agent. Figure 3 shows three example predictions produced by AutoBot-Fgo trained with c 10. We observe that the model learns to produce trajectories that are agreement with the road network and covers most possible futures directions. We also note that for each direction, AutoBot-Ego assigns different modes to different speeds to more efficiently cover the possibilitics.
We compare results on this benchmark to the state-of-the-ar in Table 1. Min ADE (5) and (10) are the average of pointwise L2 distances between the predicted trajectory and ground truth over 5 and 10 most likely predictions respectively. A prediction is classified as a miss if the maximum pointwise L2 distance between the prediction and ground truth is greater than meters. Miss Rate Top-5 (2m) and Top-10 10 (2m) the proportion of misses over all agents, where for each agent,
5

---

Published as a conference paper at ICLR 2022
Min ADE
Min ADE
Miss Rate
Miss Rate
Min FDE
Off Road
Model
(5)
(10)
Top-5 (2m)
Top-10 (2m)
(1)
Rate
Noah prediction
1.59
1.37
0.69
0.62
9.23
0.08
CXX
1.63
1.29
0.69
0.60
8,86
0.08
LISA(MHA JAM)
1,81
1.24
0.59
0.46
8,57
0.07
Trajectron++
1.88
1.51
0.70
0.57
9.52
0.25
CoverNet
2.62
1.92
0.76
0.64
11.36
0.13
Physies Oracle
3.70
3,70
0.88
0.88
9.09
0.12
WIMP
1.84
1.11
0.55
0.43
8,49
0.04
GOHOME
1,42
1.15
0.57
0.47
6,99
0.04
AutoBot- Ego (c=10)
1.43
1.05
0.66
0.45
8.66
0.03
AutoBot-Egg Ego (ensemble)
1.37
1.03
0.62
0.44
8.19
0.02
Table I: Quantitative Results on the nuScenes dataset. Other methods: LISA (Messaoud et al., 2020);
Trajectron++ (Salzmann al., 2020); CoverNet (Phan-Minhh et al., 2020); Physics Oracle (Caesar et al., 2020); WIMP (Khandelwal al., 2020); GOHOME (Gilles et al., 2021)
Figure 3: NuScenes Trajectory Prediction. Top row: birds-eye-view road network with ground truth trajectory data where given trajectory is cyan and held-out trajectory information is pink. Bottom diverse trajectories generated modes AutoBot -goo The model generates trajec- tories adhere the road network captures distinct possible futures.
Ground Truth Trajectories
the different
AutoBo:
Trajectories modes)
5 and 10 most likely predictions respectively are evaluated check if they're 're misses. Min FDE (1) is the L2 distance between the final points of the prediction and ground truth of the most likely prediction averaged over all agents. Off Road Rate the fraction predicted trajectories that are not entirely contained in the drivable area of the map.
We can see that AutoBot-Egr obtains the best performance on the Min ADE (10) metric and on Off Road Rate metric, and strong performance the other metrics. Furthermore, using ensemble of three AutoBot-Fgo models slightly improves the performance. This demonstrates our model's ability to capture multi-modal distributions over future trajectories due the latent variables. AutoBot-Ego excelled in maintaining predictions on the road, showing that the model has effectively used its capacity attend the map information. Finally, note that like all other models, AutoBot-Ego struggles to predict the probability of the correct mode, which results in poor performance of the Min FDE (1). The most interesting part of AutoBot-Ego is its computational requirements. We highlight that AutoBot-Ego was trained on a single Nvidia GTX 1080 Ti for 3 hours. This is in large part due the of the decoder seed parameters which allow ego-agent in single forward pass. Note that in Appendix C.5, we present the inference speed of AutoBot-Ego compared to an rugo-segres version. Ablation studies on various components of AutoBot-Fgg are also presented in Appendix D.1.
predict the future trajectory of the
4.2 ARGOVERSE RESULTS
Table 2 shows the results of AutoBot-Ego (with c 6) on the Argoverse test set. While our method does not achieve top scores, it uses much less computation, less data augmentation and no ensembling compared to alternatives. Further, we compare quite well to the 2020 Argoverse competition winner, a model called "Jean". addition, we note that the Argoverse challenge does not restrict the use of additional training data and that many of the top entries the leaderboard have publicly available papers or codebases. Many methods are therefore optimized for Argoverse by means of various data augmentations. For example, Scene Transformer (Ngiam et al., 2021) used data augmentation scheme that includes agent dropout and scene rotation which we have not, and the current top method and winner of the 2021 challenge (QCraft) makes of ensembling.

---

Published as a conference paper at ICLR 2022
Model
Min ADE (1)
Min FDE ())
Miss Rate (1)
DAC (1)
AutoBot-Ego (Valid Set)
0.73
1.10
0.12
AutoBot-Ego (Test Set)
0.89 (top-5)
1.41 (top-S)
0.16 (top-6)
0.9886 (top-3)
Jean (2020 winner)
0.97
1.42
0.13
0.9868
WIMP
0.90
1,42
0.17
0.9815
TNT
0.94
1.54
0.13
0.9889
LaneGCN
0.87
1.36
0.16
0,9800
TPCN
0,85
1.35
0.16
0,9884
mm Transformer
0.84
1.34
0.15
0.9842
GOHOME
0.94
1,45
0.105
0.9811
Scene Transformer
0.80
1.23
0.13
0.9899
Table 2:
Quantitative Results on the Argoverse test set. Other methods that had an associated paper:
Jean (Mereat et al., 2020); WTMP (Khandelwal et al., 2020); TNT (Zhao et al., 2020a); LaneGCN (Liang al., 2020); TPCN (Ye al., 2021); mmTransformer (Liu et al., 2021); GOHOME (Gilles al., 2021); Scene Transformer (Ngiam et al., 2021).
Ego Agent's Min ADE (.)
Number Of Collisions (1)
Scene-level
Scene-level
Model
Min ADE(6) (L)
Min FDE(6) (L)
Linear Extrapolation 0.439
2220
0.409
0.897
AutoBot-AntiSocaal
AutoBot-Ego 0.196
1827
0.316
0.632
0.098
1144
0.214
0.431
AutoBot
0.095
139
0.128
0.234
Table 3: TrajNet ablation studies for a multi-agent forecasting scenario. We investigate the impact of using the social attention in the encoder and in the decoder. We found that AutoBot is able to cause significantly fewer collisions between agents compared to variants without the social attention, and improves the prediction accuracy the scene-level metrics while maintaining strong performance the ego-agent. Metrics reported the validation
We believe that a fair comparison of the performance between neural network methods should ensure that each architecture is parameter-matee Argoverse leaderboard does not require entrants to report their parameter or FLOP count, we report available public statistics to enable comparison here. First, Scene Transformer (Ngiam al., 2021) reports that their model trains in 3 days on TPUs and TPU-v3s have peak performance of 420 TeraFL.OPS, which means their model could have consumed up to 3 days 24 hours x 420 TFLOPS, or 088,00,0000 TFLOPS. On the other hand, AutoBot-Ego hours, consuming roughly 396,000 TFLOPs, or 0.4% as much. Jean's 's method (Mercat et al., 2020) consumes roughly 8,600,00 000 TFLOPs (training method uses only 5% of that compute). While the method currently topping the leaderboard has not revealed their model details, the fact that they used ensembles very likely requires significantly more compute than our method.
and the training is FLOO-matched. While the
trained
Nvidia 1080ti for 10
4.3 TRAINET++ SYNTHETIC DATASET
In this section, we demonstrate the utility of the having social and temporal attention the encoder and decoder of AutoBot. We chose to evaluate the synthetic partition of the TrajNet++ dataset since this was specifically designed to have a high-level of interaction between scene agents (Kothari al., 2021). In this task, we are provided with the state of all agents for the past timesteps and tasked with predicting the next 12 timesteps for all agents. The scene is initially normalized around one (ego) agent such that the agent's 's final position in the input sequence aligned with the positive direction. There are a total of 54, 513 unique scenes in the dataset, which we split into 49, 062 training scenes and 5, 451 validation scenes. For this experiment, we apply our general architecture (AutoBot) to multi-agen scene forecasting, and perform an ablation study on the social attention components in the encoder and the decoder.
(0, and its heading is
Table 3 presents three variants of the model, and a baseline extrapolation model which shows that the prediction problem cannot be trivially solved. The first, MutoBot-AntiSoche corresponds to a model without the intra-set attention functions in the encoder and decoder (see Figure 2). The second variant, AutoBot-Fgo, only sees the past social context and has social attention during
7

---

Published as a conference paper at ICLR 2022
Ausbot
Araooot
LSTM
LSTM
q
q
q
q
q
Figure Omniglot qualitative results. Left: stroke completion task. We show examples of characters generated using AutoBot and an LSTM baseline. The first two columns show the ground-truth image of the character and the corresponding ground-truth strokes. In this task, the models are provided with the first half of all strokes (cyan) and are tasked with generating the other half (pink). The four other columns show the generated future strokes, one for cach latent mode. We can see that AutoBot produces more consistent and realistic characters compared to the LSTM baseline. Right: character completion task. We show two examples of characters completed using AutoBot-Egg and an LSTM baseline. The first two columns show the ground-truth image of the character as before. In this case, the models are provided with two complete strokes and are tasked stroke complete the character. We observe that AutoBot-Ege Ego generates realistic characters across all modes, given ambiguous input strokes.
with generating
the trajectory generation process. The third model corresponds to the general architecture, AutoBot. All models were trained with e = and we report scene-level error metrics as defined by Casas al. (2020). Using the social attention in the decoder significantly reduces the number of collisions between predicted agents' futures and significantly outperforms its ablated counterparts on scene-level minADE and minFDE. We expect this theoretical level since the ego-centric formulation makes the independence assumption that the scene's 's future rollout can be decomposed as the product of the individual agent's s future motion in isolation, as remarked in Casas et al. (2020). AutoBot does not make this assumption as it forecasts the future of all agents jointly. We refer the reader to Appendix D.3 for additional ablation experiments and qualitative results on this dataset.
4.4 OMNIGLOT, STROKE- BASED IMAGE GENERATION
To demonstrate the general efficacy of our approach to sequence generation, we demonstrate AutoBots on a character completion task, a diverse alternative to automotive and pedestrian trajectory prediction. We set up two different tasks on the Omniglot dataset (Lake et al., 2015): (a) stroke-completion task in which the model is presented with the first half of each stroke and has to complete all of them task which the model given several full strokes and has to draw the last stroke of the character. In task (a), the model is trained dataset (training set), in which strokes are encoded as sequential sets of points, and we are comparing the performance of AutoBot qualitatively against that of an LSTM bascline with social attention and discrete latents (similar to AutoBot but with classic RNN architecture) on a held-out test set. With the first task, we illustrate that our model is flexible enough to learn legibly complete characters from a variety of different languages. In task (b), we train AutoBot-Eggo only on characters "F", "H", and "II" demonstrate that our discrete latent variable captures plausible character completions given the context of the other strokes, e.g. character "F" can be completed as an "H" when the top stroke is missing. The results of these experiments can be found Fig. 4. Implementation details and additional results can be found in Appendix section D.5.
in parallel and (b) sharacter-completie
the entire Omniglot
5 RELATED WORK
Processing set structured data. Models of set-valued data (i.e. unordered collections) should be permutation-invarian (see Appendix for formal definitions), and be capable of processing sets with arbitrary cardinality. Canonical forms of feed forward neural networks and recurrent neural network models do not have these properties (Lee et al., 2019). One type of proposed solution is to integrate pooling mechanisms with feed forward networks process set data (Zaheer et al.. 2017: Su et al., 2015; Hartford et al., 2016). This is referred as social pooling in the context of motion prediction (Alahi et al., 2016; Deo and Trivedi, 2018; Lee et al., 2017). Attention models (Lee et al.,
8

---

Published as a conference paper at ICLR 2022
2019; Casas et al., 2020; Yuan et al., 2021; Ngiam et al., 2021: Zhao et al., 2020b) or graph- graph-based models (Zhao et al., 2020a; Messaoud et al., 2020; Li et al., 2020b; Gao et al., 2020) have performed well in set processing tasks that require modelling high-order element interactions. In contrast, provide a joint single pass attention-based approach for modelling more complex data-structures composed by sets and sequence, while also addressing how to generate diverse and heterogeneous data-structures, Although those properties exist in concurrent motion prediction papers (Tang and Salakhutdinov, 2019; Casas et al., 2020; Yuan al., 2021; Ngiam et al., 2021; Li et al., 2020a), our proposed method is targeting a more general set of problems while still getting competitive results in motion prediction benchmarks at fraction of other models' compute.
Diverse set generation. Generation of diverse samples using neural sequence models is an important part of our model and well-studied problem. For example, in Bowman et al. (2015) a variational (CVAE) approach is used to construct models that condition on attributes for dialog generation. The CVAE approach has also been integrated with transformer networks (Linn et al., 2020; Wang and Wan, 2019) for diverse response generation. For motion prediction, Lee et al. (2017) and Deo and Trivedi (2018) use a variational mechanism where each agent has their own intentions encoded. Li et al. (2020a) output Gaussian mixtures at every timestep and sample from it to generate multiple futures. The most similar strategy to ours was developed by Tang and Salakhutdinov (2019). Other recent works (Suo al., 2021; Casas et al., 2020) use continuous latent variables to model the multiple actors' intentions jointly. We built on this method by using one joint transformer conditioned on the latent variables and added maximum entropy likelihood loss that improves the variability of the generated outcomes.
Sequential state modelling. We are interested in data that can be better represented as a sequence. The main example of this is language modelling (Vaswani et al., 2017: Devlin et al., 2018), but also includes sound (Kong et al., 2020), video (Becker et al., 2018), and others (Kant et al., 2020; Yang et al., 2021). For motion prediction, prior work has largely focused on employing RNNs to model the input and/or output sequence (Alahi et al., 2016; Lee et al., 2017: Deo and Trivedi, 2018; Tang and Salakhutdinov, 2019; Messaoud et al., 2020; Casas et al., 2020; Li et al., 2020b:a). With its recent success in natural language processing, Transformers (Vaswani et al., 2017) have been adopted by recent approaches for motion forecasting. Yu et al. (2020) uses a graph encoding strategy for pedestrian motion prediction. Giuliari et al. (2021) and Liu et al. (2021) propose a transformer- based encoding to predict the future trajectories of agents separately. In contrast, AutoBots produces representations that jointly encode both the time and social dimensions by treating this structure as a sequence of sets. This provides reuse of the model capacity and lightweight approach. Concurrent to our method, Yuan et al. (2021), also jointly encodes social and time dimensions but flattens both of those axis into a single dimension.
CONCLUSION
In this paper, we propose the Latent Variable Sequential Set Transformer to model time-evolution of sequential sets using discrete latent variables. We make of the multi-head attention block to efficiently perform intra-set attention, to model the -eeeddepedee between the input and output sequence, and to predict the prior probability distribution over the latent modes. We validate our AutoBot-Fgo by achieving competitive results on the nuScenes benchmark for ego-centric motion forecasting in the domain of autonomous vehicles. We further demonstrate that AutoBot can model diverse sequences of sets that adhere to social conventions. We validate that since our model can attend the hidden state of all agents during the generative process, produces scene-consistent predictions on the +raj dataset when compared with ego-centric forecasting methods.
Our approach is limited operate in a setting where a perceptual pipeline is providing structured sets representing the other entities under consideration. If our approach is used in a robot or autonomous car, any errors or unknown entities not captured by sets used as input to the algorithm would not be modelled by our approach. Memory limitations may also limit the number of sets that can be processed at any time. In future work, we plan to extend this line of work to broader and more interpretable scene-consistency. Further, we would like to investigate the combination of a powerful sequence model like AutoBot with a reinforcement learning algorithm to create a data-efficient model-based RL agent for multi-agent settings.
Limitations:

---

Published as a conference paper at ICLR 2022
REFERENCES
Alexandre Alahi, Kratarth Goel, Vignesh Ramanathan, Alexandre Robicquet, Li Fei-Fei, and Silvio Savarese. Social Istm: Human trajectory prediction in crowded spaces. In Proceedings of the IEEE conference on computer vision and pattern recognition, pages 961-971, 2016.
Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey Hinton. Layer normalization. arXiv preprint arXiv:1607.064400 2016.
Stefan Becker, Ronny Hug. Wolfgang Hubner, and Michael Arens. Red: A simple but effective baseline predictor for the trajnet benchmark. In Proceedings of the European Conference Computer Vision (ECCV), pages 0-0, 2018.
Yuriy Biktairov, Maxim Stebelev, Irina Rudenko, Oleh Shliazhko, and Boris Yangel. PRANK: motion prediction based on ranking. CoRR, hiss/2010.12007,, 2020. URL https:://// /ttee://aaraae /ttpiv.avvivvaa///i//e..//.aa org/ abs/ 2010 12007.
C. M. Bishop. Pattern Recognition and Machine Learning. Springer; 1st ed. 2006, 2006.
Samuel R Bowman, Luke Vilnis, Oriol Vinyals, Andrew M Dai, Rafal Jozefowicz, and Samy Bengio. Generating sentences from a continuous space. arXiv preprint
2015.
H. Caesar, Varun Bankiti, A. Lang, Sourabh Vora, Venice Erin Liong, Q. Xu, A. Krishnan, Yu Pan, Giancarlo Baldan, and Oscar Beijbom. nuscenes: A multimodal dataset for autonomous driving. 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 11618- 11628, 2020.
Sergio Casas, Cole Gulino, Simon Suo, Katie Luo, Renjie Liao, and Raquel Urtasun. Implicit latent variable model for scene-consistent motion forecasting. arXiv preprini urXiv:2007. 12036, 2020.
Ming-Fang Chang, John Lambert, Patsorn Sangkloy, Jagjeet Singh, Stawomir Bąk, Andrew Hartnett, De Wang, Peter Carr, Simon Lucey, Deva Ramanan, and James Hays. Argoverse: 3d tracking and forecasting with rich maps. 2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 8740-8749, 2019.
Nachiket Deo and M. Trivedi. Convolutional social pooling for vehicle trajectory prediction. 2018 IEEE/CVF Conference Computer Vision and Pattern Recognition Workshops (CVPRW), pages 1549-15498, 2018
Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova. Bert: Pre-training of deep bidirectional transformers for language understanding. arXiv preprint arXiv: /810004800 2018.
Jiyang Gao, Chen Sun, Hang Zhao, Yi Shen, Dragomir Anguelov, Congcong Li, and Cordelia Schmid. Vectornet: Encoding hd maps and agent dynamics from vectorized representation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 11525-11533, 2020.
Thomas Gilles, Stefano Sabatini, Dzmitry Tsishkou, Bogdan Stanciulescu, and Fabien Moutarde. Gohome: Graph-oriented heatmap output forfuture motion estimation. arXiv preprint (rXiv:2109001821 2021.
Francesco Giuliari, Irtiza Hasan, M. Cristani, and Fabio Galasso. Transformer networks for trajectory forecasting. 2020 25th International Conference on Pattern Recognition (ICPR), pages 10335- 10342, 2021.
Jason Hartford, James R Wright, and Kevin Leyton-Browe Deep learning for predicting human strategic behavior. In NIPS, 2016.
Yash Kant, Dhruv Batra, Peter Anderson, Alexander Schwing, Devi Parikh, Jiasen Lu, and Harsh Agrawal. Spatially aware multimodal transformers for textvqa. In Computer Vision-ECCY 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part IX 16, pages 715-732. Springer, 2020.
10

---

Published as a conference paper at ICLR 2022
Aniruddha Kembhavi, Minjoon Seo, Dustin Schwenk, Jonghyun Choi, Ali Farhadi, and Hannaneh Hajishirzi. Are you smarter than a sixth grader? textbook question answering for multimodal machine comprchension. In 2017 IEEE Conference Computer Vision and Pattern Recognition (CVPR), pages 5376-5384, 2017. doi: 10.1109/CVPR.2217.551
Siddhesh Khandelwal, William Qi. Jagjeet Singh, Andrew Hartnett, and Deva Ramanan. What-if motion prediction for autonomous driving, 2020.
Qiugiang Kong, Yong Xu, Wenwu Wang, and Mark D Plumbley. Sound event detection of weakly labelled data with cnn-transformer and automatic threshold optimization. JEEE/ACM Transactions on Audio, Speech, and Language Processing. 28:2450-222 24600,2460.221 2020.
Parth Kothari, Sven Kreiss, and Alexandre Alahi, Human trajectory forecasting in crowds: deep learning perspective. IEEE Transactions on Intelligent Transportation Systems, 2021.
Brenden M Lake, Ruslan Salakhutdinov, and Joshua B Tenenbaum. Human-level concept learning through probabilistic program induction. Science, 350(6266):1332-13333 2015.
Juho Lee, Yoonho Lee, Jungtaek Kim, Adam Kosiorek, Seungjin Choi, and Yee Whye Teh. Set trans- former: A framework for iveetion-based permutation-invariint neural networks. In International Conference on Machine Learning, pages 3744-37 3753. PMLR, 2019.
N. Lee, W. Choi, Paul Vernaza, C. Choy, P. Torr, and Manmohan Chandraker. Desire: Distant future prediction in dynamic scenes with interacting agents. 2017 IEEE Conference Computer Vision and Puttern Recognition (CVPR), pages 2163-2174, 2017.
Jiachen Li, Fan Yang, Masayoshi Tomizuka, and Chiho Choi. Evolvegraph: Multi-agent trajectory prediction with dynamic relational reasoning. arXiv preprint MXiv:20033139244 2020a.
Lingyun Luke Li, Bin Yang, Ming Liang, Wenyuan Zeng, Mengye Ren, Sean Segal, and Raquel Urtasun. End-to-end contextual perception and prediction with interaction transformer. arXiv preprint ar0iv:2008. 059927, 2020b.
Ming Liang, Binh Yang, Rui Hu, Yun Chen, Renjie Liao, Song Feng, and Raquel Urtasun. Learning lane graph representations for motion forecasting. ArXiv, abs/2007. 2020.
Zhaojiang Lin, Genta Indra Winata, Peng Xu, Zihan Liu, and Pascale Fung. Variational transformers for diverse response generation. arXiv preprint arXiv:2003. /2738, 2020.
Yicheng Liu, Jinghuai Zhang, Liangji Fang, Qinhong Jiang, and Bolei Zhou, Multimodal motion prediction with stacked transformers. ArXiv, abs/2103. 2021.
Jean Mercat, Thomas Gilles, Nicole Zoghby, Guillaume Sandou, Dominique Beauvois, and Guillermo Gil. Multi-head attention for joint multi-modal vehicle motion forecasting. In IEEE International Conference Robotics and Automation, 2020.
Kaouther Messaoud, Nachiket Deo, Mohan M Trivedi, and Fawzi Nashashibi. Multi-head attention with joint agent-map representation for trajectory prediction in autonomous driving. arXiv, pages arXiv-2005, 2020.
Jiquan Ngiam, Benjamin Caine, Vijay Vasudevan, Zhengdong Zhang, Hao-Tien Lewis Chiang. Jeffrey Ling, Rebecca Roelofs, Alex Bewley, Chenxi Liu, Ashish Venugopal, et al. Scene transformer: A unified multi-task model for behavior prediction and planning. arXiv preprint 2021.
08417.
Tung Phan-Minh, E. Grigore. F. Boulton, Oscar Beijbom, and Eric M. Wolff. Covernet: Multimodal behavior prediction using trajectory sets. 2020 IEEF/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pages 14062-14071, 2020.
Matthew Richardson, Christopher J.C. Burges, and Erin Renshaw. MCTest: A challenge dataset for the open-domain machine comprehension of text. In Proceedings of the 2013 Conference on Empirical Methods in Natural Language Processing, pages 193-203. Seattle, Washington, USA, October 2013, Association for Computational Linguistics. URL https: //aca/acaaa aaaaaaaannthtt ahtoloo org/DD 1- 3-20
11

---

Published as a conference paper at ICLR 2022
Amir Sadeghian, Vineet Kosaraju, Agrim Gupta, Silvio Savarese, and Alexandre Alahi. Trajnet: Towards a benchmark for human trajectory prediction. arXiv preprint, 2018.
Tim Salzmann, B. Ivanovic, P. Chakravarty, and M. Pavone, Trajectron++: Multi-agent generative trajectory forecasting with heterogeneous data for control. ArXiv, abs/2001.03093, 2020.
Hang Su, Subhransu Maji, Evangelos Kalogerakis, and Erik Learned- Miller. Multi-view convolutional neural networks for 3d shape recognition. In Proceedings of the IEEE international conference on computer vision, pages 945 953,, 2015.
Simon Suo, Sebastian Regalado, Sergio Casas, and Raquel Urtasun. Trafficsim: Learning to simulate realistic multi-agent behaviors. In Proceedings of the IEEE/CVF Conference Computer Vision and Pattern Recognition, pages 10400-10409, 2021.
Charlie Tang and Russ R Salakhutdinov. Multiple futures prediction. In Advances in Neural Information Processing Systems, pages 15424-15434, 2019.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, L. Kaiser, and Illia Polosukhin. Attention is all you need. In NIPS, 2017.
T. Wang and Xiaojun Wan. T-cvae: Transformer-based conditioned variational autoencoder for story completion. IJCAI, 2019.
Zhengyuan Yang, Yijuan Lu, Jianfeng Wang. Xi Yin, Dinei Florencio, Lijuan Wang, Cha Zhang, Lei Zhang, and Jiebo Luo. Tap: Text-aware pre-training for text-vqa and text-caption. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 8751-8761, 2021.
Maosheng Ye, Tongyi Cao, and Qifeng Chen. Tpen: Temporal point cloud networks for motion forecasting, ArXiv, abs/2103.03067, 2021.
De Jong Yeong, Gustavo Adolfo elasco-Herránder: John Barry, and Joseph Walsh. Sensor and sensor fusion technology in autonomous vehicles: A review. Sensors (Basel, Switzerland), 21, 2021.
Cunjun Yu. Xiao Ma, J. Ren, H. Zhao, and Shuai Yi. Spatio-temporal graph transformer networks for pedestrian trajectory prediction. In ECCV, 2020.
Ye Yuan, Xinshuo Weng, Yanglan Ou, and Kris Kitani. Agentformer: Agent-aware transformers for socio-temporal multi-agent forecasting. In Proceedings of the IEEF/CVF International Conference on Computer Vision (ICCV), 2021.
Manzil Zaheer, Satwik Kottur, Siamak Ravanbakhsh, Barnabas Poczos, Ruslan Salakhutdinov, and Alexander Smola. Deep sets. arXiv preprint arXiv:I703.066 14, 2017.
Hang Zhao, Jiyang Gao, T. Lan, Chen Sun, Benjamin Sapp, Balakrishnan Varadarajan, Y. Shen, Yuning Chai, C. Schmid, Congcong Li, and Dragomir Anguelov. Tnt: Target-driven trajectory prediction. ArXiv, abs/2008.0822) 082944, 2020a.
Hengshuang Zhao, Li Jiang, Jiaya Jia, Philip Torr, and Vladlen Koltun. Point transformer, 2020b.
12

---

Published as a conference paper at ICLR 2022
A BACKGROUND
We now review several components of the Transformer (Vaswani et al., 2017) and Set Transformer (Lee et al., 2019) architectures, mostly following the notation found in their manuscripts.
Muadi-Head Self Attention (MHSA) as proposed in (Vaswani et al., 2017) is a function defined on ny query vectors Q E R and key-value pairs (K IR xd, and V R==xdd With single attention head, the function resolves the queries Q by performing the computation:
Attn Q K, V) = softmax (oiinaaxxKKK) QKT)V JV,
where d is the dimension of the query and key vectors and d, is the dimension of the value vector. MHSA consists of h attention heads performing the operation shown in Eq. 7 with h linear projections of keys, queries, and values. The final output is a linear projection of the concatenated output of each attention head. These computations can be expressed as
MHSA(Q, K, V) concat (headd(head1,..... head,, W
where
(8)
head; = Att KW vw!),
and where W W and W.K along with the output projection matrix Wo are the learnable projection matrices of each attention head in the MHSA. Note that W wy and W project the initial set of inputs smaller dimensionality. For example, W E Rd projects the original queries with dimension dy dM which is typically chosen be dg/ This operation has useful applications in learning a representation of an input set where all its elements interact with each other.
An interesting application of MHSA is perform self-attention on an input set (Lee al., 2019). Given an input set X, one can perform intra-set attention by using X as the queries, keys and values and having a residual connection, i.e. MHSA(X) X X+MHSA(X, X, X). This operation has useful applications in learning a representation of an input set where all its elements interact with each other.
Multi-head Attention Block Decoders (MABD) were also introduced in Vaswani et al. (2017). and were used to produce decoded sequences. Given an input matrix X and an output matrix Y representing sequences, the decoder block performs the following computations:
MABD(Y, X) = LN(H + (F))(H))
where H LN(H' + MH'S(H'' X, X))
and H' = LN(MHSA(YY)
During training, MABD be trained efficiently by shifting Y backward by one (with start token at the beginning) and computing the output with one forward pass. During testing where one does not have access to Y. the model would then generate the future autoregressively. In order to avoid conditioning on the future during training. the MHSA(Y) operation in MABD employs time-masking. This would prevent the model from accessing privileged future information during training.
B PERMUTATION EQUIVARIANCE OF THE MODEL
In this section, we prove the permutation equivariance of the Latent Variable Sequential Set Trans- former with respect to the dimension M. To do so, we analyze the model's response to permu- tation along this dimension of the input tensor. This tensor, X, is of dimension (K, M, t), where K. M, E N. The dimensions and K are ordered, i.e. if we index into the tensor along dimension M (we denote this operation as X,, where m E (1,..., MM) then we retrieve (K, (T,,t)a d)ddiennsia matrix. However, the M dimension is unordered. This implics that when we index by another dimension, for example t, we retrieve an invertible multiset of vectors, denoted X, (æI...., x where r (1..... f) and a,.. X, are K-dimensional.
Our proof that our model is permutation equivariant on M demonstrates a clear difference with the properties of similar models. For example, the model proposed in Set Transformers (Lee et al., 2019) features permutation equivariant encoder combined with a permutation invariant decoder. And while Deep Sets (Zaheer et al., 2017) provides theoretical framework for designing neural network
13

---

Published as a conference paper at ICLR 2022
models that operating sets, they do not directly address the processing of hetereogeneous data structures as we do here.
The rest of the section will proceed in the following manner:
1. Provide definitions and lemmas of the mathematical structures we will use in our proofs.
2. Show the permutation equivariance of the AutoBot encoder on M.
3. Show the permutation equivariance of the AutoBot decoder on M.
For clarity, we re-state the motivation for proving permutation equivariance in this particular model. In tasks like motion prediction, we may not be able uniquely identify moving objects (perhaps due to occlusion or similar visual features). Therefore, we may wish to represent these objects as a collection of states evolving over time. Readers may benefit from this perspective on X as multiset for
time-evolving (t) collection of M objects with K attributes. We described X, generality, as it permits repeated elements. However, in motion prediction tasks it might be safe to assume that two distinet objects will share the same attribute vector. By proving that model is equivariant with respect to M, we ensure that it will not waste capacity on spurious order-specific features. The inductive bias for this type of task differs from other settings like machine translation where word order is meaningful.
B.1 DEFINITIONS AND LEMMAS
permutation equivariant function is one that achieves the same result if you permute the set then apply the function to the permuted set, as you would achieve if you had applied the function to the set then applied the same permutation to the output set. We provide a formal description in Def B.
Definition B.1 (Permutation Equivariance). Let X be the set of n vectors {æ1,..., In.. and Sn the group of permutations of n elements π X - X. Suppose we have function f X + X, we call f Permutation Equivariant iff [ = (X))
One of the central components of the Sequential Set Transformer is the Multi-Head Self Attention function (MHSA X + X) originally introduced in (Vaswani et al., 2017). MHSA computes an updated representation for one input vector in set of input vectors by processing its relation with a collection of other vectors. Without the explicit addition of positional information to the content of these vectors, this operation is invariant to input order. MHSA is content-based attention mechanism, and updates the contents of a single input vector based on the contents of the others. However, if define fMHSA as the function applying MHSA update each input vector, then we have created a permutation equivariant function, described in Lemma B.I.
Lemma B.1 (fMAn is Permutation Equivariant on M). Let X be a set of D-dimensional vectors (II,-.., TM.) ), and fMAB RRDxM - Dx be a function that applies MAB (as defined in Equation I of the main text) to each x.m. E X. Then fAA (X,) is permutation equivariant because VT E Sm, fMAB (XX)) = (fMAAB(X))
Having discussed the preliminaries, we now show the permutation equivariance of AutoBot's encoder.
B.2 AUTOBOT ENCODER PROPERTIES
Theorem B.2 (AutoBot Encoder is Permutation Equivariant on M). Let X be a tensor of dimension (K, M, t) where K, M, E N, and where selecting on t. retrieves invertible multiset of sequences, denoted X, = { M where T E (I....,t) t} and xm E X, are K K-dimensional. The AutoBot encoder, as defined in Section 3.1, is permutation equivariant on M.
Proof. What follows is a direct proof of the permutation cquivariance of the AutoBot encoder. The encoder is a composition of three functions, rFFN: k [g xt. fMAB : [ddx X RRd. x X and MMAB dx x d X
First, we show that the embedding function rFFN is equivariant on M. Let rFFN represent the application of function R + 4x to each element of set X, where E {1, (1....,t). t). IFEN is equivariant with respect to M by Definition B. Specifically. we can see that the function IFEN satisfies the equation Vn E SM, IFFN(z(XX)) (IFFN((X.)) In the AutoBot encoder, the
Then
14

---

Published as a conference paper at ICLR 2022
function rFFN is applied to each set x, and because it is permutation equivariant on M for each, the application to each set represents a permutation equivariant transformation of the entire tensor.
Next, we show that the function MMAB. which applies Multi-Head Attention Block (MAB) to each set, (X,....,X,, X,). is equivariant on M. fMAB is permutation equivariant with respect to the M dimension of X because MAB is permutation equivariant each set X by Lemma B.
Finally, we define a function MAA which applies a positional encoding along the time dimension of X then Multi-head Attention Block (MAB) to each matrix in X, {XI..... xM)) equivariant with respect to the M dimension of X because the collection of matrices (X1,..., XM)
fMA is permutation
an unordered set, and the uniform application of a function to each element of this set satisfies permutation equvariance.
composition of equivariant functions is equivariant, so the encoder is equivariant on M.
B.3 AUTOBOT DECODER PROPERTIES
We now provide a direct proof the the permutation equivariance of the AutoBot decoder on M. The decoder is an auto regresegee function initialized with a seed set X, containing M vectors of dimension dg, each concatenated with a vector of latent variables and other context dimension 2dg. For 7 E {1,.... T I iterations, we concatenate the output of the decoder with the seed, the previous decoder output, and the latents to produce a tensor of dimension (3dx. M, T). In particular, the decoder is composition of three functions, rFFNdec: R3dxx7 - Rdx fMAB Rd, XM Rdd x and MABD Rd x M - IRxx x xr
Theorem B.3 (AutoBot Decoder is Permutation Equivariant on M). Given an invertible multiset X, with M vector-valued elements, and suto-regressive generator iterated for E $11 ....... . f). the AutoBot decoder, formed by a composition of three functions, rFFN R3dg xM - RRxx X
MAB Rd xM -> Rd and MMABD RARXXXM XTXM - RdxxxTM M. is equivariant M.
Proof. First, we must establish that the function rFFNdec: gild xM + Rdd xM is equivariant with respect to M. The AutoBot decoder applies rFFN dec each set X where E f (}. transforming the M vectors of dimension 3dx + dz. Because rFFN represents a uniform application of edee to each element of the set, we can see that satisfies the definition of permutation equivariance, specifically that VT E SM. rFFN (T (X+))
(IFFNNa.(X..))
Next, we take the function MABD. which first applies a function PE : Rdd.,,r Rd X adding position information along the temporal dimension, then applies SMABD Rdxxx - Rx to each matrix (X,,.., XM). J. fMARD is permutation equivariant with respect to the M dimension of X because the collection of matrices (X.,.... Xas) } is unordered set, and the uniform application of a function to transform each element independently does not impose an order on this set.
Similar to the final step of the previous proof, we that function fMAB applying mulli-head self- attention block (MAB) to each set (X.+1....... x is equivariant on M because MAB is permutation
equivariant on each set (see Lemma B.
The composition of equivariant functions is equivariant, so the decoder is equivariant.
C MODEL DETAILS
C.1 EGO-CENTRIC FORECASTING
AutoBots can be used to predict the future trajectories of all agents in scene. However, in some settings we may wish forecast the future trajectory of only single agent in the set. We call this variant "AutoBot-Egoo.. Referring back to the encoder presented in Section 3.1, instead of propagating the encoding of all elements the set, we can instead select the state of that element produced by the intra-set attention, identically to produce the tensor C which is an encoding of the ego element's 's history conditioned on the past sequence of all elements.
Sr. With this,
encoder proceeds with the rest of the computation
AutoBot-Ego's decoder proceeds as presented in section 3.2, with one exception. Since AutoBot-Ega only generates the future trajectory of an ego element, we do not perform intra-set attention in the
15

---

Published as a conference paper at ICLR 2022
Figure 5: AutoBots Context Encoder, Example of the birds-eye-view road net- work provided by the Nuscenes dataset. Our CNN encoder produces volume which we partition equally into modes. This allows each generated trajectory to be conditioned on a different representation of the imput image.
Encoder
Figure TrajNet Map Example. Ex- ample birds-eye-view map provided by the TrajNet benchmark. The footage was captured by drone and pedestrians and bicycles are freely moving through the scene.
decoder. As a consequence of only dealing with a single ego-agent in the future generation, we do repeat the seed parameter matrices Q; across the M dimension, where E {l,.... c}. The objective function is updated to compute the likelihood of one element's future given the entire input sequence. That is, in AutoBot-Fgo, y X 141:7
C.2 EXPECTATION MAXIMIZATION DETAILS
Our model also computes the distribution we employ e learnable vectors (one for each mode) concatenated into the matrix P E R o) which behave like seeds to compute distribution modes prior observing the future (these seed parameters should not be confused with the decoder seed parameters Q). The distribution is generated by performing the computations
of discrete random variables. To achieve this,
p( 2 X (.)) softmax (rLin(F)), where F MABD(P 1sc. C),
where rLin is a row-wise linear projection layer to vector of size c.
C.3 CONTEXT ENCODING DETAILS
In the nuScenes dataset, we are provided with a birds-eye-view 128 128 RGB image of the road network, as shown in Figure 5. In the TrajNet dataset, we are also provided with birds-eye-viens of the scene, as shown Figure 6. We process this additional context using 4 layer convolutional neural network (CNN) which encodes the map information into volume of size 7 x 7 x (7 +c) where is is the number modes. We apply a 2D dropout layer with rate 0.1 on this output volume before processing it further. As our model employs discrete latent variables, we found it helpful to divide this volume equally among all e modes, where each mode receives flattened version of the x x 7 volume. As described in Section 3.2, this context is copied across the M and T dimensions, and concatenated with the decoder seed parameters during the sequence generation process. Intuitively, each mode's generated trajectory is conditioned on different representation of the context.
C.4 IMPLEMENTATION AND TRAINING DETAILS
We implemented our model using the Pytorch open-source framework. Table 4 shows the values of parameters used in our experiments across the different datasets. The MAB encoder and decoder blocks in all parts of AutoBots use dropout. We train the model using the Adam optimizer with an initial learning rate. For the Nuscenes and TrajNet++ datasets, we anneal the learning rate every 10 epochs for the first 20 epochs by factor of followed by annealing it by factor of 1.33 every 10 epochs for the next 30 epochs. For the Argoverse dataset, we anneal the learning rate every epochs by factor 2 for the first 20 epochs. In all datasets, we found it helpful to clip the gradients
16

---

Published as a conference paper at ICLR 2022
Parameter
Description
Toy
OmniGlot
Nuscenes
Argoverse
TrajNet
TrajNet
Hidden dimension throughout all parts of the model.
dx
64
128
128
128
128
128
Learning Rate
Learning rate used in Adam Optimizer.
lc-4
5e-4
7.5c-4
7.5c-4
7.5c-4
5c-4
Batch size
Batch size used during training.
64
64
64
64
64
Number of discrete latent variables.
10
4
10
6
5
X
Entropy regularization strength.
3.0
1.0
30.0
30.0
5.0
30.0
Dropout
Amount of dropout used in MHSA functions.
0.0
0.2
0.1
0.1
0.3
0.05
Number stacked pocial/Temporal blocks the encoder.
Tene
1
2
2
2
Number of stacked social/Tempora blocks
Ldec
in the decoder.
11
Number of attention heads in all MIISAs.
8
16
16
16
Table 4: Hyperparameters used for AutoBots across all four datasets.
to maximum magnitude of 5.0. For the Nuscenes experiments, our model takes approximately 80 epochs to converge, which corresponds approximately hours of compute time on a single Nvidía Geforce GTX 1080Ti GPU, using approximately 2 GB of VRAM. We used mirroring of all trajectories and the map information as a data augmentation method since the dataset's scenes originate from Singapore and Boston. For Argoverse experiments, epochs to converge, which corresponds to approximately 10 hours on the same resources. In order improve the performance, doubling the size of the dataset. We chose agents that move by at least 10 meters. This, in-turn, increases the training time approximately 24 hours. For the TrajNet, TrajNet+++ and Omniglot datasets, our model takes approximately 100 epochs which corresponds to maximum 2 hours of compute time the same resources.
model takes approximately 30
other agents in the certain scenes
ego-agent, effectively
C.5 PERFORMANCE COMPARING SEED PARAMETER DECODER WITH AUTOREGRESSIVE DECODER
One of the main architectural contributions of our work compared to our contemporaries is the use of seed parameters in the decoder network allow for one-shot prediction of the entire future trajectory (compared to the established way of rolling out the model autoregressively timestep by timestep). To this end, created ablation of our model that predicts the future timesteps autoregressively and we compare its performance to ours across different variables (number of agents, input sequence length, output sequence length, and number of modes). Figure shows the inference rate (i.e. full trajectory predictions per second) of AutoBot-Ego (top) and AutoBot (bottom). We can see that in the single-agent case (AutoBot-Ege ), both the seed parameter version of the model as well as the autoregressive ablation are able to operate in real-time (2 30 FPS), save for when using a prediction horizon longer than 20 steps. We chose 30 FPS real-time indicator line since most cameras used in self-driving cars output at least 30 FPS (Yeong al., 2021); I); Lidar data is usually slower to retrieve 5-20 FPS. The performance difference more obvious when comparing AutoBot to its ablation. The autoregressive model is barely able perform at 30 FPS when the number of agents is low, the number of input steps is very low, the predicted future steps are under 20, and the number of modes low.
17

---

Published as a conference paper at ICLR 2022
AutoBot- Ege (Single-Agent inference)
200
Seed Parameters
Autoreg
150
100
50
00-000-
0000
6000
Agents
Input Length
Hortzon
AutoBot (Multi-Agent Joint Inference)
Framnes Socord inore- betten
1.50
Send Paramaters
Autoreg
50-
100 Number of Agants
200
100
200
100
200
100
200
Inpur Length
Prediction Horizon
Number of Modes
Figure 7: AutoBots's Inference Rate To show the impact of using a seed parameter decoder as opposed to an autoregressive decoder, we compare the inference rate of AutoBots with its autoregressive but otherwise identical ablation. We plot how the inference rate varies with increasing size of input/output variables. From left-to-right, these variables correspond agent, 3) prediction horizon, 4) number of modes. The top row corresponds AutoBot-Fge while bottom row corresponds to AutoBot.
number of agents, 2) number of observable timesteps for each
D AUTOBOTS ADDITIONAL RESULTS
D.1 ABLATION STUDY NUSCENES
Table 5 shows the results of ablating various components from AutoBot-Ego the Nuscenes dataset. As we can see, the model trained without the map information performs significantly worse across all metrics. This shows that indeed AutoBot-Ego learns attend the map information effectively, The 'No Latent Variable' version of our model is a model without the latent variable, and therefore, without the decoder seed parameters. This model then becomes auto-regressive and samples from the output bivariate Gaussian distribution at every timestep. We note that this version of the model significantly slower to train (taking approximately three times as long) due its autoregressive nature. In order to generate ten trajectories, the generation process is repeated = 10 times. This model uses the same transformer encoder and decoder blocks as described in Section As we can from its performance generating diverse enough trajectories. This is expected since the output model Φ is now tasked with accounting for both the aleatoric and epistemic uncertainties in the process. Finally, the 'No Data Augmentation' model shows that although helpful to augment the data with mirroring, even without it the model is still capable of producing strong performance.
the ADE metrics, without latent variables, the model is not capable of
Min ADE
Min ADE (10)
Miss Rate Top-5 (2m)
Miss Rate
Off Road
AutoBot-Ego
(5)
Top-10 (2m)
Rate
No Map Information
1.88
1.35
0.72
0.58
0.22
No Latent Variable
1.77
1.34
0.76
0.64
0.07
No Data Augmentation
1.47
1.08
0.66
0.47
0.03
AutoBot-Ego
1.43
1.05
0.66
0.45
0.03
Table 5: NuScenes Ablation Study..
We study the performance benefits of AutoBot-EEgoover various
ablated counterparts.
18

---

Published as a conference paper at ICLR 2022
D.2 ANALYSIS OF ENTROPY REGULARIZATION
This dataset contains the trajectories of single particle having identical pasts but diverse futures, with different turning rates and speeds, shown in Figure (top row). In the second row of Figure 8, we show that our model trained without entropy regularization, X = 0.0, can cover all modes. However, the short trajectory variants are all represented using only one of the modes (2nd row, 3rd column), as shown by the growth of the uncertainty ellipses as the trajectory moves forward. This shows that AutoBots learn increase the variance of the bivariate Gaussian distribution in order to cover the space of trajectories. The third row of Figure 8, shows the learned trajectories for this variant trained with A = 3.0. As we can see, the entropy regularization penalizes the model if the uncertainty of bivariate Gaussians is too large. By pushing the variance of the output distribution to a low magnitude, we restrict the model such that the only way it can achieve high likelihood is to minimize the distance between the mean of the bivariate distributions and ground-tree trutth-rthhhhu trajectories.
Figure 8: Influence of entropy loss on particle experiment. The input trajec- tory (cyan, only shown in top row) is iden- tical all cases. The model trained with- out an entropy loss term (middle row) cov- ers all modes with high variance, while the model with moderate entropy regu- larization (bottom row) is able to leam dearts modes with low variance and without overlap.
Catee
bage
Desa D
Ablation Study Importance of Entropy Loss Term. Table 6, we show an ablation study to evaluate the importance of the entropy component training the model for nuScenes. Again, observe that for multi-modal predictions, enforcing low entropy on the predicted bivariate distributions has great positive impact on the results, with an increased Åe resulting in better performance. This was observed on all evaluated metrics with the exception of the min FDE, which was poor regardless of the model. Furthermore, we can see from Table that, general, all prior models have difficulty obtaining strong performance on the min FDE (1) metric, with the best value having an overall eπтоr of 8.49 meters. This may be due the difficulty of predicting the correct mode when one only has access to two seconds of the past sequence.
Min ADE
Min ADE (10)
Miss Rate
Miss Rate
Min FDE
Off Road
(5)
Top-5 (2m)
Top-10 (2m)
(1)
Rate
0
1.75
1.29
0.67
0.58
9.55
0.10
1
1.71
1.22
0.65
0.57
9.03
0.08
1.70
1.14
0.62
0.53
9.11
0.05
30
1.75
1.14
0.63
0.54
9.51
0.04
40
1.72
1.11
0.60
0.51
9.01
0.04
Table 6: NuScenes Quantitative Results by Entropy Regularization. We study AutoBot-Egoos performance when varying the entropy regularization term on nuScenes, and see that increased entropy improves model performance (particularly Min ADE 10 and Off Road Rate) up a limit.
D.3 TRAINET+ ABLATION AND QUALITATIVE RESULTS
Table 7 shows the effect of increasing the number of social/temporal encoder and decoder layers in the general AutoBot architecture. We can see that increasing the number of layers in the decoder has positive impact the performance of the joint metrics. This is expected since as we increase the number of social layers, the model can use high-order interactions to model the scene. Figure shows example scene generated by each of the model variants.
D.4 TRAINET, A PEDESTRIAN MOTION PREDICTION DATASET
We further validate AutoBots on pedestrian trajectory prediction dataset TrajNet (Sadeghian et al., 2018) where are provided with the state of all agents for the past eight timesteps and are tasked
19

---

Published as a conference paper at ICLR 2022
Ego Agent's
Number Of
Scene-level
Scene-level Min FDE (1)
Model
Min 6DE (I)
Collisions ())
Min ADE( ())
Lene 1, Ldec 1
0.108
273
0.152
0.281
Lene = 2, Ldee 1
0.102
214
0.138
0.254
Lenc 3, Late
0.099
227
0.135
0.245
Lene = 1, Liec 2
0.099
164
0.132
0.248
Lene = 1, Liter = 3
0.091
123
0.121
0.225
Lone 2, Loder 2
0.095
139
0.128
0.234
Lenc 3, Ldee 3
0.090
98
0.119
0.216
Table 7: TrajNet++ ablation studies for multi-agent forecasting scenario. We investigate the impact of the number social/temporal layers the encoder the decoder.
Bround-Truth Scene
AutoBol Ego
Observed Trajectory
Trajectory Occurence
AuloBo-Anlisoia
Figure 9: TrajNet+++ qualitative results. In this example scene, five agents moving towards each other with their input trajectories (blue) and their ground-truth future trajectories (black) shown in the left figure. The different rows show three of c scene predictions made by each model variant. AutoBot-Fgo and AutoBot-Antise AntiSocial produce some modes containing collisions or socially inconsistont trajectories, while AutoBot, which has social attention in the decoder. achieves consistent trajectories.
with predicting the next 12 timesteps. In this experiment, we apply our general architecture (AutoBot) to multi-agent scene forecast situation, and ablate various the social components in the decoder only, and decoder and encoder.
Number Of Collisions
Scene-level Min ADE(5)
Scene-level Min FDE(5)
Model
AutoBot-AntiSoch
474
0.571
1.02
AutoBot-Ego
466
0.602
1.07
AutoBot
326
0.471
0.833
Table TrajNet ablation studies for multi-agent forecasting scenario. We evaluate the impact of using the extra social data as a set and the actual multi-agent modelling on this context. We found that AutoBot is able to cause fewer collisions between agents compared to variants without the social attention.
In Table present three variants of the model. The first, AutoBot-Antiociall corresponds to a model without the intra-set attention functions in the encoder and decoder (see Figure 2). The second variant, AutoBot-Fgo, only sees the past social context and has social attention during
20

---

Published as a conference paper at ICLR 2022
the trajectory generation process. The third model corresponds to the general architecture. AutoBot. All models were trained with c = 5 and we report scene-level error metrics as defined by (Casas et al., 2020). Using the social data with the multi agent modelling in the decoder greatly reduces the number of collisions between predicted agent futures. Furthermore, we can see that AutoBot significantly outperforms its ablated counterparts on generating scenes that match ground-truth, as observed by the reduced scene-level minADE and minFDE. We expect this theoretically since the ego-centric formulation makes the independence assumption that the scene's future evolution can be decomposed as the product of the individual agent's 's future motion in isolation, as remarked in (Casas et al., 2020). AutoBot does not make this assumption as we condition between agents every timestep.
Figure 10 shows an example of multi-agent trajectory generation using the three different approaches for predicting the scene evolution. We present four of the five modes generated by the model for this example input scene. One can observe that the only model achieving scene-consistent weneseetooneeeeeeeeeeeee predictions across all modes is AutoBot. Interestingly, we also observe that the different modes correspond to different directions of motion, highlighting the utility using modes, In fact, the modes AutoBot condition the entire scene future evolution, producing alternative realistic multi-agent futures. We refer the reader to Figures 11, 12, 13 and 14 where we show some additional qualitative results of AutoBots' variants compared to AutoBot effectiveness of the sequential set attention mechanism in the decoding process of our method and how important it to predict all agents in parallel compared a 1-by-1 setup.
the TrajNet dataset. These results highlight the
Ground
AutaBal-Ego (No multi agent torecast)
AuloEol-AniBosee the social attention)
AutoBol (mufti agent torecasting)
Truth
Figure 10: TrajNet qualitative results. In this example, we see two agents moving together in past (cyan) and future (pink, left). We compare how different variants of AutoBots make predictions in this social situation. AutoBot-Ego Ego and while AutoBot, which has intra-set attention during decoding, achieves consistent trajectories.
produce some modes containing collisions unrealistic trajectories.
D.5 OMNIGLOT ADDITIONAL DETAILS AND RESULTS
The LSTM baseline used in our Omniglot experiments have a hidden size of 128 and use e latent modes. The input is first projected into this space using an embedding layer, which is identical to the one used in AutoBot. All input sequences of sets are encoded independently using a shared L.TTM encoder. During the decoding phase, we autoregressively generate the next stroke step using an LSTM decoder. In order to ensure consistency between the predicted future strokes, inspired by Social LSTM (Alahi al., 2016), the LSTM baseline employs an MHSA layer at every timestep the character. We concatenate
operating on the hidden state of the different strokes making transformation of the one-hot vector representing the latent mode with the socially encoded hidden state at every timestep. The output model is identical to the one used AutoBot, generating sequence of bivariate Gaussian distributions for each stroke. We performed a hyperparameter search on the learning rate and found the optimal learning rate to be 5e - 4. Furthermore, as with all our other experiments, we found helpful to employ gradient clipping during training.
We provide additional results on the two tasks defined in Section 4.4. Figure 15 shows additional successful AutoBot results on task (completing multiple strokes) compared to an LSTM baseline equipped with social attention. These results highlight the effectiveness of sequential set transformers for generating consistent diverse futures given an input sequence of sets. Figure 16 shows examples where both models fail to correctly complete the character in task 1. Figure 17 compares AutoBot-Egr with the LSTM baseline on predicting a new stroke given the two first strokes of character (task 2). These results highlight that although not all modes predict the correct character across all modes, all
21

---

Published as a conference paper at ICLR 2022
Ground Truth
AutoBot-Ego
AutoBat-Antisoc
AutoBot
Input Traj
Output Traj (GT)
Output Traj (Pred.)
Figure 11: TrajNet qualitative results 1/4. Example scene with multiple agents moving together. These trajectories are plotted over the birds-eye-view image of the scene where we zoom into interesting trajectories. We that only AutoBot produces trajectories that are realistic in group setting across all modes.
Ground Truth
AutoBot-Ego
AutoBat-Antisoc
AutoBot
Input Traj
Output Traj (GT)
Output Traj (Pred.)
with two agents moving separately in road setting. We want to highlight this interesting scenario where some modes of AutoBots- and AutoBots-AntiSocial results trajectories that lead into the road, while AutoBot seems produces trajectories more in line with the ground-truth, and lead the agent to cross the road safely.
Figure 12: Tea qualitative results 2/4. Example
generated predictions consistent with realistic characters (i.e., confusing "F" with "II" or "H" with "II").
22

---

Published as a conference paper at ICLR 2022
Ground Truth
AutoBot-Ego
AutoBot-Antisoc
AutoBot
Input
Output Traj (GT)
Output Traj (Pred.)
Figure 13: TrajNet qualitative results 3/4. Additional example scenes with multiple agents moving together. Again, we wish to highlight the advantage of modelling the scene jointly, which is evident by the results of AutoBot.
Ground Truth
AutoBot-Ego
AutoBot-Antisoc
AutoBot
Input Traj
Output Traj (GT)
Output Traj (Pred.)
Figure 14: TrajNet qualitative results 4/4. Example scenes with two agents moving together. Again, we see that AutoBot produces trajectories consistent with the scene across all modes (e.g., not crashing into the bushes) and maintains the social distance between the walking agents.
23

---

Published as a conference paper at ICLR 2022
LSTM
MAR
AutoBot
1
LSTM
3/ 4
3/
2
AutoBot
LSTM
π
π
π
π
π
π
AutoBot
LSTM
A
A
A
X
A
A
A
AAR
A
AutoBot
7
LSTM
ita
Fta
7t
76
AutoBot
K
x
H
LSTM
J
H
AutoBot
Figure 15: Omniglot Task 1 Additional Results. These are some additional random characters from the Omniglot stroke completion task. We can see that AutoBot produces plausible characters on the test set, while the LSTM struggles to create legible ones.
24

---

Published as a conference paper at ICLR 2022
5
LSTM
T
9
9
D
AutoBot
LSTM
G
6
C
2 3
E
7
5
3
2
2
2
2
AutoBot
LSTM
b
b
A
4
A
F
1
8
AutoBot
Figure 16: Omniglot Task 1 Failure Cases. There were some characters where AutoBot failed to learn the correct character completion. We can currently only speculate why that is the case. Our first intuition was that this might occur when there are 90 degree or less angles in the trajectory that is to be predicted but in Fig. 15, we can see that there are examples where this does not seem to be problem. We believe that this might be due to a rarity of specific trajectories in the dataset but more investigation is required.
25

---

Published as a conference paper at ICLR 2022
LSTM
AutoBot
LSTM
AutoBot
J
LSTM
J
T
AutoBot
LSTM
T
T
π
T
π
T
π
π
π
π
AutoBot
T
LSTM
F
F
F
AutoBot
LSTM
7
7
7
7
7
AutoBot
Omniglot Task 2 Additional Results. These are some additional random characters from the Omniglot character completion task. Again, we can see that AutoBot produces plausible characters on test set, where different modes capture plausible variations (e.g. F to H and H to PI), while the LSTM struggles to capture valid characters in its discrete latents.
Figure 17:
26


Reviews:
[{"review_id": "RAf_y1rO4zf", "reviewer": "Reviewer_rcqg", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Pros:\n+ This proposed model is simple and effective. It applies multi-head attention modules on temporal and social axis both in encoder and decoder to capture motion and interaction among agents.\n+ The proposed model has been tested with comprehensive experiments, including on toy examples, real-world dataset such as nuScenes, Argoverse and TrajNet, as well as an image generation dataset (Omniglot). It shows comparable results compared to previous state-of-the-art models.\n+ The latent variable is used to represent the discrete future for each agents. In spirit, it is same as Multiple Future Prediciton (Tang et al 2019), however, this paper adapts it into the transformer model and make it work as well.\n+ The training of this model is very fast, as pointed out on page 6, it consumes 0.4% computation of SceneTransformer and 5% of Jean's method on Argoverse. This shows the efficiency of the proposed model.\n+ Code is provided and it's nice and clean.\n\nCons:\n- It's related to Axial Attention in Multidimensional Transformers, Jonathan Ho et al, 2019, and should cite that work.\n- In the provided code, the decoder is autoregressive instead of what is proposed in the main paper. It should be updated accordingly, and more ablation/discussion on these two type of decoder would be very helpful.\n- Minor point: for equation (5) on L(\\theta) = Q + const, the const should be KL divergence and is not a const w.r.t. \\theta.\n", "summary_of_the_paper": "This paper adapts transformer to multi-agent motion forecasting. The attention layers are applied on time and agent axis to capture motion and social information. A latent variable is introduced on the output to capture discrete motion for each agent.\nExtensive experiments are conducted on various dataset with good performance. The training time of the proposed model is significantly faster than previous methods.", "summary_of_the_review": "This paper compiles an interesting model for multi-agent motion forecasting. Although the transformer module is similar to Axial Transformer [Ho 2019] and the latent variable formulation and optimization are similar to MFP [Tang 2019], it is the first to successfully apply those two module together for motion forecasting task, with clean and concise code released. Thus I would vote for weak accept.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "IYvdcy4cF1C", "reviewer": "Reviewer_DomP", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "(+): Extensive experiments and results: The authors conducted experiments on 4 datasets, covering 3 different tasks (argo and nuscenes are both prediction for AV). These are very extensive. I do recognize that the method generalizes well to multiple different tasks.\n\n(-) Novelty: I think the novelty should be good if it wasn't for the SceneTransformer paper. According to ICLR review guidelines, Arxiv papers published 30 days prior to the ICLR deadline are considered prior work. The SceneTransformer's formulation is very similar to that proposed in this paper, e.g. using transformers to model cross-agent and cross-time information propagation. Although SceneTransformer does not have an object query that is the latent variable Z, it was still able to learn a diverse set of future predictions. It's unclear if the explicit formulation of latent variable makes a difference. It seems that on Argoverse, this paper performs very similarly than SceneTransformer, which further confirms my view that the methods are very similar.", "summary_of_the_paper": "The paper tackles the multi-agent trajectory prediction problem, primarily for autonomous driving, but experiments also include TrajNet and predicting Omniglot strokes. The authors claim their contributions in a very general sense:\n\n- Novel method on modeling sequences of structured continuous variables and capture multi-modal distributions.\n- Strong results on nuscenes, argoverse, trajnet and omniglot stroke prediction.\n\nThe key modeling novelty is the latent variable sequential set transformer which applies self-attention across different agents and across time in the scene. \n\nI mostly agree with the authors' assessment in general. I have some concerns on the claimed novelty (see below).", "summary_of_the_review": "I am giving a borderline reject for now due to the novelty concerns in light of a very similar paper in my opinion. I know that there are a lot of devils in the detail between the two works, and it would have been an obvious accept if this works outperforms SceneTransformer significantly.\n\nIn the rebuttal, may I ask the authors to further highlight the difference between their approach to SceneTransformer? And where do they see that their method is intuitively better than SceneTransformer?\n\nPOST REBUTTAL: The authors pointed out that SceneTransformer should be considered concurrent work. If that is the case, I have no issues to accept this paper.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "4: All of the claims and statements are well-supported and correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "FnwoW-4wjL", "reviewer": "Reviewer_DDxv", "summary": "", "questions": "", "limitations": "", "rating": 8, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Paper Strengths\n\n1. The proposed method using the transformer to model both social and temporal information in both the encoder and decoder makes sense. Also, the ability to generate scene-consistent predictions are verified with both qualitative and quantitative results on TrajNet\n\n2. The computation requirement of the proposed method is very appealing compared to other transformer-based motion prediction methods, as explained in section 4.2. With experience in running other similar methods, training usually takes 2-3 days on top-performing GPUs, whereas the proposed method only requires training on 1080ti for 10 hours. \n\n3. The experiments are extensive, where the proposed method has been evaluated on three motion prediction datasets and one stroke completion dataset\n\nPaper Weaknesses\n\n1. Why is the proposed method very fast compared to others? Although a summary of computing requirements is provided, it would be interesting to perform a deeper analysis as to why the proposed transformer-based method is faster than other transformer-based methods, smaller feature dimensions? The fewer number of blocks? This could be helpful to the community in the future to develop efficient and high-performance transformer-based prediction methods\n\n2. The ablation experiment can be expanded. Currently, the only ablation experiment is done in Table 3 on TrajNet with two variations of removing social modeling. However, the TrajNet challenge is a bit outdated and data annotation is imprecise. Would be more convincing to perform the ablation on the nuScenes or Argoverse dataset. Also, there are other components of the proposed method that are not ablated. For example, how is the proposed transformer-based approach compared to RNN-based or Graph-based methods? How is the map affecting the performance? How are the seed parameters useful? How many time/social blocks are needed?\n\n3. The novelty is a bit weak considering other works have similar transformer-based motion prediction methods, especially [A1] as mentioned in section 5. Even ignoring this “concurrent” work, other works including [A2, A3] also use stacked transformers to do both social and temporal modeling. It is incorrect to me that the paper claims [A2, A3] did not encode the time and social dimensions using transformers. It would be great if the differences with these works can be explained in a more detailed and precise way\n\n4. The experiments on TrajNet and Omniglot are not very convincing, where no quantitative comparison is properly made with prior work. For example, on the TrajNet challenge, there are plenty of prior motion prediction methods that have been evaluated on this dataset, so it should be easy to have a table for comparison. For stroke completion, I am not very familiar with this dataset, but only comparing with a simple LSTM baseline is not convincing. Including experiments on more datasets is great, but would be good to show convincing experiments. \n\nReferences\n\n[A1] AgentFormer: Agent-Aware Transformers for Socio-Temporal Multi-Agent Forecasting. ICCV 2021\n\n[A2] Spatio-Temporal Graph Transformer Networks for Pedestrian Trajectory Prediction. ECCV 2020\n\n[A3] Multimodal Motion Prediction with Stacked Transformers. CVPR 2021\n\n\nPost-rebuttal review\n\nAfter carefully reading other reviews and the authors’ comments, I would like to increase my rating to a score of 8 -- accept good paper (or a score of 7 if there is such an option). The change to my score is mainly because I am satisfied with most of the authors’ comments to my concerns, which include: 1) the clarification on why the proposed method is much faster; 2) the expanded ablation experiments; 4) the clarification for experiments on Omnlglot. I strongly appreciate the authors' efforts on clarifying these questions and concerns!\n\nAs a result, the main comment left on my side is 3) the limited novelty compared to prior work. This does not mean that I find the authors’ reply unconvincing. In fact, it is quite clear and honest. But essentially, the main difference to prior work seems to be just the use of seed parameters that increases performance and also runtime speed, since other points are already proposed in prior work, including joint social-temporal modeling in encoders/decoders using transformers, and scene-consistent predictions. But I would like to emphasize that this relatively incremental novelty is fine to me, and personally, I would like to have a try on the seeing parameters proposed in this work. So it would be wonderful for the community if the code can be carefully documented and released (currently I did not find it in the supp or any anonymous URL). Finally, as a minor point, it would be great if the map or context image can be overlaid on the trajectories for the provided video results. \n\nIn short, if there is no other significant downside/concern pointed by other reviewers, I would like to stick to my rating above and recommend accepting this paper\n\n", "summary_of_the_paper": "This paper proposes a transformer-based VAE model for motion prediction that can output multi-model and scene-consistent predictions. Specifically, the transformer is employed for modeling both social and temporal information (but separately, first temporal and then social and repeat). The proposed method achieves state-of-the-art performance on the nuScenes dataset and top performance on the Argoverse dataset", "summary_of_the_review": "Novelty is slightly weak but empirical performance is strong and the method might be practically useful", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "9wsV5qAPiX4", "reviewer": "Reviewer_HbLp", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths: The paper proposes a solid architecture by considering the requirements of the problem. The proposed architecture relies on Transformers and the self-attention function which is an effective tool in modeling relations between different entities. Trajectory forecasting, especially in a multi-agent setting, is highly stochastic and may result in many feasible predictions. The paper incorporates a set of learnable seed parameters into the Transformer architecture to enable multi-moded prediction where the seeds are expected to lead to different modes. The model is evaluated in established benchmarks and achieves competitive performance with a lower computational requirement.\n\nWeaknesses and Questions: The main contributions are the set assumption and the discrete latent variable. Unfortunately, the paper does not provide empirical evidence for these design choices. Ablations evaluating the proposed approaches could be insightful. \n\nCould the authors clarify why the proposed set assumption is feasible for multi-agent tasks? I can see the advantage of this design choice if the set of agents is dynamic (i.e., gets updated in time). However, this scenario is not covered in the experimental setup. It would be useful to see if the set assumption contributes to the performance. \n\nIt wasn’t clear to me how the discrete latent variable is modeled. It is partly due to the cluttered notation (see the next comment). I kindly ask the authors to provide more details. An ablation evaluating the proposed model without using the discrete latent variable could be useful. I also would like to point out alternatives in case the authors are not aware of this line of work. Latent variables in sequence models have been studied [1]. While it is reported to be effective, there are more straightforward approaches to get multi-moded predictions [2]. Have the authors experimented with continuous latent variables or mixture density outputs for multi-modal predictions?\n\nI find the notation not rigorous. Several symbols are used in multiple definitions, making the paper hard to follow. For example, \n\n1- $Z$: \nThe discrete latent variable appears in the objective function in section 3.3 with the symbol $Z$. It is also defined to be the seed as “learnable seed parameters $Z \\in \\mathcal{R}^{(d_K; M; T)}$” in section 3.2. Are they the same? \n\n2- $K$:\nSection 3.1: “... having M elements (or agents) with K attributes (e.g. x/y position).”  \nSection 4.1: “... AutoBot-Ego trained with K = 10.” or “AutoBot-Ego obtains the best overall Min ADE out of 5 predictions (K = 5).” \n\nI have the following questions:\n\n3- Could the authors clarify this statement in section 4.1:  “In order to generate ten predictions from five modes, at test time, we create combinations between the latent modes one-hot vector by summing one-hot vectors.”\n4- What is the difference between the “AutoBot-Ego (K=10)” and “AutoBot-Ego (ensemble)” entries in Tab. 1?\n\nFinally, [3] seems to be relevant and should be included in the related work discussion if the authors also agree. \n\n[1] Chung, Junyoung, et al. \"A recurrent latent variable model for sequential data.\" Advances in neural information processing systems 28 (2015): 2980-2988.\n\n[2] Graves, Alex. \"Generating sequences with recurrent neural networks.\" arXiv preprint arXiv:1308.0850 (2013).\n\n[3] Li, Jiachen, et al. \"Evolvegraph: Multi-agent trajectory prediction with dynamic relational reasoning.\" arXiv preprint arXiv:2003.13924 (2020).\n\n-- Post-rebuttal edit --\nI thank the authors for the clarifications. I read other reviews and the responses. I am glad to see that the authors addressed the questions and provided additional experimental evidence. The ablations improve the quality of the submission. In fact, the latent variable ablation can be included in the main submission. Although the technical novelty is slightly limited, the proposed approach is solid. I increase my score to 6. ", "summary_of_the_paper": "This submission presents a Transformer-based architecture for trajectory prediction tasks involving multiple agents. Such a multi-agent setting requires learning of spatio-temporal representations capturing both the long-term temporal dependencies as well as the social interactions between the agents. The paper formulates the task as modeling of the sequence of sets where every entry in the set corresponds to an agent’s observation. The proposed architecture augments the set transformer with a discrete latent variable to be able to make multiple predictions into the future. The given seed sequence is first encoded into a context representation which is later used to make future predictions with a decoder where the agents can be modeled jointly or independent from each other. ", "summary_of_the_review": "The proposed architecture is tailored for the task of multi-agent trajectory prediction task and presented to be effective in the benchmarks. However, we do not know if it is due to the superior performance of the underlying computation units (i.e., transformer blocks) or the proposed extensions (i.e., discrete latent variable/learnable seed, set assumption). I also think that the presentation could be improved. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper studies the problem of motion prediction for multiple agents in a scene using transformer-based VAE like architecture. The paper received mixed reviews initially which generally tended towards borderline acceptance. All reviews appreciated extensive experiments but had some clarifications and requests for ablations. The authors provided a strong rebuttal that addressed many of the reviewers' concerns. The paper was discussed and all the reviewers updated their reviews in the post-rebuttal phase. Reviewers unanimously agree that the paper should be accepted. AC agrees with the reviewers and suggests strong acceptance. The authors are urged to incorporate reviewers' comments in the camera-ready.
Decision: Accept (Spotlight)