Latent Variable Sequential Set Transformers for Joint Multi-Agent Motion Prediction

1 INTRODUCTION . Many problems require processing complicated hierarchical compositions of sequences and sets . For example , multiple choice questions are commonly presented as a 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 a variety of different motion forecasting tasks . Suppose we have a set of M ∈ N sequences X = { ( x1 , . . . , xK ) 1 , . . . , ( x1 , . . . , xK ) M } each withK ∈ N elements . Allowing also that X evolves across some time horizon T , we denote this sequence of sets as X = ( X1 , . . . , XT ) . 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 t steps of context . In our approach here , to better model the multi-modal distribution over possible futures , we also allow to include a discrete latent variable Z , to create sequential models of sets of the form P ( Xt+1|Xt , . . . , X1 ) = ∑ Z P ( Xt+1 , Z|Xt , . . . , X1 ) . 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 to as “ AutoBot ” ) . We evaluate AutoBot on nuScenes ( Caesar et al. , 2020 ) and Argoverse ( Chang et al. , 2019 ) , two autonomous driving trajectory prediction benchmarks , on the synthetic partition of the TrajNet++ ( Kothari et al. , 2021 ) dataset , a pedestrian trajectory forecasting benchmark , and Omniglot ( Lake et al. , 2015 ) , a dataset of hand-drawn characters that we use for predicting strokes . All Figure 1 : Overview : Our model generates multiple possible futures for agents based on several input timesteps . The input trajectories are encoded into a context tensor that captures their respective behavior and interaction with surrounding agents . While Transformer-based architectures are often autoregressive , we decode all future steps at once . This is accomplished via learnable seed parameters . Our decoder transforms these , together with the map and the context into the possible future trajectories that we call “ modes ” . of these provide empirical validation of our model ’ s ability to perform robust multi-modal prediction . Specifically , we make the following contributions : • We propose a novel approach for modelling sequences of set structured continuous variables , and extend this to a latent-variable formulation to capture multi-modal distributions . • We provide a theoretical analysis highlighting our model ’ s permutation sensitivity with respect to different components of the input data structures . • We validate our method with strong empirical results on diverse tasks , including the largescale nuScenes and Argoverse datasets for autonomous driving ( Caesar et al. , 2020 ; Chang et al. , 2019 ) , the multi-lingual Omniglot dataset of handwritten characters ( Lake et al. , 2015 ) , and the synthetic partition of the TrajNet++ pedestrian trajectory prediction task ( Sadeghian et al. , 2018 ) . With this submission , we included an implementation of our method in a Jupyter notebook for easy reproduction and to demonstrate that the architecture is performant enough to train on a desktop GPU . 2 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 a query is executed against a key-value database , returning values where the key matches the query best . While ( Vaswani et al. , 2017 ) defines MHSA on three tensors , for convenience , we input a 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 a residual connection , MHSA ( X ) = X+MHSA ( X , X , X ) . Multi-Head Attention Blocks ( MAB ) resemble the encoder proposed in 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 . 8 ( appendix ) followed by a row-wise feed-forward neural network ( rFFN ) , with residual connections and layer normalization ( LN ) ( Ba et al. , 2016 ) after each block . Given an input set X containing nx elements of dimension d , and some conditioning variables C containing nc elements of dimension d , the MAB can be described by the forward computation , MAB ( X ) = LN ( H + rFFN ( H ) ) , where H = LN ( X + MHSA ( X , X , X ) ) . ( 1 ) Multi-head Attention Block Decoders ( MABD ) were also introduced in Vaswani et al . ( 2017 ) , and were used to produce decoded sequences . Given input matrices X and Y representing sequences , the decoder block performs the following computations : MABD ( Y , X ) = LN ( H + rFFN ( H ) ) where H = LN ( H′ + MHSA ( H′ , X , X ) ) and H′ = LN ( MHSA ( Y ) ) ( 2 ) Our model , described throughout the following section , makes extensive use of these functions . 3 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 a future sequence for only one element in a scene . 3.1 ENCODER : INPUT SEQUENCE REPRESENTATION . AutoBot takes as an input a sequence of sets , X1 : t = ( X1 , . . . , Xt ) , which in the motion prediction setting may be viewed as the state of a scene evolving over t 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 a 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 et al. , 2017 ) . For this step , we view the data as a collection of matrices , { X0 , . . . , XM } , that describe the temporal evolution of agents . The encoder processes temporal relationships between sets using a MAB , performing the following operation on each set where m ∈ { 1 , . . . , M } : Sm = MAB ( rFFN ( Xm ) ) where rFFN ( . ) is an embedding layer projecting each element in the set to the size of the hidden state dK . Next we process temporal slices of S , retrieving sets of agent states Sτ at some timestep τ and is processed using Sτ = MAB ( Sτ ) . These two operations are repeated Lenc times to produce a context tensor C ∈ R ( dK , M , t ) summarizing the entire input scene , where t 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-modal data distributions . To generate c ∈ N different predictions for the same input scene , the AutoBot decoder employs c matrices of learnable seed parameters Qi ∈ R ( dK , T ) where T is the prediction horizon , where i ∈ { 1 , ... , c } . Intuitively , each matrix of learnable seed parameters corresponds to a setting of the discrete latent variable in AutoBot . Each learnable matrix Qi is then repeated across the agent dimension M times to produce the input tensor Q′i ∈ R ( dK , M , T ) Additional contextual information ( such as a rasterized image of the environment ) is encoded using a convolutional neural network to produce a vector of features mi 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 T dimensions , producing the tensor Mi ∈ R ( dK , M , T ) . Each tensor Q′i is concatenated with Mi along the dK dimension , as shown on the left side of Figure 2 ( bottom ) . This tensor is then processed using a rFFN ( . ) to produce the tensor H ∈ R ( dK , M , T ) . We begin to decode by processing the time dimension , conditioning on the encoder ’ s output C , and the encoded seed parameters and environmental information in H. The decoder processes each of m agents in H separately with an MABD layer : H′m = MABD ( Hm , 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 is socially consistent between set elements , we process each temporal slice of H′ , retrieving sets of agent states H′τ at some future timestep τ where each element h′τ ∈ 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 Ldec times to produce a final output tensor for mode i . The decoding process is repeated c times with different learnable seed parameters Qi and additional contextual information mi . The output of the decoder is a tensor O ∈ R ( dK , M , T , c ) which can then be processed element-wise using a neural network φ ( . ) to produce the desired output representation . In several of our experiments , we generate trajectories in ( x , y ) space , and as 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 a dual purpose . The first is that they account for the diversity in future prediction , where each matrix Qi ∈ R ( dk , T ) ( where i ∈ { 1 , ... , c } ) corresponds to one 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 a 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 ) ) resort to having a single output model generate the future of each agent independently . 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 .

Summary: 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:  - Novel method on modeling sequences of structured continuous variables and capture multi-modal distributions. - Strong results on nuscenes, argoverse, trajnet and omniglot stroke prediction.  The key modeling novelty is the latent variable sequential set transformer which applies self-attention across different agents and across time in the scene.   I mostly agree with the authors' assessment in general. I have some concerns on the claimed novelty (see below).