TrimSO2 speeds up training of memory-heavy molecule simulations
TrimSO(2): Fast and Memory-Efficient Training of Equivariant Interatomic Potentials
Distributed, Parallel, and Cluster Computing
Summary
Simulating molecules accurately often needs complex math that takes a lot of time and memory. The authors found a way to reorganize these calculations so the computer doesn’t have to hold large intermediate steps separately. This means training models that predict molecular behavior can run much faster and use less memory. Their method, called TrimSO(2), works well on powerful GPUs and speeds up training by two to three times while saving memory.
What this means in practice
- •For computational chemistry teams: Train molecular simulation models faster and with less GPU memory, enabling larger or more complex materials simulations.
- •For machine learning engineers: Optimize training workflows for equivariant neural networks by integrating dataflow steps to reduce memory overhead during model updates.
Authors
Chen Wang, Yuanchang Zhou, Guangming Tan, Mingzhen Li, Weile Jia
Abstract
Equivariant interatomic potentials enable accurate and efficient simulations of molecules and materials. Among them, models based on SO(2) convolutions reduce the cost of tensor products and achieve state-of-the-art accuracy, but their training remains time-consuming and memory-intensive. Existing training implementations materialize large edge tensors for representation transformations surrounding dense matrix multiplications, increasing memory usage and data movement. We observe that these transformations can be integrated into the model's existing dataflow, avoiding separate intermediate materialization. We introduce TrimSO(2), a fast and memory-efficient algorithm for SO(2)-based equivariant models. It realizes this principle through rotation on gather, recombination with activation, and inverse rotation on reduction, with support for the higher-order differentiation required for training. By avoiding intermediate materialization, TrimSO(2) delivers 2.71-16.73$\times$ operator speedups across forward, backward, and double backward and reduces model-level high-bandwidth memory traffic by 3.81-9.88$\times$. Finally, TrimSO(2) achieves 2.07-2.94$\times$ training speedups and 1.53-3.16$\times$ reductions in peak training memory over PyTorch implementations for official eSEN, UMA, and EquiformerV3 workloads on A100 and H100 GPUs.