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: 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.