Quantitative training method and device of isovariant diffusion model, equipment and storage medium

Abstract

The embodiment of the invention provides a quantitative training method and device of an isotropic diffusion model, equipment and a storage medium, and relates to the technical field of neural networks. The method comprises the steps of inputting a calibration sample into an initial isotropic diffusion model for data processing, obtaining initial output data of an attention module in the initial isotropic diffusion model, and generating a quantitative weight sequence corresponding to a floating point weight sequence based on a sample length and a time step length of the calibration sample; updating the attention module according to the quantized weight sequence to obtain an updated isovariant diffusion model, inputting the calibration sample into the updated isovariant diffusion model for data processing, obtaining quantized output data corresponding to the quantized attention module, and calculating a total loss value according to the quantized output data, the initial output data, the floating point weight sequence and the quantized weight sequence, and obtaining a trained quantitative isovariant diffusion model according to the total loss value. The quantization weight is generated based on the sample length and the time step length, and the data generation quality of the quantization model is improved on the premise of keeping the isotropy.

Description

Technical Field

[0001] The present application relates to the field of neural network technology, and in particular to a quantitative training method, device, equipment and storage medium for an equivariant diffusion model. Background Art

[0002] As a commonly used generative model, the diffusion model (DM) is widely used in many fields such as image generation, video generation, and protein structure generation in the scientific field. However, the generation process of the diffusion model requires multiple rounds of iterations at different time steps, which is time-consuming, computationally intensive, and computationally complex. In addition, compared with the common diffusion model in the field of computer vision, the equivariant diffusion model (EDM) commonly used in the scientific field also requires the generation of variable-length data, which further increases the amount of computation.

[0003] In order to solve the problem of large amount of computation in equivariant diffusion model, the related technologies proposed quantization methods to compress the model and reduce the computation time. However, the quantization algorithms in these technologies can only perceive the time step and cannot achieve targeted quantization for data of different lengths, thus affecting the quality of the generated data of equivariant diffusion model obtained after quantization training. Summary of the invention

[0004] The main purpose of the embodiments of the present application is to propose a quantized training method, device, equipment and storage medium for an equivariant diffusion model, so as to improve the quality of generated data of the quantized equivariant diffusion model.

[0005] To achieve the above object, a first aspect of an embodiment of the present application proposes a quantization training method for an equivariant diffusion model, comprising:

[0006] Acquire multiple calibration samples and an initial equivariant diffusion model, input the calibration samples into the initial equivariant diffusion model for data processing, and at least acquire initial output data of an attention module in the initial equivariant diffusion model, wherein parameter weights corresponding to the attention module constitute a floating-point weight sequence;

[0007] Based on the sample length and time step of the calibration sample, a quantized weight sequence corresponding to the floating-point weight sequence is generated, the attention module is updated according to the quantized weight sequence to obtain a quantized attention module, and the attention module is replaced by the quantized attention module to obtain an updated equivariant diffusion model;

[0008] Inputting the calibration sample into the updated equivariant diffusion model for data processing to at least obtain quantized output data corresponding to the quantized attention module;

[0009] A total loss value is calculated according to the quantized output data, the initial output data, the floating-point weight sequence and the quantized weight sequence, and a trained quantized equivariant diffusion model is obtained according to the total loss value.

[0010] In some embodiments, generating a quantized weight sequence corresponding to the floating-point weight sequence based on the sample length and training time step of the calibration sample includes:

[0011] Obtain each parameter weight in the floating-point weight sequence, and obtain the maximum value and the minimum value of the parameter in the parameter weight;

[0012] Calculate the quantization step size according to a preset number of quantization bits, the maximum value of the parameter and the minimum value of the parameter;

[0013] Determining a first intermediate parameter based on the training time step and the sample length;

[0014] The quantization weight corresponding to each of the parameter weights is calculated according to the first intermediate parameter, the quantization step size and the basic parameter, and the quantization weight sequence is obtained based on the quantization weight.

[0015] In some embodiments, the step of calculating, according to the first intermediate parameter, the quantization step size and the basic parameter, a quantization weight corresponding to each basic parameter includes:

[0016] Calculating a ratio of the basic parameter to the quantization step length, and rounding the ratio to obtain a second intermediate parameter;

[0017] The quantization weight is obtained according to the product of the first intermediate parameter, the second intermediate parameter and the quantization step size.

[0018] In some embodiments, the calculating of the total loss value according to the quantized output data, the initial output data, the floating-point weight sequence, and the quantized weight sequence, and obtaining the trained quantized equivariant diffusion model according to the total loss value includes:

[0019] Calculating a reconstruction loss value according to the quantized output data and the initial output data, measuring a quantized loss value according to the floating-point weight sequence and the quantized weight sequence, and calculating the total loss value according to the reconstruction loss value and the quantized loss value;

[0020] Adjusting the quantized attention module based on the total loss value until the trained quantized attention module is obtained;

[0021] A trained quantized equivariant diffusion model is obtained based on the quantized attention module.

[0022] In some embodiments, obtaining a plurality of calibration samples comprises:

[0023] Get multiple random samples;

[0024] Obtaining a sampling duration and a time step, and calculating the number of generation times according to the sampling duration and the time step;

[0025] A cyclic generation process consistent with the generation times is performed on each of the random samples to obtain a corresponding calibration sample. During the generation process, the random sample is input into the initial equivariant diffusion model for data generation to obtain an updated random sample.

[0026] In some embodiments, obtaining an initial equivariant diffusion model includes:

[0027] Obtain multiple initial protein structure data, time steps and total time for adding noise;

[0028] Calculate the number of noise additions according to the total noise addition duration and the time step;

[0029] Performing a cyclic noise adding process consistent with the noise adding times on the initial protein structure data to obtain a pre-trained protein sample;

[0030] An equivariant diffusion model is obtained, and the equivariant diffusion model is trained using a plurality of the pre-trained protein samples to obtain the initial equivariant diffusion model.

[0031] In some embodiments, the initial protein structure data includes a plurality of initial amino acid data, and the cyclic noise adding process includes:

[0032] In each noise adding process, obtaining noise data corresponding to each of the initial amino acid data corresponding to the initial protein structure data;

[0033] superimposing the noise data on the initial amino acid data one by one to obtain the updated initial protein structure data;

[0034] The initial protein structure data is updated at least once according to the number of times of adding noise.

[0035] In some embodiments, superimposing the noise data one by one on the initial amino acid data to obtain the updated initial protein structure data comprises:

[0036] Obtaining the three-dimensional coordinates of the amino acid, the first angle between the main carbon atom and the amino group, and the second angle between the main carbon atom and the hydroxyl group in the initial amino acid data;

[0037] Acquire three-dimensional coordinate noise data, first angle noise data, and second angle noise data corresponding to the noise data;

[0038] The three-dimensional coordinate noise data is superimposed on the amino acid three-dimensional coordinates, the first angle noise data is superimposed on the first angle, and the second angle noise data is superimposed on the second angle to obtain the updated initial protein structure data.

[0039] In some embodiments, the initial equivariant diffusion model includes at least: an embedding module, the attention module, a self-attention module and a feedforward network module which are cascaded in sequence, and the trained quantized equivariant diffusion model based on the quantized attention module includes:

[0040] Obtain a quantized embedding module corresponding to the embedding module, a quantized attention module corresponding to the attention module, a quantized self-attention module corresponding to the self-attention module, and a quantized feedforward network module corresponding to the feedforward network module;

[0041] The quantized equivariant diffusion model is obtained according to the quantized embedding module, the quantized attention module, the quantized self-attention module and the quantized feedforward network module.

[0042] In some embodiments, the step of inputting the calibration sample into the initial equivariant diffusion model for data processing to obtain at least initial output data of the attention module in the initial equivariant diffusion model includes:

[0043] Acquiring initial input data for the attention module based on the calibration sample;

[0044] Perform multiple linear transformations on the initial input data to obtain multiple intermediate vectors, and obtain three attention vectors based on the multiple intermediate vectors;

[0045] Performing weighted summation on the three attention vectors according to the attention weight of each of the attention vectors to obtain an attention output vector, wherein the floating-point weight sequence includes the attention weight;

[0046] The softmax function is performed on the attention output vector to obtain the initial output data.

[0047] To achieve the above-mentioned purpose, a second aspect of an embodiment of the present application proposes a quantization training device for an equivariant diffusion model, comprising:

[0048] Initial output acquisition module: used to obtain multiple calibration samples and an initial equivariant diffusion model, input the calibration samples into the initial equivariant diffusion model for data processing, and at least obtain the initial output data of the attention module in the initial equivariant diffusion model, and the parameter weights corresponding to the attention module constitute a floating-point weight sequence;

[0049] A weight quantization module: used to generate a quantized weight sequence corresponding to the floating-point weight sequence based on the sample length and time step of the calibration sample, update the attention module according to the quantized weight sequence to obtain a quantized attention module, and replace the attention module with the quantized attention module to obtain an updated equivariant diffusion model;

[0050] Quantized output acquisition module: used to input the calibration sample into the updated equivariant diffusion model for data processing, and at least obtain the quantized output data corresponding to the quantized attention module;

[0051] Quantization training module: used to calculate the total loss value according to the quantized output data, the initial output data, the floating-point weight sequence and the quantized weight sequence, and obtain the trained quantized equivariant diffusion model according to the total loss value.

[0052] To achieve the above objectives, a third aspect of an embodiment of the present application proposes an electronic device, which includes a memory and a processor, wherein the memory stores a computer program, and when the processor executes the computer program, the method described in the first aspect is implemented.

[0053] To achieve the above-mentioned purpose, the fourth aspect of an embodiment of the present application proposes a storage medium, which is a storage medium. The storage medium stores a computer program, and when the computer program is executed by a processor, the method described in the first aspect is implemented.

[0054] The quantitative training method, device, equipment and storage medium of the equivariant diffusion model proposed in the embodiment of the present application obtains multiple calibration samples and an initial equivariant diffusion model, inputs the calibration samples into the initial equivariant diffusion model for data processing, obtains at least the initial output data of the attention module in the initial equivariant diffusion model, generates a quantized weight sequence corresponding to the floating-point weight sequence based on the sample length and time step of the calibration samples, updates the attention module according to the quantized weight sequence to obtain the quantized attention module, replaces the attention module with the quantized attention module to obtain the updated equivariant diffusion model, inputs the calibration samples into the updated equivariant diffusion model for data processing, obtains at least the quantized output data corresponding to the quantized attention module, calculates the total loss value according to the quantized output data, the initial output data, the floating-point weight sequence and the quantized weight sequence, and obtains the trained quantized equivariant diffusion model according to the total loss value. When the embodiment of the present application performs quantization of the equivariant diffusion model, the quantized weights are generated based on the sample length and time step of the calibration samples, and then the reconstruction error and the quantization error are fused to obtain the total loss value to complete the quantization training process, which improves the effect of the quantization training as a whole, and improves the data generation quality of the equivariant diffusion model while maintaining the equivariance. BRIEF DESCRIPTION OF THE DRAWINGS

[0055] Figure 1 It is a flow chart of a quantitative training method for an equivariant diffusion model provided in an embodiment of the present application.

[0056] Figure 2 The model structure of the initial equivariant diffusion model provided in the embodiment of the present application.

[0057] Figure 3 It is a flow chart for obtaining an initial equivariant diffusion model provided in an embodiment of the present application.

[0058] Figure 4 It is a schematic diagram of generating pre-trained protein samples in the embodiment of the present application.

[0059] Figure 5 This is a flowchart of the cyclic noise addition process provided in an embodiment of the present application.

[0060] Figure 6 It is a flowchart of obtaining multiple calibration samples provided in an embodiment of the present application.

[0061] Figure 7 It is a schematic diagram of the attention module provided in an embodiment of the present application.

[0062] Figure 8 It is a flowchart of generating a quantization weight sequence corresponding to a floating-point weight sequence based on the sample length and time step of the calibration sample provided in an embodiment of the present application.

[0063] Fig. 9It is a flowchart of calculating the total loss value according to the quantized output data, the initial output data, the floating-point weight sequence and the quantized weight sequence, and obtaining the trained quantized equivariant diffusion model according to the total loss value, provided by an embodiment of the present application.

[0064] Fig.10 It is a structural block diagram of a quantization training device for an equivariant diffusion model provided in yet another embodiment of the present application.

[0065] Fig.11 It is a schematic diagram of the hardware structure of the electronic device provided in the embodiment of the present application. DETAILED DESCRIPTION

[0066] In order to make the purpose, technical solution and advantages of the present application more clearly understood, the present application is further described in detail below in conjunction with the accompanying drawings and embodiments. It should be understood that the specific embodiments described herein are only used to explain the present application and are not used to limit the present application.

[0067] It should be noted that although the functional modules are divided in the device schematic and the logical order is shown in the flowchart, in some cases, the steps shown or described may be performed in a different order than the module division in the device or the order in the flowchart.

[0068] Unless otherwise defined, all technical and scientific terms used herein have the same meaning as those commonly understood by those skilled in the art to which this application belongs. The terms used herein are only for the purpose of describing the embodiments of this application and are not intended to limit this application.

[0069] First, some nouns involved in this application are analyzed:

[0070] Artificial intelligence (AI) is a new technical science that studies and develops theories, methods, technologies and application systems for simulating, extending and expanding human intelligence. AI is a branch of computer science that attempts to understand the essence of intelligence and produce a new type of intelligent machine that can respond in a similar way to human intelligence. Research in this field includes robotics, speech recognition, image recognition, natural language processing and expert systems. AI can simulate the information process of human consciousness and thinking. AI is also a theory, method, technology and application system that uses digital computers or machines controlled by digital computers to simulate, extend and expand human intelligence, perceive the environment, acquire knowledge and use knowledge to obtain the best results.

[0071] Generative models refer to models that can randomly generate the required data. Common generative models include variational auto-encoders (VAE), generative adversarial networks (GAN), generative pre-trained models (GPT), diffusion models (DM), etc. As a commonly used generative model, the diffusion model (DM) is widely used in many fields such as image generation, video generation, and protein structure generation in the scientific field. However, the generation process of the diffusion model requires multiple rounds of iterations at different time steps, which is time-consuming, computationally intensive, and computationally complex. In addition, compared with the common diffusion models in the field of computer vision, the equivariant diffusion model (EDM) commonly used in the scientific field also requires the generation of variable-length data, which further increases the amount of computation.

[0072] In order to solve the problem of large amount of computation in equivariant diffusion model, the related technologies proposed quantization methods to compress the model and reduce the computation time. However, the quantization algorithms in these technologies can only perceive the time step and cannot achieve targeted quantization for data of different lengths, thus affecting the quality of the generated data of equivariant diffusion model obtained after quantization training.

[0073] Based on this, the embodiments of the present application provide a quantization training method, apparatus, device and storage medium for an equivariant diffusion model. When quantizing the equivariant diffusion model, quantization weights are generated based on the sample length and time step of the calibration sample, and then the reconstruction error and the quantization error are fused to obtain the total loss value to complete the quantization training process, thereby improving the effect of the quantization training as a whole and improving the data generation quality of the equivariant diffusion model while maintaining equivariance.

[0074] The embodiments of the present application provide a quantitative training method, apparatus, device and storage medium for an equivariant diffusion model, which are specifically described through the following embodiments. First, the quantitative training method for an equivariant diffusion model in the embodiments of the present application is described.

[0075] The embodiments of the present application can acquire and process relevant data based on artificial intelligence technology. Among them, artificial intelligence (AI) is the theory, method, technology and application system that uses digital computers or machines controlled by digital computers to simulate, extend and expand human intelligence, perceive the environment, acquire knowledge and use knowledge to obtain the best results. In other words, artificial intelligence is a comprehensive technology in computer science that attempts to understand the essence of intelligence and produce a new intelligent machine that can respond in a similar way to human intelligence. Artificial intelligence is to study the design principles and implementation methods of various intelligent machines so that the machines have the functions of perception, reasoning and decision-making.

[0076] Artificial intelligence technology is a comprehensive discipline that covers a wide range of fields, including both hardware-level and software-level technologies. The basic technologies of artificial intelligence generally include sensors, dedicated artificial intelligence chips, cloud computing, distributed storage, big data processing technology, operation / interaction systems, mechatronics and other technologies. Artificial intelligence software technology mainly includes computer vision technology, speech processing technology, natural language processing technology, and machine learning / deep learning.

[0077] The quantitative training method of the equivariant diffusion model provided in the embodiment of the present application relates to the field of neural network technology. The quantitative training method of the equivariant diffusion model provided in the embodiment of the present application can be applied to the terminal, can also be applied to the server side, and can also be a computer program running in the terminal or the server side. For example, the computer program can be a native program or software module in the operating system; it can be a local (Native) application (Application, APP), that is, a program that needs to be installed in the operating system to run, such as a client that supports quantitative training of the equivariant diffusion model, that is, a program that can be run only by downloading it to the browser environment; it can also be a small program that can be embedded in any APP. In short, the above-mentioned computer program can be any form of application, module or plug-in. Among them, the terminal communicates with the server via a network. The quantitative training method of the equivariant diffusion model can be executed by the terminal or the server, or by the terminal and the server in collaboration.

[0078] In some embodiments, the terminal may be a smart phone, a tablet computer, a laptop computer, a desktop computer or a smart watch, etc. The server may be an independent server, or a cloud server that provides basic cloud computing services such as cloud services, cloud databases, cloud computing, cloud functions, cloud storage, network services, cloud communications, middleware services, domain name services, security services, content delivery networks (CDN), and big data and artificial intelligence platforms; or a service node in a blockchain system, wherein each service node in the blockchain system forms a peer-to-peer (P2P) network, and the P2P protocol is an application layer protocol running on the Transmission Control Protocol (TCP) protocol. The terminal and the server may be connected via Bluetooth, Universal Serial Bus (USB) or a network, etc., which is not limited in this embodiment.

[0079] The present application can be used in many general or special computer system environments or configurations. For example: personal computers, server computers, handheld or portable devices, tablet devices, multiprocessor systems, microprocessor-based systems, set-top boxes, programmable consumer electronics, network PCs, minicomputers, mainframe computers, distributed computing environments including any of the above systems or devices, etc. The present application can be described in the general context of computer-executable instructions executed by a computer, such as program modules. Generally, program modules include routines, programs, objects, components, data structures, etc. that perform specific tasks or implement specific abstract data types. The present application can also be practiced in distributed computing environments, in which tasks are performed by remote processing devices connected through a communication network. In a distributed computing environment, program modules can be located in local and remote computer storage media including storage devices.

[0080] It should be noted that in each specific implementation of the present application, when it comes to the need to perform relevant processing based on data related to user identity or characteristics such as user information, user behavior data, user historical data, and user location information, the user's permission or consent will be obtained first, and the collection, use, and processing of these data will comply with relevant laws, regulations, and standards. In addition, when the embodiment of the present application needs to obtain the user's sensitive personal information, the user's separate permission or consent will be obtained through a pop-up window or by jumping to a confirmation page. After clearly obtaining the user's separate permission or consent, the necessary user-related data for the normal operation of the embodiment of the present application will be obtained.

[0081] The following describes a quantitative training method for an equivariant diffusion model in an embodiment of the present application.

[0082] Figure 1 is an optional flow chart of the quantization training method of the equivariant diffusion model provided in the embodiment of the present application. Figure 1 The method may include but is not limited to steps 110 to 140. It can also be understood that this embodiment Figure 1 The order of step 110 to step 140 is not specifically limited, and the order of steps can be adjusted or some steps can be reduced or added according to actual needs.

[0083] Step 110: Obtain multiple calibration samples and an initial equivariant diffusion model, input the calibration samples into the initial equivariant diffusion model for data processing, and at least obtain initial output data of the attention module in the initial equivariant diffusion model.

[0084] In one embodiment, the model structure of the initial equivariant diffusion model that needs to be trained for quantization is first described. Figure 2 , Figure 2 The model structure of the initial equivariant diffusion model provided in the embodiment of the present application. Figure 2 The initial equivariant diffusion model includes the following modules, which are cascaded in sequence: an embedding module, an attention module, a self-attention module, and a feedforward network module, wherein the feedforward network module includes a first feedforward network module and a second feedforward network module. Specifically, the input data is firstly transformed into an embedding vector x by the embedding operation of the embedding module, and then sent to the attention module to calculate the attention weight y, and then the attention weight y is sent to the self-attention module for feature extraction, after which the extracted features and the embedding vector x are simultaneously sent to the first feedforward network module for data processing, and then the second feedforward network is used for the final data processing to obtain the output data of the initial equivariant diffusion model. It can be understood that both the first feedforward network module and the second feedforward network module include fully connected layers and activation functions.

[0085] In one embodiment, the quantization training of the initial equivariant diffusion model refers to the following Figure 2Each independent module in, such as the embedding module, the attention module, the self-attention module, the first feedforward network module and the second feedforward network module, is individually quantized and trained to obtain the corresponding quantized modules. For example, the quantized embedding module corresponding to the embedding module, the quantized attention module corresponding to the attention module, the quantized self-attention module corresponding to the self-attention module, the quantized first feedforward network module corresponding to the first feedforward network module and the quantized second feedforward network module corresponding to the second feedforward network module are obtained respectively. Finally, the quantized equivariant diffusion model is obtained according to the quantized embedding module, the quantized attention module, the quantized self-attention module, the quantized first feedforward network module and the quantized second feedforward network module. In other words, the quantized embedding module is used to replace the embedding module, the quantized attention module is used to replace the attention module, the quantized self-attention module is used to replace the self-attention module, the quantized first feedforward network module is used to replace the first feedforward network module, and the quantized second feedforward network module is used to replace the second feedforward network module to obtain the quantized result of the equivariant diffusion model.

[0086] It is understandable that since it is difficult for the model to converge if the quantization training is performed as a whole, the embodiment of the present application does not use the initial equivariant diffusion model as a whole for quantization training, but quantizes each independent module. The module is finally replaced to obtain the quantized equivariant diffusion model. In this way, the convergence speed of the quantization training is improved. The following embodiment illustrates the process of quantizing the attention module to obtain the quantized attention module, and other independent modules can also be quantized in the same way. Training.

[0087] In addition, it can be understood that the equivariant diffusion model of the embodiment of the present application is also widely used in many fields such as image generation, video generation, and protein structure generation in the scientific field. The following embodiments are explained by taking protein structure generation as an example, which does not limit the application field.

[0088] Next, the construction process of the initial equivariant diffusion model in the embodiment of the present application is described by taking the generation scenario of protein structure data as an example. Figure 3 , Figure 3 is a flowchart of obtaining an initial equivariant diffusion model provided in an embodiment of the present application, which specifically includes the following steps:

[0089] Step 310: Obtain a plurality of initial protein structure data, time steps and total time for adding noise.

[0090] In one embodiment, multiple initial protein structure data can be downloaded from the Protein Data Bank (PDB), which includes at least 224,931 initial protein structure data. Proteins are composed of multiple amino acids in a "dehydration condensation" manner to form a peptide chain (protein sequence), and the peptide chain obtains a 3D structure through dynamic folding. The longest peptide chain length L=512 is defined in the embodiment of the present application.

[0091] In one embodiment, referring to Figure 4 , Figure 4 It is a schematic diagram of generating pre-trained protein samples in the embodiment of the present application. Figure 4 The initial protein structure data is shown in . The protein structure includes amino acids that are consistent with the length of the peptide chain, and the initial protein structure data also includes initial amino acid data corresponding to each amino acid. Specifically, each initial amino acid data includes the three-dimensional coordinates of the amino acid, the first angle between the main carbon atom and the amino group, and the second angle between the main carbon atom and the hydroxyl group.

[0092] Therefore, the i-th initial protein structure data is expressed as:

[0093] P i ={(x j ,y j ,z j ,α j ,β j )|j∈[1,512]}

[0094] Among them, (x j ,y j ,z j ) represents the three-dimensional coordinates of the amino acid in the jth initial amino acid data, α j is the first angle between the main carbon atom and the amino group in the jth initial amino acid data, β j Represents the second angle between the main carbon atom and the hydroxyl group in the jth initial amino acid data.

[0095] In one embodiment, the time step Δt=0.02 is set, and a number is randomly selected in the interval [0, 1] as the total noise adding time t, for example, t=0.64.

[0096] Step 320: Calculate the number of noise additions according to the total noise addition duration and the time step.

[0097] In one embodiment, the ratio of the time step to the total duration of noise addition is used as the number of noise additions. Therefore, assuming Δt=0.02 and t=0.64, the number of noise additions n is expressed as: n=t / Δt=32.

[0098] Step 330: Perform a cyclic denoising process on the initial protein structure data for a number of times that the denoising is performed to obtain a pre-trained protein sample.

[0099] In one embodiment, adding noise refers to adding noise, and the cyclic adding noise process refers to the adding noise process corresponding to the number of times of adding noise for an input data. The input data of the first adding noise process is the initial protein structure data, and the output data is the initial protein structure data after adding noise. The input data of the subsequent adding noise process are all the output data of the previous adding noise process, and the final output data is obtained after adding noise n times. In other words, the initial protein structure data is updated n times.

[0100] In one embodiment, referring to Figure 5 , Figure 5 : is a flowchart of the cyclic noise adding process provided in the embodiment of the present application, which specifically includes the following steps:

[0101] Step 510: In each noise adding process, noise data corresponding to each initial amino acid data corresponding to the initial protein structure data is obtained.

[0102] In one embodiment, the noise data corresponding to the j-th initial amino acid data is expressed as: (xn j ,yn j ,zn j ,αn j ,βn j ), where (xn j ,yn j ,zn j ) represents the coordinate noise corresponding to the three-dimensional coordinates of the j-th amino acid, αn j is the first angle noise corresponding to the jth first angle, β j represents the second angle noise corresponding to the jth second angle. And the value of each noise is in the Gaussian distribution Random sampling in

[0103] Step 520: Superimpose the noise data on the initial amino acid data one by one to obtain updated initial protein structure data.

[0104] In one embodiment, after the noise data is superimposed on the initial amino acid data, for the i-th initial protein structure data, the updated data is expressed as:

[0105] P i ={(x j +xn j ,y j +yn j ,z j +zn j ,α j+αn j ,β j +βn j )|j∈[1,512]}

[0106] Step 530: Update the initial protein structure data at least once according to the number of noise additions.

[0107] In one embodiment, for each initial protein structure data, the denoising process is cycled n times, and each denoising process includes independent denoising of all initial amino acid data.

[0108] In one embodiment, in combination Figure 4 For the initial protein structure data x0, a noise addition process is performed at each time step, and the cycle is executed n times, and finally the pre-training protein sample x is obtained. T It is understandable that each piece of initial protein structure data can be used as a pre-training protein sample after the above-mentioned cyclic noise addition. For example, N = 10,000 pieces of initial protein structure data can be selected, and after cyclic noise addition, 10,000 corresponding pre-training protein samples can be obtained to form a pre-training data set. It is expressed as: Understandably, Figure 4 It is illustrated in the figure that the process from the initial protein structure data to the pre-trained protein samples is the noise addition process, while the process from the pre-trained protein samples to the initial protein structure data is the generation process.

[0109] Step 340: Acquire an equivariant diffusion model, and train the equivariant diffusion model using a plurality of pre-trained protein samples to obtain an initial equivariant diffusion model.

[0110] In one embodiment, based on Figure 2 The structure shown in the figure generates an equivariant diffusion model, and then uses the pre-training dataset The equivariant diffusion model is trained using multiple pre-trained protein samples to obtain an initial equivariant diffusion model.

[0111] The optimizer used in the training process is defined as ADAM optimizer. The training process has a total of iter = 100000 steps, and the number of training data for each iteration is batch_size = 8. The initial learning rate is lr = 0.001, and the learning rate adopts cosine annealing, which will be reduced to 0 in the last training iteration. After iter times of training, the equivariant diffusion model can reach convergence, thus obtaining the initial equivariant diffusion model.

[0112] After obtaining the initial equivariant diffusion model through the above process, it is necessary to generate a calibration data set, complete the quantitative training with a small amount of calibration data set, and improve the data generation speed of the quantitative model. Figure 6 , Figure 6 This is a flowchart of obtaining multiple calibration samples provided by an embodiment of the present application, which specifically includes the following steps:

[0113] Step 610: Obtain multiple random samples.

[0114] In one embodiment, the calibration samples are obtained by using random samples, and each random sample corresponds to a calibration sample. At this time, it is assumed that the calibration data set includes CM=2000 calibration samples in total. It is expressed as: At the same time, the sample length of all calibration samples in the calibration data set is defined to obey The distribution of the samples is truncated to ensure that the sample length of each calibration sample does not exceed the peptide chain length L = 512, that is, length = max(100, min(length, L)).

[0115] Among them, multiple amino acids in the random sample, take the i-th random sample as an example, assuming that its length = 350, at this time the random sample N i It is expressed as:

[0116]

[0117] Correspondingly, (x′ j ,y′ j ,z′ j ) represents a random sample N i The random three-dimensional coordinates of the amino acid corresponding to the jth amino acid in , α j ′ represents the random first angle between the main carbon atom and the amino group in the jth amino acid, β j ' represents the random second angle between the main carbon atom and the hydroxyl group in the jth amino acid. The random three-dimensional coordinates of the amino acid, the random first angle and the random second angle are all random numbers between 0 and 1.

[0118] Step 620: Obtain the sampling duration and time step, and calculate the number of generation times according to the sampling duration and time step.

[0119] In one embodiment, the sampling duration corresponding to the random sample is defined as And through the truncation operation to ensure that st i ~(0,1), that is, st i =max(0,min(st i ,1), at this time, the time step is Δt, then the number of times the random sample is generated is gi It is expressed as:

[0120]

[0121] Step 630: Execute a cyclic generation process consistent with the number of generation times for each random sample to obtain a corresponding calibration sample. During the generation process, the random sample is input into the initial equivariant diffusion model for data generation to obtain an updated random sample.

[0122] In one embodiment, each random sample has a corresponding number of times it is generated. i As an example, the random sample is input into the initial equivariant diffusion model for data generation, and the output of the initial equivariant diffusion model is obtained as the updated random sample. Figure 4 It can be seen that the purpose of the generation process is to generate protein structure data. Next, multiple cycles are performed according to the number of generation times, and the updated random samples are input into the initial equivariant diffusion model again for updating. In the last generation process, random samples N are obtained. i The corresponding calibration sample S i .

[0123] According to the above method, the corresponding CM calibration samples are obtained according to the CM random samples, thereby obtaining the calibration data set

[0124] Next, the initial equivariant diffusion model can be quantized based on multiple calibration samples. In one embodiment, since each module needs to be independently quantized and trained later, the parameter weights corresponding to each module in the initial equivariant diffusion model are first obtained. Since each module contains more parameters, the parameter weights can form a weight sequence. Taking the attention module as an example, since the parameters in the initial equivariant diffusion model are floating point precision, their parameter weights form a floating point weight sequence, expressed as: FW = {fw k |k∈[1,K]}, where K represents the total number of parameters of the attention module, fw k represents the kth parameter weight, and FW represents a floating-point weight sequence.

[0125] In one embodiment, the calibration samples need to be input into the initial equivariant diffusion model for data processing, and the input data and output data corresponding to each module, such as the embedding module, the attention module, the self-attention module, the first feedforward network module, and the second feedforward network module, are obtained in the data stream to form data pairs corresponding to each module. If there are CM calibration samples, each module has CM data pairs.

[0126] In one embodiment, taking the attention module as an example, a data pair of the attention module in the initial equivariant diffusion model is obtained, and the data pair is expressed as [initial input data, initial output data].

[0127] Reference Figure 7 , Figure 7 is a schematic diagram of an attention module provided in an embodiment of the present application, Figure 7 As can be seen from the figure, for the attention module, the initial input data of the attention module is first obtained based on the calibration sample, and then the initial input data is linearly transformed multiple times to obtain multiple intermediate vectors, represented as: V, Q, K, Q_pt, K_pt, V_pt.

[0128] Then, three attention vectors are obtained based on multiple intermediate vectors. Specifically, vector Q and vector K are dot-multiplied to obtain vector QK, and vector QK and vector V are dot-multiplied to obtain attention vector Comp1. Vector Q_pt and vector K_pt are dot-multiplied to obtain vector QK_pt, and vector QK_pt and vector V_pt are dot-multiplied to obtain attention vector Comp2. Vector Q and vector K_pt are dot-multiplied to obtain vector QK_pt', and vector QK_pt' and vector V_pt are dot-multiplied to obtain attention vector Comp3.

[0129] Next, the three attention vectors are weighted and summed according to the attention weight of each attention vector to obtain the attention output vector. In this process, the corresponding attention weight is added to each attention vector to measure its contribution to the result. The attention weight is contained in the floating-point weight sequence corresponding to the attention module and is a trainable parameter. Among them, the attention output vector is expressed as:

[0130] w L1 Comp1+w L2 Comp2+w L3 Comp3

[0131] Among them, w L1 represents the attention weight corresponding to the attention vector Comp1, w L2 represents the attention weight corresponding to the attention vector Comp2, w L3 Represents the attention weight corresponding to the attention vector Comp3.

[0132] Finally, the softmax function is calculated on the attention output vector to obtain the initial output data, so the initial output data attn new It is expressed as:

[0133] attn new =softmax(w L1 Comp1+w L2 Comp2+wL3 Comp3)

[0134] Step 120: Based on the sample length and time step of the calibration sample, generate a quantized weight sequence corresponding to the floating-point weight sequence, update the attention module according to the quantized weight sequence to obtain a quantized attention module, and replace the attention module with the quantized attention module to obtain an updated equivariant diffusion model.

[0135] In one embodiment, during the quantization training process, a quantizer is used to quantize the parameters of each module. The quantizer designed in this embodiment can take into account the sample length and the time step. Figure 8 , Figure 8 This is a flowchart of generating a quantized weight sequence corresponding to a floating-point weight sequence based on the sample length and time step of a calibration sample provided in an embodiment of the present application, which specifically includes the following steps:

[0136] Step 810: Obtain each parameter weight in the floating-point weight sequence, and obtain the maximum parameter value and the minimum parameter value in the parameter weight.

[0137] In one embodiment, for a floating point weight sequence FW={fw k |k∈[1,K]}, from which we can get the maximum and minimum values ​​of the parameter weights, and get the maximum value of the parameter max(fw k ) and the minimum value of the parameter min(fw k ).

[0138] Step 820: Calculate the quantization step size according to the preset quantization bit number, the maximum parameter value and the minimum parameter value.

[0139] In one embodiment, the quantization bit b is a preset value, and the value range is [0,2 b -1], the quantization step size is Δw k It is expressed as:

[0140]

[0141] It is understandable that if the kth parameter weight corresponds to the weight of a group of parameters, each parameter weight can have a corresponding quantization step size, which can improve the accuracy of the quantization process. At the same time, all parameter weights can also use the same quantization step size Δw, which is not limited in this embodiment.

[0142] Step 830: Determine a first intermediate parameter based on the training time step and the sample length.

[0143] In one embodiment, the first intermediate parameter μ k It is expressed as:

[0144]

[0145] Among them, t represents the training time step, which is set according to the actual training situation, and length represents the sample length. It can be seen that the embodiment of the present application realizes the quantization process under the condition of variable length and variable time step through the first intermediate parameter, which can solve the problem that the relevant technology cannot realize targeted quantization for data of different lengths, thereby affecting the quality of the generated data of the equivariant diffusion model obtained after quantization training.

[0146] In one embodiment, the calculation description of the first intermediate parameter is: when the training time step is in the first interval [0, 0.3) and the sample length is in the second interval [100, 240), the first intermediate parameter is the first value μ k1 ; When the training time step is in the third interval [0.3, 0.7) and the sample length is in the fourth interval [240, 380), the first intermediate parameter is the second value μ k2 ; When the training time step is in the fifth interval [0.7, 1.0) and the sample length is in the sixth interval [380, 512), the first intermediate parameter is the third value μ k3 It is understandable that the above intervals can be adjusted according to actual conditions, and the specific value of the first intermediate parameter is set according to actual conditions. This embodiment does not limit this.

[0147] Step 840: Calculate the quantization weight corresponding to each parameter weight according to the first intermediate parameter, the quantization step size and the basic parameter, and obtain a quantization weight sequence based on the quantization weight.

[0148] In one embodiment, the specific calculation process of the quantization weight is described as: calculating the ratio of the basic parameter and the quantization step, rounding the ratio to obtain the second intermediate parameter, and then obtaining the quantization weight according to the product of the first intermediate parameter, the second intermediate parameter and the quantization step.

[0149] Therefore, for the kth parameter weight fw k The corresponding quantization weight Q k It is expressed as:

[0150]

[0151] Among them, int represents the rounding operation. Represents the second intermediate parameter.

[0152] Therefore, the corresponding quantization weights are calculated for all parameter weights, and the obtained quantization weight sequence is expressed as: QW = {Q k |k∈[1,K]}.

[0153] After obtaining the quantized weight sequence, the attention module is updated according to the quantized weight sequence to obtain the quantized attention module, that is, the corresponding parameter weights in the attention module are all changed into the corresponding quantized weights. Next, the quantized attention module is used to replace the attention module in the initial equivariant diffusion model to obtain the updated equivariant diffusion model.

[0154] Step 130: Input the calibration sample into the updated equivariant diffusion model for data processing, and at least obtain the quantized output data corresponding to the quantized attention module.

[0155] In one embodiment, the calibration sample is input into the updated equivariant diffusion model for data processing, and the input and output data pairs corresponding to the quantized attention module are obtained according to the above process, and the data pairs are expressed as [quantized input data, quantized output data].

[0156] Step 140: Calculate a total loss value according to the quantized output data, the initial output data, the floating-point weight sequence, and the quantized weight sequence, and obtain a trained quantized equivariant diffusion model according to the total loss value.

[0157] In one embodiment, in order to improve the training effect, the reconstruction error and the quantization error are integrated into the total loss value. Fig. 9 , Fig. 9 The flowchart of calculating the total loss value according to the quantized output data, the initial output data, the floating-point weight sequence and the quantized weight sequence provided in the embodiment of the present application, and obtaining the trained quantized equivariant diffusion model according to the total loss value specifically includes the following steps:

[0158] Step 910: Calculate the reconstruction loss value according to the quantized output data and the initial output data, quantize the loss value according to the floating-point weight sequence and the quantized weight sequence, and calculate the total loss value according to the reconstruction loss value and the quantized loss value.

[0159] In one embodiment, the reconstruction loss value L rec It is expressed as:

[0160]

[0161] Among them, predict j represents the quantized output data corresponding to the jth calibration sample, GT j represents the initial output data corresponding to the jth calibration sample.

[0162] Furthermore, the quantified loss value L quant It is expressed as:

[0163]

[0164] Among them, W fi Represents the i-th parameter weight, which is the aforementioned fwk , W qi Represents the quantized weight corresponding to the i-th parameter weight, which is the aforementioned Q k .

[0165] The total loss value L is expressed as:

[0166]

[0167] Step 920: Adjust the quantized attention module based on the total loss value until a trained quantized attention module is obtained.

[0168] In one embodiment, the training optimizer is defined as an AdamW optimizer, the initial learning rate is 0.0002, and the quantization training of the attention module is completed through 5000 iterations to obtain the trained quantized attention module.

[0169] Step 930: Obtain a trained quantized equivariant diffusion model based on the quantized attention module.

[0170] In one embodiment, according to the above process, a quantized embedding module corresponding to the embedding module, a quantized attention module corresponding to the attention module, a quantized self-attention module corresponding to the self-attention module, a quantized first feedforward network module corresponding to the first feedforward network module, and a quantized second feedforward network module corresponding to the second feedforward network module are obtained respectively. Finally, a quantized equivariant diffusion model is obtained according to the quantized embedding module, the quantized attention module, the quantized self-attention module, the quantized first feedforward network module, and the quantized second feedforward network module. That is to say, the quantized embedding module is replaced with the embedding module, the quantized attention module is replaced with the attention module, the quantized self-attention module is replaced with the self-attention module, the quantized first feedforward network module is replaced with the first feedforward network module, and the quantized second feedforward network module is replaced with the second feedforward network module, to obtain a quantized equivariant diffusion model of the equivariant diffusion model. Among them, the quantized equivariant diffusion model QW is expressed as:

[0171] QW={qw m |m∈[1,M]}

[0172] Among them, qw m represents the model parameters in the quantized equivariant diffusion model, and M represents the total number of parameters in the quantized equivariant diffusion model.

[0173] In one embodiment, a verification and evaluation of a quantization-equivariant diffusion model is carried out. Specifically, test samples Test are generated in the same way as the calibration samples, and then the test samples are input into the initial equivariant diffusion model FW and the quantization-equivariant diffusion model QW respectively. Through the generation process, 100 iterations are performed to generate the output data corresponding to the two models respectively. Among them, the output data corresponding to the initial equivariant diffusion model is Output_f, and the output data corresponding to the quantization-equivariant diffusion model is Output_q. And the time required for the two models to complete the generation process is counted. The time corresponding to the initial equivariant diffusion model is time_f, and the time corresponding to the quantization-equivariant diffusion model is time_q.

[0174] Next, compare the magnitudes of time_q and time_f. When time_q < time_f, it is considered that the quantization-equivariant diffusion model has improved the data generation speed. At the same time, Output_f and Output_q are respectively input into the ProteinMPNN+ESMFold model, and the output corresponding to Output_f is Gen_f and the output corresponding to Output_q is Gen_q. Gen_f and Gen_q are respectively compared with all protein results in the protein database and the similarity is calculated, and the maximum similarity value is taken. Among them, the maximum similarity corresponding to Gen_f is score_f, and the maximum similarity corresponding to Gen_q is score_q. When score_q > score_f, it is considered that the data generated by the quantization-equivariant diffusion model has higher quality.

[0175] Finally, 10 test samples with different lengths are generated, and the above verification process is respectively executed to compare the quality and speed of data generation. According to the verification results, after quantization training, the obtained quantization-equivariant diffusion model can complete fast and high-quality data generation for different input data lengths.

[0176] As can be seen from the above, the quantization-equivariant diffusion model obtained by the quantization training method of the equivariant diffusion model provided in the embodiment of the present application can realize length perception and time step perception in the generation process when facing variable-length data generation tasks, and improve the quantization training effect. And through the proposed quantization evaluation method, when facing variable-length data generation tasks, it can accurately evaluate the data generation speed and the quality of the generated data of the quantization-equivariant diffusion model after quantization training. Through evaluation, it can be known that the quantization-equivariant diffusion model of the embodiment of the present application can improve the quality and generation speed of the generated data while ensuring equivariance.

[0177] The technical solution provided by the embodiment of the present application is to obtain multiple calibration samples and an initial equivariant diffusion model, input the calibration samples into the initial equivariant diffusion model for data processing, obtain at least the initial output data of the attention module in the initial equivariant diffusion model, generate a quantized weight sequence corresponding to the floating-point weight sequence based on the sample length and time step of the calibration samples, update the attention module according to the quantized weight sequence to obtain the quantized attention module, replace the attention module with the quantized attention module to obtain an updated equivariant diffusion model, input the calibration samples into the updated equivariant diffusion model for data processing, obtain at least the quantized output data corresponding to the quantized attention module, calculate the total loss value according to the quantized output data, the initial output data, the floating-point weight sequence and the quantized weight sequence, and obtain the trained quantized equivariant diffusion model according to the total loss value. The embodiment of the present application first constructs an equivariant diffusion model, and uses the pre-trained protein sample to train the equivariant diffusion model to obtain the initial equivariant diffusion model. Then, based on the initial equivariant diffusion model, a calibration data set required for quantization training of a specified distribution is generated, and when the equivariant diffusion model is quantized, quantized weights are generated based on the sample length and time step of the calibration sample. Then, the reconstruction error and quantization error are fused to obtain the total loss value to complete the quantization training process, which improves the effect of quantization training as a whole and improves the data generation quality of the equivariant diffusion model while maintaining equivariance.

[0178] The embodiment of the present application also provides a quantization training device for an equivariant diffusion model, which can implement the quantization training method for the equivariant diffusion model. Fig.10 , the device comprises:

[0179] Initial output acquisition module 1010: used to obtain multiple calibration samples and an initial equivariant diffusion model, input the calibration samples into the initial equivariant diffusion model for data processing, and at least obtain the initial output data of the attention module in the initial equivariant diffusion model. The parameter weights corresponding to the attention module constitute a floating-point weight sequence.

[0180] Weight quantization module 1020: used to generate a quantized weight sequence corresponding to the floating-point weight sequence based on the sample length and time step of the calibration sample, update the attention module according to the quantized weight sequence to obtain a quantized attention module, and replace the attention module with the quantized attention module to obtain an updated equivariant diffusion model.

[0181] Quantized output acquisition module 1030: used to input the calibration sample into the updated equivariant diffusion model for data processing, and at least obtain the quantized output data corresponding to the quantized attention module.

[0182] Quantization training module 1040: used to calculate the total loss value according to the quantized output data, the initial output data, the floating-point weight sequence and the quantized weight sequence, and obtain the trained quantized equivariant diffusion model according to the total loss value.

[0183] The specific implementation of the quantization training device for the equivariant diffusion model of this embodiment is basically the same as the specific implementation of the quantization training method for the equivariant diffusion model described above, and will not be described in detail here.

[0184] The present application also provides an electronic device, including:

[0185] at least one memory;

[0186] at least one processor;

[0187] at least one program;

[0188] The program is stored in the memory, and the processor executes the at least one program to implement the quantitative training method of the equivariant diffusion model implemented in the present application. The electronic device can be any intelligent terminal including a mobile phone, a tablet computer, a personal digital assistant (PDA), a car computer, etc.

[0189] See also Fig.11 , Fig.11 The hardware structure of an electronic device of another embodiment is illustrated, and the electronic device includes:

[0190] The processor 1101 may be implemented by a general-purpose central processing unit (CPU), a microprocessor, an application-specific integrated circuit (ASIC), or one or more integrated circuits, and is used to execute relevant programs to implement the technical solutions provided in the embodiments of the present application;

[0191] The memory 1102 may be implemented in the form of a read-only memory (ROM), a static storage device, a dynamic storage device, or a random access memory (RAM). The memory 1102 may store an operating system and other application programs. When the technical solution provided in the embodiment of this specification is implemented by software or firmware, the relevant program code is stored in the memory 1102, and the processor 1101 calls and executes the quantization training method of the equivariant diffusion model in the embodiment of the present application;

[0192] Input / output interface 1103, used to implement information input and output;

[0193] Communication interface 1104, used to realize communication interaction between the device and other devices, which can be realized by wired means (such as USB, network cable, etc.) or wireless means (such as mobile network, WIFI, Bluetooth, etc.); and

[0194] A bus 1105 that transmits information between various components of the device (e.g., the processor 1101, the memory 1102, the input / output interface 1103, and the communication interface 1104);

[0195] The processor 1101 , the memory 1102 , the input / output interface 1103 and the communication interface 1104 are connected to each other in communication within the device via the bus 1105 .

[0196] An embodiment of the present application further provides a storage medium, which is a storage medium storing a computer program. When the computer program is executed by a processor, the quantization training method of the above-mentioned equivariant diffusion model is implemented.

[0197] As a non-transitory storage medium, the memory can be used to store non-transitory software programs and non-transitory computer executable programs. In addition, the memory may include a high-speed random access memory, and may also include a non-transitory memory, such as at least one disk storage device, a flash memory device, or other non-transitory solid-state storage device. In some embodiments, the memory may optionally include a memory remotely arranged relative to the processor, and these remote memories may be connected to the processor via a network. Examples of the above-mentioned network include, but are not limited to, the Internet, an intranet, a local area network, a mobile communication network, and combinations thereof.

[0198] The quantitative training method, device, equipment and storage medium of the equivariant diffusion model proposed in the embodiment of the present application, by obtaining multiple calibration samples and an initial equivariant diffusion model, inputting the calibration samples into the initial equivariant diffusion model for data processing, at least obtaining the initial output data of the attention module in the initial equivariant diffusion model, generating a quantized weight sequence corresponding to the floating-point weight sequence based on the sample length and time step of the calibration samples, updating the attention module according to the quantized weight sequence to obtain the quantized attention module, replacing the attention module with the quantized attention module to obtain an updated equivariant diffusion model, inputting the calibration samples into the updated equivariant diffusion model for data processing, at least obtaining the quantized output data corresponding to the quantized attention module, calculating the total loss value according to the quantized output data, the initial output data, the floating-point weight sequence and the quantized weight sequence, and obtaining the trained quantized equivariant diffusion model according to the total loss value. The embodiment of the present application first constructs an equivariant diffusion model, and uses the pre-trained protein sample to train the equivariant diffusion model to obtain the initial equivariant diffusion model. Then, based on the initial equivariant diffusion model, a calibration data set required for quantized training of a specified distribution is generated, and when the equivariant diffusion model is quantized, quantized weights are generated based on the sample length and time step of the calibration sample. Subsequently, the reconstruction error and quantization error are fused to obtain the total loss value to complete the quantization training process, which improves the overall effect of quantization training and improves the data generation quality of the equivariant diffusion model while maintaining equivariance.

[0199] The embodiments described in the embodiments of the present application are intended to more clearly illustrate the technical solutions of the embodiments of the present application and do not constitute a limitation on the technical solutions provided in the embodiments of the present application. Those skilled in the art will appreciate that with the evolution of technology and the emergence of new application scenarios, the technical solutions provided in the embodiments of the present application are also applicable to similar technical problems.

[0200] Those skilled in the art will appreciate that the technical solutions shown in the figures do not constitute a limitation on the embodiments of the present application, and may include more or fewer steps than shown in the figures, or a combination of certain steps, or different steps.

[0201] The device embodiments described above are merely illustrative, and the units described as separate components may or may not be physically separated, that is, they may be located in one place or distributed on multiple network units. Some or all of the modules may be selected according to actual needs to achieve the purpose of the solution of this embodiment.

[0202] Those skilled in the art will appreciate that all or some of the steps in the methods disclosed above, and the functional modules / units in the systems and devices may be implemented as software, firmware, hardware, or a suitable combination thereof.

[0203] The terms "first", "second", "third", "fourth", etc. (if any) in the specification of the present application and the above-mentioned drawings are used to distinguish similar objects, and are not necessarily used to describe a specific order or sequence. It should be understood that the data used in this way can be interchangeable where appropriate, so that the embodiments of the present application described herein can be implemented in an order other than those illustrated or described herein. In addition, the terms "including" and "having" and any of their variations are intended to cover non-exclusive inclusions, for example, a process, method, system, product or device comprising a series of steps or units is not necessarily limited to those steps or units clearly listed, but may include other steps or units that are not clearly listed or inherent to these processes, methods, products or devices.

[0204] It should be understood that in the present application, "at least one (item)" means one or more, and "plurality" means two or more. "And / or" is used to describe the association relationship of associated objects, indicating that three relationships may exist. For example, "A and / or B" can mean: only A exists, only B exists, and A and B exist at the same time, where A and B can be singular or plural. The character " / " generally indicates that the objects associated before and after are in an "or" relationship. "At least one of the following" or similar expressions refers to any combination of these items, including any combination of single or plural items. For example, at least one of a, b or c can mean: a, b, c, "a and b", "a and c", "b and c", or "a and b and c", where a, b, c can be single or multiple.

[0205] In the several embodiments provided in the present application, it should be understood that the disclosed devices and methods can be implemented in other ways. For example, the device embodiments described above are only schematic. For example, the division of the above units is only a logical function division. There may be other division methods in actual implementation, such as multiple units or components can be combined or integrated into another system, or some features can be ignored or not executed. Another point is that the mutual coupling or direct coupling or communication connection shown or discussed can be through some interfaces, indirect coupling or communication connection of devices or units, which can be electrical, mechanical or other forms.

[0206] The units described above as separate components may or may not be physically separated, and the components shown as units may or may not be physical units, that is, they may be located in one place or distributed on multiple network units. Some or all of the units may be selected according to actual needs to achieve the purpose of the solution of this embodiment.

[0207] In addition, each functional unit in each embodiment of the present application may be integrated into one processing unit, or each unit may exist physically separately, or two or more units may be integrated into one unit. The above-mentioned integrated unit may be implemented in the form of hardware or in the form of software functional units.

[0208] If the integrated unit is implemented in the form of a software functional unit and sold or used as an independent product, it can be stored in a computer-readable storage medium. Based on this understanding, the technical solution of the present application is essentially or the part that contributes to the prior art or all or part of the technical solution can be embodied in the form of a software product, and the computer software product is stored in a storage medium, including multiple instructions to enable a computer device (which can be a personal computer, server, or network device, etc.) to execute all or part of the steps of the methods of various embodiments of the present application. The aforementioned storage medium includes: U disk, mobile hard disk, read-only memory (ROM), random access memory (RAM), disk or optical disk and other media that can store programs.

[0209] The preferred embodiments of the present invention are described above with reference to the accompanying drawings, but the scope of the rights of the present invention is not limited thereto. Any modification, equivalent substitution and improvement made by a person skilled in the art without departing from the scope and essence of the present invention should be within the scope of the rights of the present invention.

Claims

1. A quantitative training method for an equivariant diffusion model, characterized in that: include: Acquire multiple calibration samples and an initial equivariant diffusion model, input the calibration samples into the initial equivariant diffusion model for data processing, and at least acquire initial output data of an attention module in the initial equivariant diffusion model, wherein parameter weights corresponding to the attention module constitute a floating-point weight sequence; Based on the sample length and time step of the calibration sample, a quantized weight sequence corresponding to the floating-point weight sequence is generated, the attention module is updated according to the quantized weight sequence to obtain a quantized attention module, and the attention module is replaced by the quantized attention module to obtain an updated equivariant diffusion model; Inputting the calibration sample into the updated equivariant diffusion model for data processing to at least obtain quantized output data corresponding to the quantized attention module; A total loss value is calculated according to the quantized output data, the initial output data, the floating-point weight sequence and the quantized weight sequence, and a trained quantized equivariant diffusion model is obtained according to the total loss value.

2. The quantitative training method of the equivariant diffusion model according to claim 1, characterized in that: The step of generating a quantized weight sequence corresponding to the floating-point weight sequence based on the sample length and training time step of the calibration sample comprises: Obtain each parameter weight in the floating-point weight sequence, and obtain the maximum value and the minimum value of the parameter in the parameter weight; Calculate the quantization step size according to a preset number of quantization bits, the maximum value of the parameter and the minimum value of the parameter; Determining a first intermediate parameter based on the training time step and the sample length; The quantization weight corresponding to each of the parameter weights is calculated according to the first intermediate parameter, the quantization step size and the basic parameter, and the quantization weight sequence is obtained based on the quantization weight.

3. The quantitative training method of the equivariant diffusion model according to claim 2, characterized in that: The step of calculating, according to the first intermediate parameter, the quantization step size and the basic parameter, a quantization weight corresponding to each basic parameter comprises: Calculating a ratio of the basic parameter to the quantization step length, and rounding the ratio to obtain a second intermediate parameter; The quantization weight is obtained according to the product of the first intermediate parameter, the second intermediate parameter and the quantization step size.

4. The quantitative training method of the equivariant diffusion model according to claim 1, characterized in that: The calculating a total loss value according to the quantized output data, the initial output data, the floating-point weight sequence and the quantized weight sequence, and obtaining a trained quantized equivariant diffusion model according to the total loss value, comprises: Calculating a reconstruction loss value according to the quantized output data and the initial output data, measuring a quantized loss value according to the floating-point weight sequence and the quantized weight sequence, and calculating the total loss value according to the reconstruction loss value and the quantized loss value; Adjusting the quantized attention module based on the total loss value until the trained quantized attention module is obtained; A trained quantized equivariant diffusion model is obtained based on the quantized attention module.

5. The quantitative training method of the equivariant diffusion model according to claim 1, characterized in that: The obtaining of a plurality of calibration samples comprises: Get multiple random samples; Obtaining a sampling duration and a time step, and calculating the number of generation times according to the sampling duration and the time step; A cyclic generation process consistent with the generation times is performed on each of the random samples to obtain a corresponding calibration sample. During the generation process, the random sample is input into the initial equivariant diffusion model for data generation to obtain an updated random sample.

6. The quantitative training method of the equivariant diffusion model according to claim 1, characterized in that: Get the initial equivariant diffusion model, including: Obtain multiple initial protein structure data, time steps and total time for adding noise; Calculate the number of noise additions according to the total noise addition duration and the time step; Performing a cyclic noise adding process consistent with the noise adding times on the initial protein structure data to obtain a pre-trained protein sample; An equivariant diffusion model is obtained, and the equivariant diffusion model is trained using a plurality of the pre-trained protein samples to obtain the initial equivariant diffusion model.

7. The quantitative training method of the equivariant diffusion model according to claim 6, characterized in that: The initial protein structure data includes a plurality of initial amino acid data, and the cyclic noise adding process includes: In each noise adding process, obtaining noise data corresponding to each of the initial amino acid data corresponding to the initial protein structure data; superimposing the noise data on the initial amino acid data one by one to obtain the updated initial protein structure data; The initial protein structure data is updated at least once according to the number of times of adding noise.

8. The quantitative training method of the equivariant diffusion model according to claim 7, characterized in that: The step of superimposing the noise data on the initial amino acid data one by one to obtain the updated initial protein structure data comprises: Obtaining the three-dimensional coordinates of the amino acid, the first angle between the main carbon atom and the amino group, and the second angle between the main carbon atom and the hydroxyl group in the initial amino acid data; Acquire three-dimensional coordinate noise data, first angle noise data, and second angle noise data corresponding to the noise data; The three-dimensional coordinate noise data is superimposed on the amino acid three-dimensional coordinates, the first angle noise data is superimposed on the first angle, and the second angle noise data is superimposed on the second angle to obtain the updated initial protein structure data.

9. The quantitative training method of the equivariant diffusion model according to claim 1, characterized in that: The initial equivariant diffusion model at least includes: an embedding module, an attention module, a self-attention module and a feedforward network module which are cascaded in sequence. The trained quantized equivariant diffusion model based on the quantized attention module includes: Obtain a quantized embedding module corresponding to the embedding module, a quantized attention module corresponding to the attention module, a quantized self-attention module corresponding to the self-attention module, and a quantized feedforward network module corresponding to the feedforward network module; The quantized equivariant diffusion model is obtained according to the quantized embedding module, the quantized attention module, the quantized self-attention module and the quantized feedforward network module.

10. The quantitative training method of the equivariant diffusion model according to any one of claims 1 to 9, characterized in that: The step of inputting the calibration sample into the initial equivariant diffusion model for data processing to obtain at least initial output data of the attention module in the initial equivariant diffusion model includes: Acquiring initial input data for the attention module based on the calibration sample; Perform multiple linear transformations on the initial input data to obtain multiple intermediate vectors, and obtain three attention vectors based on the multiple intermediate vectors; Performing weighted summation on the three attention vectors according to the attention weight of each of the attention vectors to obtain an attention output vector, wherein the floating-point weight sequence includes the attention weight; The softmax function is performed on the attention output vector to obtain the initial output data.

11. A quantitative training device for an equivariant diffusion model, characterized in that: include: Initial output acquisition module: used to obtain multiple calibration samples and an initial equivariant diffusion model, input the calibration samples into the initial equivariant diffusion model for data processing, and at least obtain the initial output data of the attention module in the initial equivariant diffusion model, and the parameter weights corresponding to the attention module constitute a floating-point weight sequence; A weight quantization module: used to generate a quantized weight sequence corresponding to the floating-point weight sequence based on the sample length and time step of the calibration sample, update the attention module according to the quantized weight sequence to obtain a quantized attention module, and replace the attention module with the quantized attention module to obtain an updated equivariant diffusion model; Quantized output acquisition module: used to input the calibration sample into the updated equivariant diffusion model for data processing, and at least obtain the quantized output data corresponding to the quantized attention module; Quantization training module: used to calculate the total loss value according to the quantized output data, the initial output data, the floating-point weight sequence and the quantized weight sequence, and obtain the trained quantized equivariant diffusion model according to the total loss value.

12. An electronic device, characterized in that: The electronic device includes a memory and a processor, the memory stores a computer program, and the processor implements the quantitative training method of the equivariant diffusion model according to any one of claims 1 to 10 when executing the computer program.

13. A storage medium storing a computer program, characterized in that: When the computer program is executed by a processor, the quantization training method of the equivariant diffusion model according to any one of claims 1 to 10 is implemented.

Images