SiT: Simulation Transformer for Particle-based Physics Simulation
Authors: Yidi Shao
Chen Change Loy
Bo Dai
Venue: ICLR 2022 Submitted (2022)

Abstract: Most existing particle-based simulators adopt graph convolutional networks (GCNs) 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 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, SiT generates interaction tokens 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 systems 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.

Under review as a conference paper at ICLR 2022
SIT: SIMULATION TRANSFORMER FOR PARTICLE-BASED PHYSICS SIMULATION
Anonymous authors Paper under double-blind review
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 un- derlying physies of particles. However, they force particles 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 unsat- isfactory performance, especially in generalization. This paper proposes Simula- tion 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 gen- erates 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 abili- ties 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.
INTRODUCTION
Particle-based physics simulation is classic and important topic in computer science. It not only fa- cilitates 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. Dif- ferent from traditional simulators, such as grid-based (Guo et al.] 2016) and mesh-based (Bronstein) 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-wis 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 system 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 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 simulate a waterfall and a river.
Recent particle-base basicle-bbaeddeee simulators (Battaglia et al.) 2016 Schenck & Fox 2018 Mrowca et al. 2018 Li et al. 2019. Sanchez-Gonzalee et al.. 2020; Ummenhoter et al. (2200) often view a system as 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 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 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 selective mechanism, which leads to computational

---

Under review as a conference paper at ICLR 2022
FluidFall
FluidStuke
RiceGrip
BoxBath
SIT DIIANE Grous
Time
(8)
(a), Samples of the datasets. FluidFall contains two drops of water. FluidShake simulates block of water in moving box. RiceGrip has deformable object squeezed by two grippers. BoxBath contains rigid cubic washed by water. (b). Rollouts from generalizations of RiceGrip, where we increase the number of particles. SiT with abstract tokens, marked by +, can still maintain the object's shape, while DPI-Net (Li et al. 2019) fail to simulate the deformable object.
Figure 1;
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 semanties, such as the shape of 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 novel Transformer-based framework, referred to as Simulation Trans- former (SiT), for particle-based physics simulation. The model inherits the powerful multi-head self-attention mechanism in Transformer (Vaswani et al. 2017) to capture partiele interactions. To further encourage efficient modeling of complex interactions, instead of treating particle interac- tions as attention weights obtained by dot-product, we introduce the notion of interaction tokens, which are high-dimensional representations for interactions, 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-speeific semantics, SiT further learns a high- dimensional 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 cap- turing particle interactions explicitly with interaction tokens and allowing dynamic inter-token at- tention, SiT dynamically focuses on essential particle interactions and reduces the computations for redundant and noisy ones. This is crucial for particle-base simulation, especially for domains containing hundreds and thousands of densely placed particles, where the modeling of all particle interactions redundant, expensive and noisy in practice. Second, thanks to the trainable abstract tokens that disentangle intrinsic material characteristics from domain-specift semantics, reuse them to apply SiT to unseen domains of the same materials without retraining. As shown in our experiments, after training yields fairly faithful simulations when the cubic is replaced with a ball or a bunny comparing with previous work.
domain consisting of fluid water and rigid cubic, SiT still
To show the effectiveness of SiT, we perform extensive evaluations on four standard environments commonly used in the literature (Li al., 2019 Sanchez-Gonzalez 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.
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 al. 2020) adopt convolutional architectures for learning high- dinensinnnional physical system, while mesh-base based simulations (Bronstein et al., 2017 Luo al. 2020 Hanocka et al. 2019; Nash et al. 2020 Qiao
ct al. 2020 Weng et al. 2021 Pfaff et al. 2021) typically simulate objects with continuous surfaces, such clothes, rigid objects, surfaces of water and on.
2

---

Under review as a conference paper at ICLR 2022
Many works (Battagla et al., 2016; Schenck & Fox 2018 Mrowca et al. 2018 Li et al. 2019
Sanchez-Gonzalee et al. 20200 Ummenhofer et al.) (020) simulate physics on particle-based sys- tems, 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 (SP-
Nets) (Schenck & Fox (018)) implemented fluid dynamics using position-based fluids (Macklin &
Müller 2013). 2013). Hierarchical Relation Network (HRN) (Mrowca et al., 2018 predicted physical dy-
namics based on hierarchical graph convolution. Dynamic Particle Interaction Networks (DPI-Net)
at 2019) combined dynamic graphs, multi-step spatial propagation, and hierarchical structure simulate particles. CConv (Ummenhofer et al. 2020 used spatial convolutions to simulate fluid
particles. Graph Network-basse Simulators (GNS) (Sanchez-Gonzalez et al. 2020) computed dy-
namics via learned message-passingg Similar to particle-based systems, COPINGNet (Shao et al.
202 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 inter- act with all its nearby particles without selective mechanism. In contrast, our SiT employs both particle and interaction tokens and selectively focus necessary particle interactions through atten- tion 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-ar performance in many natural langruage processing tasks (Devlin et al. 2019 Radford et al. 20199) Brown et al. 2020). Recently, Transformer starts to show great expand- ability and applicability in many other fields, such as computer vision (Wang et al. 2018 Carion itt al. 2020, Dosovitskiy et al., [2021; Wang et al.. [202]] et al.) 202 I) and graph representations
(Zhou al.] |020] Zhang et al. [020 Dwivedi & Bresson
20200 To our knowledge, no attempt
has been made to apply Transformer on physics simulation.
Our SiT models the interactions between partieles 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 (020)) which we refer as GraphTrans for short. However, there are differences in 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 scalar for fur- ther 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 to- weighted scores any more. We adopt GraphTrans (Dwivedi & Bresson, (020)) in particle-based simulation and compare it with SiT in experiments. The quantitative results show that SiT achieves better results than GraphTrans.
kens will attend to each other. The interaction tokens
METHODOLOGY
3.1 PROBLEM FORMULATION
For particle-basee based system composed N particles, we use X {2 )ー= to denote the system state at time step t. where æ denotes the state of i-th particle. Specifically, x = [p], q a.], where p of E 3R3 refer to position and velocity, and a, E &Rd represents fixed particle attributes such as its material type. The goal of simulator to learn model Φ(-) from previous rollouts of system to causally predict a rollout trajectory in a specific time period conditioned on the initial system state X The prediction is run in a recursive manner, where the simulator will predict
the state p = <(x') at time step + 1 based on the state X = {æf] time step In practice, we will predict the velocities of particles Q = {q] }, and obtain their positions via P a+1 = pf + Al g where 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, they indicate the energy transition of system when constrained by
3

---

Under review as a conference paper at ICLR 2022
MLP
Atlention
Encode
Decode
Network
X
Leanablle emhellding
Particle States
Internction Tokens
Abstract Tokens
Particlo Tokens
Particle Velocities
Figure 2: Overview of our transformer- based model with abstract tokens, which optional. SiT is composed an encoding layer, stacks of blocks extended from Transformer encoder blocks, and decoding layer. Given the particles states at time t, the encoding layer will first encode the particles' states and learnable embeddings into particle tokens and abstract tokens, and initialize interaction tokens by particles' states. In each blocks, the sub-network first updates interaction tokens given particle tokens. Then all tokens will dynamically attend each other via uniform attention and update particle and abstract tokens. An MI.P attached after attention. Finally, decoding layer will predict the particles' velocities time + given particle tokens.
material characteristics and physical laws. However, it is infeasible to know a priori how should particles interact with each other. Thus, 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, is naturally a good backbone of building an efficient particle- based simulator. To apply vanilla Transformer in particle-basee simulation, first encode the states of particles into corresponding particle tokens = {wa} by
(1) where of E Rd is a d, dimensional vector and (v)((.) is an encoding layer implemented as a MLP. Subsequently, particle interactions achieved by T. blocks of self-attention modules, where in the 1-th block, particle tokens will attend to each other selectively as:
of -V))
Σ o
exp(wi;)
ddn E, exp(wx)) Toje
w
(4)
Since system usually contains hundreds of particles and the interactions among particles oceur within the neighbors in our settings, considering all possible interactions, of which the number quadratically increases with respect the number of particles, is computationally redundant and in- efficient. Therefore, we follow previous literature Li et al. (2019) [Sancee Sanchez-Gonzalez et al. (2020) to assume that interactions of distant particles can be omitted, which realized by window function:
(Ilp) l/1l2 < R),
where ((condition) is a indicator function that returns if the condition is satisfied, and 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, decoding layer is applied to the updated token of i-th partiele to obtain its velocity:
c c+1
( (ottt))
(6)
where fele(.)) is implemented by another MLP.
3.3 SIMULATION TRANSFORMER (SIT)
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, vanilla
4

---

Under review as a conference paper at ICLR 2022
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 partiele 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 par- ticles, domain-specific semanties, such as the shape of rigid cubic, are mixed with domain-agnosio semantics, such as the stickiness attribute of materials. In such cases, the former may be misinter- preted as the latter, leading to poor generalization ability. To overcome the limitations of vanilla Transformer, we propose novel Transformer-basee former (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
simulator, referred to as Simulation Trans-
Interaction Tokens. To capture the rich semantics of particle interactions, instead a scalar, SiT assigns an interaction token to each potential particle interaction. Specifically, the token for the interaction between particles ær, and ær is initialized by
u
f7 (x),,)))))
where u/y E dR is a da dimensional vector and f?ne() (-) is the encoding function for interactions.
Subsequently, in each self-attention block SiT will update the interaction tokens
1+1, + (of² o/"" uuly s(pp.p))
where (-)-) is learnable sub-network shared for all interaction tokens. Inspired by GNN-based methods (Mrowca et al. 2018 Li et al.] 2019); Manchez-Gonzalle et al., 2020) SST adopts the learnable sub-network to enlarge its capacity for handling complex particle interactions. In practice, the sub-network is implemented as
v (o .0 uij u.f,,(plspp)) s(p(.))) = MLP([U)* , ,.fi, s(P..))).
As is shown in Figure 6 the update of particle tokens in the self-attention module is thus achieved by
a
Σ
w + 3
ca 1,
sk
(10)
3
k
exp( (wx))
(11)
Vdn. (aE, (esp(w") + B Ex exp(wx))
exp(wxt
w
(12)
Jhh a [a exp(urf) + B BBE, exp(wxk)
w
(v) 1. o To
(13)
w (of uxx 1+1.
(14)
where a, 8 E {0,1) 1) are the gate controllers, j and are the indexes of particles within the recep- tion window defined by Equation 5) Equation 13 computes the attention scores between particle tokens, and Equation [4 computes the attention scores between partiele and interaction tokens. Sub- sequently, these attention scores are normalized by Equation T1 and Equation [2] They are then used by Equation 10 to aggregate semantics of particle and interaction tokens, where the gate controllers
a, are hyper-parameters that work as balancing coefficients. When 1, 3 0, Equation TO 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
worth noting that SiT provides a more flexible way for further extensions when compared to vanilla Transformer. Specifically, we can change the value of a and 33 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.
5

---

Under review as a conference paper at ICLR 2022
Abstract Tokens. To improve the generalization ability of SiT and disentangle domain- specific mantics from its domain- domain-agnostic counterparts, we further apply SiT with anterial-specife abstract tokens.
For N, types of materials, SiT adopts Na abstract tokens A {a | each of which a learn- able vector of the same length of particle tokens. Ideally, the abstract token ag should capture the aomain-agnostic semantics of k-th material. They act as particle tokens for additional abstract par- ticles, following the same update formula as particle tokens. Therefore with N, abstract tokens SiT will have N + Na particle tokens in total: {a),... ONN,: #1. x ). 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 general- izing to domains that have same materials but vary in particle amount and configuration.
3.4 TRAINING OBJECTIVE
To train SiT with previous rollouts, the standard mean square error (MSE) loss is applied on the output of our SIT:
M ,10H)) = N Σ g
9, g.-11111 1122,
(15)
where Q++1 = [(] 1==1 is the estimated velocity, Q++ = {q M==1 is the ground truth, and II- 1// 1.2 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. such cases, the MSE loss will be biased to the material with more particles. Therefore, we further apply material weighted MSE (WMSE) loss to reduce the effect of imbalance:
WMSE(Q**1. ,Q*!!)
1 N
Σ
1
Σ llg
L 112
(16)
where K is the number of material types, Ng is the number of particles belonging to the k-th material ) ground truth for k-th material, and Qu+1 1Q Jk=1 Q111= (Qt+ )
+++ and = Σx Nr. Q++1 = ( k Jizi and Q++11 = {g N₂ denote the estimated velocity and
[Q
EXPERIMENTS
4.1
BASE ENVIRONMENTS
We adopt the environments commonly used in the literature (i et al. 2019 Sanchez-Gonzaler Gonzalez et al., 2020 Ummenhoter et 2020). Samples are displayed in Figure 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 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.3
We compare SIT with 4 recent approaches: DPI-Net (Li et al. 2019), CConv (Ummenhofer et al.) 2020), GNS (Sanchez Gonzzlezz 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
all environments with fewer model parameters. The effec- tiveness of abstract tokens are more obvious for RiceGrip and BoxBath, which are more complex. Quantitative results show that SiT can better simulate the dynamics different environments.
Comparison with DPI-Net and GNS. DPI-Net Li al. 2019) and GNS (Sanchez-Gomme Gonzalez et al.) 2020) adopt cassage-passing graph networks for particle- based simulation. As shown in the results,

---

Under review as a conference paper at ICLR 2022
Quantitative results and model parameters four enviromments. SiT achieves superior and rea- sonable performance with model parameters on all environments. When adding trainable abstract tokens, which introduce a few more parameters, SiT, marked by +, further improves performance in complex envi- ronments. As CConv is designed for fluid dynamics, we report the results on BoxBath for reference of the simulation fluid parts, which is marked by
Table 1:
FluidFall
FluidShake
RiceGrip
BoxBath
Methods
MSE
Params
MSE
Params
MSE
Params
MSE
WMSE
Params
DPI-Net
0.12=00.06
0.62M
1.43=0.52
1.02M
0.11=0.21
2.17M
1.91
1.52:+.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+0087
2.62=0.87
1.59M
GraphTrans
0.12-10.03
0.49M
8.77113997
0.96M
0.36=0.27
0.96M
3.17:10.07
2.33
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.40M
1.08-+0.39
0.77M
0.07=0.07
0.82M
1.57-10.06
0.39
0.77M
Ground Truth
DPI-Net
CConv
GNS
GripTTans
SIT
SIT
Phsdiall (e-51)
DooDails (x-1)
Figure 3: Qualitative results FluidFall and BoxBath. When the two drops getting closer in FluidFall hut close enough merge, previous methods are likely to mix them together due the newly added neighbors from different drops and incorrect interactions. But SiT can still predict faithful rollouts.
SiT achieves better performance, as the selective mechanism can eliminate the unnecessary interac- tions, which further demonstrated by the heat map in Appendix 6.2.1 SiT with abstract tokens further improves the performances especially ness of abstract tokens modeling complex deformations and multi-materials interactions. For FluidFall, though the MSEs are close 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.
RiceGrip and BoxBath, suggesting the effective-
Comparison with CConv. CConv (Ummenhofer al. 2020 designs convolutional layers carefully tailored to modeling fluid dynamics, such as an MPnaMhnaghan 1992 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, only report result on BoxBath for reference off simulation on fluid parts.
Comparison with Graph Trans (Dwivedi & Bresson 2020). The GraphTrans is also a Transformer- based method and updates the interactions using element- wise product from node representations. We represent particles by nodes and adopt layers in Graph same hidden dimensions and number of blocks. Other settings are completely the same as SiT for better comparison. As is shown in experiments, Graph fails simulate particles 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.
backbone, which and the
4.2 GENERALIAZTIONS
To challenge the robustness of model more complex settings, add more particles for RichGrip and FluidShake, and change the size and shape of rigid object for BoxBath. 5 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 all environments. Quantitative results are summarized in Table while qualitative results are depicted in Figure 4 and Figure (b). As is shown, SiT with abstract tokens better simulate particle dy-

---

Under review as a conference paper at ICLR 2022
Table 2: MSEs generalizations. The lists of numbers in FluidShake and RiceGrip the range of particles, while tuples BoxBath denotes number fluid particles, number rigid particles, and shape of rigid objects respectively. Training settings are marked by +. We compare SiT with abstract tokens and DPI-Net. SiT+ achieves the best results on all cases.
FluidShake [450,627]*
RiceGrip [570,980]*
Methods
[720,900]
[924,1080]
[1104,1368]
[1060,1345]
[1347,1640]
DPI-Net
2.14
2.78=00.58
3.53+10771
2.33++30.64
0,87+t4.00
SiT+ (Ours)
1.54+0037
1.90=0.63
2.30:10.69
0.14+0.10
0.17=00.15
BoxBath (
Methods
(1280,64,cubic)
9600,41bunny
(960,125,cubic)
(960,136,ball)
(960,120,cuboid)
DPI-Net
4.28++0.11
2.49:10.23
2.75++0.42
2.71++0035
2.97:+00338
SiT+ (Ours)
0.64
1.72:+0.08
1.744-0.14
1.67:+0.13
0.32
Ground Iruth (I-123)
DPI-Net
SIT
Ground truth (1-60)
DPI-Net
FluidShake
Figure 4:
Rendered rollouts on generalized FluidShake and BoxBath. More particles participate in simulations. SiT with abstract tokens (SiT+) can still faithfully predict the wave of fluid in FluidShake and the positions of rigid cubic in BoxBath.
namics 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.223
4.3 ABLATION STUDIES
We comprehensively analyze SiT and explore the effectiveness of our model in the following ways: (a) the necessity interaction tokens; (b) the architectures sub-network, which generates interaction tokens; (c) the values for gate controllers a, 3; (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 FluidShake for (a) to (e). We verify (f) on BoxBath and RiceGrip. The quantitative results are in Table and Table
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 one dimension. As shown in Table 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 self- attention blocks, where q 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-paranters for a, 3. We set a = 1, F, and = 0, 3 = in configuration H, Consequently. both configuration E and H achieve good results, suggesting that the participants interaction tokens are the key to boost the performance. configuration H for SiT which shows the best results.
in configuration E, a = 1, 3 in configuration
In practice,
Multi-head uniform attention. We apply multi-head mechanism configuration G and set head number 8. The results show that multi-head mechanism bring some improvement. But it also introduces more parameters and consumes lots of memory when computing.
Sensitiveness to R. Quantitative results reported on FluidShake. As shown in Table 4 SIT is more robust when varying the radius R, suggesting the effectiveness selective mechanism and robustness of SiT.
8

---

Under review as a conference paper at ICLR 2022
Table 3: Ablation studies. We comprehensively explore the effectiveness of SiT, including the necessity of interaction tokens, the structure sub-network, the value of a, 3 and the usage of multi-hcad in uniform attention. We report (SSs(1e-2) on FluidShake, which is complex environments involving outer forces.
Configurations
Interaction modeling
B
Multi-head
FluidShake
A A(Transformer)
Scalar by dot product
0
21.03++10.54
B
Scalar by MLP
0
1
300.43:17..1
Tokens by element-wise product
0
1
17.70:16.99
Tokens by self attention
3.1710.81
Tokens by MLP
0
14.27+42.00
F
Tokens by MLP
1
1.19=0.44
G H(SiT)
Tokens by MLP
0 1 1
v
0.04
Tokens by MLP
1.36
Table 4: Ablation studies on sensitiveness to radius R and abstract tokens. The left parts are MSFs(1es 1e-2) on FluidShake, which suggest that SiT is more robust within small range of R. Our default setting all domains marked by *. The right parts MSEs(Ic-2) RiceGrip and BoxBath. We replace abstract tokens with dummy tokens, which are fixed value vectors have same connectivities abstract tokens.
Methods 0.07
Aº 0.08
R = 0.09
Methods
RiceGrip
BoxBath
DPI-Net
2.60=0.56
1.43:0.52
1.66:10.48
SiT w dummy tokens
2.12:10:46
3.98:10.09
SIT
1.38:10.36
1.08:1.36
1.37:10.35
SIT +
0.07+0.00
1.57
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 41 SiT with dummy tokens fails on both domains, suggesting the abstract tokens are able to learn the materials semantics and boost SiT's performances.
4.4 FEW SHOTS LEARNING
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 I 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-wee 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.
As shown in Figure 5
Grand Trcn(t 200)
Predictiong 20% data)
Profictio(411% data)
545 Care
Predictina((tr dana
Prediction 50% data)
Profiction (00% deo)
30
(a)
(b)
Figure 5: Few shots learning on FluidShake. Figure (a) shows the rollouts. Figure (h) draws the MSE curves as the number of training examples increases. Our model are robust even with only 60% of training examples.
CONCLUSION
In this paper, 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 semanties. Finally, SiT makes successful attempt to apply Transformer into physics simulation and achieve superior performances over existing methods.

---

Under review as a conference paper at ICLR 2022
REFERENCES
Peter W. Battaglia, Razvan Pascanu, Matthew Lai, Danilo Jimenez Rezende, and Koray Kavukcuoglu. Interaction networks for learning about objects, relations and physics. In Ad- vances in Neural Information P Processing Systems 29: Annual Conference on Neural Information Processing Systems 2016, December 5-10, 2016, Barcelona, Spain, 2016.
Michael M. Bronstein, Joan Bruna, Yann LeCun, Arthur Szlam, and Pierre Vandergheynst. Geo- metric deep learning: Going beyond euclidean data. IEEE Signal Process. Mag., 2017.
Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhari- wal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herberrt-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler. Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei. Language models are few-shot learners. In Advances in Neu- ral Information Processing Systems 33: Annual Conference on Neural Information Processing Systemy 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020.
Nicolas Carion, Francisco Massa, Gabriel Synnaeve, Nicolas Usunier, Alexander Kirillov, and Sergey Zagoruyko. End-to-end object detection with transformers. In Computer Vision ECCV 2020 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part 2020.
Jacob Devlin, Ming-Weo Mei Chang, Kenton Lee, and Kristina Toutanova. BERT: pre-training of deep bidirectional transformers for language understanding. In Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies. NAACL-HLT 2019, Minneapolis, MN, USA, June 2-7, 2019, Volume I (Long and Short Papers), 2019.
Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszko- reit, and Neil Houlsby. An image is worth 16x 16 words: Transformers for image recognition scale. In 9th International Conference on Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021, 2021.
Vijay Prakash Dwivedi and Xavier Bresson. generalization of transformer networks to graphs. CoRR, 2020.
Xiaoxiao Guo, Wei Li, and Francesco Torio. Convolutional neural networks for steady flow ap- proximation. In Proceedings of the 22nd ACM SIGKDD International Conference on Knowledge Discovery and Data Mining, San Francisco, CA, USA, August 13-17, 2016, 2016.
Rana Hanocka, Amir Hertz, Noa Fish, Raja Giryes, Shachar Fleishman, and Daniel Cohen-Or. MeshCNN: network with an edge. ACM Trans. Graph., 2019.
Thomas N. Kipf and Max Welling. Semi-supervised classification with graph convolutional net- works. CoRR, 2016.
Sangseung Lee and Donghyun You. Data-driven prediction of unsteady flow over a circular cylinder using deep learning. Journal of Fluid Mechanics, 2019.
Yunzhu Li, Jiajun Wu, Russ Tedrake, Joshua B. Tenenbaum, and Antonio Torralba. Learning particle dynamics for manipulating rigid bodies, deformable objects, and fluids. In 7th International Conference on Learning Representations, ICLR 2019, New Orleans, LA, USA, May 6-9, 2019, 2019.
Ze Liu, Yutong Lin, Yue Cao, Han Hu, Yixuan Wei, Zheng Zhang, Stephen Lin, and Baining Guo. Swin Transformer: Hierarchical vision transformer using shifted windows. CoRR, 2021.
Ran Luo, Tianjia Shao, Huamin Wang, Weiwei Xu, Xiang Chen, Kun Zhou, and Yin Yang. NNWarp: Neural network-based nonlinear deformation. IEEE Trans. Vis. Comput. Graph., 2020.
Miles Macklin and Matthias Müller. Position based fluids. ACM Trans. Graph., 2013.
10

---

Under review as a conference paper at ICLR 2022
Joe J Monaghan. Smoothed particle hydrodynamics. Annual review of astronomy and astrophysics, 1992.
Damian Mrowca, Chengxu Zhuang, Elias Wang, Nick Haber, Li Fei-Fei, Josh Tenenbaum, and Daniel L. Yamins. Flexible neural representation for physics prediction. In Advances in Neu- ral Information Processing Systems 31: Annual Conference on Neural Information Processing Systems 2018, NeurlPS 2018, December 3-8, 2018, Montréul, Canada, 2018.
Charlie Nash, Yaroslav Ganin, S. M. Ali Eslami, and Peter W. Battaglia. PolyGen: An autoregressive generative model of 3d meshes. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, 2020.
Tobias Pfaff, Meire Fortunato, Alvaro Sanzhez-Gonzalez, and Peter W. Battaglia. Learning mesh- based simulation with graph networks. In 9th International Conference on Learning Representa- tions, ICLR 2021, Virtual Event, Austria, May 3-7, 2021, 2021.
Yi-Ling Qiao, Junbang Liang, Vladlen Koltun, and Ming C. Lin. Scalable differentiable physics for learning and control. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, 2020.
Alec Radford, Jeffrey Wu, Rewon Child, David Luan, Dario Amodei, Ilya Sutskever, et al. Language models are unsupervised multitask learners. OpenAl blog, 2019.
Alvaro Sanchez-Gonzalez, Jonathan Godwin, Tobias Pfaff, Rex Ying, Jure Leskovec, and Peter W. Battaglia. Learning to simulate complex physics with graph networks. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, 2020.
Connor Schenck and Dieter Fox. SPNets: Differentiable fluid dynamics for deep neural networks. In 2nd Annual Conference on Robot Learning, CoRL 2018, Zürich, Switzerland, 29-31 October 2018, Proceedings, 2018.
Han Shao, Tassilo Kugelstadt, Torsten Hädrich, Wojtek Palubicki, Jan Bender, Soren Pirk, and Do- minik Michels. Accurately solving rod dynamics with graph learning. In Thirty-Fifth Conference on Neural Information Processing Systems, 2021.
Nils Thuerey, Konstantin Weißennw, Lukas Prantl, and Xiangyu Hu. Deep learning methods for reynolds-averaged navier-stokes simulations of airfoil flows. AIAA Journal, 2020.
Benjamin Ummenhofer, Lukas Prantl, Nils Thuerey, and Vladlen Koltun. Lagrangian fluid simula- tion with continuous convolutions. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020, 2020.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. In Advances Neural Infor- mation Processing Systems 30: Annual Conference on Neural Information Processing Systems 2017, December 4-9, 2017, Long Beach, CA, USA, 2017.
Huiyu Wang, Yukun Zhu, Hartwig Adam, Alan Yuille, and Liang-Chieh Chen. MaX-Deeplab: End-to-end panoptic segmentation with mask transformers. In IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2021, virtual, June 19-25, 2021, 2021.
Rui Wang, Karthik Kashinath, Mustafa Mustafa, Adrian Albert, and Rose Yu. Towards physics- informed deep learning for turbulent flow prediction. In KDD 200 The 26th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, Virtual Event, CA, USA, August 23-27, 2020, 2020.
Xiaolong Wang, Ross B. Girshick, Abhinav Gupta, and Kaiming He. Non-local neural networks. In 2018 IEEE Conference on Computer Vision and Pattern Recognition, CVPR 2018, Salt Lake City, UT, USA, June 18-22, 2018, 2018
Zehang Weng, Fabian Paus, Anastasiaa Varava, Hang Yin, Tamim Asfour, and Danica Kragic. Graph-based task-specific prediction models for interactions between deformable and rigid ob- jects. CoRR, 2021.
11

---

Under review as a conference paper at ICLR 2022
Jiawei Zhang, Haopeng Zhang, Congying Xia, and Li Sun. Graph-Bert: Only attention is needed for learning graph representations. CoRR, 2020.
Dawei Zhou, Lecheng Zheng, Jiawei Han, and Jingrui He. A data-driven graph generative model for temporal interaction networks. In KDD '0: The 26th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, Viriual Event, CA, USA, August 23-27, 2020, 2020.
12

---

Under review as a conference paper at ICLR 2022
APPENDIX
6.1 MODEL DETAILS
6.1.1 UNIFORM ATTENTION
SoftMax
Scale
Abstrett
Interaction Tokens
Figure 6:
Overview of our uniform attention. The abstract tokens are optional. The tokens will
attend to each other and update accordingly.
Figure illustrates how our uniform attention works. Our uniform attention extend self-attention (aawwani et al.) (017) 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.
6.1.2 MULTI-HEAD UNIFORM ATTENTION
The multi-head version of uniform attention is formulated by
W [SA1, SAz, .SAM]
(17)
SAm
a: Σ
m. Wij
mm 1, + 3 Σ
With
mm (Vm (.
(18)
k
exp(umm..)
m,,
(19)
Vdh aa E, exp(wx"")) + B Ex (
exp(wik m.u)
(20)
Vdn. (a E, exp(wx" """"""""))) + 3 Mn exp(w m
m,o
'Qmlp of 1-1,1 (Kmlp)) 10 111))
(21)
sj
Wikk
m,u
(Q (Qm.l 1-1.) (K ()
L.
(22)
where m E .... M is the index of head, W E Rdd xd2 is the output weight for the concate- nation of M heads; Qu.. K Km.,, ym., are the projection weights of query, key, value for head block L. The rest symbols are the same with those in uniform attention.
6.1.3 IMPLEMENTATION DETAILS
Inputs and outputs details, For FluidFall, FluidShake, and BoxBath, we only use particles' states at time inputs and output the velocities at time t+ For RiceGrip, we concatenate particles states
13

---

Under review as a conference paper at ICLR 2022
from to t 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 as inputs and output the velocities at time f 1 except RiceGrip, where we concatenate particles states from t 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 blocks shown in Figure 2 for FluidFall and blocks for the other three environments. We set = 0, 8 = 1 in our model for default, which show the hest 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 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 [6 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 mainly compare the fluid simulations. The model-related typer-parameters for baselines are the same as the original papers, such as the number of message passing is 10 for GNS. Other
for baselines are first chosen the same as their original papers, and then fine-tuned within small range of changes. For example, in terms of the initial learning rate, 0,0001 works better on GNS while others adopt 0.0008.
6.2 EXPERIMENTS
6.2.1 ADDTIONAL ABLATION STUDIES
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 (b).
6.2.2 BASE ENVIRONMENTS
Dataset details. We use the same setting for 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 randomly ini- tialized. 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 mov- ing 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 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-
14

---

Under review as a conference paper at ICLR 2022
(a) Sample in BoxBath at time = 42.
0.10
1.0
0.30
1.0
0.10
1.0
0.10
1.0
0.05
0.05
Y 0.05
Y 0.05
0.45
0.50
0.0
0.45
0.50
0.0
0.45
0.50
0.0
0.45
0.50
0.0
(b) Attention scores (c) Attention scores (d) Attention scores (e) Attention scores block
block
block
block
Figure 7: (a) is a sample frame on BoxBath. We select one particle in the center of red circle to show its attention scores interactions from neighbors. (b)(c)d)(e) show the heat maps from uniform attention in each block of SiT. The attention scores on interactions are assigned to the corresponding neighbors, which shown by different colors. The darker the color, the lower attention score particles.
ticles with 151 frames for each rollout. There 2700 rollouts in training set and 300 rollouts in validation set.
Results. We display some rollouts for all experiments in this section. The results RiceGrip for CConv are not available because CConv was designed specifically for fluid dynamics. The results on BoxBath for CConv only used as reference of the simulation fluid parts. And compare CCony with SiT mainly on the fluid simulations.
15

---

Under review as a conference paper at ICLR 2022
Ground Truth
t= 30
t 38
t 44
t= 54
t= 62
DPI-Net
CConv
GNS
Graph Trans
SiT
S
Figure 8: Rollouts FluidFall. Here display the settings where the larger drop is above the smaller drop. DPI-Net and GNS are likely mix the drops together, while CConv alters the shape the upper drop before falling on the floor. GraphTrans has worse performance maintaining shape water drops. SiT changes little for the shape of the upper drop. When using abstract tokens, which is marked by +, SiT can achieve hetter results which are closer to the ground truth.
16

---

Under review as a conference paper at ICLR 2022
DPI-Net Ground Truth
t= 63
t 92
t= 146
t= 260
t= 296
CConv
GNS
GraphTran rrans
SiT
SiT+
Figure 9: Rollouts FluidShake. DPI-Net still has some artifacts in predicting the wave of water, such as the
rollout at 63. CConv, which is designed specifically for fluid dynamics, rollouts fairly good results, such surface of water, this domain. GNS has artifacts when predicting waves 63) and predicts overly smooth surfaces of fluid (t 90). GraphTrans fails this domain, suggesting the insufficiency of capturing the complex interactions from outer forces. SiT can also achieve compatible results with less parameters.
17

---

Under review as a conference paper at ICLR 2022
Ground Truth
8
t= 14
t= 22
t= =030
t= 38
DPI-Net
GNS
Gruph Trans
SiT
SiT+
Figure 10: Rollouts on RiceGrip. For the surface the rice, we can see at t 38 that DPI-Net seems to put more pressure to the rice; GNS and GraphTrans has difficulties in maintaining the smooth surfaces of the rice. SiT is capable of predicting faithful rollouts, while SiT with abstract tokens can better simulate the deformable object.
18

---

Under review as a conference paper at ICLR 2022
t= 24
t 43
t= 52
t= 78
t= 95
DPI-Net Ground Truth
CConv
GNS
Graph Trams
SiT
S
Rollouts BoxBath. As CConv is designed for fluid dynamics, it treats all particles fluid and fails this domain. GNS able to handle rotations rigid box, has difficulties in modeling motions of rigid box when pushed by waves fluid, suggesting the insufficient abilities dealing with interactions between different materials. SiT with abstract tokens achieves more faithful results in terms of the rotation of rigid cubics, the interactions between fluid particles.
Figure 11:
19

---

Under review as a conference paper at ICLR 2022
Table 5: MSE and WMSE results from models on generalizations. The lists of numbers in FluidShake and RiceGrip are the range of particles, while the tuples in BoxBarh denotes number of fluid particles, number of rigid particles, and shape of rigid objects respectively. Training settings are marked by
FluidShake [450,627]*
RiceGrip [570,980]*
Methods
[720,900]
[924,1080]
[1104,1368]
[1060,1345]
[1347.1640]
DPI-Net
2.14:10.36
2.7811.58
3.5310.71
2.33:130.64
0.87:1:000
CConv
1.86.10.42
2.23
2.50:10.77
N/A
N/A
GNS
2.06++0.75
2.82:+1.93
2.98+1.22
0.56++0.31
0.57=t0.29
GraphTrans
10.24:33.5
12.27+4.80
12.72++3755
0.63:::028
0.68++0.37
SiT+ (Ours)
1.54+0037
1.90=0.63
2.30++0.69
0.14+00.10
0.17=0015
BoxBath (960,64,cubie)*
Methods
cuubi)
(960,41.bnny))
(960,125,cubic)
(960,136,ball)
(960,120,cuboid)
DPI-Net
4.28:1:.111
2.49:10.23
2.75+:0.42
2.71
2.97:10.38
CConv
4.30
0.03+
3,214:0.82
3.80++0.22
3.05+t
GNS
2.87-10.19
2.22:10.28
1.971:0288
2.08:0022
181110.18
GraphTrans
3.87:10.08
3.55
3.92:10.15
3.86:10.17
3.56:10.13
SiT+ (Ours)
2.6410.111
1.72:1.08
1.74:10.14
1.67-0.113
1.92:10.32
BoxBath WMSE (
Methods
(1280,64.cubic)
(960,41.bunny)
1960,125,cubic)
(960,136,ball)
(960,120,cuboid)
DPI-Net
3.14
3.65
3,42:+0.90
3.13::0.68
3,96
CConv
4.30:11.09
3.59
4.09110.22
4.43:10.35
3.43:10.90
GNS
4.55
8.09.13.32
3.40:1:1.0
3.60:10.74
2.18:10.66
GraphTrans
3.37-1-0775
3.59
5.31-10.73
5.38:10.69
3.81-1:0775
SiT+ (Ours)
1.82:10.21
1.81
1.9410.511
1.76:10.35
1.86
Table 6: Details of generalization settings.
Environments
Number of training particles
Number of generalization particles
FluidShake Large
[450, 720]
[720, 1500]
RiceGrip Large
[558, 1060]
[1060, 1798]
BoxBath Large Fluid
Fluid: 960.
Rigid: 64
Fluid: 1280.
Rigid 64
BoxBath Bunny
Fluid: 960.
Rigid: 64
Fluid: 960.
Rigid: 41
BoxBath Large Cubic
Fluid: 960.
Rigid: 64
Fluid: 960.
Rigid: 125
BoxBath Large Ball
Fluid: 960.
Rigid: 64
Fluid: 960.
Rigid: 136
BoxBath Large Cuboid
Fluid: 960.
Rigid: 64
Fluid: 960.
Rigid: 120
6.2.3 GENERALIZATION DETAILS
We report WMSE results BoxBath and other models' generalization performance for references in 5able5 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 980; the generalized FluidShake contains particles from 720 to 1368, while 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 rollouts on environments.
20

---

Under review as a conference paper at ICLR 2022
DPI-Net Ground Truth
t= 64
t= 123
t= 179
t= 215
t= 296
CConv
GNS
SiT+ GraphTrans
Figure 12: Rollouts on generalized FluidShake. DPI-Net achieves fairly good performance, while it still has difficulties to correctly predict the waves, such as the rollout at = 123. CConv and GNS have similar problems, such as the rollout at f = 215 and I = 123 respectively. GraphTrans still fails in this environment. Our S can still achieve hetter performance.
21

---

Under review as a conference paper at ICLR 2022
Ground Truth
t= 20
t= 23
t= 26
t= 30
t= 38
DPI-Net
GNS
Graph Trans
SiT+
Rollouts on generalized RiceGrip. DPI-Net fails and cannot maintain the shape of the deformable object, as is shown in the prediction at t = 38. GNS tends to rotate the corner of the rice, as is shown at t 38, the left lower comer is twisted when the grippers do not have contact with the rice's surface. Our S has better rollouts than GraphTrans, suggesting the effectiveness and robustness of our method.
Figure 13:
22

---

Under review as a conference paper at ICLR 2022
t= 35
t= 41
t= 60
t= 96
t= 118
Ground Truth
DPI-Net
CConv
GNS
GraphTrans Trans
SiT+
Rollouts on generalized BoxBath, where we add more fluid particles. DPI-Net misses the cubic at the end. While GNS is able to simulate the rotations of rigid box, it fails to make further predictions of the rigid after falls into the water. While SiT+ predicts wrong motions of several fluid particles at the beginning, it still achieve better results compared with other models.
Figure 14:
23

---

Under review as a conference paper at ICLR 2022
46
60
74
t= 88
t= 102
Ground Truth
DPI-Net
CConv
GNS
GraphTran Trans
S
Figure 15: Rollouts generalized BoxBath, where add rigid particles. DPI-Net starts predict
wrong rollouts for fluid particles at the beginning, worse results for rigid cubic. GNS seems overfit the rotations of rigid box. SiT+ achieves better results comparing with other models.
24

---

Under review as a conference paper at ICLR 2022
29
38
t= 107
t= 144
Ground Truth
DPI-Net
CConv
GNS
GraphTran Trans
S
Figure 16: Rollouts generalized BoxBath, where change the shape from cubic bunny. The shape of
bunny is more challenging, as ears the bunny do not have direct connections the body. While DPI-Net
fails in this scene, the bunny is submerged in the rollouts of GNS. SiT+ predicts rollouts as correctly as possible.
25

---

Under review as a conference paper at ICLR 2022
38
44
58
t= 68
t= 113
Ground Truth
DPI-Net
CConv
GNS
Graph Trans
S
Figure 17: Rollouts generalized BoxBath, where rigid ball. For simple comparison, the positions rigid ball suggest models' abilities simulation. SiT+ predicts more faithful positions ball overall shapes of the fluid.
26

---

Under review as a conference paper at ICLR 2022
t=
62
72
t= 119
Ground Truth
DPI-Net
CConv
GNS
Graph Trans
SiT+
Rollouts generalized BoxBath, where rigid cubiod. It seems that the cuboid is harder move. DPI-Net fails this scenes by predicting wrong rotations cuboid, GNS overfits to the rotations of rigid box. SiT+ still achieves better results.
Figure 18:
27


Reviews:
[{"review_id": "x7Rm3cDfiT_", "reviewer": "Reviewer_DCt1", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "SiT uses the Transformer architecture to model particle dynamics. In contrast to a vanilla Transformer, SiT incorporates interaction tokens and abstract tokens to better model particle interactions and material properties, respectively. SiT is evaluated in four simulated environments: FluidFall, FluidShake, BoxBath, RiceGrip, though it would be nice to see environments with more non-liquid multi-object interactions. ", "summary_of_the_paper": "The paper proses the Simulation Transformer (SiT) to simulate particle dynamics, using the Tranformers' attention mechanism to attend to critical particle interactions. They demonstrate qualitative and quantitative improvements over prior work with various architectures in multiple environments. ", "summary_of_the_review": "SiT provides a novel Transformer-based architecture for predicting particle dynamics that outperforms previous graph-convolutional and Transformer-based architectures.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}, {"review_id": "r9DBSyCsVn", "reviewer": "Reviewer_dYcg", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "- In the abstract the sentence 'most existing particle-based simulators adopt graph convolutional networks to model the underlying physics of particles' seems off. For most learning-based methods this might be true, but there of course exists a lot of simulators that are not learning-based. \n- In the abstract it is not clear what is meant by 'material semantics', 'particle tokens', 'interaction tokens'.\n- S1: 'Different from traditional simulators ...'. Here it would help to provide a few examples and references. \n- S1: 'Consequently, they follow the same pipeline ...'. It is not clear which pipeline is meant by 'the same pipeline'. Does it mean the same pipeline is used to model different material properties? Please clarify. \n- S1: 'where they can robustly estimate the dynamics of a system with varying number and configuration of particles'. This is a strong claim, please provide a reference or tone-down this statement.\n- S1 'rich semantics'. Please define what is meant by this. \n- S1 'deploys an abstract token'. Not clear what this means. Please clarify. \n- S1 The second to last paragraph about the contributions of the proposed paper remains too vague. The text refers to several concepts (e.g. like interaction tokens, trainable abstract tokens, intrinsic material characteristics, system-specific semantics, good balance between efficiency and complexity, particle interactions etc.) that have all not been introduced up to this point and it is not clear what the contributions of the proposed paper are compared to existing concepts. \n- S3.1 'The goal of a simulator is to learn a model ...' Does this refer to a learning-based simulator? In other literature simulators are commonly not described to 'learn a model'. It might be useful to more carefully define what is meant by a simulator in the context of this paper. \n- S3.2 'a selective mechanism is thus in great need to help the simulator'. This sentence is not clear. Is this supposed to mean 'a selective mechanism is needed to help the simulator'?\n- 'is naturally a good backbone of building'. Why? It would help to add a few more details here. Are transformers 'good' because particles can be interpreted as discrete tokens? \n- S3.2 'considering all possible interactions is computationally redundant' This is not clear. Why would that be the case?\n- S3.3 'rich semantics': up to this point the text has mentioned the term 'semantics' several times without clearly defining what is meant by it. \n- S3.3 'the semantics of a system is scattered in the tokens'. Without a clear definition of what is meant by 'semantics', sentences like these are difficult to comprehend. \n- S3.3 'To improve generalization ability of SiT and disentangle system-specific semantics from its system-agnostic counterparts'. This sentence is not clear. What is meant by 'disentangle system-specific semantics'? What is meant by 'system-agnostic counterparts'? Please provide more details. \n- S3.3 'Once abstract tokens capture system-agnostic semantics'? At this point 'system semantics' has not be defined, therefore this sentence is not clear. \n\nW.r.t the related work, I consider the following recent work on simulating rod-dynamics with GNNs missing: \n\nH. Shao, T. Kugelstadt, W. Pałubicki, J. Bender, S. Pirk, D. L. Michels, Accurately Solving Rod Dynamics with Graph Learning, Conference on Neural Information Processing Systems (NeurIPS), 2021", "summary_of_the_paper": "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.  ", "summary_of_the_review": "\nOverall, I am very supportive of the topic this paper addresses and generally also of using transformers for solving physical systems. Using transformers for solving physics is an interesting research direction. However, in its current form I am not enthusiastic about this paper which is due to the following main reason: \n\n(1) While the paper is generally well-written, several sentences are left to vague and unclear. In the introduction several concepts and terms are used without carefully introducing them. This makes understanding the text difficult and blurs what the contribution is. Also, some sentences almost read like a collection of buzzwords which makes is very difficult to comprehend what the actual meaning is. Consequently, while the experiments and results are technically sound, I got the impression more care should be taken in carefully explaining and formalizing the concepts that are aimed to be discussed in the paper. \n\n(2) With 12 pages the Appendix provides a significant amount of additional material, some of which can be considered quite relevant for the main contribution of the paper (e.g. Table 4). This indicates that the contribution of this paper may better be discussed in a format that has less restrictive page constraints. There are reasons why papers have page limits and I have the impression that for this work the authors are trying to circumvent them by providing a lot of material in the Appendix. \n\n(3) The paper does not discuss any limitations of the introduced method. \n\nTherefore, I don't think the paper is ready for publication. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "oTPVGjihtcf", "reviewer": "Reviewer_yAKg", "summary": "", "questions": "", "limitations": "", "rating": 6, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "**Strengths:**\n1. Paper is well written, well organized, and narrative is coherent.\n\n2. Experimentation for comparisons with previous state of the art models are convincing as authors compare with 4 recent related approaches and demonstrate that SiT improves performance over other models for fluid simulation in different contexts using standard datasets. \n\n3. Specifically, an interesting contribution of the paper is the use of explicit (disentangled) \"abstract tokens\" that encode material specific properties. The authors demonstrate (through multi-material simulations like a solid cube floating in a fluid) that the SiT method is able to leverage the material specific properties learned in the \"abstract tokens\" to better simulate the properties of the solid interacting with the fluid. For example: other methods (without abstract tokens) which don't explicitly differentiate between the solid and the fluid simulate the solid disintegrating and merging with the flowing fluid but SiT maintains the solid shape and is able to simulate the motion of the solid over the moving fluid.\n\n4. Demonstrated SiT models are much more light-weight (i.e., far fewer parameters). Training methodology further explicitly prioritizes learning properties for different material types by adopting a weighted MSE loss.\n\n\n**Weaknesses:**\n\n1. Some of the initial claims made by the authors of disadvantages of GCNs (as they presuppose particle interaction neighborhoods where interaction occurrs with all particles in neighborhood) seem to be part of the proposed method as well? For example, even in the SiT model, the particle token interactions are defined by a window function governed by a radius `R`, isn't the particle token interacting with all other particles within radius `R`? \n\n2. Incomplete ablation analysis, authors do not demonstrate the importance of abstract tokens.\n\n3. In the \"few shot generalization\" case, it is unclear what the goal of the experiment is? Also, in in Fig. 5b what does \"robust\" mean? i.e., What does it mean to claim that the SiT model (i.e., without abstract tokens) is \"robust even at 60% training data\" ? Secondly, the claim of \"few-shot\" learning while the model uses 60% of training data is questionable. \n\n\n**Questions:**\n\n1. How is the parameter governing the extent of the particle window to inform particle connections (`R`) selected? \n\n2. How sensitive is SiT to the selection of `R` ?\n\n3. Also, in the case of pair-wise interaction components i.e., particle interaction tokens, how is setting allowable connections (using Eq 5) any different from setting connections in the GCN works (that authors describe as inferior) earlier in their paper?\n\n4. Authors claim that GCN based mechanisms are unable to provide selection mechanisms? Would employing attention in the GCN context not serve a a selection mechanism? \n\n5. All abstract tokens of a particular material interact with each other? How scalable is this operation for large fluid simulations?", "summary_of_the_paper": "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.", "summary_of_the_review": "The paper makes a contribution of being the first work to employ transformer to model particle based simulations where the model explicitly accounts for particle, particle interaction as well as material properties during the simulation. The authors have performed a rigorous comparison with state of the art models and the achieved improvements and the explanations rendered for the improvements are sensible. Also, authors demonstrate also through qualitative and quantitative comparisons, the generalizability of the SiT model in various multi-material settings which are also convincing. ", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "3: Some of the paper’s claims have minor issues. A few statements are not well-supported, or require small changes to be made correct.", "technical_novelty_and_significance": "2: The contributions are only marginally significant or novel.", "empirical_novelty_and_significance": "2: The contributions are only marginally significant or novel."}, "extra_text": {}}, {"review_id": "SzA_cAhlOiF", "reviewer": "Reviewer_ZsKn", "summary": "", "questions": "", "limitations": "", "rating": 3, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "Strengths:\n\nS1: Testing the ability of transformers on physical simulation, as opposed to message passing, is indeed an interesting question that needed to be asked. I suspect many groups have probably tried this, and generally found worse performance (Consistent with the findings for the GraphTrans baseline in this paper), but it is good to actually show some published results.\n\nS2: Given that a pure transformer does not seem to work very well, there is another question worth asking: Is there a hybrid between message passing and transformer that could get the best of both worlds?\n\nWhile I am very supportive of studying those two questions, I am afraid the current experimental setup is not doing a great job to support the claims that are being made due to some weaknesses:\n\nW1: The proposed model, (equation 10) contains alpha and beta terms, and actually the conclusion from the authors is that alpha=0 beta=1 is the best working configuration. Now, if the alpha term is ignored, all that is left is the beta term, and the beta term is mostly a weighted aggregation of equation (8), what the authors call \"Interaction Tokens\". Looking closely at this function, this is pretty much a high dimensional message function, as used by the GNN baselines (DPI-Net and GNS), the only difference is that in this case this message plays two roles, the role of a high dimensional message, and the role of an edge specific query, which is combined with a receiver node specific key to build an additional attention weight for the aggregation. First of all I would probably not call this a transformer, because the mechanism is neither cross attention, or even self attention, since each edge key is used for only one edge, and never shared for more than one key/query product. Also, the mechanism does not offer the main advantage of transformers which is not having to require to compute specialized edge level embeddings, which are expensive. Instead, it pretty much is a standard message passing model, with an additional scalar attention component.\n\nW2: Now, even if we leave aside whether the model should be called a transformer or not, then the remaining question is: how important is the contribution of the softmaxed \"\\omega\" attention weights in the right term of equation (10). This this is the only major difference between this SiT and the other message passing baselines (e.g. DPI-Net and GNS), the easiest way to study that difference in a clean way would have been to take one of the most competitive message passing baselines (e.g. GNS), and added the proposed attention mechanism (e.g. omega) to the very same implementation, while keeping the rest exactly the same. Still, I tried to gave it my best shot at using evidence presented in the paper to actually find an answer to that question by looking at the comparisons to the closest best performing message passing baseline, but several experimental choices made this very hard (See W3, W4, W5).\n\nW3: It is hard to take much from the comparisons in the BoxBath environments. Specifically, because, as highlighted by comparing the MSE and WMSE losses, the main differences in performance are just depending on whether the model keeps the box together or not. However predictions for the box particles are made with completely different approaches. In DPI-Net, GraphTrans, SiT, and SiT+, a single solid translation and rotation prediction is made for the whole box (it is by design impossible for the box to break apart), while for CConv and GNS the model is asked to produce independent predictions for each particle in the box (so the box can actually break). This difference seems to be the main dominant factor in the comparison between those two groups of models, and specifically between GNS and SiT, however, the explanation of this difference in the text does not seem clear enough. Furthermore, comparison about how to make predictions for the box is completely orthogonal to the topic in the paper, and just makes the results harder to interpret. There does not seem to be a good reason not to use the same approach (either model the box as a whole, or make per particle predictions) for all baselines. Otherwise the results presented in Table 1, and Figure 3 don’t actually tell much about how the extra attention weight contributes to the predictions, and the usage of the WMSE metric seems to just highlight this unfair advantage that some models have over the others when modeling the box. This also seems to directly oppose the main claim at the end of the abstract “Without bells and whistles…achieves superior performance … than existing methods”, as this is precisely a “bell and whistle” that is only used for a subset of the models.\n\n\nW4: Comparison in FluidFall is also hard to interpret. “though the MSE results are close to each other. SiT can better maintain the shape of drops while DPI-Net and GNS fail due to the wrong prediction of viscosity”. This statement seems to make a quite general claim of why the DPI-Net and GNS fail on this domain, but there is not actual evidence supporting this explanation. First in Figure 3 frames differences are not actually that obvious, but instead we can look at Supplementary Figure 8, where the problem is more apparent. Second, in the animations from the original DPI-Net paper (http://dpi.csail.mit.edu/, “Rollout from our model”), the model predictions for this same domain look qualitatively much better than the results shown on the same domain for DPI-Net. Similarly, the GNS paper also models goop like materials with viscosity (https://sites.google.com/corp/view/learning-to-simulate#h.p_JUaMhhhLYDkV), and this particular artifact that not seem to happen. So not sure if there is any difference in the implementations that could explain this. Third, from my experience, these sort of artifacts appear when (1) the training distribution does not have a lot of variance (like it seems to be the case for this dataset) and (2) a model that has a lot of capacity per message passing layer does not have enough depth to learn effective long range interactions/correlations, so it learns bad short range correlations that exploit the biases in the training dataset and make the model overfit. The reason this happens for DPI-Net and GNS, but not SiT, may be dependent on the choice of hyperparameters so it is unclear, if for a different choice of hyperparameters, results would have told a different story. Note, these biases / failure modes are also more typical of models that are not translation equivariant and try to exploit biases on the absolute positions, but not sure whether the implementations of the models are translation equivariant or not. \n\nW5: GNS is not implemented for FluidShake or Rice Grip. From FluidFall and BoxBath, GNS seems to be the most competitive model with SiT, so it is unclear why no comparison is provided for FluidShake and for RiceGrip. This would probably be ok if it was not because of W3, and W4, but it seems crucially important considering the potential issues with those comparisons.\n\nW6: The role of the abstract tokens in the model is unclear. The results clearly show that SiT+ is better than SiT performance-wise, however, it is unclear why. The authors intent is that they “abstract tokens capture system-agnostic semantics, they can be reused by SiT when generalizing to systems that have same materials but vary in particle amount and configuration.”,  however I don’t think this hypothesis is tested explicitly. Abstract tokens seem to help on FluidFall (which has a single material), and the training distribution on box, batch, for which the whole training datasets consists of 4x4x4 boxes, which cannot be explained by that hypothesis. On RiceGrip, the difference between SiT and SiT+ is larger, however, actually in that domain SiT that worse than DPI-Net, which I find suspicious. On the other hand if I understand abstract tokens correctly, they actually connect to all nodes of the same type, in a bidirectional way, and if this is the case they actually are sort of a global node that allow communication between any two pairs of particles of the same material in two hops, regardless of how far they are. How do we know that the differences observed are not just due to this extra range of communication? To test this, an ablation where “dummy abstract nodes” without the material embedding but with the same connectivity are added should be included. Authors also say “baselines … may regard the latter as part of the intrinsic material semantics and fail to generalize to systems with the same materials but different particle amounts and configurations.”, but this is not backed by evidence, and in fact several of the baselines in the paper have already shown generalization to different shapes at test time.\n\nW7: (Referring to the baselines) “First, it forces each particle to interact with all its nearby particles without providing a selection mechanism, which leads to computational redundancy and prevents the discovery of inherent patterns of particle interaction.” \n“They require the particles to interact with all nearby neighbors without selections, which will aggregate redundant and irrelevant neighbors”. This is slightly misleading. Currently, all methods in the paper are using a fixed radius of connectivity to select a number of candidate neighbors (if this claim is True, this should not be necessary for SiT). Second, the SiT model still has to consider all interactions within the radius of connectivity, and learn to discard some by downweighting them. This is something that the baselines could also learn to do, by just outputting messages that are 0 for unnecessary interactions, and a similar plot like Sup. Fig. 7 (b) could be made by showing the norm of the messages. It is true that the explicit scalar weight is a nicer way to disable full edges, rather than disabling all element wise message components, but for the other reasons indicated above, there is no evidence showing that SiT is actually empirically better than the baselines at doing this.\n\n\nOther observations comments:\n\nO1: I would recommend using the more general term GNN (Graph Neural Network) instead of GCN (Graph Convolutional Network) to refer to the corresponding baselines, as most readers will think of node update rules as done “Kipf et. al” model when reading the term GCN.\n\nO2: The transformer description (eqs. 2, 3, 4), seems to be missing something, as none of these equations that are part of the iterative update of the nodes actually contain learnable weights. I think the computation of separate keys, queries, and values may be missing. Not sure if this is a bug in the math, or if it is actually implemented that way. Perhaps this is also missing in some of the equations (10, 11, 12, 13, 14). \n\nO3: More explanation is needed about specifics of the model, for example, does the encoder uses absolute positions, or does it (like several or the baselines) use relative positions in the encoder, to become translation equivariant? There is also some explanation in the supplementation about number of blocks, number of MLP layers etc, but this is ambiguous and the main hyperparameters (such as number of message passing steps), should be more prominent, and probably in the main text. Also some explanation about how the model and baselines were fine tunes should be included.\n\nO4: Not sure if the definition of WMSE is correct (eq 16). My understanding is that the goal is is to increase the relative weight of particles with materials that appear less often when there is an imbalance of number of particles of different materials. Currently, if you replace equation (15), into equation (16), if is unclear what happens with the N in equation (15). If it gets replaced by an N_k, then it would cancel out with the N_k in equation (16), which means no rebalancing happens, since K is just a constant (number of material types). If it remains as N, then you get N_k/N and the weighting goes into the opposite direction, materials with less particles, get even less weight on the sum on average. It would be useful to just fully expand equation (15).\n\nO5: Not sure if there is something wrong with the top row of figure 3, since the camera position seems to be slightly different for each frame.\n\nO6: “Notice that CConv is not suitable for rigid box, we only report result on BoxBath for reference of simulation on fluid parts.” Why?\n\nO7: In the related work it may be worth mentioning “CConv” before “GNS” as it was published earlier.\n", "summary_of_the_paper": "The paper studies the use of transformers for particle-based physics simulation. The main contributions are as follows.\n\nC1. the paper proposes a specific transformer-inspired form of message passing, which, unlike transformer, still does explicit message computation.\n\nC2. the paper proposes a very specific way to encode material type information as abstract tokens.\n\nC3. the paper compares to several non-transfomer based baselines, as well as a pure transformer.\n\nC4. 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.\n", "summary_of_the_review": "The intent of the paper is good, and the results could bring interesting insights about whether attention inspired architectures could bring to the simulation domain. However the empirical evaluation makes it very hard to know which improvements are a result of the specific technical additions added to the model (extra scalar attention for the explicit messages and abstract tokens), and which improvements are just due to other small, but possibly very relevant, differences (including different input encoders/decoders, different featurizations, different outputs, etc.) between SiT and the baselines. My recommendation would be to perform the same investigations by adding the extra attention to the existing baselines one by one (or at least to the most promising ones) while keeping other parameters identical and test if the form or proposed attention helps universally. Similarly, the idea of adding extra abstract nodes for material types, and its impact on generalization, could be applied to existing baselines too and tested and studied orthogonally to the attention contribution. This would make the paper conclusions much more insightful and useful than in their current form.\n\nEDIT POST REBUTTAL: I want to thank the authors for the additional experiments and extensive replies. While many interesting ideas are introduced in the paper, the way the comparisons with baselines is performed makes it very hard to really understand which aspects of the proposed architectures is making the model perform better than baselines on some settings. \n\nSince the authors are proposing a new neural network layer, it would be more useful if, rather than compare to different baselines each with their original hyper-parameters from their papers, etc, the authors took the approach of integrating the different parts of the newly proposed layer into existing architectures (possibly including non-simulation settings), and try to understand better that way how the new layer may help in a more apples-to-apples comparison.", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "", "extra_scores": {"correctness": "2: Several of the paper’s claims are incorrect or not well-supported.", "technical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work.", "empirical_novelty_and_significance": "3: The contributions are significant and somewhat new. Aspects of the contributions exist in prior work."}, "extra_text": {}}]

Author rebuttal:

Meta-review: The submission received split reviews: two reviewers recommended weak accepts, and the other two weak rejects.  The AC went through the reviews, responses, and discussions carefully.  The AC agrees that this paper is well-written and has demonstrated the possibility of using transformers for particle-based physical simulation.  The AC also believes that the authors have addressed the concerns of reviewer dYcg, despite that the reviewer didn't engage in the discussions.  The contributions are however not most exciting, and none of the reviewers would like to champion the submission.

Further, the AC agrees with the knowledgeable and responsible reviewer ZsKn that the presentation and experiments can be better positioned to highlight the key contribution. As reviewer ZsKn has summarized, it's recommended that "the authors took the approach of integrating the different parts of the newly proposed layer into existing architectures (possibly including non-simulation settings), and try to understand better that way how the new layer may help in a more apples-to-apples comparison."

The recommendation is reject, and the authors are encouraged to revise the paper for the next venue.
Decision: Reject