Molecular dynamics reasoning acceleration method based on time coarsening generative model and multi-scale machine learning potential and engine system using method

Abstract

The invention discloses a molecular dynamics reasoning acceleration method based on a time coarsening generation model and a multi-scale machine learning potential and an engine system using the method, and the method comprises the following steps: S1, data preparation: obtaining an initial state of a to-be-simulated system; s2, training and reasoning a time coarsening model; s3, performing multi-time-step integration and multi-scale machine learning potential; and S4 and S5, carrying out confidence evaluation and active sampling. And step S6, hardware collaborative scheduling. According to the method, the effective time scale of molecular dynamics simulation can be greatly expanded, and rapid reasoning from microsecond to millisecond and even longer time scale is realized. While the quantum chemical precision is maintained, the method reduces the demand for a large number of continuous small time steps, and achieves acceleration of more than two orders of magnitude. An active sampling and online correction mechanism ensures the robustness of the model in cross-system and cross-time scale applications. Besides, according to the hardware scheduling scheme designed by the invention, the engine can be deployed on a GPU or a customized ASIC, hardware resources are efficiently utilized, and the method is suitable for scientific research and industrial environments.

Description

Technical Field

[0001] The invention relates to the fields of molecular dynamics simulation, deep learning and hardware acceleration technology, specifically to a method and system for accelerating molecular dynamics inference by combining a time-coarsened generative model with a multi-scale machine learning potential, which can be used to efficiently infer the dynamic behavior of complex molecular systems over long time scales. Background Technology

[0002] Molecular dynamics (MD) simulations are crucial tools for understanding the behavior of complex molecular systems. However, classical MD is limited by time-step stability, typically requiring femtosecond time steps, making simulations spanning milliseconds or even seconds extremely time-consuming. Many critical dynamic processes, such as protein folding, crystal defect migration, material phase transitions, and chemical reactions, involve sparse transitions spanning long periods, which traditional MD struggles to capture within a reasonable computational timeframe. Furthermore, traditional multiscale modeling methods rely on empirical or analytical potentials, often failing to simultaneously balance accuracy and efficiency, and struggling to handle couplings between interactions across different timescales.

[0003] With the development of deep learning, machine learning potential models have been introduced to replace traditional potential functions, achieving force predictions with near-quantum chemical accuracy. However, existing machine learning potentials are still based on integration over fixed time steps, requiring numerous consecutive calls and failing to overcome the time step limitation. Meanwhile, generative models (such as variational autoencoders and normalized flows) demonstrate the ability to capture distributional characteristics in simulated data, but their application in simulating long-scale dynamic transitions remains exploratory. Therefore, a new method is urgently needed to combine generative models with machine learning potentials, learning dynamic transitions through time coarsening, significantly expanding the simulation timescale and improving inference efficiency without sacrificing accuracy. Summary of the Invention

[0004] The purpose of this invention is to propose a method and system for accelerating molecular dynamics reasoning that can solve the above-mentioned problems.

[0005] To address the above problems, this invention provides a method for accelerating molecular dynamics inference based on a time-coarsened generative model and multi-scale machine learning potential, characterized by the following steps: Step S1: Obtain the initial molecular state of the system to be simulated and construct a molecular diagram containing information on atom type, coordinates, and optional velocities; Step S2: Use the pre-trained time coarsening generative model to learn the historical molecular dynamics trajectory, infer the transitions of the system on time scales much larger than the basic time step, and obtain the approximate state of the next macro time step. Step S3: Construct two collaborative machine learning potential models based on a multi-time-step integration strategy. The small-scale potential is responsible for quickly calculating local interactions and is called in most time steps, while the large-scale potential is responsible for slower-changing long-range interactions and is called in interval time steps. Step S4: At each macro time step, first use the time coarsening generation model to generate candidate states, then use the small-scale potential to calculate local interactions and evaluate the confidence level, and determine whether to call the large-scale potential for correction based on the confidence level. Step S5: When the confidence level is low or the prediction error exceeds the threshold, active sampling and online correction are triggered. The simulation results are corrected using a high-precision model or quantum chemical calculations, and the corrected data is fed back to update the generation model and machine learning potential. Step S6: The generated model, small-scale potential, and large-scale potential are mapped to different execution units of GPU or ASIC through a hardware scheduler. Pipeline scheduling and quantization optimization techniques are used to achieve efficient parallelism of the inference process.

[0006] As a further improvement of the present invention, the time coarsening generation model is a generation model based on normalized flow, which uses a reversible network to learn the probability distribution of historical molecular dynamics trajectories, generates a state sequence spanning hundreds or thousands of time steps through time warp mapping, and has the ability to migrate across different chemical systems.

[0007] As a further improvement of the present invention, the small-scale machine learning potential uses a quantized model running at INT8 or FP16 precision to quickly estimate local atomic interactions in most time steps, while the large-scale machine learning potential uses a high-precision model running at FP32 or higher precision to calculate long-range interactions over longer time intervals and correct for the cumulative error of the small-scale potential.

[0008] As a further improvement of the present invention, it also includes an active sampling and online correction step based on confidence level, wherein the confidence level is jointly evaluated by the generating model and the small-scale potential. When the prediction uncertainty exceeds a preset threshold, a high-precision model or a small amount of quantum chemical calculations are automatically selected for correction to ensure the accuracy of long-term simulation.

[0009] As a further improvement of the present invention, the multi-time-step integration strategy adaptively adjusts the calling frequency of small-scale potential and large-scale potential according to the local rate of change of the system; in the stage of drastic system change, the calling frequency of large-scale potential is increased, and in the stage of gradual system change, the calling frequency of large-scale potential is reduced to save computing resources.

[0010] As a further improvement of the present invention, the hardware scheduler divides the operation into operators such as matrix multiplication, convolution, and activation based on the computation graph of the generative model and the machine learning potential. Through operator fusion, dynamic quantization and tensor parallelism, the throughput of GPU or ASIC is improved and the latency is reduced.

[0011] As a further improvement of the present invention, the method also includes joint training and optimization of the generative model, machine learning potential and hardware scheduler, so that the inference engine remains efficient and stable under different systems and time scales by continuously iterating and updating the model parameters and scheduling strategy.

[0012] Another subject of the invention is a molecular dynamics inference acceleration engine system for implementing the method of any one of claims 1–7, comprising: The data acquisition module is used to read the initial structural data of the molecular system and generate molecular diagrams; The time coarsening generation module is used to learn historical trajectories and predict system states at macro time steps; The small-scale potential module is used to quickly compute local interactions at fine time steps; Large-scale potential module, used to periodically calculate long-range interactions and correct system states; The confidence assessment module is used to monitor the prediction uncertainty of the generative model and small-scale potential, and to decide whether to initiate active sampling. The active sampling and online correction module is used to schedule high-precision models or quantum computing for correction when the confidence level is insufficient, and to feed back the correction results; The hardware scheduler module is used to allocate the computing tasks of the above modules to GPU or ASIC execution units, and achieve efficient parallelism through pipeline scheduling and quantization optimization.

[0013] As a further improvement of the present invention, the confidence evaluation module uses a Bayesian neural network or ensemble learning method to quantify the prediction variance of the time coarsening generation model and the small-scale potential in the current state, and combines historical error statistics to determine whether it is necessary to call the large-scale potential or high-precision model.

[0014] As a further improvement of the present invention, the hardware scheduler module includes execution queues for different types of operators. By analyzing the operator dependencies of the generative model and the machine learning potential, it optimizes the execution order using deep pipeline and graph scheduling algorithms, reduces memory transfer and idle time, and improves overall utilization.

[0015] The beneficial effects of this invention are that it significantly expands the effective timescale of molecular dynamics simulations, enabling rapid inference from microseconds to milliseconds and even longer timescales. While maintaining quantum chemical-level accuracy, this method reduces the need for a large number of consecutive small time steps, achieving an acceleration of approximately two orders of magnitude. Active sampling and online calibration mechanisms ensure the robustness of the model in cross-system and cross-timescale applications. Furthermore, the hardware scheduling scheme designed in this invention allows the engine to be deployed on GPUs or custom ASICs, efficiently utilizing hardware resources and making it suitable for research and industrial environments. Detailed Implementation

[0016] The technical solution of the present invention will be further described below with reference to the accompanying drawings and specific embodiments.

[0017] This embodiment includes the following steps: Step S1: Data Preparation. Obtain the initial state of the system to be simulated, record the atomic composition, three-dimensional coordinates, and optional initial velocity information, and construct a molecular diagram. Nodes in the diagram represent atoms, and edges represent chemical bonds or spatial proximity relationships.

[0018] Step S2: Temporal Coarsening Model Training and Inference. Using existing high-precision MD trajectories as the training set, the normalized flow model is trained to learn the probability flow and time-warped mapping of the system state evolution over time. After training, during the inference phase, starting from the current system state, a generative model is used to directly generate candidate states for macro-time steps spanning multiple basic time steps. This process significantly reduces the number of time steps requiring explicit integration.

[0019] Step S3: Multi-time-step integration and multi-scale machine learning potential. Within each macro-time step, a multi-time-step integration strategy is used to invoke the machine learning potential: small-scale potentials are run at fine-time steps to quickly estimate local interactions; large-scale potentials are invoked at predetermined or adaptive time intervals to calculate long-range or slowly varying interactions and correct for accumulated errors in the small-scale potentials. Small-scale potentials use quantized models to run with low precision, suitable for deployment on mobile or embedded hardware; large-scale potentials use higher-precision models to ensure accuracy.

[0020] Steps S4 and S5: Confidence Assessment and Active Sampling. The prediction confidence of the generative model and small-scale potential is monitored in real time by the confidence assessment module. For example, a Bayesian neural network or ensemble learning can be used to assess the prediction variance. If the confidence is low or deviates significantly from the training data distribution, active sampling and online correction are triggered: high-precision model or quantum chemical calculations are used to obtain real force or energy data to correct the current state, and the corrected samples are added to the training set to update the generative model and machine learning potential.

[0021] Step S6: Hardware Co-scheduling. At the implementation level, the computation of the generative model, small-scale potential, and large-scale potential in the engine are mapped to different execution units of the GPU or ASIC, respectively. The hardware scheduler analyzes operator dependencies and performs operator fusion and pipeline scheduling for operations such as matrix multiplication, convolution, and activation. By combining techniques such as dynamic quantization and tensor parallelism, memory access and computation order are optimized to maximize hardware throughput and reduce inference latency.

[0022] Through the above steps, this invention enables rapid inference of the dynamics of complex molecular systems while ensuring physical accuracy, making it particularly suitable for scenarios requiring long-scale simulations. Active sampling and online calibration mechanisms ensure the robustness of the model under different chemical systems, temperatures, and pressures; hardware co-scheduling allows the engine to run efficiently on GPUs, FPGAs, or dedicated ASICs, providing a scalable and portable dynamics simulation solution for scientific research and industrial applications.

[0023] 1. System initialization and molecular map construction First, the system reads the initial structural data (such as a PDB file) of the system to be simulated through the data acquisition module, extracts atom type, three-dimensional coordinates, and optional velocity information, and constructs a molecular graph. In the graph (not shown), nodes represent atoms, and edges represent chemical bonds or spatial proximity relationships (generated based on the cutoff radius). To be compatible with systems of different scales, a gradient-environment subgraph strategy is adopted: for systems with more than 1000 atoms, a gradient subgraph (not shown, used for loss calculation) and an environment subgraph (not shown, used as model input) are divided according to a radius threshold, with random atoms as the center, to balance computational efficiency and physical consistency.

[0024] 2. Training and Inference of Temporally Coarsened Generative Models The temporal coarsening generation module employs a generative model based on a normalized flow or stochastic interpolation framework (such as UniSim), and its training and inference process includes: Pre-training phase: Multi-head pre-training is performed using large-scale molecular datasets (such as PCQM4Mv2, ANI1x, PepMD, etc.) to learn a unified atomic representation. Atoms in the periodic table are expanded into multiple discrete patterns through atomic embedding extension techniques to form an extended vocabulary to capture features of different chemical environments.

[0025] Inference Phase: Starting from the current molecular state, the vector field of state transition is learned using a pre-trained model to generate candidate states spanning hundreds or thousands of time steps. The generation process is discretized into a multi-step Markov process through stochastic differential equations (SDEs) to achieve autoregressive trajectory generation.

[0026] 3. Co-calculation of multi-time-step integration and multi-scale potential In the multi-timestep integration strategy, the system constructs two collaborative machine learning potential models: Small-scale potential module: Employs a quantization model (INT8 / FP16 precision) to rapidly calculate local atomic interactions such as bond lengths and bond angles at fine time steps. This module is called frequently (e.g., every 1–10 femtoseconds) and is suitable for deployment on dedicated computing units of GPUs or ASICs.

[0027] Large-scale potential module: Employs a high-precision model (FP32 or higher) called at interval time steps (e.g., every 100 femtoseconds) to calculate long-range interactions such as electrostatics and van der Waals, and corrects for accumulated errors in small-scale potentials. The potential function uses an equivariant graph neural network (e.g., TorchMD-NET) to ensure rotation and translation invariance.

[0028] 4. Confidence assessment and active sampling mechanism The confidence assessment module monitors the prediction uncertainty of the generative model and small-scale potential in real time. Evaluation method: Use Bayesian neural networks or ensemble learning methods to quantify the prediction variance, and set a threshold by combining historical error statistics.

[0029] Triggering conditions: When the confidence level falls below a threshold or the prediction error exceeds the limit, the active sampling and online correction module schedules a high-precision model (such as quantum chemical calculations) to correct the current state. The corrected data is fed back to the training set to update the generative model and machine learning potential parameters online, enhancing generalization ability.

[0030] 5. Hardware Co-scheduling and Optimization The hardware scheduler module maps the generated model, small-scale potential, and large-scale potential to different execution units of the GPU or ASIC: Operator-level optimization: By analyzing the operator dependencies in the computation graph, operator fusion and pipeline scheduling are performed on operations such as matrix multiplication, convolution, and activation to reduce memory transfer overhead.

[0031] Quantization and parallel techniques: Dynamic quantization (such as INT8 precision) is used for small-scale potentials, combined with tensor parallelism to improve throughput. The scheduler contains multiple execution queues and uses graph scheduling algorithms to optimize task ordering and maximize hardware utilization.

[0032] 6. System workflow integration The entire engine system achieves end-to-end acceleration through modular collaboration: Input the initial state of the molecule to generate a molecular diagram.

[0033] The time coarsening generation module predicts candidate states for macro time steps.

[0034] The multi-scale potential modules integrate collaboratively within the macro time step, and the confidence level assessment determines whether to trigger correction.

[0035] The hardware scheduler dynamically allocates computing resources to ensure low-latency inference. This method improves the simulation timescale by tens of times while maintaining physical accuracy (e.g., free energy landscape error <0.5 kcal / mol), and is applicable to a variety of systems from small molecules to protein complexes.

[0036] The technical principles of the present invention have been described above with reference to specific embodiments. These descriptions are merely for explaining the principles of the invention and should not be construed as limiting the scope of protection of the invention in any way. Based on this explanation, those skilled in the art can readily conceive of other specific embodiments of the invention without inventive effort, and these embodiments will all fall within the scope of protection of the present invention.

Claims

1. A method for accelerating molecular dynamics inference based on a time-coarsened generative model and multi-scale machine learning potential, characterized in that, Includes the following steps: Step S1: Obtain the initial molecular state of the system to be simulated and construct a molecular diagram containing information on atom type, coordinates, and optional velocities; Step S2: Use the pre-trained time coarsening generative model to learn the historical molecular dynamics trajectory, infer the transitions of the system on time scales much larger than the basic time step, and obtain the approximate state of the next macro time step. Step S3: Construct two collaborative machine learning potential models based on a multi-time-step integration strategy. The small-scale potential is responsible for quickly calculating local interactions and is called in most time steps, while the large-scale potential is responsible for slower-changing long-range interactions and is called in interval time steps. Step S4: At each macro time step, first use the time coarsening generation model to generate candidate states, then use the small-scale potential to calculate local interactions and evaluate the confidence level, and determine whether to call the large-scale potential for correction based on the confidence level. Step S5: When the confidence level is low or the prediction error exceeds the threshold, active sampling and online correction are triggered. The simulation results are corrected using a high-precision model or quantum chemical calculations, and the corrected data is fed back to update the generation model and machine learning potential. Step S6: The generated model, small-scale potential, and large-scale potential are mapped to different execution units of GPU or ASIC through a hardware scheduler. Pipeline scheduling and quantization optimization techniques are used to achieve efficient parallelism of the inference process.

2. The method for accelerating molecular dynamics inference based on a time-coarsened generative model and multi-scale machine learning potential as described in claim 1, characterized in that, The time-coarsening generation model is a normalized flow-based generation model that uses a reversible network to learn the probability distribution of historical molecular dynamics trajectories. It generates state sequences spanning hundreds or thousands of time steps through time-warped mapping and has the ability to migrate across different chemical systems.

3. The method for accelerating molecular dynamics inference based on a time-coarsened generative model and multi-scale machine learning potential as described in claim 1 or 2, characterized in that, The small-scale machine learning potential employs a quantized model running at INT8 or FP16 precision to rapidly estimate local atomic interactions at most time steps, while the large-scale machine learning potential employs a high-precision model running at FP32 or higher precision to calculate long-range interactions over longer time intervals and correct for the accumulated errors of the small-scale potential.

4. The method for accelerating molecular dynamics inference based on a time-coarsened generative model and multi-scale machine learning potential as described in claim 3, characterized in that, It also includes a confidence-based active sampling and online correction step, where the confidence level is jointly evaluated by the generative model and the small-scale potential. When the prediction uncertainty exceeds a preset threshold, a high-precision model or a small amount of quantum chemical calculations are automatically selected for correction to ensure the accuracy of long-term simulation.

5. The method for accelerating molecular dynamics inference based on a time-coarsened generative model and multi-scale machine learning potential according to claim 4, characterized in that, The multi-time-step integration strategy adaptively adjusts the calling frequency of small-scale and large-scale potentials based on the local rate of change of the system. During periods of rapid system change, the calling frequency of large-scale potentials is increased, while during periods of gradual system change, the calling frequency of large-scale potentials is decreased to save computational resources.

6. The method for accelerating molecular dynamics inference based on a time-coarsened generative model and multi-scale machine learning potential according to claim 5, characterized in that, The hardware scheduler divides operations into operators such as matrix multiplication, convolution, and activation based on the computation graph of the generative model and the machine learning potential. It optimizes the process by using techniques such as operator fusion, dynamic quantization, and tensor parallelism to improve the throughput of GPUs or ASICs and reduce latency.

7. The method for accelerating molecular dynamics inference based on a time-coarsened generative model and multi-scale machine learning potential as described in claim 6, characterized in that, The method also includes joint training and optimization of the generative model, machine learning potential, and hardware scheduler. By continuously iterating and updating model parameters and scheduling strategies, the inference engine can maintain high efficiency and stability under different systems and time scales.

8. A molecular dynamics inference acceleration engine system for implementing the method of any one of claims 1–7, characterized in that, include: The data acquisition module is used to read the initial structural data of the molecular system and generate molecular diagrams; The time coarsening generation module is used to learn historical trajectories and predict system states at macro time steps; The small-scale potential module is used to quickly compute local interactions at fine time steps; Large-scale potential module, used to periodically calculate long-range interactions and correct system states; The confidence assessment module is used to monitor the prediction uncertainty of the generative model and small-scale potential, and to decide whether to initiate active sampling. The active sampling and online correction module is used to schedule high-precision models or quantum computing for correction when the confidence level is insufficient, and to feed back the correction results; The hardware scheduler module is used to allocate the computing tasks of the above modules to GPU or ASIC execution units, and achieve efficient parallelism through pipeline scheduling and quantization optimization.

9. A molecular dynamics inference acceleration engine system according to claim 8, characterized in that, The confidence assessment module uses a Bayesian neural network or ensemble learning method to quantify the prediction variance of the time coarsening generative model and the small-scale potential in the current state, and combines historical error statistics to determine whether it is necessary to call the large-scale potential or high-precision model.

10. A molecular dynamics inference acceleration engine system according to claim 8 or 9, characterized in that, The hardware scheduler module includes execution queues for different types of operators. By analyzing the operator dependencies of the generative model and machine learning potential, it optimizes the execution order using deep pipeline and graph scheduling algorithms, reducing memory transfer and idle time, and improving overall utilization.