baller2vec: A Multi-Entity Transformer For Multi-Agent Spatiotemporal Modeling

Multi-agent spatiotemporal modeling is a challenging task from both an algorithmic1 design and computational complexity perspective . Recent work has explored the2 efficacy of traditional deep sequential models in this domain , but these architectures3 are slow and cumbersome to train , particularly as model size increases . Further,4 prior attempts to model interactions between agents across time have limitations,5 such as imposing an order on the agents , or making assumptions about their relation-6 ships . In this paper , we introduce baller2vec1 , a multi-entity generalization of7 the standard Transformer that can , with minimal assumptions , simultaneously and8 efficiently integrate information across entities and time . We test the effectiveness9 of baller2vec for multi-agent spatiotemporal modeling by training it to perform10 two different basketball-related tasks : ( 1 ) simultaneously modeling the trajectories11 of all players on the court and ( 2 ) modeling the trajectory of the ball . Not only does12 baller2vec learn to perform these tasks well ( outperforming a graph recurrent13 neural network with a similar number of parameters by a wide margin ) , it also14 appears to “ understand ” the game of basketball , encoding idiosyncratic qualities of15 players in its embeddings , and performing basketball-relevant functions with its16 attention heads.17 1 Introduction18 Whether it is a defender anticipating where the point guard will make a pass in a game of basketball,19 a marketing professional guessing the next trending topic on a social media platform , or a theme20 park manager forecasting the flow of visitor traffic , humans frequently attempt to predict phenomena21 arising from processes involving multiple entities interacting through time . When designing learning22 algorithms to perform such tasks , researchers face two main challenges:23 1 . Given that entities lack a natural ordering , how do you effectively model interactions between24 entities across time ? 25 2 . How do you efficiently learn from the large , high-dimensional inputs inherent to such26 sequential data ? 27 Prior work in athlete trajectory modeling , a widely studied application of multi-agent spatiotemporal28 modeling ( MASM ; where entities are agents moving through space ) , has attempted to model player29 interactions through “ role-alignment ” preprocessing steps ( i.e. , imposing an order on the players ) 30 [ 1 , 2 ] or graph neural networks [ 3 ] , but these approaches may destroy identity information in the31 former case ( see Section 4.2 ) or limit personalization in the latter case ( see Section 5.1 ) . Recently,32 researchers have experimented with variational recurrent neural networks ( VRNNs ) [ 4 ] to model33 1All data and code for the paper are available at : < anonymized > . Submitted to 35th Conference on Neural Information Processing Systems ( NeurIPS 2021 ) . Do not distribute . the temporal aspects of player trajectory data [ 3 , 2 ] , but the inherently sequential design of this34 architecture limits the size of models that can feasibly be trained in experiments.35 Transformers [ 5 ] were designed to circumvent the computational constraints imposed by other36 sequential models , and they have achieved state-of-the-art results in a wide variety of sequence37 learning tasks , both in natural language processing ( NLP ) , e.g. , GPT-3 [ 6 ] , and computer vision , e.g.,38 Vision Transformers [ 7 ] . While Transformers have successfully been applied to static multi-entity39 data , e.g. , graphs [ 8 ] , the only published work we are aware of that attempts to model multi-entity40 sequential data with Transformers uses four different Transformers to separately process information41 temporally and spatially before merging the sub-Transformer outputs [ 9 ] .42 t = 4 t = 8 t = 13 Figure 1 : After solely being trained to model the trajectory of the ball ( ) given the locations of the players and the ball on the court through time , a self-attention ( SA ) head in baller2vec learned to anticipate passes . When the ball handler ( ) is driving towards the basket at t = 4 , SA assigns near-zero weights ( black ) to all players , suggesting no passes will be made . Indeed , the ball handler did not pass and dribbled into the lane ( t = 8 ) . SA then assigns a high weight ( white ) to a teammate ( ) , which correctly identifies the recipient of the pass at t = 13 . In this paper , we introduce a multi-43 entity Transformer that , with minimal44 assumptions , is capable of simultane-45 ously integrating information across46 agents and time , which gives it pow-47 erful representational capabilities . We48 adapt the original Transformer archi-49 tecture to suit multi-entity sequen-50 tial data by converting the standard51 self-attention mask matrix used in52 NLP tasks into a novel self-attention53 mask tensor . To test the effectiveness54 of our multi-entity Transformer for55 MASM , we train it to perform two dif-56 ferent basketball-related tasks ( hence57 the name baller2vec ) : ( 1 ) simulta-58 neously modeling the trajectories of59 all players on the court ( Task P ) and60 ( 2 ) modeling the trajectory of the ball61 ( Task B ) . Further , we convert these tasks into classification problems by binning the Euclidean62 trajectory space , which allows baller2vec to learn complex , multimodal trajectory distributions via63 strictly maximizing the likelihood of the data ( in contrast to variational approaches , which maximize64 the evidence lower bound and thus require priors over the latent variables ) . We find that:65 1. baller2vec is an effective learning algorithm for MASM , obtaining a perplexity of 1.64 on66 Task P ( compared to 15.72 when simply using the label distribution from the training set ) and67 13.44 on Task B ( vs. 316.05 ) ( Section 4.1 ) . Further , compared to a graph recurrent neural68 network ( GRNN ) with similar capacity , baller2vec is ∼3.8 times faster and achieves a69 10.5 % lower average negative log-likelihood ( NLL ) on Task P ( Section 4.1 ) .70 2. baller2vec demonstrably integrates information across both agents and time to achieve71 these results , as evidenced by ablation experiments ( Section 4.2 ) .72 3 . The identity embeddings learned by baller2vec capture idiosyncratic qualities of players,73 indicative of the model ’ s deep personalization capabilities ( Section 4.3 ) .74 4. baller2vec ’ s trajectory bin distributions depend on both the historical and current context75 ( Section 4.4 ) , and several attention heads appear to perform different basketball-relevant76 functions ( Figure 1 ; Section 4.5 ) , which suggests the model learned to “ understand ” the77 sport.78 2 Methods79 2.1 Multi-entity sequences80 Let A = { 1 , 2 , . . . , B } be a set indexing B entities and P = { p1 , p2 , . . . , pK } ⊂ A be the K entities81 involved in a particular sequence . Further , let Zt = { zt,1 , zt,2 , . . . , zt , K } be an unordered set of K82 feature vectors such that zt , k is the feature vector at time step t for entity pk . Z = ( Z1 , Z2 , . . . , ZT ) is83 thus an ordered sequence of sets of feature vectors over T time steps . When K = 1 , Z is a sequence84 of individual feature vectors , which is the underlying data structure for many NLP problems.85 We now consider two different tasks : ( 1 ) sequential entity labeling , where each entity has its own label86 at each time step ( which is conceptually similar to word-level language modeling ) , and ( 2 ) sequential87 labeling , where each time step has a single label ( see Figure 3 ) . For ( 1 ) , let V = ( V1 , V2 , . . . , VT ) be88 a sequence of sets of labels corresponding to Z such that Vt = { vt,1 , vt,2 , . . . , vt , K } and vt , k is the89 label at time step t for the entity indexed by k. For ( 2 ) , let W = ( w1 , w2 , . . . , wT ) be a sequence of90 labels corresponding to Z where wt is the label at time step t. The goal is then to learn a function f91 that maps a set of entities and their time-dependent feature vectors Z to a probability distribution92 over either ( 1 ) the entities ’ time-dependent labels V or ( 2 ) the sequence of labels W .93 2.2 Multi-agent spatiotemporal modeling94 In the MASM setting , P is a set of K different agents and95 Ct = { ( xt,1 , yt,1 ) , ( xt,2 , yt,2 ) , . . . , ( xt , K , yt , K ) } is an un-96 ordered set of K coordinate pairs such that ( xt , k , yt , k ) are97 the coordinates for agent pk at time step t. The ordered98 sequence of sets of coordinates C = ( C1 , C2 , . . . , CT ) ,99 together with P , thus defines the trajectories for the100 K agents over T time steps . We then define zt , k as:101 zt , k = g ( [ e ( pk ) , xt , k , yt , k , ht , k ] ) , where g is a multilayer102 perceptron ( MLP ) , e is an agent embedding layer , and ht , k103 is a vector of optional contextual features for agent pk at104 time step t. The trajectory for agent pk at time step t is105 defined as ( xt+1 , k−xt , k , yt+1 , k−yt , k ) . Similar to Zheng106 et al . [ 10 ] , to fully capture the multimodal nature of the107 trajectory distributions , we binned the 2D Euclidean space108 into an n× n grid ( Figure 2 ) and treated the problem as a classification task . Therefore , Z has a cor-109 responding sequence of sets of trajectory labels ( i.e. , vt , k = Bin ( ∆xt , k , ∆yt , k ) , so vt , k is an integer110 from one to n2 ) , and the loss for each sample in Task P is : L = ∑T t=1 ∑K k=1− ln ( f ( Z ) t , k [ vt , k ] ) ,111 where f ( Z ) t , k [ vt , k ] is the probability assigned to the trajectory label for agent pk at time step t by f ; 112 i.e. , the loss is the NLL of the data according to the model.113 For Task B , the loss for each sample is : L = ∑T t=1− ln ( f ( Z ) t [ wt ] ) , where f ( Z ) t [ wt ] is the114 probability assigned to the trajectory label for the ball at time step t by f , and the labels correspond115 to a binned 3D Euclidean space ( i.e. , wt = Bin ( ∆xt , ∆yt , ∆zt ) , so wt is an integer from one to n3 ) .116 2.3 The multi-entity Transformer117 We now describe our multi-entity118 Transformer , baller2vec ( Figure 3 ) .119 For NLP tasks , the Transformer self-120 attention mask M takes the form of a121 T × T matrix ( Figure 4 ) where T is122 the length of the sequence . The ele-123 ment at Mt1 , t2 thus indicates whether124 or not the model can “ look ” at time125 step t2 when processing time step126 t1 . Here , we generalize the stan-127 dard Transformer to the multi-entity128 setting by employing a T × K ×129 T × K mask tensor where element130 Mt1 , k1 , t2 , k2 indicates whether or not131 the model can “ look ” at agent pk2 at132 time step t2 when processing agent133 pk1 at time step t1 . Here , we mask134 all elements where t2 > t1 and leave135 all remaining elements unmasked , i.e.,136 baller2vec is a “ causal ” model.137 In practice , to be compatible with Transformer implementations in major deep learning libraries,138 we reshape M into a TK × TK matrix ( Figure 4 ) , and the input to the Transformer is a matrix139 with shape TK × F where F is the dimension of each zt , k . Irie et al . [ 11 ] observed that positional140 encoding [ 5 ] is not only unnecessary , but detrimental for Transformers that use a causal attention141 mask , so we do not use positional encoding with baller2vec . The remaining computations are142 identical to the standard Transformer ( see code ) .143 3 Experiments144 3.1 Dataset145 We trained baller2vec on a publicly available dataset of player and ball trajectories recorded from146 631 National Basketball Association ( NBA ) games from the 2015-2016 season.2 All 30 NBA teams147 and 450 different players were represented . Because transition sequences are a strategically important148 part of basketball , unlike prior work , e.g. , Felsen et al . [ 1 ] , Yeh et al . [ 3 ] , Zhan et al . [ 2 ] , we did not149 terminate sequences on a change of possession , nor did we constrain ourselves to a fixed subset of150 sequences . Instead , each training sample was generated on the fly by first randomly sampling a game,151 and then randomly sampling a starting time from that game . The following four seconds of data were152 downsampled to 5 Hz from the original 25 Hz and used as the input.153 Because we did not terminate sequences on a change of possession , we could not normalize the154 direction of the court as was done in prior work [ 1 , 3 , 2 ] . Instead , for each sampled sequence,155 we randomly ( with a probability of 0.5 ) rotated the court 180◦ ( because the court ’ s direction is156 arbitrary ) , doubling the size of the dataset . We used a training/validation/test split of 569/30/32157 games , respectively ( i.e. , 5 % of the games were used for testing , and 5 % of the remaining 95 % 158 of games were used for validation ) . As a result , we had access to ∼82 million different ( albeit159 overlapping ) training sequences ( 569 games × 4 periods per game × 12 minutes per period × 60160 seconds per minute × 25 Hz × 2 rotations ) , ∼800x the number of sequences used in prior work.161 For both the validation and test sets , ∼1,000 different , non-overlapping sequences were selected for162 evaluation by dividing each game into d 1,000N e non-overlapping chunks ( where N is the number of163 games ) , and using the starting four seconds from each chunk as the evaluation sequence.164 3.2 Model165 We trained separate models for Task P and Task B . For all experiments , we used a single Transformer166 architecture that was nearly identical to the original model described in Vaswani et al . [ 5 ] , with167 dmodel = 512 ( the dimension of the input and output of each Transformer layer ) , eight attention heads,168 2https : //github.com/linouk23/NBA-Player-Movements dff = 2048 ( the dimension of the inner feedforward layers ) , and six layers , although we did not use169 dropout . For both Task P and Task B , the players and the ball were included in the input , and both170 the players and the ball were embedded to 20-dimensional vectors . The input features for each player171 consisted of his identity , his ( x , y ) coordinates on the court at each time step in the sequence , and a172 binary variable indicating the side of his frontcourt ( i.e. , the direction of his team ’ s hoop ) .3 The input173 features for the ball were its ( x , y , z ) coordinates at each time step.174 The input features for the players and the ball were processed by separate , three-layer MLPs before175 being fed into the Transformer . Each MLP had 128 , 256 , and 512 nodes in its three layers , respectively,176 and a ReLU nonlinearity following each of the first two layers . For classification , a single linear layer177 was applied to the Transformer output followed by a softmax . For players , we binned an 11 ft× 11 ft178 2D Euclidean trajectory space into an 11 × 11 grid of 1 ft × 1 ft squares for a total of 121 player179 trajectory labels . Similarly , for the ball , we binned a 19 ft × 19 ft × 19 ft 3D Euclidean trajectory180 space into a 19× 19× 19 grid of 1 ft× 1 ft× 1 ft cubes for a total of 6,859 ball trajectory labels.181 We used the Adam optimizer [ 12 ] with an initial learning rate of 10−6 , β1 = 0.9 , β2 = 0.999 , and182 = 10−9 to update the model ’ s parameters , of which there were ∼19/23 million for Task P/Task B,183 respectively . The learning rate was reduced to 10−7 after 20 consecutive epochs of the validation loss184 not improving . Models were implemented in PyTorch and trained on a single NVIDIA GTX 1080 Ti185 GPU for seven days ( ∼650 epochs ) where each epoch consisted of 20,000 training samples , and the186 validation set was used for early stopping.187 3.3 Baselines188 As our naive baseline , we used the marginal distribution of189 the trajectory bins from the training set for all predictions.190 For our strong baseline , we implemented a baller2vec-191 like graph recurrent neural network ( GRNN ) and trained192 it on Task P ( code is available in the baller2vec repos-193 itory ) .4 Specifically , at each time step , the player and ball194 inputs were first processed using MLPs as in baller2vec,195 and these inputs were then fed into a graph neural net-196 work ( GNN ) similar to Yeh et al . [ 3 ] . The node and197 edge functions of the GNN were each a Transformer-198 like feedforward network ( TFF ) , i.e. , TFF ( x ) = LN ( x+199 W2ReLU ( W1x+ b1 ) + b2 ) , where LN is Layer Normal-200 ization [ 13 ] , W1 and W2 are weight matrices , b1 and b2201 are bias vectors , and ReLU is the rectifier activation function . For our RNN , we used a gated recurrent202 unit ( GRU ) RNN [ 14 ] in which we replaced each of the six weight matrices of the GRU with a TFF.203 Each TFF had the same dimensions as the Transformer layers used in baller2vec . Our GRNN had204 ∼18M parameters , which is comparable to the ∼19M in baller2vec . We also trained our GRNN205 for seven days ( ∼175 epochs ) .206 3.4 Ablation studies207 To assess the impacts of the multi-entity design and player embeddings of baller2vec on model208 performance , we trained three variations of our Task P model using : ( 1 ) one player in the input209 without player identity , ( 2 ) all 10 players in the input without player identity , and ( 3 ) all 10 players in210 the input with player identity . In experiments where player identity was not used , a single generic211 player embedding was used in place of the player identity embeddings . We also trained two variations212 of our Task B model : one with player identity and one without . Lastly , to determine the extent to213 which baller2vec uses historical information in its predictions , we compared the performance of214 our best Task P model on the full sequence test set with its performance on the test set when only215 predicting the trajectories for the first frame ( i.e. , we applied the same model to only the first frames216 of the test set ) .217 3We did not include team identity as an input variable because teams are collections of players and a coach , and coaches did not vary in the dataset because we only had access to half of one season of data ; however , with additional seasons of data , we would include the coach as an input variable . 4We chose to implement our own strong baseline because baller2vec has far more parameters than models from prior work ( e.g. , ∼70x Felsen et al . [ 1 ] ) . 4 Results218 4.1 baller2vec is an effective learning algorithm for multi-agent spatiotemporal modeling.219 The average NLL on the test set220 for our best Task P model was221 0.492 , while the average NLL for222 our best Task B model was 2.598.223 In NLP , model performance is often224 expressed in terms of the perplexity225 per word , which , intuitively , is the226 number of faces on a fair die that has227 the same amount of uncertainty as228 the model per word ( i.e. , a uniform229 distribution over M labels has a per-230 plexity of M , so a model with a per231 word perplexity of six has the same232 average uncertainty as rolling a fair233 six-sided die ) . In our case , we con-234 sider the perplexity per trajectory bin , defined as : PP = e 1 NTK ∑N n=1 ∑T t=1 ∑K k=1− ln ( p ( vn , t , k ) ) , where235 N is the number of sequences . Our best Task P model achieved a PP of 1.64 , i.e. , baller2vec236 was , on average , as uncertain as rolling a 1.64-sided fair die ( better than a coin flip ) when predicting237 player trajectory bins ( Table 1 ) . For comparison , when using the distribution of the player trajectory238 bins in the training set as the predicted probabilities , the PP on the test set was 15.72 . Our best Task239 B model achieved a PP of 13.44 ( compared to 316.05 when using the training set distribution ) .240 Compared to our GRNN , baller2vec was ∼3.8 times faster and had a 10.5 % lower average NLL241 when given an equal amount of training time ( Table 2 ) . Even when only given half as much training242 time as our GRNN , baller2vec had a 9.1 % lower average NLL.243 4.2 baller2vec uses information about all players on the court through time , in addition to244 player identity , to model spatiotemporal dynamics.245 Results for our ablation experiments can be seen in246 Table 3 . Including all 10 players in the input dra-247 matically improved the performance of our Task P248 model by 18.0 % vs. only including a single player.249 Including player identity improved the model ’ s per-250 formance a further 4.4 % . This stands in contrast to251 Felsen et al . [ 1 ] where the inclusion of player identity252 led to slightly worse model performance ; a coun-253 terintuitive result given the range of skills among254 NBA players , but possibly a side effect of their role-255 alignment procedure . Additionally , when replacing256 the players in each test set sequence with random257 players , the performance of our best Task P model de-258 teriorated by 6.2 % from 0.492 to 0.522 . Interestingly,259 including player identity only improved our Task B260 model ’ s performance by 2.7 % . Lastly , our best Task261 P model ’ s performance on the full sequence test set262 ( 0.492 ) was 70.6 % better than its performance on the single frame test set ( 1.67 ) , i.e. , baller2vec is263 clearly using historical information to model the spatiotemporal dynamics of basketball.264 4.3 baller2vec ’ s learned player embeddings encode individual attributes.265 Neural language models are widely known for their ability to encode semantic relationships between266 words and phrases as geometric relationships between embeddings—see , e.g. , Mikolov et al . [ 16,267 17 ] , Le and Mikolov [ 18 ] , Sutskever et al . [ 19 ] . Alcorn [ 20 ] observed a similar phenomenon in a268 baseball setting , where batters and pitchers with similar skills were found next to each other in the269 embedding space learned by a neural network trained to predict the outcome of an at-bat . A 2D270 UMAP [ 21 ] of the player embeddings learned by baller2vec for Task B can be seen in Figure 5.271 Like ( batter|pitcher ) 2vec [ 20 ] , baller2vec seems to encode skills and physical attributes in272 its player embeddings.273 4.4 baller2vec ’ s predicted285 trajectory bin distributions depend286 on both the historical and current context.287 Because baller2vec explicitly models the dis-288 tribution of the player trajectories ( unlike varia-289 tional methods ) , we can easily visualize how290 its predicted trajectory bin distributions shift291 in different situations . As can be seen in Fig-292 ure 6 , baller2vec ’ s predicted trajectory bin293 distributions depend on both the historical and294 current context . When provided with limited295 historical information , baller2vec tends to296 be less certain about where the players might297 go . baller2vec also tends to be more certain298 when predicting trajectory bins at “ easy ” mo-299 ments ( e.g. , a player moving into open space ) vs.300 “ hard ” moments ( e.g. , an offensive player choos-301 ing which direction to move around a defender ) .302 4.5 Attention heads in baller2vec appear to perform basketball-relevant functions.303 focus on teammates of the ball handler ( ) . Middle and right : attention head 6-2 seems to predict ( middle ; ) who the ball handler will be in a future frame ( right ) . Players are shaded according to the sum of the attention weights assigned to the players through time with reference to the ball in the current frame ( recall that each player occurs multiple times in the input ) . Higher attention weights are lighter . For both of these attention heads , the sum of the attention weights assigned to the ball through time was small ( 0.01 for both the left and middle frames where the maximum is 1.00 ) . Additional examples can be found in Figures S2 and S3 . One intriguing property of the attention mecha-304 nism [ 22–25 ] is how , when visualized , the atten-305 tion weights often seem to reveal how a model306 is “ thinking ” . For example , Vaswani et al . [ 5 ] 307 discovered examples of attention heads in their308 Transformer that appear to be performing var-309 ious language understanding subtasks , such as310 anaphora resolution . As can be seen in Figure311 7 , some of the attention heads in baller2vec312 seem to be performing basketball understand-313 ing subtasks , such as keeping track of the ball314 handler ’ s teammates , and anticipating who the315 ball handler will pass to , which , intuitively , help316 with our task of predicting the ball ’ s trajectory.317 5 Related Work318 5.1 Trajectory modeling in sports319 There is a rich literature on MASM , particularly320 in the context of sports , e.g. , Kim et al . [ 26 ] ,321 Zheng et al . [ 10 ] , Le et al . [ 27 , 28 ] , Qi et al.322 [ 29 ] , Zhan et al . [ 30 ] . Most relevant to our work323 is Yeh et al . [ 3 ] , who used a variational recurrent324 neural network combined with a graph neural325 network to forecast trajectories in a multi-agent setting . Like their approach , our model is permutation326 equivariant with regard to the ordering of the agents ; however , we use a multi-head attention327 mechanism to achieve this permutation equivariance while the permutation equivariance in Yeh et al.328 [ 3 ] is provided by the graph neural network . Specifically , Yeh et al . [ 3 ] define : v → e : ei , j =329 fe ( [ vi , vj , ti , j ] ) and e→ v : oi = fv ( ∑ j∈Ni [ ei , j , ti ] ) , where vi is the initial state of agent i , ti , j is330 an embedding for the edge between agents i and j , ei , j is the representation for edge ( i , j ) , Ni is the331 neighborhood for agent i , ti is a node embedding for agent i , oi is the output state for agent i , and fe332 and fv are deep neural networks.333 Assuming each individual player is a different “ type ” in fe ( i.e. , attempting to maximize the level of334 personalization ) would require 4502 = 202,500 ( i.e. , B2 ) different ti , j edge embeddings , many of335 which would never be used during training and thus inevitably lead to poor out-of-sample performance.336 Reducing the number of type embeddings requires making assumptions about the nature of the337 relationships between nodes . By using a multi-head attention mechanism , baller2vec learns to338 integrate information about different agents in a highly flexible manner that is both agent and time-339 dependent , and can generalize to unseen agent combinations . The attention heads in baller2vec340 are somewhat analogous to edge types , but , importantly , they do not require a priori knowledge about341 the relationships between the players.342 Additionally , unlike recent works that use variational methods to train their generative models [ 3 , 1 , 2 ] ,343 we translate the multi-agent trajectory modeling problem into a classification task , which allows us to344 train our model by strictly maximizing the likelihood of the data . As a result , we do not make any345 assumptions about the distributions of the trajectories nor do we need to set any priors over latent346 variables . Zheng et al . [ 10 ] also predicted binned trajectories , but they used a recurrent convolutional347 neural network to predict the trajectory for a single player trajectory at a time at each time step.348 5.2 Transformers for multi-agent spatiotemporal modeling349 Giuliari et al . [ 31 ] used a Transformer to forecast the trajectories of individual pedestrians , i.e. , the350 model does not consider interactions between individuals . Yu et al . [ 9 ] used separate temporal and351 spatial Transformers to forecast the trajectories of multiple , interacting pedestrians . Specifically , the352 temporal Transformer processes the coordinates of each pedestrian independently ( i.e. , it does not353 model interactions ) , while the spatial Transformer , which is inspired by Graph Attention Networks354 [ 8 ] , processes the pedestrians independently at each time step . Sanford et al . [ 32 ] used a Transformer355 to classify on-the-ball events from sequences in soccer games ; however , only the coordinates of the356 K-nearest players to the ball were included in the input ( along with the ball ’ s coordinates ) . Further,357 the order of the included players was based on their average distance from the ball for a given358 temporal window , which can lead to specific players changing position in the input between temporal359 windows . As far as we are aware , baller2vec is the first Transformer capable of processing all360 agents simultaneously across time without imposing an order on the agents.361 6 Limitations362 While baller2vec does not have a mechanism for handling unseen players , a number of different363 solutions exist depending on the data available . For example , similar to what was proposed in Alcorn364 [ 20 ] , a model could be trained to map a vector of ( e.g. , NCAA ) statistics and physical measurements365 to baller2vec embeddings . Alternatively , if tracking data is available for the other league , a single366 baller2vec model could be jointly trained on all the data.367 At least two different factors may explain why including player identity as an input to baller2vec368 only led to relatively small performance improvements . First , both player and ball trajectories are369 fairly generic—players tend to move into open space , defenders tend to move towards their man or370 the ball , point guards tend to pass to their teammates , and so on . Further , the location of a player371 on the court is often indicative of their position , and players playing the same position tend to have372 similar skills and physical attributes . As a result , we might expect baller2vec to be able to make373 reasonable guesses about a player ’ s/ball ’ s trajectory just given the location of the players and the ball374 on the court.375 Second , baller2vec may be able to infer the identities of the players directly from the spatiotemporal376 data . Unlike ( batter|pitcher ) 2vec [ 20 ] , which was trained on several seasons of Major League377 Baseball data , baller2vec only had access to one half of one season ’ s worth of NBA data for378 training . As a result , player identity may be entangled with season-specific factors ( e.g. , certain379 rosters or coaches ) that are actually exogenous to the player ’ s intrinsic qualities , i.e. , baller2vec380 may be overfitting to the season . To provide an example , the Golden State Warriors ran a very specific381 kind of offense in the 2015-2016 season—breaking the previous record for most three-pointers made382 in the regular season by 15.4 % —and many basketball fans could probably recognize them from a383 bird ’ s eye view ( i.e. , without access to any identifying information ) . Given additional seasons of data,384 baller2vec would no longer be able to exploit the implicit identifying information contained in385 static lineups and coaching strategies , so including player identity in the input would likely be more386 beneficial in that case.387 7 Conclusion388 In this paper , we introduced baller2vec , a generalization of the standard Transformer that can389 model sequential data consisting of multiple , unordered entities at each time step . As an architecture390 that both is computationally efficient and has powerful representational capabilities , we believe391 baller2vec represents an exciting new direction for MASM . As discussed in Section 6 , training392 baller2vec on more training data may allow the model to more accurately factor players away393 from season-specific patterns . With additional data , more contextual information about agents ( e.g.,394 a player ’ s age , injury history , or minutes played in the game ) and the game ( e.g. , the time left in395 the period or the score difference ) could be included as input , which might allow baller2vec to396 learn an even more complete model of the game of basketball . Although we only experimented with397 static , fully connected graphs here , baller2vec can easily be applied to more complex inputs—for398 example , a sequence of graphs with changing nodes and edges—by adapting the self-attention mask399 tensor as appropriate . Lastly , as a generative model ( see Alcorn and Nguyen [ 33 ] for a full derivation ) ,400 baller2vec could be used for counterfactual simulations ( e.g. , assessing the impact of different401 rosters ) , or combined with a controller to discover optimal play designs through reinforcement402 learning.403 References404 [ 1 ] Panna Felsen , Patrick Lucey , and Sujoy Ganguly . Where will they go ? predicting fine-grained405 adversarial multi-agent motion using conditional variational autoencoders . In Proceedings of406 the European Conference on Computer Vision ( ECCV ) , pages 732–747 , 2018.407 [ 2 ] Eric Zhan , Stephan Zheng , Yisong Yue , Long Sha , and Patrick Lucey . Generating multi-agent408 trajectories using programmatic weak supervision . In International Conference on Learning409 Representations , 2019 . URL https : //openreview.net/forum ? id=rkxw-hAcFQ.410 [ 3 ] Raymond A Yeh , Alexander G Schwing , Jonathan Huang , and Kevin Murphy . Diverse genera-411 tion for multi-agent sports games . In Proceedings of the IEEE Conference on Computer Vision412 and Pattern Recognition , pages 4610–4619 , 2019.413 [ 4 ] Junyoung Chung , Kyle Kastner , Laurent Dinh , Kratarth Goel , Aaron Courville , and Yoshua414 Bengio . A recurrent latent variable model for sequential data . In Advances in Neural Information415 Processing Systems , 2015.416 [ 5 ] Ashish Vaswani , Noam Shazeer , Niki Parmar , Jakob Uszkoreit , Llion Jones , Aidan N Gomez,417 Łukasz Kaiser , and Illia Polosukhin . Attention is all you need . In Advances in Neural Informa-418 tion Processing Systems , pages 5998–6008 , 2017.419 [ 6 ] Tom B Brown , Benjamin Mann , Nick Ryder , Melanie Subbiah , Jared Kaplan , Prafulla Dhariwal,420 Arvind Neelakantan , Pranav Shyam , Girish Sastry , Amanda Askell , et al . Language models are421 few-shot learners . arXiv preprint arXiv:2005.14165 , 2020.422 [ 7 ] Alexey Dosovitskiy , Lucas Beyer , Alexander Kolesnikov , Dirk Weissenborn , Xiaohua Zhai,423 Thomas Unterthiner , Mostafa Dehghani , Matthias Minderer , Georg Heigold , Sylvain Gelly,424 Jakob Uszkoreit , and Neil Houlsby . An image is worth 16x16 words : Transformers for image425 recognition at scale . In International Conference on Learning Representations , 2021 . URL426 https : //openreview.net/forum ? id=YicbFdNTTy.427 [ 8 ] Petar Veličković , Guillem Cucurull , Arantxa Casanova , Adriana Romero , Pietro Lio , and Yoshua428 Bengio . Graph attention networks . In International Conference on Learning Representations,429 2018.430 [ 9 ] Cunjun Yu , Xiao Ma , Jiawei Ren , Haiyu Zhao , and Shuai Yi . Spatio-temporal graph transformer431 networks for pedestrian trajectory prediction . In Proceedings of the European Conference on432 Computer Vision ( ECCV ) , August 2020.433 [ 10 ] Stephan Zheng , Yisong Yue , and Jennifer Hobbs . Generating long-term trajectories using deep434 hierarchical networks . Advances in Neural Information Processing Systems , 29:1543–1551,435 2016.436 [ 11 ] Kazuki Irie , Albert Zeyer , Ralf Schlüter , and Hermann Ney . Language modeling with deep437 transformers . In Proc . Interspeech 2019 , pages 3905–3909 , 2019. doi : 10.21437/Interspeech.438 2019-2225 . URL http : //dx.doi.org/10.21437/Interspeech.2019-2225.439 [ 12 ] Diederik P Kingma and Jimmy Ba . Adam : A method for stochastic optimization . In Interna-440 tional Conference on Learning Representations , 2015.441 [ 13 ] Jimmy Lei Ba , Jamie Ryan Kiros , and Geoffrey E Hinton . Layer normalization . In Advances in442 Neural Information Processing Systems , 2016.443 [ 14 ] Kyunghyun Cho , Bart van Merriënboer , Caglar Gulcehre , Dzmitry Bahdanau , Fethi Bougares,444 Holger Schwenk , and Yoshua Bengio . Learning phrase representations using RNN encoder–445 decoder for statistical machine translation . In Proceedings of the 2014 Conference on Empirical446 Methods in Natural Language Processing ( EMNLP ) , pages 1724–1734 , Doha , Qatar , October447 2014 . Association for Computational Linguistics . doi : 10.3115/v1/D14-1179 . URL https:448 //www.aclweb.org/anthology/D14-1179.449 [ 15 ] Basketball-Reference.com . 2015-16 nba player stats : Totals , February 2021 . URL https:450 //www.basketball-reference.com/leagues/NBA_2016_totals.html.451 [ 16 ] Tomas Mikolov , Ilya Sutskever , Kai Chen , Greg Corrado , and Jeffrey Dean . Distributed452 representations of words and phrases and their compositionality . In Advances in Neural453 Information Processing Systems , 2013.454 [ 17 ] Tomas Mikolov , Kai Chen , Greg Corrado , and Jeffrey Dean . Efficient estimation of word455 representations in vector space . arXiv preprint arXiv:1301.3781 , 2013.456 [ 18 ] Quoc Le and Tomas Mikolov . Distributed representations of sentences and documents . In457 International conference on machine learning , pages 1188–1196 . PMLR , 2014.458 [ 19 ] Ilya Sutskever , Oriol Vinyals , and Quoc V Le . Sequence to sequence learning with neural459 networks . In Advances in Neural Information Processing Systems , 2014.460 [ 20 ] Michael A Alcorn . ( batter|pitcher ) 2vec : Statistic-free talent modeling with neural player461 embeddings . In MIT Sloan Sports Analytics Conference , 2018.462 [ 21 ] Leland McInnes , John Healy , and James Melville . Umap : Uniform manifold approximation463 and projection for dimension reduction . arXiv preprint arXiv:1802.03426 , 2018.464 [ 22 ] Alex Graves . Generating sequences with recurrent neural networks . arXiv preprint465 arXiv:1308.0850 , 2013.466 [ 23 ] Alex Graves , Greg Wayne , and Ivo Danihelka . Neural turing machines . arXiv preprint467 arXiv:1410.5401 , 2014.468 [ 24 ] Jason Weston , Sumit Chopra , and Antoine Bordes . Memory networks . In International469 Conference on Learning Representations , 2015.470 [ 25 ] Dzmitry Bahdanau , Kyunghyun Cho , and Yoshua Bengio . Neural machine translation by jointly471 learning to align and translate . In International Conference on Learning Representations , 2015.472 [ 26 ] Kihwan Kim , Matthias Grundmann , Ariel Shamir , Iain Matthews , Jessica Hodgins , and Irfan473 Essa . Motion fields to predict play evolution in dynamic sport scenes . In 2010 IEEE Computer474 Society Conference on Computer Vision and Pattern Recognition , pages 840–847 . IEEE , 2010.475 [ 27 ] Hoang M Le , Yisong Yue , Peter Carr , and Patrick Lucey . Coordinated multi-agent imitation476 learning . In International Conference on Machine Learning , volume 70 , pages 1995–2003,477 2017.478 [ 28 ] Hoang M Le , Peter Carr , Yisong Yue , and Patrick Lucey . Data-driven ghosting using deep479 imitation learning . In MIT Sloan Sports Analytics Conference , 2017.480 [ 29 ] Mengshi Qi , Jie Qin , Yu Wu , and Yi Yang . Imitative non-autoregressive modeling for trajectory481 forecasting and imputation . In Proceedings of the IEEE/CVF Conference on Computer Vision482 and Pattern Recognition , pages 12736–12745 , 2020.483 [ 30 ] Eric Zhan , Albert Tseng , Yisong Yue , Adith Swaminathan , and Matthew Hausknecht . Learning484 calibratable policies using programmatic style-consistency . In International Conference on485 Machine Learning , pages 11001–11011 . PMLR , 2020.486 [ 31 ] Francesco Giuliari , Irtiza Hasan , Marco Cristani , and Fabio Galasso . Transformer networks for487 trajectory forecasting . In International Conference on Pattern Recognition , 2020.488 [ 32 ] Ryan Sanford , Siavash Gorji , Luiz G Hafemann , Bahareh Pourbabaee , and Mehrsan Javan.489 Group activity detection from trajectory and video data in soccer . In Proceedings of the490 IEEE/CVF Conference on Computer Vision and Pattern Recognition Workshops , pages 898–491 899 , 2020.492 [ 33 ] Michael A. Alcorn and Anh Nguyen . baller2vec++ : A look-ahead multi-entity transformer493 for modeling coordinated agents . arXiv preprint arXiv:2104.11980 , 2021.494 Checklist495 1 . For all authors ... 496 ( a ) Do the main claims made in the abstract and introduction accurately reflect the paper ’ s497 contributions and scope ? [ Yes ] 498 ( b ) Did you describe the limitations of your work ? [ Yes ] See Section 6.499 ( c ) Did you discuss any potential negative societal impacts of your work ? [ No ] Our work500 does not introduce new ethical challenges.501 ( d ) Have you read the ethics review guidelines and ensured that your paper conforms to502 them ? [ Yes ] 503 2 . If you are including theoretical results ... 504 ( a ) Did you state the full set of assumptions of all theoretical results ? [ Yes ] 505 ( b ) Did you include complete proofs of all theoretical results ? [ Yes ] 506 3 . If you ran experiments ... 507 ( a ) Did you include the code , data , and instructions needed to reproduce the main experi-508 mental results ( either in the supplemental material or as a URL ) ? [ Yes ] 509 ( b ) Did you specify all the training details ( e.g. , data splits , hyperparameters , how they510 were chosen ) ? [ Yes ] 511 ( c ) Did you report error bars ( e.g. , with respect to the random seed after running experi-512 ments multiple times ) ? [ N/A ] 513 ( d ) Did you include the total amount of compute and the type of resources used ( e.g. , type514 of GPUs , internal cluster , or cloud provider ) ? [ Yes ] 515 4 . If you are using existing assets ( e.g. , code , data , models ) or curating/releasing new assets ... 516 ( a ) If your work uses existing assets , did you cite the creators ? [ Yes ] 517 ( b ) Did you mention the license of the assets ? [ No ] We link directly to the dataset.518 ( c ) Did you include any new assets either in the supplemental material or as a URL ? [ Yes ] 519 ( d ) Did you discuss whether and how consent was obtained from people whose data you ’ re520 using/curating ? [ N/A ] 521 ( e ) Did you discuss whether the data you are using/curating contains personally identifiable522 information or offensive content ? [ N/A ] 523 5 . If you used crowdsourcing or conducted research with human subjects ... 524 ( a ) Did you include the full text of instructions given to participants and screenshots , if525 applicable ? [ N/A ] 526 ( b ) Did you describe any potential participant risks , with links to Institutional Review527 Board ( IRB ) approvals , if applicable ? [ N/A ] 528 ( c ) Did you include the estimated hourly wage paid to participants and the total amount529 spent on participant compensation ? [ N/A ] 530

Summary: The manuscript presents a modification to the original transformer architecture that allows modeling (un-)ordered sequences of unordered sets. The TxT self-attention masks for sequences of T inputs in the original transformer are replaced by TxKxTxK mask tensors for a sequence of T input sets with K unordered entities each. The architecture is used  in an autoregressive fashion for modeling the trajectories of basketball players and the ball. At each time-step the input consists of a set of players, represented by an identity embedding, raw x and y coordinates, and an optional context feature which indicates the direction of the player’s frontcourt. In task P, the set of target labels indicate the 2d motion corresponding to each input player. In Task B, the output labels indicate the 3d motion of the ball.