Accurate and efficient molecular dynamics based on machine learning and non von Neumann architecture

The study targets the long-standing trade-off in molecular dynamics between accuracy and efficiency. AIMD delivers high accuracy by evaluating PES via DFT but is computationally prohibitive for large systems, whereas CMD is efficient but can suffer from large PES errors. Recent machine-learning MD (MLMD) approaches reduce this gap, achieving AIMD-level accuracy with improved speed, yet still lag CMD by around two orders of magnitude. Additionally, most MD runs on von Neumann (vN) architectures that suffer from a memory wall, where data shuttling between processing and memory dominates time and energy. The authors propose a paradigm shift to a non-von Neumann (NvN) architecture combined with a tailored ML potential to achieve both AIMD-level accuracy and CMD-level efficiency, addressing the memory bottleneck and enabling scalable, energy-efficient MD.

The paper reviews: (1) AIMD’s accuracy versus cost limitations; (2) CMD’s efficiency versus limited accuracy with classical force fields; (3) MLMD (e.g., DeepMD) that fits PES with neural networks to approach AIMD accuracy while being orders faster, yet still significantly slower than CMD; (4) the vN memory wall problem where data movement dominates computing time/energy; and (5) prior special-purpose MD computers (e.g., Anton) that accelerate CMD but inherit accuracy limitations of classical force fields and are application-specific. This motivates a special-purpose, ML-driven, non-vN MD machine to overcome both accuracy and efficiency constraints.

The authors propose NVNMD, a special-purpose MD system coupling a modified DeepMD-style ML potential with a non-von Neumann (NvN) processing-in-memory (PIM) hardware design. - Workflow: (i) Train an ML model to reproduce PES on standard vN hardware (CPU/GPU) with TensorFlow-based open-source code; (ii) quantize and deploy the trained model to the NVNMD hardware and run MD through a LAMMPS-like interface substituting the force field with the QNN model. - Training: Two-stage: first a continuous neural network (CNN) is trained (e.g., ~1e5 epochs, higher learning rate); then quantized neural network (QNN) fine-tuning (~1e4 epochs, lower learning rate) minimizes quantization error. Training data come from public datasets (e.g., MD17 for molecules) or generated via DFT/AIMD. - Architecture: Heterogeneous system with a Master Processing Unit (MPU, CPU) and a Slave Processing Unit (SPU, FPGA) connected by a high-speed interface (PCIe 3.0 x16, DMA, full-duplex memory channels). The MPU handles neighbor-list building, data encoding/decoding, and time integration; the SPU performs PES evaluation and force/virial computation with PIM to minimize off-chip data movement. - PIM and pipeline: Model weights/activations reside in on-chip memories (BRAM/URAM), enabling pipelined forward/backward computations without frequent DRAM access. Modules M1–M6 implement forward energy evaluation and backprop for forces/virials; FIFOs carry intermediate results on-chip. - QNN and hardware-friendly NN: Weights and activations are quantized (e.g., y≈13 bits) to reduce memory and computation. Multiplications are replaced with bit-shifts and additions (multiplication-less neural network) by representing weights as sums of powers-of-two (with limited K terms, e.g., K=3) to reduce DSP usage. Trigonometric or costly activation functions (e.g., tanh) are replaced by a lightweight, piecewise-linear, shift-friendly activation φ closely approximating tanh in value and first derivative, facilitating both vN training and NvN inference. - Software integration: A modified LAMMPS interface exposes ensemble control (NVE/NVT/NPT), timestep, and thermostats; the force provider is the uploaded QNN. Atomic data are packed into compact formats for MPU↔SPU transfer. - Implementation: Prototype SPU on Xilinx UltraScale+ FPGA (e.g., xcvu9p) at ~250 MHz; MPU is an Intel i7-10700K CPU. PCIe 3.0 x16 with DMA supports high-throughput communication. Resource utilization includes LUTs, FFs, BRAM/URAM, and DSP blocks; NVNMD keeps intermediate data on-chip to avoid off-chip bottlenecks. - Test systems: Three molecules (benzene, naphthalene, aspirin) using MD17; three solids (Sb/Si, GeTe, Li10GeP2S12) from literature datasets. Accuracy assessed via energy RMSE and force errors; properties (RDF, ADF, coordination) validated; dynamic processes (GeTe melt–quench–anneal; Li-ion diffusion in LGPS) examined. Performance measured as time per step per atom and energy per step per atom, compared with MLMD and CMD.

- Accuracy (energies): For benzene, naphthalene, aspirin, Sb/Si, GeTe, and Li-Ge-P-S systems, NVNMD PES RMSEs are ~0.19, 0.39, 0.32, 0.14, 0.09, and 0.14 kcal mol^-1, all below the chemical accuracy threshold (1.0 kcal mol^-1). On identical solid datasets (meV atom^-1): NVNMD (QNN) vs MLMD (DeepMD) RMSEs are Sb 6.2 vs 2.2; GeTe 3.7 vs 4.1; Li-Ge-P-S 6.1 vs 1.3, indicating small quantization-induced differences while retaining high accuracy. - Accuracy (forces and properties): Force RMSEs closely match MLMD (e.g., water: |ΔF| NVNMD 21.4 meV Å^-1 vs MLMD 20.4 meV Å^-1). Structural and dynamical properties (RDF/ADF/coordination in amorphous GeTe) computed by NVNMD closely agree with MLMD/AIMD. GeTe phase transformations (crystal→liquid→amorphous→recrystallization) are reproduced in NVT. For Li10GeP2S12 (900 atoms, NVE 100 ps), the Li diffusion coefficient is 2.03×10^-10 m^2 s^-1, matching literature (≈2.00×10^-10 m^2 s^-1), and anisotropic diffusion (z > x,y) is observed. - Time efficiency: NVNMD achieves ≈2.0×10^-7 s step^-1 atom^-1 across tested systems on a single FPGA, roughly two orders of magnitude faster than representative MLMD GPU implementations (typically 10^-5–10^-6 s step^-1 atom^-1 per atom) and comparable to CMD benchmarks, thereby reaching CMD-like speed with MLMD/AIMD-level accuracy. - Energy efficiency: With total system power ~108 W, the per-step-per-atom energy is ≈2.1×10^-5 J, yielding 2–3 orders of magnitude better energy efficiency than reported MLMD CPU+GPU configurations at similar accuracy. Gains stem from avoiding repeated off-chip data movement via PIM and quantized, multiplication-less inference on FPGA. - Generality and availability: The approach is demonstrated on diverse molecular and bulk systems and is deployed on an accessible in-house FPGA server (links provided) with open-source training code.

The prototype demonstrates that a non-von Neumann, PIM-based special-purpose MD system can deliver AIMD/MLMD-level accuracy with CMD-like speed and markedly improved energy efficiency. Despite being implemented on a relatively low-end FPGA (e.g., Xilinx UltraScale+ at ~250 MHz) with limited on-chip resources, NVNMD achieves strong performance. The authors argue that transitioning to an ASIC could provide further substantial gains: higher clock rates (GHz range) may yield ~10× speed-up; vastly greater transistor budgets could enable 100–1000× resource scaling, improving parallelism and throughput. Compared to prior special-purpose MD machines like Anton, which target classical force fields for biomolecular simulations, NVNMD focuses on ML-based interatomic potentials to maintain quantum-level accuracy while leveraging specialized hardware for efficiency. The results indicate that removing the vN memory bottleneck via PIM and co-designing algorithms (QNN, multiplication-less ops, custom activations) with hardware can overcome the traditional accuracy-efficiency trade-off in MD.

The work introduces NVNMD, a co-designed ML algorithm and non-von Neumann hardware system for molecular dynamics that achieves AIMD-level accuracy and CMD-level efficiency. By quantizing and restructuring DeepMD-like models for PIM execution (multiplication-less arithmetic and hardware-friendly activations) on FPGA, NVNMD reduces data movement and energy use while maintaining high fidelity in energies, forces, structural properties, and dynamic behaviors across molecules and solids. The system attains ~2×10^-7 s step^-1 atom^-1 and 2–3 orders of magnitude better energy efficiency than typical MLMD GPU implementations. The authors provide open-source training tools and a public inference server. Future work includes migrating to ASIC for higher clock rates and larger on-chip resources, scaling to larger models/systems, expanding materials coverage, and refining quantization/training to further narrow residual accuracy gaps.

- Prototype hardware: Implemented on a lower-frequency, resource-limited FPGA (~250 MHz), which constrains throughput and model size relative to potential ASIC implementations. - Quantization effects: Although small, QNN introduces accuracy differences compared to CNN/DeepMD (e.g., meV atom^-1 deviations), which may be system-dependent. - Dataset heterogeneity: Some cross-method accuracy comparisons (e.g., Table 1) draw from literature datasets rather than identical training/testing sets, limiting strict head-to-head comparisons; refined analyses are provided for select systems only. - Training dependency: Model training still relies on vN CPU/GPU resources and high-quality DFT/AIMD data; generalization depends on dataset quality and coverage. - Specialized infrastructure: Requires access to the NVNMD hardware/service; portability to other platforms may need reimplementation or hardware support for PIM/quantization schemes. - Communication overheads: Although minimized via PCIe, DMA, and compact data formats, MPU–SPU communication still exists and could impact scaling in certain regimes.