TRACE: 用于颗粒动力学的时空接触记忆图网络模拟器

# TRACE：面向颗粒动力学的时空接触记忆图网络模拟器  
来源：https://arxiv.org/html/2609.02991 \\setheadertext D&CG Lab\\correspondingemail‡通讯作者\. \\emailiconchangjianz@student\.unimelb\.edu\.au Data\-Driven\-Computational\-Geotechnics/TRACE (https://github.com/Data-Driven-Computational-Geotechnics/TRACE)\\setheadertitleTRACE：面向颗粒动力学的时空接触记忆图网络模拟器 Changjian Zhou, Negin Yousefpour, Jie Qi, Junfeng Fang, Guillermo A\. Narsilio, Hans Petter Jostad所属机构：墨尔本大学工程与信息技术学院 新加坡国立大学计算学院 挪威岩土工程研究所  

#### 摘要  
学习型图模拟器为颗粒动力学提供了一种高效的替代方案，可替代高保真度求解器。然而，颗粒运动强烈依赖于颗粒间接触历史，当颗粒接触形成、断裂和重组时，该历史信息难以保持。现有模拟器主要通过节点特征或节点级记忆存储时序信息。本文引入TRACE，一种图网络模拟器，其直接在接触边上存储交互历史。每条边通过基于注意力的消息传递和门控循环单元维护一个持久记忆，并借助边标识字典在接触图变化时保存此记忆。物理结构化解码器预测颗粒间法向和切向接触力、施加库仑摩擦限制，并应用等大反向内力。模型采用单步预训练后接自回归展开微调的训练策略。我们在二维和三维颗粒柱坍塌基准测试上评估了TRACE。结果表明，TRACE均能生成稳定且物理一致的长期展开轨迹，精确复现最终沉积几何形态及坍塌过程中释放的动能。与图网络模拟器（GNS）和节点记忆图神经模拟器（NMGNS）相比，TRACE在两个基准测试中将长期展开位置误差降低31–62%，最终沉积误差降低58–89%，同时使用更少参数并保持近乎零的颗粒互穿。TRACE在二维和三维场景中分别实现了相较于物质点法（MPM）参考解算器12.2倍和8.9倍的加速。代码已开源于https://github.com/Data-Driven-Computational-Geotechnics/TRACE。  

关键词：图神经网络；学习型物理模拟；科学机器学习；颗粒材料；接触边记忆；颗粒柱坍塌。  

## 符号列表  
符号 描述  
N, d 颗粒数量与空间维度  
i, j 接触对的颗粒索引，排序满足i<j  
δ_ij = (r_i + r_j) - ‖x_i - x_j‖ > 0时，两颗粒沿接触法线方向各向半重叠距离位移：  
x_i ← x_i + ½δ_ij n_ij, x_j ← x_j - ½δ_ij n_ij (28)  
其中n_ij为单位接触法线。该基于位置的（雅可比式）投影并行应用于所有重叠颗粒对，并在每步迭代固定次数。  
(2) *边界投影*。颗粒被限制在模拟域内。任何中心穿透底板或侧壁的颗粒将被投影回边界，且指向墙内的速度分量设为零，使颗粒既不穿越也不粘附边界。重叠投影仅调整位置；墙壁投影同时将向内速度分量置零。两者仅在展开阶段生效。结合解码器内置的力级保证，它们确保模拟轨迹无互穿和域外状态，同时不限制学习的力模型。  

### 2.7训练策略  
第一阶段采用教师强制法。每时间步模型以参考状态为输入独立预测单步加速度，仅接触记忆跨时间窗口传播。第二阶段使用自回归展开：从单一参考状态开始，所有后续状态均由模型自身预测生成。无梯度漂移段先产生偏离轨迹的状态以反映模型累积误差，随后应用多步位置监督以抑制误差进一步增长。第一阶段旨在学习精确的局部动力学，第二阶段则提升模型自身展开分布下的长期稳定性。  

#### 2.7.1单步预训练  
模型首先被训练预测每个颗粒的单步加速度。同时，循环接触记忆跨连续帧传播，使其能学习与未来预测相关的历史交互。该阶段结合归一化加速度回归、输入噪声和截断时间反向传播。在时间步t，目标加速度通过有限差分参考轨迹中的颗粒速度获得：  
a_i^{⋆,t} = (v_i^{t+1} - v_i^t)/Δt (29)  
其中a_i^{⋆,t}为颗粒i的目标加速度，v_i^t和v_i^{t+1}为其在两个连续步的参考速度，Δt为模拟时间步。预测加速度与目标加速度在归一化空间中比较。单步损失为：  
ℓ^t = (1/N) ∑_{i=1}^N Huber(norm(a_i^t), norm(a_i^{⋆,t})) (30)  
其中ℓ^t为单步损失，N为颗粒数量，a_i^t为解码器预测的加速度（见第2.5节），norm(·)为基于训练集拟合的通道归一化器。Huber损失对小残差呈二次性，对大残差呈线性，使训练对偶尔的加速度离群值不敏感。同时向输入位置和速度添加小高斯扰动以提升对展开期间遇到的不完美状态的鲁棒性。  

由于接触记忆循环演化（见第2.3节），训练使用截断时间反向传播，窗口包含T_W个连续帧。在每个窗口内，接触记忆和标识字典按推理方式精确传播，每步均由方程(29)中的加速度目标监督。窗口损失为：  
ℒ = (1/T_W) ∑_{k=0}^{T_W-1} ℓ^{t_0+k} (31)  
其中t_0表示窗口起始。梯度通过窗口内的记忆递归传播，使接触状态保留历史信息以提升后续预测。  

#### 2.7.2展开微调  
教师强制预训练主要使模型接触参考状态。然而在自回归推理中，每个预测状态成为下一步输入，导致误差累积并逐渐使展开偏离训练分布。输入噪声仅部分缓解此不匹配，因其引入小的独立扰动而非模型自身产生的时空相关误差。  

为减少训练-测试不匹配，我们遵循Brandstetter等人的前推原则，通过自回归展开微调模型。每个训练样本包含无梯度漂移段和监督展开段。漂移段从参考状态开始，自回归推进模型随机采样步数（最多至最大漂移深度H_max）。几何约束（见第2.6节）严格按推理方式应用。该段无监督或梯度追踪执行，目的是生成现实的偏离轨迹颗粒状态及相应接触记忆。由此状态初始化T_ft步的监督展开。由于模型状态已偏离参考轨迹，方程(29)定义的单步加速度目标不再与当前模型状态一致，故监督直接应用于颗粒位置：  
ℒ_ft = (1/T_ft) ∑_{k=1}^{T_ft} (1/N) ∑_{i=1}^N ‖x_i^{t_0+k} - ĥx_i^{t_0+k}‖_2^2 (32)  
其中t_0为监督段起始，x_i^{t_0+k}为约束展开生成的颗粒i位置，ĥx_i^{t_0+k}为相应参考位置。损失对所有颗粒和监督展开步的平方位置误差取平均。通过优化模型自身生成状态的预测，该目标提升了模型限制长期模拟误差累积的能力。  

## 3数值实验  
### 3.1数据集  
我们在Sanchez-Gonzalez等人引入的沙子基准测试上评估方法，考虑其二维（2D-Sand）和三维（3D-Sand）设置。每个场景中，一个或多个致密颗粒体在封闭矩形容器内随机位置初始化并在重力作用下释放。随着物料撞击边界、扩散并与相邻物体相互作用，接触在整个模拟过程中持续出现、滑动、分离和重组。详细基准规范见附录A.1。  

两个数据集均使用MPM生成，这是一种用于大变形颗粒动力学的高保真连续体求解器，已成为学习型颗粒模拟器的标准基准。其主要特征汇总于表1。  

表1：本文所用颗粒沙子基准数据集统计  
属性 | 2D-Sand | 3D-Sand  
---|---|---  
空间维度 | 2 | 3  
真实求解器 | MPM | MPM  
训练/验证/测试轨迹数 | 1000/30/30 | 1000/100/100  
每场景颗粒数（最小-最大） | 107-1976 | 4586-19762  
每场景平均颗粒数 | ≈1159 | ≈10000  
每轨迹步数 | 320 | 350  
时间步Δt | 2.5×10⁻³ | 2.5×10⁻³  
模拟域 | [0.1,0.9]² | [0.1,0.9]³ → [0.2,0.8]³  
边界条件 | 刚性围壁 | 刚性围壁  

注：两个数据集均使用MPM生成。一个或多个致密颗粒体在封闭容器内随机初始化并在重力作用下释放，经历自由落体、撞击、扩散及壁/物体间碰撞后达到稳定沉积。颗粒速度和加速度通过位置有限差分重构。无真实接触力数据。图构建时，颗粒在固定邻域半径内连接（2D为0.015，3D为0.025）。由于MPM数据仅包含点位置，每个颗粒被分配均匀半径用于接触推理，该半径取最近邻中位间距的一半左右。全程中，轨迹指参考解算器产生的一次模拟序列。每个测试轨迹由其初始颗粒配置定义，模型通过从该初始状态展开进行评估。每轨迹仅存储颗粒位置，速度和加速度通过有限差分重构，加速度目标使用训练集统计信息对每个通道独立归一化。重要的是，两个数据集均未提供真实接触力。因此，TRACE预测的接触法向力和切向力完全在无力监督下学习，要求模型仅从颗粒运动学推断路径依赖的摩擦相互作用。  

### 3.2计算环境  
本文所有实验（包括训练、展开评估和计时测量）均在单个工作站上进行（表2）。该机器配备AMD Ryzen Threadripper PRO 5995WX处理器（64核/128线程）、512GB系统内存和四块NVIDIA GeForce RTX 4090 GPU（各24GB显存）。软件栈为Ubuntu 22.04 LTS、PyTorch 2.5.1（CUDA 12.1）和PyTorch Geometric 2.8.0。所有模型均在单精度下训练和评估。2D-Sand训练使用单GPU，3D-Sand训练使用四GPU数据并行，而所有3D-Sand评估在单GPU上运行。所有模型和基线在同一机器上评估以确保两个案例研究的准确性和运行时间公平比较。所有报告时间均在热身阶段后测量并启用显式设备同步。完整硬件和软件配置如下：  

表2：硬件与软件配置  
组件 | 规格  
---|---  
CPU | AMD Ryzen Threadripper PRO 5995WX，64核/128线程  
系统内存 | 512GB  
GPU | 4×NVIDIA GeForce RTX 4090，各24GB GDDR6X  
操作系统 | Ubuntu 22.04.3 LTS，Linux内核6.8  
深度学习框架 | PyTorch 2.5.1（CUDA 12.1，cuDNN 9.1）  
图库 | PyTorch Geometric 2.8.0  
Python | 3.13  
数值精度 | float32  
训练并行 | 1 GPU（2D-Sand）；4-GPU DistributedDataParallel（3D-Sand）  
推理 | 所有报告展开和计时均使用1 GPU  

### 3.3评估协议  
测试时，模拟器完全按部署方式应用。给定某步状态，预测颗粒加速度，积分器推进状态，结果作为下一步输入。第3.3.1节定义了执行此循环的两种展开协议。性能通过位置均方根误差（RMSE）量化运动精度（第3.3.2节）、归一化沉积误差评估沉积形态（第3.3.3节）、动能诊断评估物理一致性（第3.3.4节）。  

#### 3.3.1短程与长程展开协议  
两种协议仅在自回归循环校正状态前的运行时长上不同。  

**短程展开**。模拟器以20步为窗口自回归运行。每窗口起始时，状态从真实值重新初始化。