S I T: S IMULATION T RANSFORMER FOR P ARTICLE - BASED P HYSICS S IMULATION
Anonymous authors
Paper under double-blind review
A BSTRACT
In this paper, we focus on learning-based methods for simulation.
Most existing particle-based simulators adopt graph neural networks (GNNs) to model the underlying physics of particles.
However, they force particles to interact with all neighbors without selection, and they fall short in capturing material semantics, such as viscosity or plastic deformations, for different particles, leading to unsatisfactory performance, especially in generalization.
This paper proposes Simulation Transformer (SiT) to simulate particle dynamics with more careful modeling of particle states, interactions, and their intrinsic properties.
Speciﬁcally, besides the particle tokens, which encode the states of particles into latent space, SiT generates interaction tokens to extract high-level representations for interactions and selectively focuses on essential interactions by allowing both tokens to attend to each other.
In addition, SiT learns material-aware representations by learnable abstract tokens, which will participate in the attention mechanism and boost the generalization capability further.
We evaluate our model on diverse environments, including ﬂuid, rigid, and deformable objects, which cover domains of different complexity and materials.
Without bells and whistles, SiT shows strong abilities to simulate particles of different materials and achieves superior performance and generalization across these environments with fewer parameters than existing methods.
Codes and models will be released.
I NTRODUCTION
Particle-based physics simulation is a classic and important topic in computer science.
It not only facilitates the exploration of underlying principles in physics, chemistry and biology, but also enables the creation of vivid visual effects such as explosion and ﬂuid dynamic in ﬁlms and games.
Different from traditional simulators, such as grid-based (Guo et al., 2016) and mesh-based (Bronstein et al., 2017) methods, particle-based simulators view a system, which is an example in one domain, as a composition of particles and imitate system changes over time by predicting the changes of particle-wise states according to current particle states and particle interactions, of which the latter represents the inﬂuence 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-speciﬁc physical priors.
Moreover, since in particle-based simulators the dynamics of a system is modeled by the states of particles and their interactions, they also have the potential to possess a strong generalization ability, where they can estimate the dynamics of a system with varying number and conﬁguration of particles in a more robust manner.
After learning the dynamics of ﬂuid water in a sandbox, the same particle-based simulator can be used to simulate a waterfall and a river.
Recent particle-based simulators (Battaglia et al., 2016; Schenck & Fox, 2018; Mrowca et al., 2018;
Li et al., 2019; Sanchez-Gonzalez et al., 2020; Ummenhofer et al., 2020) often view a system as a graph, and adopt graph neural network (GNN) (Kipf & Welling, 2016) as the basic network structure.
In these attempts, each particle is treated as a node in the graph, with edges linking it to all its neighboring particles, assuming interactions mainly occur between particles that are close to each other.
Subsequently, state updates of particles are achieved by combining node features with the summation of edge features.
While such a GNN-based formulation obtains satisfying simulation results, it faces two issues that affect efﬁciency and generality.
First, it forces each particle to interact with all its nearby particles without providing a selective mechanism, which leads to computational redundancy and prevents the discovery of inherent patterns of particle interaction.
Second, the
GNN-based formulation uses particle-wise attributes to capture both material characteristics, such as viscosity or plastic deformations, and domain-speciﬁc semantics, such as the shape of a rigid material.
Therefore, it may regard the latter as part of the intrinsic material semantics and fail to generalize to domains with the same materials but different particle amounts and conﬁgurations.
In this paper, we propose a novel Transformer-based framework, referred to as Simulation Transformer (SiT), for particle-based physics simulation.
The model inherits the powerful multi-head self-attention mechanism in Transformer (Vaswani et al., 2017) to capture particle interactions.
To further encourage efﬁcient modeling of complex interactions, instead of treating particle interactions as attention weights obtained by dot-product , we introduce the notion of interaction tokens, which are high-dimensional representations for interactions, to model the rich semantics of particle interactions, such as how the particle is restored after deformations.
In addition, to disentangle local material-speciﬁc characteristics from global domain-speciﬁc semantics, SiT further learns a highdimensional abstract token for each type of material to capture material semantics, forcing particles of the same material to interact with their corresponding abstract token.
The proposed SiT is more appealing than previous methods in several aspects.
First, through capturing particle interactions explicitly with interaction tokens and allowing dynamic inter-token attention, SiT dynamically focuses on essential particle interactions and reduces the computations for redundant and noisy ones.
This is crucial for particle-based simulation, especially for domains containing hundreds and thousands of densely placed particles, where the modeling of all particle interactions is redundant, expensive and noisy in practice.
Second, thanks to the trainable abstract tokens that disentangle intrinsic material characteristics from domain-speciﬁc semantics, we can reuse them to apply SiT to unseen domains of the same materials without retraining.
As shown in our experiments, after training on one domain consisting of ﬂuid water and a rigid cubic, SiT still yields fairly faithful simulations when the cubic is replaced with a ball or a bunny comparing with previous work.
To show the effectiveness of SiT, we perform extensive evaluations on four standard environments commonly used in the literature (Li et al., 2019; Sanchez-Gonzalez et al., 2020), covering domains of different complexity and materials.
The proposed method achieves superior performance across all these environments with fewer parameters compared to existing methods.
We further demonstrate the generalization ability of SiT by adjusting the environments to create new domains and applying
SiT to these domains without retraining.
In all cases, SiT obtains more realistic simulation results than previous methods, which tend to overﬁt to the training domains.
R ELATED W ORK
Physics simulation by neural networks .
There are many different kind of representations for physics simulations.
Grid-based methods (Lee & You, 2019; Thuerey et al., 2020;
Wang et al., 2020) adopt convolutional architectures for learning high-dimensional physical system, while mesh-based simulations (Bronstein et al., 2017; Luo et al., 2020;
Hanocka et al., 2019; Nash et al., 2020; Qiao et al., 2020; Weng et al., 2021; Pfaff et al., 2021) typically simulate objects with continuous surfaces, such as clothes, rigid objects, surfaces of water and so on.
Many works (Battaglia et al., 2016;
Schenck & Fox, 2018; Mrowca et al., 2018; Li et al., 2019;
Sanchez-Gonzalez et al., 2020; Ummenhofer et al., 2020) simulate physics on particle-based systems, where all objects are represented by groups of particles.
Speciﬁcally, Interaction Network (IN) (Battaglia et al., 2016) simulated interactions in object-level.
Smooth Particle Networks (SPNets)
(Schenck & Fox, 2018) implemented ﬂuid dynamics using position-based ﬂuids (Macklin &
M¨uller, 2013).
Hierarchical Relation Network (HRN) (Mrowca et al., 2018) predicted physical dynamics based on hierarchical graph convolution.
Dynamic Particle Interaction Networks (DPI-Net) (Li et al., 2019) combined dynamic graphs, multi-step spatial propagation, and hierarchical structure to simulate particles.
CConv (Ummenhofer et al., 2020) used spatial convolutions to simulate ﬂuid particles.
Graph Network-based Simulators (GNS) (Sanchez-Gonzalez et al., 2020) computed dynamics via learned message-passing.
Similar to particle-based systems, COPINGNet (Shao et al., 2021) applies graph networks to simulate rod dynamics, where the discretized rod is the basic unit similar to particle.
Previous work mostly adopted graph networks for simulations and required each particle to interact with all its nearby particles without selective mechanism.
In contrast, our SiT employs both particle and interaction tokens and selectively focus on necessary particle interactions through attention mechanism.
Experiments show that SiT surpasses existing GNN-based methods and has more robust performances in generalizations.
Transformer .
Transformer (Vaswani et al., 2017) was designed for machine translation and achieved state-of-the-art performance in many natural langruage processing tasks (Devlin et al., 2019; Radford et al., 2019; Brown et al., 2020).
Recently, Transformer starts to show great expandability and applicability in many other ﬁelds, such as computer vision (Wang et al., 2018; Carion et al., 2020;
Dosovitskiy et al., 2021; Wang et al., 2021; Liu et al., 2021), and graph representations (Zhou et al., 2020; Zhang et al., 2020; Dwivedi & Bresson, 2020).
To our knowledge, no attempt has been made to apply Transformer on physics simulation.
Our SiT models the interactions between particles by trainable sub-network given corresponding particle tokens.
The same notion of extracting potential semantics between nodes is also applied in
Graph Transformer (Dwivedi & Bresson, 2020), which we refer as GraphTrans for short.
However, there are differences in our formulations.
Speciﬁcally, GraphTrans adopts element-wise product between node representations followed by multi-layer perceptron (MLP) to update the interaction embeddings, which store the attention scores in each dimension and are reduced to a scalar for further attention mechanism to update node embeddings.
In contrast, our model learns semantic tokens for interactions through sub-network instead of element-wise product.
When updating the particle tokens, which are referred as node representations in GraphTrans, both particle and interaction tokens will attend to each other.
The interaction tokens are not weighted scores any more.
We adopt GraphTrans (Dwivedi & Bresson, 2020) in particle-based simulation and compare it with SiT in experiments.
The quantitative results show that SiT achieves better results than GraphTrans.
M ETHODOLOGY 3. P ROBLEM F ORMULATION
For a particle-based system composed of N particles, we use X t = { x ti } Ni =1 to denote the system state at time step t , where x ti denotes the state of i -th particle.
Speciﬁcally, x ti =
[ p ti , q ti , a i ] , where p ti , q ti ∈ R 3 refer to position and velocity, and a i ∈ R d a represents ﬁxed particle attributes such as its material type.
The goal of a simulator is to learn a model φ ( · ) from previous rollouts of a system to causally predict a rollout trajectory in a speciﬁc time period conditioned on the initial system state X 0 .
The prediction is run in a recursive manner, where the simulator will predict the state ˆ X t +
= φ ( X t ) at time step t + 1 based on the state X t = { x ti } at time step t .
In practice, we will predict the velocities of particles ˆ Q t +
=
{ ˆ q t +1 i } , and obtain their positions via ˆ p t
+1 i
= p t i + ∆ t · ˆ q t +1 i , where ∆ t is a domain-speciﬁc constant.
3. S IMULATION VIA V ANILLA T RANSFORMER
To accurately simulate the changes of a system over time, it is crucial to effectively model the interactions among particles, as they indicate the energy transition of a system when constrained by material characteristics and physical laws.
However, it is infeasible to know a priori how should particles interact with each other. Thus, a selective mechanism is needed to help the simulator focus only on necessary interactions.
Since Transformer (Vaswani et al., 2017) is capable of modeling the dynamical attention scores between tokens via the self-attention module, we can regard the attention weights as the intensity of connectivities and the importance of interactions, where the larger the attention weight the more important the interaction is. Thus, it is naturally a good backbone of building an efﬁcient particlebased simulator.
To apply vanilla Transformer in particle-based simulation, we ﬁrst encode the states of particles into corresponding particle tokens V
= { v ti } by
[Equation]
where v ti ∈ R
d
h is a d h dimensional vector and f enc V ( · ) is an encoding layer implemented as a MLP.
Subsequently, particle interactions are achieved by L blocks of self-attention modules, where in the l -th block, particle tokens will attend to each other selectively as:
[Equation]
Since a system usually contains hundreds of particles and the interactions among particles occur within the neighbors in our settings, considering all possible interactions, of which the number quadratically increases with respect to the number of particles, is computationally redundant and inefﬁcient.
Therefore, we follow previous literature Li et al.
Sanchez-Gonzalez et al.
(2020) to assume that interactions of distant particles can be omitted, which is realized by a window function:
[Equation]
where I ( condition ) is a indicator function that returns 1 if the condition is satisﬁed, and R deﬁnes the extent of the window.
This window function will generate a mask to retain only interactions between neighboring particles as potential candidates in the self-attention modules.
To predict particle states in the next time step, a decoding layer is applied to the updated token of i -th particle to obtain its velocity:
[Equation]
where f dec V ( · ) is implemented by another MLP.
3. S IMULATION T RANSFORMER (S I T)
Although the vanilla Transformer provides a ﬂexible approach for particle-based simulation, directly applying it leads to inferior simulation results as shown in our experiments.
In particular, a vanilla
Transformer uses attention weights, which are scalars obtained via dot-product, to represent particle interactions.
A single scalar is insufﬁcient to reﬂect the rich semantics of particle interactions.
For example, in the case of ﬂuid water, when two particles get closer and closer, their interaction will at ﬁrst act as surface tension to push them towards each other.
And when they are sufﬁciently close, the repulsion between them will dominate the interaction and stop them from further moving towards each other.
In addition, since the semantics of one domain is scattered in the tokens of different particles, domain-speciﬁc semantics, such as the shape of rigid cubic, are mixed with domain-agnostic semantics, such as the stickiness attribute of materials.
In such cases, the former may be misinterpreted as the latter, leading to poor generalization ability.
To overcome the limitations of vanilla Transformer, we propose a novel Transformer-based simulator, referred to as Simulation Transformer (SiT), that extends vanilla Transformer with two new types of tokens, namely interactions tokens and abstract tokens.
An overview of SiT is included in Figure 2.
Interaction Tokens.
To capture the rich semantics of particle interactions, instead of a scalar, SiT assigns an interaction token to each potential particle interaction.
Speciﬁcally, the token for the interaction between particles x ti and x tj is initialized by
[Equation]
where u tij ∈ R
d h is a d h dimensional vector and f enc I ( · ) is the encoding function for interactions.
Subsequently, in each self-attention block SiT will update the interaction tokens as
[Equation]
[Equation]
where ψ ( · ) is a learnable sub-network shared for all interaction tokens.
Inspired by GNN-based methods (Mrowca et al., 2018; Li et al., 2019; Sanchez-Gonzalez et al., 2020) , SiT adopts the learnable sub-network to enlarge its capacity for handling complex particle interactions.
In practice, the sub-network is implemented as
[Equation]
As is shown in Figure [Equation] the update of particle tokens in the self-attention module is thus achieved by
[Equation]
where α, β ∈ { 0 , 1 } are the gate controllers, j and k are the indexes of particles within the reception window deﬁned by Equation 5.
Equation 13 computes the attention scores between particle tokens [Equation] and Equation 14 computes the attention scores between particle and interaction tokens.
Subsequently, these attention scores are normalized by Equation 11 and Equation 12.
They are then used by Equation 10 to aggregate semantics of particle and interaction tokens, where the gate controllers α, β are hyper-parameters that work as balancing coefﬁcients.
When α = 1 , β = 0 , Equationbecomes the original self-attention formula in vanilla Transformer.
The multi-head attention version of Equation 10 can be found in Appendix 6.1.2.
It is worth noting that SiT provides a more ﬂexible way for further extensions when compared to vanilla Transformer.
Speciﬁcally, we can change the value of α and β in Equation 10 to adjust the signiﬁcance of different tokens.
And changing the window function in Equation 5 allows SiT to consider different interaction patterns.
Finally, by adjusting the learnable sub-network, we can also generalize SiT to consider more complex particle interactions such as interactions beyond a pair of particles and interactions of interactions.
Abstract Tokens.
To improve the generalization ability of SiT and disentangle domain-speciﬁc semantics from its domain-agnostic counterparts, we further apply SiT with material-speciﬁc abstract tokens.
For N a types of materials, SiT adopts N a abstract tokens A
= { a k } N a k
=1 , each of which is a learnable vector of the same length of particle tokens.
Ideally, the abstract token a k should capture the domain-agnostic semantics of k -th material.
They act as particle tokens for additional abstract particles, following the same update formula as particle tokens.
Therefore with N a abstract tokens SiT will have N + N a particle tokens in total: { a 1 , · · · , a N a , x t 1 , · · · , x tN } , 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 ﬁeld of abstract tokens to the particles of the same material.
Once abstract tokens capture domain-agnostic semantics, they can be reused by SiT when generalizing to domains that have same materials but vary in particle amount and conﬁguration.
3. T RAINING O BJECTIVE
To train SiT with previous rollouts, the standard mean square error (MSE) loss is applied on the output of our SiT:
[Equation]
i [Equation] (15) where ˆ Q t +1 = { ˆ q t {++
+}
i Ni =1 is the estimated velocity , Q t +1 = { q t +1 i } Ni =1 is the ground truth, and || · || 2 is L2 norm.
While the MSE loss works well for cases with relatively low complexity, in practice some domains often contain multiple types of materials with imbalance numbers of particles.
In such cases, the MSE loss will be biased to the material with more particles.
Therefore, we further apply a material weighted MSE (WMSE) loss to reduce the effect of imbalance:
[Equation]
where K is the number of material types, N k is the number of particles belonging to the k -th material and N = (cid:80) k N k , ˆ Q t +1 k = { ˆ q t +1 i,k } N k
i =1 and Q t +1 k
= { q t
+1 i,k } N k i = 1 denote the estimated velocity and ground truth for k -th {+material, and ˆ Q t +
+}
= { ˆ
[Equation]
t +
[Equation]
E XPERIMENTS 4. B ASE E NVIRONMENTS
We adopt the environments commonly used in the literature (Li et al., 2019; Sanchez-Gonzalez et al., 2020; Ummenhofer et al., 2020).
Samples are displayed in Figure 1.
There are four environments in total: FluidFall is a basic simulation for two drops of water; FluidShake is more complex and simulate the water in a randomly moving box; BoxBath simulates the water washing a rigid cubic in ﬁxed box; RiceGrip simulates the interactions between deformable rice and two rigid grippers.
More details of data generalization can be found in Appendix 6.1.3.
We compare SiT with 4 recent approaches: DPI-Net (Li et al., 2019),
CConv (Ummenhofer et al., 2020), GNS (Sanchez-Gonzalez et al., 2020), and GraphTrans (Dwivedi & Bresson, 2020).
For fair comparison, we adopt similar training schedules.
Quantitative results and model parameters are in
Table 5.
Qualitative results are in Figure 3.
More details can be found in Appendix 6.1.3 and 6.2.
SiT achieves superior performance on all environments with fewer model parameters.
The effectiveness of abstract tokens are more obvious for RiceGrip and BoxBath , which are more complex.
Quantitative results show that SiT can better simulate the dynamics on different environments.
Comparison with DPI-Net and GNS .
DPI-Net
(Li et al., 2019) and GNS (Sanchez-Gonzalez et al., 2020)
adopt message-passing graph networks for particle-based simulation.
As shown in the results,
[Table]
getting closer in FluidFall but not close enough to merge, previous methods are likely to mix them together due to 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 interactions, which is further demonstrated by the heat map in Appendix 6.2.1.
SiT with abstract tokens further improves the performances especially on RiceGrip and BoxBath , suggesting the effectiveness of abstract tokens in modeling complex deformations and multi-materials interactions.
For
FluidFall , though the MSEs are close to each other, SiT can better maintain the shape of drops while DPI-Net and GNS fail because they do not have selective mechanism and equally treat the incorrect interactions, which are introduced by the newly added neighbors from different drops when they get closer but not close enough to merge.
Comparison with CConv .
CConv (Ummenhofer et al., 2020) designs convolutional layers carefully tailored to modeling ﬂuid dynamics, such as an SPH(Monaghan, 1992)-like local kernel, different sub-networks for ﬂuid and boundary particles.
Without modeling interactions, CConv propagates information only among particles weighted by distances.
On the other hand, SiT assigns high-level tokens for interactions and achieves better results, suggesting the necessity of interaction tokens.
Notice that CConv is not suitable for rigid box, we only report result on BoxBath for reference of simulation on ﬂuid parts.
Comparison with GraphTrans (Dwivedi & Bresson, 2020).
The GraphTrans is also a Transformerbased method and updates the interactions using element-wise product from node representations.
We represent particles by nodes and adopt layers in GraphTrans as our backbone, which and use the same hidden dimensions and number of blocks.
Other settings are completely the same as SiT for better comparison.
As is shown in experiments, GraphTrans fails to simulate particles on all cases with much higher errors.
This is because our SiT directly models the interactions by interaction tokens generated by sub-network, which is necessary to capture high-level semantics.
4. G ENERALIAZTIONS
To challenge the robustness of our model on more complex settings, we add more particles for
RichGrip and FluidShake , and change the size and shape of rigid object for BoxBath .
Table 5 shows the upper and lower bound for the number of particles in training set and generalization validation set, as well as the shapes of rigid object.
We mainly compare SiT with DPI-Net, which has fairly good and reasonable performance on all environments.
Quantitative results are summarized in Table 5, while qualitative results are depicted in Figure 4 and Figure 1 (b).
As is shown, SiT with abstract tokens can better simulate particle dy- [Table] namics in generalized environment.
suggesting that the abstract tokens and the selective mechanism involving interaction tokens can provide more stable and robust performance.
More details about the results and rollouts can be found in Appendix 6.2.3.
4. A BLATION STUDIES
We comprehensively analyze our SiT and explore the effectiveness of our model in the following ways: (a) the necessity of interaction tokens; (b) the architectures of sub-network, which generates interaction tokens; (c) the values for gate controllers α, β ; (d) the usage of multi-head in uniform attention; (e) the sensitiveness of SiT to radius R ; (f) the effectiveness of abstract tokens.
We conduct our experiments on FluidShake for (a) to (e). We verify (f) on BoxBath and RiceGrip . The quantitative results are in Table 3 and Table 4.
Necessity of interaction tokens .
We apply vanilla Transformer encoder by conﬁguration A, which uses the same hidden dimension and number of blocks as SiT, and SiT by conﬁguration B, which only changes the output of sub-network to one dimension.
As shown in Table 3, both conﬁguration A and B have low performance, suggesting the scalars are insufﬁcient to describe complex interactions.
Architectures of sub-network .
Conﬁguration C adopts element-wise product between particle tokens, which follows the GraphTrans (Dwivedi & Bresson, 2020); conﬁguration D adopts selfattention blocks, where q is the interaction token from previous layer, k and v are the neighbor particle tokens.
As shown in Table 3, conﬁguration C and D have worse performance than SiT, which suggests that the element-wise product is insufﬁcient to model complex interactions, while self-attention block introduces more parameters and complexities, making it harder to train.
Hyper-paramters for α, β .
We set α = 1 , β = 0 in conﬁguration E, α = 1 , β = 1 in conﬁguration
F, and α = 0 , β = 1 in conﬁguration H, Consequently, both conﬁguration E and H achieve good results, suggesting that the participants of interaction tokens are the key to boost the performance.
In practice, we use conﬁguration H for SiT which shows the best results.
Multi-head uniform attention .
We apply multi-head mechanism in conﬁguration G and set head number as 8.
The results show that multi-head mechanism can bring us some improvement.
But it also introduces more parameters and consumes lots of memory when computing.
Sensitiveness to R . Quantitative results are reported on FluidShake . As shown in Table 4, SiT is more robust when varying the radius R , suggesting the effectiveness of selective mechanism and robustness of SiT. [Table]
[Table]
Effectiveness of abstract tokens . We replace the abstract tokens with dummy tokens, which are randomly initialized vectors with ﬁxed values but have same connectivities as abstract tokens. As shown in 4, SiT with dummy tokens fails on both domains, suggesting the abstract tokens are able to learn the materials semantics and boost SiT’s performances. 4. F EW S HOTS L EARNING
As shown in Figure 5, we conduct few shots learning experiments on FluidShake using SiT without abstract tokens.
Even when trained on 60% of data, SiT achieves similar MSEs comparing with DPI-Net trained on 100% of data.
When trained with further less data, SiT learns not well enough.
As the number of particles differs in each rollout, SiT does not have enough training on all possible cases. On the other hand, the weights are not shared in SiT and would not learn enough on less data comparing with the shared-weight blocks in DPI-Net. SiT is more robust with no less than 60% of the whole training data, but is sensitive to extremely less data.
C ONCLUSION
In this paper, we propose Simulation Transformer (SiT) and introduce interaction tokens and abstract tokens to simulate domains of different complexity and materials, including hundreds and thousands of particles.
Our experimental results show the necessity and effectiveness of interaction and abstract tokens as well as selective mechanism in attention for particle-based simulation.
Generalization results further suggest the robustness of SiT. SiT is also ﬂexible for extensions and applications, such as more complex and efﬁcient architectures for sub-network, changing the particle dynamics by attending to alternative abstract tokens with different materials semantics.
Finally, SiT makes a successful attempt to apply Transformer into physics simulation and achieve superior performances over existing methods.
6.1. M ULTI -H EAD U NIFORM A TTENTION
The multi-head version of uniform attention is formulated by
[Equation]
where m ∈ 1 , 2 , · · · , M is the index of head, W l ∈ R d h ×
d h is the output weight for the concatenation of M heads; Q m,l , K m,l , V m, ; are the projection weights of query, key, value for head m at block l .
The rest symbols are the same with those in uniform attention.
6.1. I MPLEMENTATION D ETAILS
Inputs and outputs details .
For FluidFall , FluidShake , and BoxBath , we only use particles’ states at time t as inputs and output the velocities at time t + 1 .
For RiceGrip , we concatenate particles states from t − 2 to t as inputs and output 6-dim vector for the velocity of the current observed position and the resting position, which is the same setting as DPI-Net for better comparison.
For BoxBath , we output 7-dim vectors, where 3 dimensions for the predicted velocities, and 4 dimensions for rotation constrains.
The rotation constraints are applied only on rigid particles which predict the rotation velocities.
We use particles’ states at time t as inputs and output the velocities at time t + 1 except
RiceGrip , where we concatenate particles states from t − 2 to t as inputs for better comparison with
DPI-Net (Li et al., 2019).
All states of particles, such as the positions, velocities, and accelerations, are ﬁrst 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 arefor default.
The MLP after uniform attention has two layers with dimensions 512 and 128.
We use 2 blocks shown in Figure 2 for FluidFall and 4 blocks for the other three environments.
We set α = 0 , β = 1 in our model for default, which show the best results in ablation studies.
The learnable embeddings for abstract tokens have the same dimensions as inputs.
Training .
We train four models independently on these four environments.
For common settings, we use Adam optimizer with initial learning rate 0.0008 and plateau scheduler with patience 3 and decreasing factor 0.8.
We choose batch size 16 for all environments and train the models forepochs on FluidShake , 5 epochs for BoxBath , 13 epochs for FluidFall , and 20 epochs for RiceGrip .
No augmentation is involved, such as adding noises to particles states during training.
We adopt
WMSE loss on BoxBath and MSE loss on the rest environments for training.
Evaluation .
We evaluate the positions between predictions and ground truths by MSE as is shown in equation 15.
For BoxBath , we also report the WMSE results by equation 16 to better show the performance of simulating different materials.
Baseline details .
For fair comparison, the following settings are the same with SiT: inputs for models, number of training epochs on different environments, learning rate schedules, and training loss on velocities.
On BoxBath , all baselines adopt the rigid constraints except CConv, which is designed speciﬁcally for ﬂuid dynamics and we use it to mainly compare the ﬂuid simulations. The model-related hyper-parameters for baselines are the same as the original papers, such as the number of message passing is 10 for GNS. Other hyper-parameters for baselines are ﬁrst chosen the same as their original papers, and then ﬁne-tuned within a small range of changes. For example, in terms of the initial learning rate, 0.0001 works better on GNS while others adopt 0.0008.
6. E XPERIMENTS 6.2. A DDTIONAL A BLATION S TUDIES
We also visualize the attention scores in each block for a speciﬁc 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.
Speciﬁcally, the particle will focus mainly on closer neightbors in the ﬁrst layer as is shown in Figure 7 (b).
6.2. B ASE E NVIRONMENTS
Dataset details .
We use the same setting for our datasets as mentioned in previous work (Li et al., 2019).
FluidFall contains two ﬂuid drops with different sizes.
The size for each drop are randomly generated with one drop larger than the other.
Position and viscosity for drops are randomly initialized.
This environment contains 189 particles with 121 frames for each rollout.
There arerollouts in training set and 300 rollouts in validation set.
FluidShake simulates the water in a moving box.
The speed of the box is randomly generated at each timestamp.
In addition, the size of the box and the number of particles are various for different rollouts.
In basic training and validation sets, the number of particles changes from 450 to 627.
This environment has 301 frames for each rollout.
There are 1800 rollouts in training set and 200 rollouts in validation set.
RiceGrip contains two grippers and a sticky rice.
The grippers’ positions and orientation are randomly initialized.
The number of particles for rice changes from 570 to 980 with 41 frames for each rollout in training and validation sets.
There are 4500 rollouts in training set and 500 rollouts in validation set.
BoxBath simulates a rigid cubic washed by water in a ﬁxed container.
The initial position of ﬂuid block and rigid cubic are randomly initialized.
This environment contains 960 ﬂuid particles and 64 rigid par- [Figure] ticles with 151 frames for each rollout.
There are 2700 rollouts in training set and 300 rollouts in validation set.
Results .
We display some rollouts for all experiments in this section.
The results on RiceGrip for
CConv are not available because CConv was designed speciﬁcally for ﬂuid dynamics.
The results on BoxBath for CConv are only used as reference of the simulation on ﬂuid parts.
And we compare
CConv with SiT mainly on the ﬂuid simulations.
[Table] [Table]
Rigid:
BoxBath Large Ball
Fluid: 960.
Rigid:
Fluid: 960.
Rigid:
BoxBath Large
Fluid: 960.
Rigid: Fluid: 960.
Rigid:
6.2. G ENERALIZATION DETAILS
We report WMSE results on BoxBath and other models’ generalization performance for references in Table 5.
SiT with abstract tokens achieves reasonable and superior performance.
We release the details of generalization settings in Table 6.
For the generalizations of FluidShake and RiceGrip , we display the range of particle numbers.
For the generalizations of BoxBath , we change the shape of rigid box into bunny, ball, cuboid, and larger cubic.
Speciﬁcally, the generalized
RiceGrip contains particles from 1060 to 1640, while the training set contains particles fromto 980; the generalized FluidShake contains particles from 720 to 1368, while the training set has 450 to 627 particles; For BoxBath , the training set has 960 ﬂuid particles and 64 rigid particles.
We ﬁrst separately increase the ﬂuid 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, andrespectively.
The followings are the rollouts on all environments.