FlashSchNet: Fast and Accurate Coarse-Grained Neural Network Molecular Dynamics

Paper: 2602.13140 Authors: Pingzhi Li, Hongxuan Li, Zirui Liu, Xingcheng Lin, Tianlong Chen Categories: cs.LG, cs.CE

Abstract

Graph neural network (GNN) potentials like SchNet have revolutionized molecular dynamics (MD) simulations by learning complex many-body interactions, offering superior accuracy and transferability compared to classical force fields. However, their computational efficiency has been hampered by fragmented kernel operations and memory-bound pipelines that fail to fully leverage GPU capabilities. FlashSchNet addresses this bottleneck by introducing an IO-aware design philosophy that carefully manages data movement between GPU high-bandwidth memory (HBM) and on-chip SRAM. Through four key optimizations—flash radial basis, flash message passing, flash aggregation, and channel-wise quantization—FlashSchNet achieves 1000 ns/day throughput on coarse-grained protein simulations with 64 parallel replicas, delivering 6.5x speedup over baseline CGSchNet while reducing peak memory by 80%.

Key Contributions

• IO-Aware Architecture: First GNN-MD framework explicitly designed around GPU memory hierarchy, minimizing expensive HBM reads/writes
• Flash Radial Basis: Fuses distance computation, Gaussian basis expansion, and cosine envelope into a single tiled kernel, computing each pairwise distance once and reusing across all basis functions
• Flash Message Passing: Eliminates materialization of edge tensors in HBM by fusing cutoff, neighbor gather, filter multiplication, and reduction operations
• Flash Aggregation: Reformulates scatter-add as CSR segment reduce, reducing atomic writes by a factor equal to feature dimension and enabling contention-free accumulation
• Channel-Wise Quantization: Exploits low per-channel dynamic range in SchNet MLP weights for 16-bit quantization with negligible accuracy loss

Technical Methodology

The core innovation of FlashSchNet lies in recognizing that GNN-MD performance is fundamentally limited by memory bandwidth rather than compute throughput. Traditional SchNet implementations suffer from three major inefficiencies:

Memory Bottleneck Analysis: The authors identify that naive SchNet implementations perform redundant distance computations (once for cutoff, again for basis expansion), materialize large intermediate edge tensors in HBM, and use atomic scatter operations that create memory contention. For a system with $N$ atoms and average degree $d$, this results in $O(Nd \cdot F)$ atomic writes where $F$ is the feature dimension.

Flash Radial Basis: Instead of computing pairwise distances $r_{ij} = \|\mathbf{r}_i - \mathbf{r}_j\|$ multiple times, FlashSchNet tiles the computation and immediately applies all $K$ Gaussian basis functions $e^{-\gamma(r_{ij} - \mu_k)^2}$ and the cosine cutoff $0.5[\cos(\pi r_{ij}/r_c) + 1]$ before writing to HBM. This reduces memory traffic from $O(Nd \cdot K)$ to $O(Nd)$ for distance computation.

Flash Message Passing: The message passing layer traditionally computes edge features $\mathbf{m}_{ij} = \mathbf{W}\phi(r_{ij}) \odot \mathbf{h}_j$, materializing all $Nd$ edge messages in HBM before aggregation. FlashSchNet fuses this with the subsequent reduction: $\mathbf{h}_i' = \sum_{j \in \mathcal{N}(i)} \mathbf{m}_{ij}$, streaming edge computations directly into accumulation without intermediate storage.

Flash Aggregation: By representing the neighbor list in CSR (Compressed Sparse Row) format, scatter-add becomes a segmented reduction that can be parallelized without atomic operations. This transforms $Nd \cdot F$ atomic writes into $N \cdot F$ contention-free writes, dramatically improving both forward and backward pass efficiency.

Performance Results

FlashSchNet demonstrates substantial improvements across multiple metrics:

Throughput: Achieves 1000 ns/day aggregate throughput over 64 parallel replicas on a coarse-grained protein system (269 beads) using a single NVIDIA RTX PRO 6000 GPU. This represents a 6.5x speedup compared to the CGSchNet baseline while maintaining equivalent accuracy.

Memory Efficiency: Reduces peak memory consumption by 80% through kernel fusion and elimination of intermediate tensors. This enables larger batch sizes and more parallel replicas on the same hardware.

Accuracy Preservation: Channel-wise 16-bit quantization maintains SchNet-level accuracy on standard benchmarks. The low per-channel dynamic range in MLP weights allows aggressive quantization without significant precision loss.

Comparison to Classical Force Fields: FlashSchNet surpasses the throughput of classical coarse-grained force fields like MARTINI while retaining the superior accuracy and transferability of learned GNN potentials. This bridges the long-standing gap between ML-based and classical MD methods.

Implications and Future Directions

FlashSchNet represents a paradigm shift in how we approach GNN-based scientific computing. The IO-aware design principles are broadly applicable beyond molecular dynamics:

Scalability: The techniques enable GNN-MD to scale to biologically relevant timescales (microseconds to milliseconds) that were previously accessible only to classical force fields, opening new research avenues in protein folding, drug discovery, and materials science.

Hardware Efficiency: By explicitly accounting for memory hierarchy, FlashSchNet demonstrates that ML models can achieve classical force field-level performance without sacrificing accuracy. This challenges the assumption that learned potentials must be slower than analytical ones.

Generalization: The flash kernel design patterns—fusing operations, avoiding materialization, reformulating scatter operations—are applicable to other GNN architectures (e.g., DimeNet, GemNet) and domains (e.g., computational chemistry, materials modeling).

Future Work: Potential extensions include multi-GPU scaling with communication-aware partitioning, integration with enhanced sampling methods, and application to all-atom simulations where the accuracy-speed tradeoff becomes even more critical.

Takeaways

1. Memory bandwidth, not compute, is the primary bottleneck in GNN-based molecular dynamics—FlashSchNet’s IO-aware design achieves 6.5x speedup by minimizing HBM traffic
2. Kernel fusion is essential for GNN efficiency: combining distance computation, basis expansion, and message passing into single-pass operations eliminates redundant memory accesses
3. CSR-based aggregation eliminates atomic contention: reformulating scatter-add as segmented reduction reduces atomic writes by a factor equal to feature dimension
4. Channel-wise quantization exploits MLP weight structure: 16-bit precision is sufficient for SchNet due to low per-channel dynamic range, enabling further throughput gains
5. Learned potentials can match classical force field speed: FlashSchNet surpasses MARTINI throughput while maintaining superior accuracy, bridging the performance gap in computational chemistry