SiT: Simulation Transformer for Particle-based Physics Simulation

1 INTRODUCTION . Particle-based physics simulation is a classic and important topic in computer science . It not only facilitates the exploration of underlying principles in physics , chemistry and biology , but also enables the creation of vivid visual effects such as explosion and fluid dynamic in films and games . Different from traditional simulators , such as grid-based ( Guo et al. , 2016 ) and mesh-based ( Bronstein et al. , 2017 ) methods , particle-based simulators view a system , which is an example in one domain , as a composition of particles and imitate system changes over time by predicting the changes of particle-wise states according to current particle states and particle interactions , of which the latter represents the influence of action-reaction forces , such as the collisions . Consequently , they follow the same forward process without separately considering different constraints to simulate different domains with varying materials , requiring no domain-specific physical priors . Moreover , since in particle-based simulators the dynamics of a system is modeled by the states of particles and their interactions , they also have the potential to possess a strong generalization ability , where they can estimate the dynamics of a system with varying number and configuration of particles in a more robust manner . After learning the dynamics of fluid water in a sandbox , the same particle-based simulator can be used to simulate a waterfall and a river . Recent particle-based simulators ( Battaglia et al. , 2016 ; Schenck & Fox , 2018 ; Mrowca et al. , 2018 ; Li et al. , 2019 ; Sanchez-Gonzalez et al. , 2020 ; Ummenhofer et al. , 2020 ) often view a system as a graph , and adopt graph neural network ( GNN ) ( Kipf & Welling , 2016 ) as the basic network structure . In these attempts , each particle is treated as a node in the graph , with edges linking it to all its neighboring particles , assuming interactions mainly occur between particles that are close to each other . Subsequently , state updates of particles are achieved by combining node features with the summation of edge features . While such a GNN-based formulation obtains satisfying simulation results , it faces two issues that affect efficiency and generality . First , it forces each particle to interact with all its nearby particles without providing a selective mechanism , which leads to computational redundancy and prevents the discovery of inherent patterns of particle interaction . Second , the GNN-based formulation uses particle-wise attributes to capture both material characteristics , such as viscosity or plastic deformations , and domain-specific semantics , such as the shape of a rigid material . Therefore , it may regard the latter as part of the intrinsic material semantics and fail to generalize to domains with the same materials but different particle amounts and configurations . In this paper , we propose a novel Transformer-based framework , referred to as Simulation Transformer ( SiT ) , for particle-based physics simulation . The model inherits the powerful multi-head self-attention mechanism in Transformer ( Vaswani et al. , 2017 ) to capture particle interactions . To further encourage efficient modeling of complex interactions , instead of treating particle interactions as attention weights obtained by dot-product , we introduce the notion of interaction tokens , which are high-dimensional representations for interactions , to model the rich semantics of particle interactions , such as how the particle is restored after deformations . In addition , to disentangle local material-specific characteristics from global domain-specific semantics , SiT further learns a highdimensional abstract token for each type of material to capture material semantics , forcing particles of the same material to interact with their corresponding abstract token . The proposed SiT is more appealing than previous methods in several aspects . First , through capturing particle interactions explicitly with interaction tokens and allowing dynamic inter-token attention , SiT dynamically focuses on essential particle interactions and reduces the computations for redundant and noisy ones . This is crucial for particle-based simulation , especially for domains containing hundreds and thousands of densely placed particles , where the modeling of all particle interactions is redundant , expensive and noisy in practice . Second , thanks to the trainable abstract tokens that disentangle intrinsic material characteristics from domain-specific semantics , we can reuse them to apply SiT to unseen domains of the same materials without retraining . As shown in our experiments , after training on one domain consisting of fluid water and a rigid cubic , SiT still yields fairly faithful simulations when the cubic is replaced with a ball or a bunny comparing with previous work . To show the effectiveness of SiT , we perform extensive evaluations on four standard environments commonly used in the literature ( Li et al. , 2019 ; Sanchez-Gonzalez et al. , 2020 ) , covering domains of different complexity and materials . The proposed method achieves superior performance across all these environments with fewer parameters compared to existing methods . We further demonstrate the generalization ability of SiT by adjusting the environments to create new domains and applying SiT to these domains without retraining . In all cases , SiT obtains more realistic simulation results than previous methods , which tend to overfit to the training domains . 2 RELATED WORK . Physics simulation by neural networks . There are many different kind of representations for physics simulations . Grid-based methods ( Lee & You , 2019 ; Thuerey et al. , 2020 ; Wang et al. , 2020 ) adopt convolutional architectures for learning high-dimensional physical system , while mesh-based simulations ( Bronstein et al. , 2017 ; Luo et al. , 2020 ; Hanocka et al. , 2019 ; Nash et al. , 2020 ; Qiao et al. , 2020 ; Weng et al. , 2021 ; Pfaff et al. , 2021 ) typically simulate objects with continuous surfaces , such as clothes , rigid objects , surfaces of water and so on . Many works ( Battaglia et al. , 2016 ; Schenck & Fox , 2018 ; Mrowca et al. , 2018 ; Li et al. , 2019 ; Sanchez-Gonzalez et al. , 2020 ; Ummenhofer et al. , 2020 ) simulate physics on particle-based systems , where all objects are represented by groups of particles . Specifically , Interaction Network ( IN ) ( Battaglia et al. , 2016 ) simulated interactions in object-level . Smooth Particle Networks ( SPNets ) ( Schenck & Fox , 2018 ) implemented fluid dynamics using position-based fluids ( Macklin & Müller , 2013 ) . Hierarchical Relation Network ( HRN ) ( Mrowca et al. , 2018 ) predicted physical dynamics based on hierarchical graph convolution . Dynamic Particle Interaction Networks ( DPI-Net ) ( Li et al. , 2019 ) combined dynamic graphs , multi-step spatial propagation , and hierarchical structure to simulate particles . CConv ( Ummenhofer et al. , 2020 ) used spatial convolutions to simulate fluid particles . Graph Network-based Simulators ( GNS ) ( Sanchez-Gonzalez et al. , 2020 ) computed dynamics via learned message-passing . Similar to particle-based systems , COPINGNet ( Shao et al. , 2021 ) applies graph networks to simulate rod dynamics , where the discretized rod is the basic unit similar to particle . Previous work mostly adopted graph networks for simulations and required each particle to interact with all its nearby particles without selective mechanism . In contrast , our SiT employs both particle and interaction tokens and selectively focus on necessary particle interactions through attention mechanism . Experiments show that SiT surpasses existing GNN-based methods and has more robust performances in generalizations . Transformer . Transformer ( Vaswani et al. , 2017 ) was designed for machine translation and achieved state-of-the-art performance in many natural langruage processing tasks ( Devlin et al. , 2019 ; Radford et al. , 2019 ; Brown et al. , 2020 ) . Recently , Transformer starts to show great expandability and applicability in many other fields , such as computer vision ( Wang et al. , 2018 ; Carion et al. , 2020 ; Dosovitskiy et al. , 2021 ; Wang et al. , 2021 ; Liu et al. , 2021 ) , and graph representations ( Zhou et al. , 2020 ; Zhang et al. , 2020 ; Dwivedi & Bresson , 2020 ) . To our knowledge , no attempt has been made to apply Transformer on physics simulation . Our SiT models the interactions between particles by trainable sub-network given corresponding particle tokens . The same notion of extracting potential semantics between nodes is also applied in Graph Transformer ( Dwivedi & Bresson , 2020 ) , which we refer as GraphTrans for short . However , there are differences in our formulations . Specifically , GraphTrans adopts element-wise product between node representations followed by multi-layer perceptron ( MLP ) to update the interaction embeddings , which store the attention scores in each dimension and are reduced to a scalar for further attention mechanism to update node embeddings . In contrast , our model learns semantic tokens for interactions through sub-network instead of element-wise product . When updating the particle tokens , which are referred as node representations in GraphTrans , both particle and interaction tokens will attend to each other . The interaction tokens are not weighted scores any more . We adopt GraphTrans ( Dwivedi & Bresson , 2020 ) in particle-based simulation and compare it with SiT in experiments . The quantitative results show that SiT achieves better results than GraphTrans . 3 METHODOLOGY . 3.1 PROBLEM FORMULATION . For a particle-based system composed of N particles , we use X t = { xti } Ni=1 to denote the system state at time step t , where xti denotes the state of i-th particle . Specifically , x t i = [ p t i , q t i , ai ] , where pti , q t i ∈ R3 refer to position and velocity , and ai ∈ Rda represents fixed particle attributes such as its material type . The goal of a simulator is to learn a model φ ( · ) from previous rollouts of a system to causally predict a rollout trajectory in a specific time period conditioned on the initial system state X 0 . The prediction is run in a recursive manner , where the simulator will predict the state X̂ t+1 = φ ( X t ) at time step t + 1 based on the state X t = { xti } at time step t. In practice , we will predict the velocities of particles Q̂t+1 = { q̂t+1i } , and obtain their positions via p̂t+1i = p t i + ∆t · q̂ t+1 i , where ∆t is a domain-specific constant . 3.2 SIMULATION VIA VANILLA TRANSFORMER . To accurately simulate the changes of a system over time , it is crucial to effectively model the interactions among particles , as they indicate the energy transition of a system when constrained by material characteristics and physical laws . However , it is infeasible to know a priori how should particles interact with each other . Thus , a selective mechanism is needed to help the simulator focus only on necessary interactions . Since Transformer ( Vaswani et al. , 2017 ) is capable of modeling the dynamical attention scores between tokens via the self-attention module , we can regard the attention weights as the intensity of connectivities and the importance of interactions , where the larger the attention weight the more important the interaction is . Thus , it is naturally a good backbone of building an efficient particlebased simulator . To apply vanilla Transformer in particle-based simulation , we first encode the states of particles into corresponding particle tokens V = { vti } by vti = f enc V ( x t i ) , ( 1 ) where vti ∈ Rdh is a dh dimensional vector and f encV ( · ) is an encoding layer implemented as a MLP . Subsequently , particle interactions are achieved by L blocks of self-attention modules , where in the l-th block , particle tokens will attend to each other selectively as : vl+1 , ti = ∑ j ŵvijv l , t j , ( 2 ) ŵvij = exp ( wvij ) √ dh · ∑ j exp ( w v ij ) , ( 3 ) wvij = ( v l , t i ) > vl , tj . ( 4 ) Since a system usually contains hundreds of particles and the interactions among particles occur within the neighbors in our settings , considering all possible interactions , of which the number quadratically increases with respect to the number of particles , is computationally redundant and inefficient . Therefore , we follow previous literature Li et al . ( 2019 ) ; Sanchez-Gonzalez et al . ( 2020 ) to assume that interactions of distant particles can be omitted , which is realized by a window function : g ( pti , p t j ) = I ( ||pti − ptj ||2 < R ) , ( 5 ) where I ( condition ) is a indicator function that returns 1 if the condition is satisfied , and R defines the extent of the window . This window function will generate a mask to retain only interactions between neighboring particles as potential candidates in the self-attention modules . To predict particle states in the next time step , a decoding layer is applied to the updated token of i-th particle to obtain its velocity : q̂t+1i = f dec V ( v L , t i ) , ( 6 ) where fdecV ( · ) is implemented by another MLP .

Summary: This paper proposes to learn particle dynamics of physical systems by the means of a transformer. In particular, the paper investigates using a vanilla transformer as well as a more customized variant called 'Simulation Transformer'. Both networks are validated based on variety of experiments of simulations of fluids, rigid and deformable objects, as well as through comparisons to existing approaches.