A method for solving molecular system steady state based on deep equilibrium model

Abstract

The application discloses a kind of solving molecular system steady-state method based on depth balanced model, comprising: constructing data set, the data in the data set is the paired data of initial conformation-steady-state conformation;The data in data set is divided into training sample and test sample;The training sample is first time sampled, and then sampled according to the path probability of Brownian bridge, to obtain the preprocessed training sample;SE (3) isomorphic neural network model based on graph Transform or spherical harmonic function representation is constructed;SE (3) isomorphic neural network model is trained using the preprocessed training sample, until SE (3) isomorphic neural network model converges, to obtain the trained SE (3) isomorphic graph neural network model;First, sample from initial conformation distribution ∏0, generate a rough conformation;Then sample from the following SDE from joint distribution Π 0,1 X1 of t=1 is the required steady-state conformation.

Description

Technical Field

[0001] This invention belongs to the field of artificial intelligence technology, specifically a method for solving the steady state of molecular systems based on a deep equilibrium model. Background Technology

[0002] Capturing the evolution of a system's geometric state is crucial for multiple scientific fields. Typically, a geometric state describes the position of each object in the system at a given time. As the system evolves, certain geometric states of interest, such as equilibrium states, can be determined. Predicting these specific states from the initial state is essential for revealing the system's behavior and function, thus providing valuable insights into important problems in practice, such as drug discovery, reaction modeling, and catalyst analysis.

[0003] While crucial, accurately predicting the future geometric state of interest from the initial state of a complex system is no easy task. Experimental methods, though valuable for direct observation, often face obstacles under stringent environmental conditions and instrumental limitations. Computational methods attempt to simulate dynamic evolution by solving Newton's equations of motion and obtain equilibrium states through iterative geometric optimization. While offering greater flexibility, such calculations are typically driven by first-principles methods or empirical laws, either requiring significant computational cost or sacrificing accuracy.

[0004] In recent years, deep learning has become a key tool for scientific discovery in many fields, offering new avenues for capturing the challenges of geometric state evolution. One direct approach is to train neural networks to predict the target geometric state directly from the initial state. Equivariant networks are typically employed, with end-to-end supervised training on paired data of the initial and target states, by appropriately encoding the inherent symmetries of the geometric state. Once the model is trained, its single forward propagation can significantly improve prediction efficiency compared to traditional methods. However, this paradigm also forces neural networks to compress iterative evolution into a single step, making the learning process difficult and potentially leading to reduced accuracy. Another research approach involves predicting forces based on the system's current geometric state using machine learning force fields (MLFFs). Trained MLFFs can simulate evolution and facilitate geometric optimization to determine the target geometric state, showing a better efficiency-accuracy balance than first-principles methods. However, MLFFs heavily rely on the availability and quality of intermediate labels (such as potential energy) to ensure stability and accuracy, limiting their application in complex scenarios. Summary of the Invention

[0005] Purpose of the Invention: To address the difficulties in the learning process and low accuracy of deep learning in capturing the evolution of geometric states, and to solve the problems of heavy dependence on the availability and quality of intermediate labels (such as potential energy) and limited application scenarios when predicting forces based on the current geometric state of a system using machine learning force fields (MLFFs), this invention proposes a method for solving the steady state of molecular systems based on a deep equilibrium model. It utilizes a modified bridge-matching framework to generate molecular steady-state conformations that satisfy specific physical symmetry constraints, named the GDB (Geometric Diffusion Bridge) framework, for accurate and efficient solving of molecular steady-state conformations and trajectory generation.

[0006] Technical solution: A method for solving the steady state of molecular systems based on a deep equilibrium model, comprising the following steps:

[0007] Construct a dataset containing paired data of initial and steady-state conformations; divide the dataset into training samples and test samples;

[0008] The training samples are first sampled over time t ~ U(0,1), where U(0,1) represents a uniform distribution in the interval (0,1). Then, the path probabilities X of the Brownian bridge are used. t Sampling is performed on ~N(X0+(X1-X0)t,σt(1-t)) to obtain preprocessed training samples; where X t ∈R d σ represents the conformation at time t in the dynamic evolution, σ represents the noise scale, and X0 and X1 are random variables representing the initial conformation and the steady-state conformation, respectively.

[0009] Construct an SE(3) equivariant neural network model based on graph Transformer or spherical harmonic function representation;

[0010] The SE(3) equivariant neural network model was trained using the preprocessed training samples until the SE(3) equivariant neural network model converged, and the trained SE(3) equivariant graph neural network model was obtained.

[0011] First, sample from the initial conformational distribution ∏0 to generate a rough conformation; then, sample from the joint distribution Π0,1 by simulating the following SDE to obtain X1 at t=1, which is the desired steady-state conformation:

[0012] dX t =v θ (X t ,X0,t)dt+σdB t ,X0~Π0

[0013] In the formula, v θ (X tX0,t) represents a time-dependent vector field of a trained SE(3) isomorphic graph neural network model, Π0 represents the initial conformational distribution, and X0,t represents the initial conformational distribution. t ∈R d B represents the conformation at time t in the dynamic evolution. t This is a standard Brownian motion.

[0014] Furthermore, the training process of the SE(3) isovariant neural network model is equivalent to minimizing the following objective function L. GDB (θ):

[0015]

[0016] Among them, P 1|t Indicates a given X t The conditional distribution of X1, Indicates X t Find the gradient.

[0017] Furthermore, the Brownian bridge is represented as:

[0018]

[0019] This invention also discloses a method for solving the steady state of a molecular system based on a deep equilibrium model, comprising the following steps:

[0020] Construct a dataset, wherein the data in the dataset consists of trajectory data composed of several frames, and the trajectory data includes N frames (X0, X...). 1 / N ,...,X1), X i ∈R 3n ;

[0021] Construct an SE(3) equivariant neural network model based on graph Transformer or spherical harmonic function representation;

[0022] The SE(3) equivariant neural network model is trained in each frame using the following loss function;

[0023]

[0024] After training, the following SDE is simulated frame-by-frame from the joint distribution Π 0,1 By sampling from the middle, the obtained X1 is the desired steady-state conformation:

[0025]

[0026] Where [Nt] represents taking the integer part.

[0027] Beneficial effects: Compared with the prior art, the present invention has the following advantages:

[0028] (1) This invention introduces a novel and systematic generative modeling method for bridging geometric states. From a probabilistic perspective, the evolution of geometric states actually reflects the transmission of their distributions. This invention establishes a GDB framework, which represents a predetermined diffusion process, with initial and target geometric states as fixed endpoints, and is controlled by a rotation-translation isovariant transformation kernel, ensuring that the evolution follows inherent symmetry. This is achieved through the Doobh-transformation, simulating the evolution of geometric states and providing a new approach for accurately and efficiently bridging geometric states;

[0029] (2) The GDB framework proposed in this invention improves performance by modeling the joint state distribution of geometric states, accurately connecting geometric states, and fully utilizing available trajectories as a fine description of dynamics. Mathematically, this invention has demonstrated that the joint distribution of geometric states across different time steps can be completely preserved through the (chain-like) isovariant diffusion bridge technique, confirming its expressive power in connecting geometric states and highlighting the importance of design choices within the proposed GDB framework. Furthermore, under mild and practical assumptions, this invention has demonstrated that, in terms of convergence, the proposed GDB framework can approximate the potential dynamics controlling the evolution of geometric state trajectories with negligible errors, illustrating the completeness and practicality of the proposed GDB framework in different scenarios. These advantages demonstrate the superiority of the proposed GDB over existing methods.

[0030] (3) To verify its effectiveness and generality, this invention conducted extensive experiments covering different data modalities (simple molecules and adsorbent-catalyst complexes), scales (small, medium, and large scales), and scenarios (with and without trajectory guidance). Numerical results show that the proposed GDB framework consistently outperforms state-of-the-art machine learning methods. In particular, the proposed method even surpasses a robust machine learning force field baseline with training data that is more than 10 times larger in the challenging structural relaxation task at OC22, while trajectory guidance further enhances the performance of the proposed method. The significantly superior performance demonstrates the powerful ability of the proposed GDB framework to capture the complex evolutionary dynamics of geometric states and to identify valuable and critical geometric states of interest in key real-world challenges. Attached Figure Description

[0031] Figure 1 This is a schematic diagram of the neural network architecture used in this invention. Detailed Implementation

[0032] The technical solution of the present invention will now be further described in conjunction with the accompanying drawings and embodiments.

[0033] This embodiment proposes a method for solving the steady state of molecular systems based on a deep equilibrium model. For a readily obtainable coarse molecular conformation, an isomorphic graph neural network is trained by constructing a bridge. The trained isomorphic graph neural network can iteratively repair this coarse molecular conformation, thereby generating a molecular steady-state conformation that satisfies specific physical symmetry constraints. This framework is named GDB (Geometric Diffusion Bridge). The generated molecular steady state can be used for efficient molecular steady-state solving and trajectory generation.

[0034] The bridge matching framework is constructed by creating a path probability Π = Π0Π1P (·|0,1) The associated stochastic differential equation (SDE), where P (·|0,1) Let Π represent a bridging process, where Π0 represents the initial coarse conformational distribution, also known as the initial distribution; and Π1 represents the steady-state conformational distribution to be sampled, also known as the terminal distribution. Note that in Π, the initial and terminal distributions are decoupled. Consider the coupled data (X0, X1) ~ Π. 0,1 Where X0 and X1 represent the initial and steady-state conformations, respectively, implying that the initial and terminal distributions are not independent. Therefore, the bridge matching framework should construct a path probability Π = Π 0,1 P (·|0,1) Related SDEs. Previous studies have found that bridge-matching frameworks cannot maintain the invariance of coupling distributions. That is, the joint distribution of (X0, X1) is usually not Π 0,1 To address this issue, the GDB framework in this embodiment constructs a path probability Π by fixing the initial point X0. (·|0) =Π (1|0) P (·|0,1) The associated stochastic differential equation, where Π (·|0) Π represents the conditional probability path distribution given an initial point. (1|0) This represents the conditional distribution of the steady-state conformation given an initial point.

[0035] Assume the path probability P is given by the following SDE:

[0036] dX t =f(X) t ,t)dt+σ(t)dB t

[0037] Among them, X t ∈R d B represents the conformation at time t in the dynamic evolution. t This is a standard Brownian motion.

[0038] According to the well-known Doob h-transform in probability theory, the path probability P given a starting point and an ending point can be obtained. (·|0,1) It is given by the following SDE:

[0039]

[0040] Given X0, X t function under conditions Conditional expectation, Indicates X t Find the gradient.

[0041] The path probability P can be designed by ourselves; in this embodiment, we use f(X) t The design, t) = 0, uses a neural network v θ (X t to approximate (X0,t) The training process of a neural network is equivalent to minimizing the following objective function:

[0042]

[0043] Among them, P 1|t Indicates a given X t The conditional distribution of X1.

[0044] The optimal value of a neural network satisfies:

[0045]

[0046] After the training process is complete, the following SDE can be used to simulate the joint distribution Π. 0,1 By sampling from the middle, the obtained X1 is the desired steady-state conformation:

[0047] dX t =v θ (X t ,X0,t)dt+σdB t ,X0~Π0

[0048] Neural networks are the v in SDE. θ (X t After training, X0, t), we first need to sample from the initial distribution Π0, i.e., generate a rough conformation. Then we simulate this SDE, and finally, the X1 obtained at time t=1 is a steady-state conformation. Therefore, the sampling (or simulating the SDE) process is the process of generating a steady-state conformation.

[0049] Furthermore, the symmetry of the system should also be considered in neural networks. Clearly, the likelihood of a molecule should not change under rotation or translation, so its distribution is SE(3) invariant. To satisfy this property, the vector field should satisfy some constraints. In fact, if the prior distribution Π0 is SE(3) invariant, and the transition kernel Π t|0 (X t If |X0) is equivariant to SE(3), then the marginal distribution Π at any time... t It is also unchanged for SE(3).

[0050] If the dataset used to train the neural network model contains trajectory information, then you can choose to use or not use this information during training. Using trajectory information will generally result in better training performance. However, because trajectory information is more expensive, many datasets do not contain it.

[0051] For methods of generating molecular steady-state conformations that do not utilize trajectory information, the following steps are included:

[0052] Step 1: Construct a dataset containing pairs of coarse and fine conformations; divide the constructed dataset into training and testing samples.

[0053] Step 2: First, perform time sampling on the training samples, t ~ U(0,1), where U(0,1) represents a uniform distribution on the interval (0,1). Then, follow the path probability X of the Brown Bridge. t Sampling is performed on ~N(X0+(X1-X0)t,σt(1-t)), where X t ∈R d This represents the conformation at time t in the dynamic evolution;

[0054] Step 3: Construct a neural network based on graph Transformer or spherical harmonic function representation, such as... Figure 1 As shown, the neural network in this embodiment is v in SDE. θ (X t ,X0,t) terms.

[0055] Step 4: Use the backpropagation algorithm and the AdamW optimization algorithm to train the neural network.

[0056] The Brownian bridge used in this embodiment is represented as follows:

[0057]

[0058] The loss function is expressed as:

[0059]

[0060] Among them, v θ (X tLet P(X0,t) be a time-dependent vector field parameterized by a neural network θ. 1,t X represents t The joint distribution with X0.

[0061] Step 5: Using the Rdkit tool, generate a rough three-dimensional conformation for a molecule's 2D structure or chemical formula. Then, simulate the following SDE to generate the steady-state conformation; the resulting X1 is the desired steady-state conformation:

[0062] dX t =v θ (X t ,X0,t)dt+σdB t ,X0~Π0

[0063] Among them, v θ (X t ,X0,t) is a time-dependent vector field parameterized by θ of a trained neural network, where Π0 represents the initial conformation distribution.

[0064] The above process can be simplified to:

[0065] The training modes of Algorithm2 that do not utilize trajectory information are as follows:

[0066] repeat;

[0067] (z0,z1)~Π 0,1 (X0,X1);

[0068] t~U(0,1);

[0069] ε~N(0,I);

[0070]

[0071] Gradient descent optimization step

[0072]

[0073] A method for generating molecular steady-state conformations using trajectory information includes the following steps:

[0074] The dataset provides trajectory data consisting of several frames, including N frames (X0, X...). 1 / N ,…,X1),X i ∈R 3n Assume the true path measure is a Markov path probability. Therefore, the true path probability can be decomposed as:

[0075]

[0076] Construct an SE(3) equivariant neural network model based on graph Transformer or spherical harmonic function representation;

[0077] The SE(3) equivariant neural network model is trained in each frame using the following loss function;

[0078]

[0079] After training, the following SDE is simulated frame-by-frame from the joint distribution Π 0,1 By sampling from the middle, the obtained X1 is the desired steady-state conformation:

[0080]

[0081] Where [Nt] represents taking the integer part.

[0082] The above process can be simplified to:

[0083] Algorithm3 uses the following training mode with trajectory information:

[0084] repeat;

[0085] (z0,…,z N )~Π 0,…,1 (X0,…,X1);

[0086]

[0087] ε~N(0,I);

[0088]

[0089] Gradient descent optimization step

[0090]

[0091] The parameters required for the experiment are briefly provided below. For the Molecule3D and QM9 datasets, the bridge SDE is parameterized by extending the graph Transformer-based equivariant network to simultaneously encode the conditions of the time steps and the initial geometric state.

[0092] During training, AdamW was used as the optimizer, with hyperparameters ∈ set to 1e-8 and (β1,β2) set to (0.9,0.999). The gradient clipping norm was set to 5.0. The peak learning rate was set to 1e-4. The batch size was set to 512. The weight decay was set to 0.0. The model was trained for 500k steps, including a 30k-step warm-up phase. After the warm-up phase, the learning rate decayed linearly to zero. The noise scale σ was set to 0.5. During inference, 10 time steps were used, and an Euler solver was employed. All models were trained on 16 NVIDIA V100 GPUs. The training process iterated for 300 epochs with a batch size of 128 (128 images per batch). The initial learning rate was 0.1, decreasing to 0.01 from epochs 151 to 225, and then to 0.001 from epochs 226 to 300.

[0093] For the OC22 dataset, GemNet-OC was used to parameterize the bridge SDE, demonstrating the compatibility of the proposed framework with different base models. For training, AdamW was used as the optimizer, with hyperparameters ∈ 1e-8 and (β1,β2) set to (0.9,0.999). The gradient clipping norm was set to 10.0. The peak learning rate was set to 5e-4. The batch size was set to 64. The weight decay was set to 0.0. The model was trained for 200k steps. After a warm-up phase, the learning rate linearly decayed to zero. The noise scale σ was set to 0.5. The trajectory length was set to N = 10. During inference, 10 time steps were used, and an Euler solver was employed. All models were trained on eight NVIDIA A100 GPUs.

[0094] The following are experimental results on the Molecule3D, QM9, and OC22 datasets:

[0095]

[0096]

[0097] Model ADwT index (%) ↑ IS baseline 44.77 SpinConv iterative solution 54.53 GemNet-dT iterative solution 59.68 GemNet-OC Iterative Solution 60.69 GemNet-OC direct prediction 60.44 GDB does not utilize trajectory information 62.4 GDB utilizes trajectory information 63.0

Claims

1. A method for solving the steady-state problem of a molecular system based on a deep equilibrium model, characterized in that: Includes the following steps: Construct a dataset containing paired data of initial and steady-state conformations; divide the dataset into training samples and test samples; The training samples are first sampled over time t ~ U(0,1), where U(0,1) represents a uniform distribution in the interval (0,1). Then, the path probabilities X of the Brownian bridge are used. t Sampling is performed on ~N(X0+(X1-X0)t,σt(1-t)) to obtain preprocessed training samples; where X t ∈R d σ represents the conformation at time t in the dynamic evolution, σ represents the noise scale, and X0 and X1 are random variables representing the initial conformation and the steady-state conformation, respectively. Construct an SE(3) equivariant neural network model based on graph Transformer or spherical harmonic function representation; The SE(3) equivariant neural network model was trained using the preprocessed training samples until the SE(3) equivariant neural network model converged, and the trained SE(3) equivariant graph neural network model was obtained. First, a rough conformation is generated by sampling from the initial conformational distribution Π0; then, a coarse conformation is obtained from the joint distribution Π0 by simulating the following SDE. 0,1 By sampling from the middle, X1 obtained at t=1 is the desired steady-state conformation: dX t =v θ (X t ,X0,t)dt+σdB t ,X0~Π0 In the formula, v θ (X t X0,t) represents a time-dependent vector field of a trained SE(3) isomorphic graph neural network model, Π0 represents the initial conformational distribution, and X0,t represents the initial conformational distribution. t ∈R d B represents the conformation at time t in the dynamic evolution. t This is a standard Brownian motion.

2. The method for solving the steady state of a molecular system based on a deep equilibrium model according to claim 1, characterized in that: The training process of the SE(3) isovariant neural network model is equivalent to minimizing the following objective function L GDB (θ): Among them, P 1|t Indicates a given X t The conditional distribution of X1, Indicates X t Find the gradient.

3. The method for solving the steady state of a molecular system based on a deep equilibrium model according to claim 1, characterized in that: The Brownian bridge is represented as:

4. A method for solving the steady-state problem of a molecular system based on a deep equilibrium model, characterized in that: Includes the following steps: Construct a dataset, wherein the data in the dataset consists of trajectory data composed of several frames, and the trajectory data includes N frames (X0, X...). 1 / N ,...,X1), X i ∈R 3n ; Construct an SE(3) equivariant neural network model based on graph Transformer or spherical harmonic function representation; The SE(3) equivariant neural network model is trained in each frame using the following loss function; After training, the following SDE is simulated frame-by-frame from the joint distribution Π 0,1 By sampling from the middle, the obtained X1 is the desired steady-state conformation: Where [Nt] represents taking the integer part.

Images