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: The paper studies the use of transformers for particle-based physics simulation. The main contributions are as follows.  C1. the paper proposes a specific transformer-inspired form of message passing, which, unlike transformer, still does explicit message computation.  C2. the paper proposes a very specific way to encode material type information as abstract tokens.  C3. the paper compares to several non-transfomer based baselines, as well as a pure transformer.  C4. comparisons are performed on a subset of 4 environments including some generalization settings. As presented in the paper, the proposed model seems to perform better than the baselines in most settings.