Year: 2021

Abstract:
In this paper, we focus on learning-based methods for simulation. Most existing particle-based simulators adopt graph neural networks (GNNs) to model the underlying physics of particles. However, they force particles to interact with all neighbors without selection, and they fall short in capturing material semantics, such as viscosity or plastic deformations, for different particles, leading to unsatisfactory performance, especially in generalization. This paper proposes Simulation Transformer (SiT) to simulate particle dynamics with more careful modeling of particle states, interactions, and their intrinsic properties. Specifically, besides the particle tokens, which encode the states of particles into latent space, SiT generates interaction tokens to extract high-level representations for interactions and selectively focuses on essential interactions by allowing both tokens to attend to each other. In addition, SiT learns material-aware representations by learnable abstract tokens, which will participate in the attention mechanism and boost the generalization capability further. We evaluate our model on diverse environments, including fluid, rigid, and deformable objects, which cover domains of different complexity and materials. Without bells and whistles, SiT shows strong abilities to simulate particles of different materials and achieves superior performance and generalization across these environments with fewer parameters than existing methods. Codes and models will be released.

heading: 1 INTRODUCTION
text: 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.

heading: 2 RELATED WORK
text: 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.

heading: 3 METHODOLOGY

heading: 3.1 PROBLEM FORMULATION
text: 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.

heading: 3.2 SIMULATION VIA VANILLA TRANSFORMER
text: 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.

heading: 3.3 SIMULATION TRANSFORMER (SIT)
text: Although the vanilla Transformer provides a flexible approach for particle-based simulation, directly applying it leads to inferior simulation results as shown in our experiments. In particular, a vanilla
Transformer uses attention weights, which are scalars obtained via dot-product, to represent particle interactions. A single scalar is insufficient to reflect the rich semantics of particle interactions. For example, in the case of fluid water, when two particles get closer and closer, their interaction will at first act as surface tension to push them towards each other. And when they are sufficiently close, the repulsion between them will dominate the interaction and stop them from further moving towards each other. In addition, since the semantics of one domain is scattered in the tokens of different particles, domain-specific semantics, such as the shape of rigid cubic, are mixed with domain-agnostic semantics, such as the stickiness attribute of materials. In such cases, the former may be misinterpreted as the latter, leading to poor generalization ability. To overcome the limitations of vanilla Transformer, we propose a novel Transformer-based simulator, referred to as Simulation Transformer (SiT), that extends vanilla Transformer with two new types of tokens, namely interactions tokens and abstract tokens. An overview of SiT is included in Figure 2.
Interaction Tokens. To capture the rich semantics of particle interactions, instead of a scalar, SiT assigns an interaction token to each potential particle interaction. Specifically, the token for the interaction between particles xti and x t j is initialized by
utij = f enc I (x t i,x t j), (7)
where utij ∈ Rdh is a dh dimensional vector and f encI (·) is the encoding function for interactions. Subsequently, in each self-attention block SiT will update the interaction tokens as
ul+1,tij = ψ ( vl,ti ,v l,t j ,u l,t ij , g(p t i,p t j) ) , (8)
where ψ(·) is a learnable sub-network shared for all interaction tokens. Inspired by GNN-based methods (Mrowca et al., 2018; Li et al., 2019; Sanchez-Gonzalez et al., 2020) , SiT adopts the learnable sub-network to enlarge its capacity for handling complex particle interactions. In practice, the sub-network is implemented as
ψ ( vl,ti ,v l,t j ,u l,t ij , g(p t i,p t j) ) = MLP([vl,ti ,v l,t j ,u l,t ij ]) · g(p t i,p t j). (9)
As is shown in Figure 6, the update of particle tokens in the self-attention module is thus achieved by
vl+1,ti = α ∑ j ŵvijv l,t j + β ∑ k ŵuiku l+1,t ik , (10)
ŵvij = exp(wvj )
√ dh · ( α ∑ j exp(w v j ) + β ∑ k exp(w u ik) ) , (11)
ŵuik = exp(wuik)√ dh · ( α ∑ j exp(w v j ) + β ∑ k exp(w u ik) ) , (12)
wvij = (v l,t i ) >vl,tj , (13)
wuik = (v l,t i ) >ul+1,tik , (14)
where α, β ∈ {0, 1} are the gate controllers, j and k are the indexes of particles within the reception window defined by Equation 5. Equation 13 computes the attention scores between particle tokens, and Equation 14 computes the attention scores between particle and interaction tokens. Subsequently, these attention scores are normalized by Equation 11 and Equation 12. They are then used by Equation 10 to aggregate semantics of particle and interaction tokens, where the gate controllers α, β are hyper-parameters that work as balancing coefficients. When α = 1, β = 0, Equation 10 becomes the original self-attention formula in vanilla Transformer. The multi-head attention version of Equation 10 can be found in Appendix 6.1.2.
It is worth noting that SiT provides a more flexible way for further extensions when compared to vanilla Transformer. Specifically, we can change the value of α and β in Equation 10 to adjust the significance of different tokens. And changing the window function in Equation 5 allows SiT to consider different interaction patterns. Finally, by adjusting the learnable sub-network, we can also generalize SiT to consider more complex particle interactions such as interactions beyond a pair of particles and interactions of interactions.
Abstract Tokens. To improve the generalization ability of SiT and disentangle domain-specific semantics from its domain-agnostic counterparts, we further apply SiT with material-specific abstract tokens.
For Na types of materials, SiT adopts Na abstract tokens A = {ak}Nak=1, each of which is a learnable vector of the same length of particle tokens. Ideally, the abstract token ak should capture the domain-agnostic semantics of k-th material. They act as particle tokens for additional abstract particles, following the same update formula as particle tokens. Therefore with Na abstract tokens SiT will have N + Na particle tokens in total: {a1, · · · ,aNa ,xt1, · · · ,xtN}, as is show in Figure 2. Although abstract tokens share the same update formula as particle tokens, unlike particle tokens that decide their potential interactions via the window function in Equation 5, abstract tokens are forced to interact with all particle tokens of its corresponding material, which is achieved by setting the reception field of abstract tokens to the particles of the same material.
Once abstract tokens capture domain-agnostic semantics, they can be reused by SiT when generalizing to domains that have same materials but vary in particle amount and configuration.

heading: 3.4 TRAINING OBJECTIVE
text: To train SiT with previous rollouts, the standard mean square error (MSE) loss is applied on the output of our SiT:
MSE(Q̂t+1, Qt+1) = 1
N ∑ i ||q̂t+1i − q t+1 i ||2 2, (15)
where Q̂t+1 = {q̂t+1i }Ni=1 is the estimated velocity , Qt+1 = {q t+1 i }Ni=1 is the ground truth, and || · ||2 is L2 norm. While the MSE loss works well for cases with relatively low complexity, in practice some domains often contain multiple types of materials with imbalance numbers of particles. In such cases, the MSE loss will be biased to the material with more particles. Therefore, we further apply a material weighted MSE (WMSE) loss to reduce the effect of imbalance:
WMSE(Q̂t+1, Qt+1) = 1
K ∑ k 1 Nk ∑ i ||q̂t+1i,k − q t+1 i,k || 2 2, (16)
whereK is the number of material types,Nk is the number of particles belonging to the k-th material and N = ∑ kNk, Q̂ t+1 k = {q̂ t+1 i,k } Nk i=1 and Q t+1 k = {q t+1 i,k } Nk i=1 denote the estimated velocity and ground truth for k-th material, and Q̂t+1 = {Q̂t+1k }Kk=1, Qt+1 = {Q t+1 k }Kk=1,

heading: 4 EXPERIMENTS

heading: 4.1 BASE ENVIRONMENTS
text: We adopt the environments commonly used in the literature (Li et al., 2019; Sanchez-Gonzalez et al., 2020; Ummenhofer et al., 2020). Samples are displayed in Figure 1. There are four environments in total: FluidFall is a basic simulation for two drops of water; FluidShake is more complex and simulate the water in a randomly moving box; BoxBath simulates the water washing a rigid cubic in fixed box; RiceGrip simulates the interactions between deformable rice and two rigid grippers. More details of data generalization can be found in Appendix 6.1.3.
We compare SiT with 4 recent approaches: DPI-Net (Li et al., 2019), CConv (Ummenhofer et al., 2020), GNS (Sanchez-Gonzalez et al., 2020), and GraphTrans (Dwivedi & Bresson, 2020). For fair comparison, we adopt similar training schedules. Quantitative results and model parameters are in Table 5. Qualitative results are in Figure 3. More details can be found in Appendix 6.1.3 and 6.2.2
SiT achieves superior performance on all environments with fewer model parameters. The effectiveness of abstract tokens are more obvious for RiceGrip and BoxBath, which are more complex. Quantitative results show that SiT can better simulate the dynamics on different environments.
Comparison with DPI-Net and GNS. DPI-Net (Li et al., 2019) and GNS (Sanchez-Gonzalez et al., 2020) adopt message-passing graph networks for particle-based simulation. As shown in the results,
DPI-Net 0.12±0.06 0.62M 1.43±0.52 1.02M 0.11±0.21 2.17M 1.91±0.08 1.52±0.33 1.98M CConv 0.08±0.02 0.84M 1.34±0.45 0.84M N/A N/A 1.98 ±0.11* 3.06±0.47* 0.84M GNS 0.05±0.02 1.59M 1.45±0.55 1.59M 0.38±0.25 1.60M 1.77±0.87 2.62±0.87 1.59M GraphTrans 0.12±0.03 0.49M 8.77±3.97 0.96M 0.36±0.27 0.96M 3.17±0.07 2.33±0.47 0.97M SiT (Ours) 0.05±0.02 0.40M 1.08±0.36 0.77M 0.15±0.12 0.82M 1.74±0.08 1.42±0.29 0.77M SiT+ (Ours) 0.04±0.01 0.40M 1.08±0.39 0.77M 0.07±0.07 0.82M 1.57±0.06 1.39±0.31 0.77M
SiT achieves better performance, as the selective mechanism can eliminate the unnecessary interactions, which is further demonstrated by the heat map in Appendix 6.2.1. SiT with abstract tokens further improves the performances especially on RiceGrip and BoxBath, suggesting the effectiveness of abstract tokens in modeling complex deformations and multi-materials interactions. For FluidFall, though the MSEs are close to each other, SiT can better maintain the shape of drops while DPI-Net and GNS fail because they do not have selective mechanism and equally treat the incorrect interactions, which are introduced by the newly added neighbors from different drops when they get closer but not close enough to merge.
Comparison with CConv. CConv (Ummenhofer et al., 2020) designs convolutional layers carefully tailored to modeling fluid dynamics, such as an SPH(Monaghan, 1992)-like local kernel, different sub-networks for fluid and boundary particles. Without modeling interactions, CConv propagates information only among particles weighted by distances. On the other hand, SiT assigns high-level tokens for interactions and achieves better results, suggesting the necessity of interaction tokens. Notice that CConv is not suitable for rigid box, we only report result on BoxBath for reference of simulation on fluid parts.
Comparison with GraphTrans (Dwivedi & Bresson, 2020). The GraphTrans is also a Transformerbased method and updates the interactions using element-wise product from node representations. We represent particles by nodes and adopt layers in GraphTrans as our backbone, which and use the same hidden dimensions and number of blocks. Other settings are completely the same as SiT for better comparison. As is shown in experiments, GraphTrans fails to simulate particles on all cases with much higher errors. This is because our SiT directly models the interactions by interaction tokens generated by sub-network, which is necessary to capture high-level semantics.

heading: 4.2 GENERALIAZTIONS
text: To challenge the robustness of our model on more complex settings, we add more particles for RichGrip and FluidShake, and change the size and shape of rigid object for BoxBath. Table 5 shows the upper and lower bound for the number of particles in training set and generalization validation set, as well as the shapes of rigid object.
We mainly compare SiT with DPI-Net, which has fairly good and reasonable performance on all environments. Quantitative results are summarized in Table 5, while qualitative results are depicted in Figure 4 and Figure 1 (b). As is shown, SiT with abstract tokens can better simulate particle dy-
namics in generalized environment. suggesting that the abstract tokens and the selective mechanism involving interaction tokens can provide more stable and robust performance. More details about the results and rollouts can be found in Appendix 6.2.3.

heading: 4.3 ABLATION STUDIES
text: We comprehensively analyze our SiT and explore the effectiveness of our model in the following ways: (a) the necessity of interaction tokens; (b) the architectures of sub-network, which generates interaction tokens; (c) the values for gate controllers α, β; (d) the usage of multi-head in uniform attention; (e) the sensitiveness of SiT to radius R; (f) the effectiveness of abstract tokens. We conduct our experiments on FluidShake for (a) to (e). We verify (f) on BoxBath and RiceGrip. The quantitative results are in Table 3 and Table 4.
Necessity of interaction tokens. We apply vanilla Transformer encoder by configuration A, which uses the same hidden dimension and number of blocks as SiT, and SiT by configuration B, which only changes the output of sub-network to one dimension. As shown in Table 3, both configuration A and B have low performance, suggesting the scalars are insufficient to describe complex interactions.
Architectures of sub-network. Configuration C adopts element-wise product between particle tokens, which follows the GraphTrans (Dwivedi & Bresson, 2020); configuration D adopts selfattention blocks, where q is the interaction token from previous layer, k and v are the neighbor particle tokens. As shown in Table 3, configuration C and D have worse performance than SiT, which suggests that the element-wise product is insufficient to model complex interactions, while self-attention block introduces more parameters and complexities, making it harder to train.
Hyper-paramters for α, β. We set α = 1, β = 0 in configuration E, α = 1, β = 1 in configuration F, and α = 0, β = 1 in configuration H, Consequently, both configuration E and H achieve good results, suggesting that the participants of interaction tokens are the key to boost the performance. In practice, we use configuration H for SiT which shows the best results.
Multi-head uniform attention. We apply multi-head mechanism in configuration G and set head number as 8. The results show that multi-head mechanism can bring us some improvement. But it also introduces more parameters and consumes lots of memory when computing.
Sensitiveness to R. Quantitative results are reported on FluidShake. As shown in Table 4, SiT is more robust when varying the radius R, suggesting the effectiveness of selective mechanism and robustness of SiT.
Effectiveness of abstract tokens. We replace the abstract tokens with dummy tokens, which are randomly initialized vectors with fixed values but have same connectivities as abstract tokens. As shown in 4, SiT with dummy tokens fails on both domains, suggesting the abstract tokens are able to learn the materials semantics and boost SiT’s performances.

heading: 4.4 FEW SHOTS LEARNING
text: As shown in Figure 5, we conduct few shots learning experiments on FluidShake using SiT without abstract tokens. Even when trained on 60% of data, SiT achieves similar MSEs comparing with DPI-Net trained on 100% of data. When trained with further less data, SiT learns not well enough. As the number of particles differs in each rollout, SiT does not have enough training on all possible cases. On the other hand, the weights are not shared in SiT and would not learn enough on less data comparing with the shared-weight blocks in DPI-Net. SiT is more robust with no less than 60% of the whole training data, but is sensitive to extremely less data.

heading: 5 CONCLUSION
text: In this paper, we propose Simulation Transformer (SiT) and introduce interaction tokens and abstract tokens to simulate domains of different complexity and materials, including hundreds and thousands of particles. Our experimental results show the necessity and effectiveness of interaction and abstract tokens as well as selective mechanism in attention for particle-based simulation. Generalization results further suggest the robustness of SiT. SiT is also flexible for extensions and applications, such as more complex and efficient architectures for sub-network, changing the particle dynamics by attending to alternative abstract tokens with different materials semantics. Finally, SiT makes a successful attempt to apply Transformer into physics simulation and achieve superior performances over existing methods.

heading: 6 APPENDIX

heading: 6.1 MODEL DETAILS

heading: 6.1.1 UNIFORM ATTENTION
text: Figure 6 illustrates how our uniform attention works. Our uniform attention extend self-attention (Vaswani et al., 2017) by additional tokens. The queries are the abstract tokens, which are optional, and particle tokens, while the keys and values are all three types of tokens. Each token will attend to tokens within its reception fields defined by the window function 5. For the abstract tokens, they are forced to attend only to the particles belonging to the same material. The output of uniform attention is the updated tokens except interaction tokens.

heading: 6.1.2 MULTI-HEAD UNIFORM ATTENTION
text: The multi-head version of uniform attention is formulated by
vl,ti = W l [SA1,SA2, · · · ,SAM ]> , (17) SAm = α ∑ j ŵm,vij (V m,lvl−1,tj ) + β ∑ k ŵm,uik (V m,lul,tik ), (18)
ŵm,vij = exp(wm,vj )
√ dh · ( α ∑ j exp(w m,v j ) + β ∑ k exp(w m,u ik ) ) , (19) ŵm,uik =
exp(wm,uik )√ dh · ( α ∑ j exp(w m,v j ) + β ∑ k exp(w m,u ik ) ) , (20) wm,vij = (Q m,lvl−1,ti ) >(Km,lvl−1,tj ), (21) wm,uik = (Q m,lvl−1,ti ) >(Km,lul,tik ), (22)
where m ∈ 1, 2, · · · ,M is the index of head, W l ∈ Rdh×dh is the output weight for the concatenation of M heads; Qm,l,Km,l, V m,; are the projection weights of query, key, value for head m at block l. The rest symbols are the same with those in uniform attention.

heading: 6.1.3 IMPLEMENTATION DETAILS
text: Inputs and outputs details. For FluidFall, FluidShake, and BoxBath, we only use particles’ states at time t as inputs and output the velocities at time t+ 1. For RiceGrip, we concatenate particles states
from t−2 to t as inputs and output 6-dim vector for the velocity of the current observed position and the resting position, which is the same setting as DPI-Net for better comparison. For BoxBath, we output 7-dim vectors, where 3 dimensions for the predicted velocities, and 4 dimensions for rotation constrains. The rotation constraints are applied only on rigid particles which predict the rotation velocities. We use particles’ states at time t as inputs and output the velocities at time t + 1 except RiceGrip, where we concatenate particles states from t− 2 to t as inputs for better comparison with DPI-Net (Li et al., 2019). All states of particles, such as the positions, velocities, and accelerations, are first normalized by mean and standard deviations calculated on corresponding training set.
SiT details. We set R = 0.08 across all environments. All hidden dimensions in SiT are 128 for default. The MLP after uniform attention has two layers with dimensions 512 and 128. We use 2 blocks shown in Figure 2 for FluidFall and 4 blocks for the other three environments. We set α = 0, β = 1 in our model for default, which show the best results in ablation studies. The learnable embeddings for abstract tokens have the same dimensions as inputs.
Training. We train four models independently on these four environments. For common settings, we use Adam optimizer with initial learning rate 0.0008 and plateau scheduler with patience 3 and decreasing factor 0.8. We choose batch size 16 for all environments and train the models for 5 epochs on FluidShake, 5 epochs for BoxBath, 13 epochs for FluidFall, and 20 epochs for RiceGrip. No augmentation is involved, such as adding noises to particles states during training. We adopt WMSE loss on BoxBath and MSE loss on the rest environments for training.
Evaluation. We evaluate the positions between predictions and ground truths by MSE as is shown in equation 15. For BoxBath, we also report the WMSE results by equation 16 to better show the performance of simulating different materials.
Baseline details. For fair comparison, the following settings are the same with SiT: inputs for models, number of training epochs on different environments, learning rate schedules, and training loss on velocities. On BoxBath, all baselines adopt the rigid constraints except CConv, which is designed specifically for fluid dynamics and we use it to mainly compare the fluid simulations. The model-related hyper-parameters for baselines are the same as the original papers, such as the number of message passing is 10 for GNS. Other hyper-parameters for baselines are first chosen the same as their original papers, and then fine-tuned within a small range of changes. For example, in terms of the initial learning rate, 0.0001 works better on GNS while others adopt 0.0008.

heading: 6.2 EXPERIMENTS

heading: 6.2.1 ADDTIONAL ABLATION STUDIES
text: We also visualize the attention scores in each block for a specific particle on BoxBath in Figure 7. The red dot in the center is the given particle. Other dots are the neighors. The red particle will attend selectively to its neighbors rather than treating them equally. Different layers focus on different aspect of semantics. Specifically, the particle will focus mainly on closer neightbors in the first layer as is shown in Figure 7 (b).

heading: 6.2.2 BASE ENVIRONMENTS
text: Dataset details. We use the same setting for our datasets as mentioned in previous work (Li et al., 2019). FluidFall contains two fluid drops with different sizes. The size for each drop are randomly generated with one drop larger than the other. Position and viscosity for drops are randomly initialized. This environment contains 189 particles with 121 frames for each rollout. There are 2700 rollouts in training set and 300 rollouts in validation set. FluidShake simulates the water in a moving box. The speed of the box is randomly generated at each timestamp. In addition, the size of the box and the number of particles are various for different rollouts. In basic training and validation sets, the number of particles changes from 450 to 627. This environment has 301 frames for each rollout. There are 1800 rollouts in training set and 200 rollouts in validation set. RiceGrip contains two grippers and a sticky rice. The grippers’ positions and orientation are randomly initialized. The number of particles for rice changes from 570 to 980 with 41 frames for each rollout in training and validation sets. There are 4500 rollouts in training set and 500 rollouts in validation set. BoxBath simulates a rigid cubic washed by water in a fixed container. The initial position of fluid block and rigid cubic are randomly initialized. This environment contains 960 fluid particles and 64 rigid par-
ticles with 151 frames for each rollout. There are 2700 rollouts in training set and 300 rollouts in validation set.
Results. We display some rollouts for all experiments in this section. The results on RiceGrip for CConv are not available because CConv was designed specifically for fluid dynamics. The results on BoxBath for CConv are only used as reference of the simulation on fluid parts. And we compare CConv with SiT mainly on the fluid simulations.

heading: 6.2.3 GENERALIZATION DETAILS
text: We report WMSE results on BoxBath and other models’ generalization performance for references in Table 5. SiT with abstract tokens achieves reasonable and superior performance.
We release the details of generalization settings in Table 6. For the generalizations of FluidShake and RiceGrip, we display the range of particle numbers. For the generalizations of BoxBath, we change the shape of rigid box into bunny, ball, cuboid, and larger cubic. Specifically, the generalized RiceGrip contains particles from 1060 to 1640, while the training set contains particles from 570 to 980; the generalized FluidShake contains particles from 720 to 1368, while the training set has 450 to 627 particles; For BoxBath, the training set has 960 fluid particles and 64 rigid particles. We first separately increase the fluid particles and rigid particles to 1280 and 125 respectively. Then, we change the shape of rigid object into bunny, ball, and cuboid with particles 41, 136, and 120 respectively. The followings are the rollouts on all environments.

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

In this paper, authors propose Simulation Transformer (SiT), a transformer based approach for particle-based fluid simulations (in contrast to all the existing approaches which are overwhelmingly based on Graph Convolutional Networks). Specifically, in their paper the authors augment the vanilla transformers by incorporating sub-networks for richer modeling particle interactions (i.e., as opposed to particle interactions being modeled as dot products and being reduced to a single number, sub-networks allow richer modeling of interactions), as well as material specific properties. Authors evaluate the proposed SiT model on diverse environments and compare against several state of the art fluid simulation models and also demonstrate generalization across different materials.

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.