Transfer learning in Scalable Graph Neural Network for Improved Physical Simulation

TL;DR

The paper tackles the data-intensity barrier in GNN-based physics simulation by introducing a pre-training and transfer-learning framework for a scalable graph U-Net (SGUNET) that employs DFS pooling to handle multi-resolution meshes. A dedicated ABCD pre-training dataset (~$20{,}000$ simulations) enables effective knowledge transfer to downstream quasi-static tasks, with parameter-sharing strategies (Uniform and First-N) and a Frobenius-norm weight-regularization term enhancing generalization. Empirically, SGUNET outperforms the MeshGraphNet baseline on pre-training and demonstrates data-efficient improvements on Deformable Plate and Deforming Plate when fine-tuned with reduced data, including an $11.05\%$ RMSE improvement on 2D deformable plates with $1/16$ data. The work offers a scalable, data-efficient path for transferring physics priors across mesh configurations, with practical implications for rapid prototyping and deployment of GNN-based simulators.

Abstract

In recent years, Graph Neural Network (GNN) based models have shown promising results in simulating physics of complex systems. However, training dedicated graph network based physics simulators can be costly, as most models are confined to fully supervised training, which requires extensive data generated from traditional physics simulators. To date, how transfer learning could improve the model performance and training efficiency has remained unexplored. In this work, we introduce a pre-training and transfer learning paradigm for graph network simulators. We propose the scalable graph U-net (SGUNET). Incorporating an innovative depth-first search (DFS) pooling, the SGUNET is adaptable to different mesh sizes and resolutions for various simulation tasks. To enable the transfer learning between differently configured SGUNETs, we propose a set of mapping functions to align the parameters between the pre-trained model and the target model. An extra normalization term is also added into the loss to constrain the difference between the pre-trained weights and target model weights for better generalization performance. To pre-train our physics simulator we created a dataset which includes 20,000 physical simulations of randomly selected 3D shapes from the open source A Big CAD (ABC) dataset. We show that our proposed transfer learning methods allow the model to perform even better when fine-tuned with small amounts of training data than when it is trained from scratch with full extensive dataset. On the 2D Deformable Plate benchmark dataset, our pre-trained model fine-tuned on 1/16 of the training data achieved an 11.05\% improvement in position RMSE compared to the model trained from scratch.

Figures (12)

• Figure 1: (a) An illustration of the composition of mesh data $\mathcal{M}$ and its representation as a heterogeneous graph $\mathcal{G}^{hetero}$. (b) Example of down-sampled graphs with pooling ratios as 3 and 2.
• Figure 2: (a) A detailed depiction of our proposed model, SguNet, which includes four primary modules: the Processor $Pr_i$ for information propagation; the Encoder for data transformation, the GUnet for graph pooling, and the Decoder for downstream tasks. (b) & (c) Mapping functions for GUNet stages and GNBs for the case where pre-trained model has more stages and per-processor GNBs than the fine-tuned model.
• Figure 3: Randomized FEA simulation dataset using geometry from ABC dataset.
• Figure 4: The FEA simulation results using ABC CAD dataset highlight various deformation modes, including compression with associated tension around a hole, as well as plate and beam bending.
• Figure 5: Simulated meshes at various stages ($t$=30 at the top row, $t$=50 at the bottom row) for different models: MGN, MGN-FT (fine-tuned with Uniform and First-N strategies), SguNet, SguNet-FT (fine-tuned with Uniform and First-N strategies), and the ground truth. All models are trained on 1/8 of the original training size. The colors indicate displacement magnitude.