Transformer with Implicit Edges for Particle-based Physics Simulation

For a particle-based system composed of $N$ particles, we use $\mathcal{X}^{t}=\{{\bm{x}}^{t}_{i}\}_{i=1}^{N}$ to denote the system state at time step $t$, where ${\bm{x}}^{t}_{i}$ denotes the state of $i$-th particle. Specifically, ${\bm{x}}^{t}_{i}=[{\bm{p}}^{t}_{i},{\bm{q}}^{t}_{i},{\bm{a}}_{i}]$, where ${\bm{p}}^{t}_{i},{\bm{q}}^{t}_{i}\in\mathbb{R}^{3}$ refer to position and velocity, and ${\bm{a}}_{i}\in\mathbb{R}^{d_{a}}$ represents fixed particle attributes such as its material type. The goal of a simulator is to learn a model $\phi(\cdot)$ from previous rollouts of a system to causally predict a rollout trajectory in a specific time period conditioned on the initial system state $\mathcal{X}^{0}$. The prediction runs in a recursive manner, where the simulator will predict the state $\hat{\mathcal{X}}^{t+1}=\phi(\mathcal{X}^{t})$ at time step $t+1$ based on the state $\mathcal{X}^{t}=\{x^{t}_{i}\}$ at time step $t$. In practice, we will predict the velocities of particles $\hat{Q}^{t+1}=\{\hat{{\bm{q}}}^{t+1}_{i}\}$, and obtain their positions via $\hat{{\bm{p}}}^{t+1}_{i}={\bm{p}}^{t}_{i}+\Delta t\cdot\hat{{\bm{q}}}^{t+1}_{i}$, where $\Delta t$ is a domain-specific constant. In the following discussion, the time-step $t$ is omitted to avoid verbose notations.

As particle-based physics systems can be naturally viewed as directed graphs, a straightforward solution for particle-based physics simulation is applying graph neural network (GNN) [2, 23, 17, 12, 22]. Specifically, we can regard particles in the system as graph nodes, and interactions between pairs of particles as directed edges. Given the states of particles $\mathcal{X}=\{{\bm{x}}_{i}\}_{i=1}^{N}$ at some time-step, to predict the velocities of particles in the next time-step, GNN will at first obtain the initial node features and edge features following:

	$\displaystyle{\bm{v}}_{i}^{(0)}$	$\displaystyle=f^{\mathrm{enc}}_{V}({\bm{x}}_{i}),$		(1)
	$\displaystyle{\bm{e}}_{ij}^{(0)}$	$\displaystyle=f^{\mathrm{enc}}_{E}({\bm{x}}_{i},{\bm{x}}_{j}),$		(2)

where ${\bm{v}}_{i},{\bm{e}}_{ij}\in\mathbb{R}^{d_{h}}$ are $d_{h}$ dimensional vectors, and $f^{\mathrm{enc}}_{V}(\cdot),f^{\mathrm{enc}}_{E}(\cdot)$ are respectively the node and edge encoders. Subsequently, GNN will conduct $L$ rounds of message-passing, and obtain the velocities of particles as:

	$\displaystyle{\bm{e}}^{(l+1)}_{ij}$	$\displaystyle=f^{\mathrm{prop}}_{E}({\bm{v}}^{(l)}_{i},{\bm{v}}^{(l)}_{j},{\bm{e}}^{(l)}_{ij}),$		(3)
	$\displaystyle{\bm{v}}^{(l+1)}_{i}$	$\displaystyle=f^{\mathrm{prop}}_{V}({\bm{v}}^{(l)}_{i},\sum_{j\in\mathcal{N}_{i}}{\bm{e}}^{(l+1)}_{ij}),$		(4)
	$\displaystyle\hat{{\bm{q}}}_{i}$	$\displaystyle=f^{\mathrm{dec}}_{V}({\bm{v}}^{(L)}_{i}),$		(5)

where $\mathcal{N}_{i}$ indicates the set of neighbors of $i$-th particle, and $f^{\mathrm{prop}}_{E}(\cdot),f^{\mathrm{prop}}_{V}(\cdot)$ and $f^{\mathrm{dec}}_{V}(\cdot)$ are respectively the node propagation module, the edge propagation module as well as the node decoder. In practice, $f^{\mathrm{enc}}_{V}(\cdot),f^{\mathrm{enc}}_{E}(\cdot),f^{\mathrm{prop}}_{E}(\cdot),f^{\mathrm{prop}}_{V}(\cdot)$ and $f^{\mathrm{dec}}_{V}(\cdot)$ are often implemented as multi-layer perceptrons (MLPs). Moreover, a window function $g$ is commonly used to filter out interactions between distant particles and reduce computational complexity:

$\displaystyle g(i,j)$

$\displaystyle=\mathbf{1}\left(\|{\bm{p}}_{i}-{\bm{p}}_{j}\|_{2}<R\right),$

(6)

where $\mathbf{1}(\cdot)$ is the indicator function and $R$ is a pre-defined threshould.

To accurately simulate the changes of a system over time, it is crucial to exploit the rich semantics conveyed by the interactions among particles, such as the energy transition of a system when constrained by material characteristics and physical laws. While GNN achieves this by explicitly modeling particle interactions as graph edges, such a treatment also leads to substantial computational overhead. Since a particle-based system contains hundreds even thousands of particles, and particles of a system are densely clustered together, this issue significantly limits the efficiency of GNN-based approaches.

Inspired by recent successes of Transformer [26] that applies computational efficient self-attention operations to model the communication among different tokens, in this paper we propose a Transformer-based method, which we refer to as Transformer with Implicit Edges, TIE, for particle-based physics simulation. We first describe how to apply a vanilla Transformer in this task. Specifically, we assign a token to each particle of the system, and therefore particle interactions are naturally achieved by $L$ blocks of multi-head self-attention modules. While the token features are initialized according to Equation 1, they will be updated in the $l$-th block as:

	$\displaystyle\omega_{ij}$	$\displaystyle=$	$\displaystyle(W^{(l)}_{Q}{\bm{v}}^{(l)}_{i})^{\top}\cdot(W^{(l)}_{K}{\bm{v}}^{(l)}_{j}),$		(7)
	$\displaystyle{\bm{v}}^{(l+1)}_{i}$	$\displaystyle=$	$\displaystyle\sum_{j}\frac{\exp(\omega_{ij}{/\sqrt{d}})}{\sum_{k}\exp(\omega_{i{k}}{/\sqrt{d}})}\cdot(W^{(l)}_{V}{\bm{v}}^{(l)}_{j}),$		(8)

where $d$ is the dimension of features, and $W^{(l)}_{Q},W^{(l)}_{K},W^{(l)}_{V}$ are weight matrices for queries, keys, and values. Following the standard practice, a mask is generated according to Equation 6 to mask out distant particles when computing the attention. And finally, the prediction of velocities follows Equation 5.

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, the vanilla Transformer uses attention weights that are scalars obtained via dot-product, to represent particle interactions, which are insufficient to reflect the rich semantics of particle interactions. To combine the merits of GNN and Transformer, TIE modify the self-attention operation in the vanilla Transformer to implicitly include edge features as in GNN in an edge-free manner. In Figure 2 we include the comparison between our proposed implicit edges and the explicit edges in GNN. Specifically, since $f^{\mathrm{prop}}_{E}$ in Equation 3 and $f^{\mathrm{enc}}_{E}$ in Equation 2 are both implemented as an MLP in practice, by expanding Equation 3 recursively and grouping terms respectively for $i$-th and $j$-th particle we can obtain:

	$\displaystyle{\bm{r}}^{(0)}_{i}$	$\displaystyle=W^{(0)}_{r}{\bm{x}}_{i},\qquad{\bm{s}}^{(0)}_{j}=W^{(0)}_{s}{\bm{x}}_{j},$		(9)
	$\displaystyle{\bm{r}}^{(l)}_{i}$	$\displaystyle=W^{(l)}_{r}{\bm{v}}^{(l)}_{i}+W^{(l)}_{m}{\bm{r}}^{(l-1)}_{i},$		(10)
	$\displaystyle{\bm{s}}^{(l)}_{j}$	$\displaystyle=W^{(l)}_{s}{\bm{v}}^{(l)}_{j}+W^{(l)}_{m}{\bm{s}}^{(l-1)}_{j},$		(11)
	$\displaystyle{\bm{e}}^{(l+1)}_{ij}$	$\displaystyle={\bm{r}}^{(l)}_{i}+{\bm{s}}^{(l)}_{j},$		(12)

where the effect of an explicit edge can be effectively achieved by two additional tokens, which we refer to as the receiver token ${\bm{r}}_{i}$ and the sender token ${\bm{s}}_{j}$. The detained expansion of Equation 3 can be found in the supplemental material. Following the above expansion, TIE thus assigns three tokens to each particle of the system, namely a receiver token ${\bm{r}}_{i}$, a sender token ${\bm{s}}_{i}$, and a state token ${\bm{v}}_{i}$. The state token is similar to the particle token in the vanilla Transformer, and its update formula combines the node update formula in GNN and the self-attention formula in Transformer:

	$\displaystyle\omega^{\prime}_{ij}$	$\displaystyle=$	$\displaystyle(W^{(l)}_{Q}{\bm{v}}^{(l)}_{i})^{\top}{\bm{r}}^{(l)}_{i}+(W^{(l)}_{Q}{\bm{v}}^{(l)}_{i})^{\top}{\bm{s}}^{(l)}_{j},$		(13)
	$\displaystyle{\bm{v}}^{(l+1)}_{i}$	$\displaystyle=$	$\displaystyle{\bm{r}}^{(l)}_{i}+\sum_{j}\frac{\exp(\omega^{\prime}_{ij}{/\sqrt{d}})}{\sum_{k}\exp(\omega^{\prime}_{i{k}}{/\sqrt{d}})}\cdot{\bm{s}}^{(l)}_{j}.$		(14)

We refer to Equation 14 as an implicit way to incorporate the rich semantics of particle interactions, since TIE approximates graph edges in GNN with two additional tokens per particle, and more importantly these two tokens can be updated, along with the original token, separately for each particle, avoiding the significant computational overhead. To interpret the modified self-attention in TIE, from the perspective of graph edges, we decompose them into the receiver tokens and the sender tokens, maintaining two extra paths in the Transformer’s self-attention module. As for the perspective of self-attention, the receiver tokens and the sender tokens respectively replace the original keys and values.

In practice, since GNN-based methods usually incorporate LayerNorm [1] in their network architectures that computes the mean and std of edge features to improve their performance and training speed, we can further modify the self-attention in Equation 13 and Equation 14 to include the effect of normalization as well:

	$\displaystyle{\left(\sigma^{(l)}_{ij}\right)^{2}}$	$\displaystyle=$	$\displaystyle\frac{1}{d}({\bm{r}}^{(l)}_{i})^{\top}{\bm{r}}^{(l)}_{i}+\frac{1}{d}({\bm{s}}^{(l)}_{j})^{\top}{\bm{s}}^{(l)}_{j}+\frac{2}{d}({\bm{r}}^{(l)}_{i})^{\top}{\bm{s}}^{(l)}_{j}-(\mu^{(l)}_{{r_{i}}}+\mu^{(l)}_{{s_{j}}})^{2},$		(15)
	$\displaystyle\omega^{\prime\prime}_{ij}$	$\displaystyle=$	$\displaystyle\frac{(W^{(l)}_{Q}{\bm{v}}^{(l)}_{i})^{\top}({\bm{r}}^{(l)}_{i}-\mu^{(l)}_{{r_{i}}})+(W^{(l)}_{Q}{\bm{v}}^{(l)}_{i})^{\top}({\bm{s}}^{(l)}_{j}-\mu^{(l)}_{{s_{j}}})}{{\sigma^{(l)}_{ij}}},$		(16)
	$\displaystyle{\bm{v}}^{(l+1)}_{i}$	$\displaystyle=$	$\displaystyle\sum_{j}\frac{\exp(\omega^{\prime\prime}_{ij}{/\sqrt{d}})}{\sum_{k}\exp(\omega^{\prime\prime}_{i{k}}{/\sqrt{d}})}\cdot{\frac{({\bm{r}}^{(l)}_{i}-\mu^{(l)}_{r_{i}})+({\bm{s}}^{(l)}_{j}-\mu^{(l)}_{s_{j}})}{\sigma^{(l)}_{ij}}},$		(17)

where $\mu^{(l)}_{{r_{i}}}$ and $\mu^{(l)}_{{s_{j}}}$ are respectively the mean of receiver tokens and sender tokens after $l$-th block. Detailed deduction can be found in the supplemental material.

To further improve the generalization ability of TIE and disentangle global material-specific semantics from local particle-wise semantics, we further equip TIE with material-specific abstract particles.

For $N_{a}$ types of materials, TIE respectively adopts $N_{a}$ abstract particles $A=\{{\bm{a}}_{k}\}_{k=1}^{N_{a}}$, each of which is a virtual particle with a learnable state token. Ideally, the abstract particle ${\bm{a}}_{k}$ should capture the material-specific semantics of $k$-th material. They act as additional particles in the system, and their update formulas are the same as normal particles. Unlike normal particles that only interact with neighboring particles, each abstract particle is forced to interact with all particles belonging to its corresponding material. Therefore with $N_{a}$ abstract particles TIE will have $N+N_{a}$ particles in total: $\{{\bm{a}}_{1},\cdots,{\bm{a}}_{N_{a}},{\bm{x}}_{1},\cdots,{\bm{x}}_{N}\}$. Once TIE is trained, abstract particles can be reused when generalizing TIE to unseen domains that have same materials but vary in particle amont and configuration.

To train TIE with existing rollouts of a domain, the standard mean square error (MSE) loss is applied to the output of TIE:

$\displaystyle\text{MSE}(\hat{Q},Q)$

$\displaystyle=$

$\displaystyle\frac{1}{N}\sum_{i}\|\hat{{\bm{q}}}_{i}-{\bm{q}}_{i}\|_{2}^{2},$

(18)

where $\hat{Q}=\{\hat{{\bm{q}}}_{i}\}_{i=1}^{N}$ and $Q=\{{\bm{q}}_{i}\}_{i=1}^{N}$ are respectively the estimation and the ground truth, and $\|\cdot\|_{2}$ is the L2 norm.

In terms of evaluation metric, since a system usually contains multiple types of materials with imbalanced numbers of particles, to better reflect the estimation accuracy, we apply the Mean of Material-wise MSE ($\text{M}^{3}\text{SE}$) for evaluation:

$\displaystyle\text{M}^{3}\text{SE}(\hat{Q},Q)$

$\displaystyle=$

$\displaystyle\frac{1}{K}\sum_{k}\frac{1}{N_{k}}\sum_{i}\|\hat{{\bm{q}}}_{i,k}-{\bm{q}}_{i,k}\|_{2}^{2},$

(19)

where $K$ is the number of material types, $N_{k}$ is the number of particles belonging to the $k$-th material. $\text{M}^{3}\text{SE}$ is equivalent to the standard MSE when $K=1$.

Quantitative results are provided in Table 1, while qualitative results are shown in Figure 4. TIE achieves superior performances on all domains. The effectiveness of abstract particles are more obvious for RiceGrip and BoxBath, which involve complex materials or multi-material interactions.

As shown in Table 2, we generate more complex domains to challenge the robustness of our full model TIE+. Specifically, we add more particles for FluidShake and RiceGrip, which we refer to as L-FluidShake and L-RiceGrip respectively. The L-FluidShake includes 720 to 1368 particles, while L-RiceGrip contains 1062 to 1642 particles. On BoxBath, we change the size and shape of rigid object. Specifically, we add more fluid particles in Lfluid-BoxBath to 1280 fluid particles, while we enlarge the rigid cube in L-BoxBath to 125 particles. We also change the shape of the rigid object into ball and bunny, which we refer to BallBox and BunnyBath respectively. Details of generalized environments settings and results can be found in supplementary materials.

Quantitative results are summarized in Table 2, while qualitative results are depicted in Figure 5. As shown in Table 2, TIE+ achieves lower M³SEs on most domains, while having more faithful rollouts in Figure 5. On L-FluidShake, TIE+ maintains the block of fluid in the air and predicts faithful wave on the surface. On L-RiceGrip, while DPI-Net and GNS have difficulties in maintaining the shape, the rice predicted by GraphTrans is compressed more densely only in the center areas where the grips have reached, the left side and right side of the rice does not deform properly compared with the ground truth. In contrast, TIE+ is able to maintains the shape of the large rice and faithfully deform the whole rice after compressed. On generalized BoxBath, TIE+ is able to predict faithful rollout when the fluid particles flood the rigid objects into the air or when the wave of the fluid particles starts to push the rigid object after the collision. Even when the rigid object changes to bunny with more complex surfaces, TIE+ generates more accurate predictions for both fluid particles and rigid particles.

We comprehensively analyze our TIE and explore the effectiveness of our model in the following aspects: (a) with and without implicit edges; (b) with and without normalization effects in attention; (c) with and without abstract particles; and (d) the sensitiveness to $R$. The experiments for (a), (b), and (d) are conducted on FluidShake and L-FluidShake, while experiment (c) is conducted on BoxBath. The quantitative results are in Table 3 and Table 4.