Efficient molecular docking method based on geometric perception attention model

Abstract

The present application belongs to the technical field of molecular docking, and particularly relates to an efficient molecular docking method based on a geometric perception attention model. A dynamic memory reduction strategy uses iterative evaluation and block technology to realize efficient parallel calculation of an attention matrix, and with the help of shared memory, communication of intermediate results between different thread bundles is completed; a hierarchical fusion strategy designs two high-performance fusion kernels for triangular attention, so that all the data generated by each thread block is saved in registers and shared memory; tensor tiling data moves through a cross-dimension data carrying mode to reorganize the mapping relationship between threads and data, realizing a global merging memory access mode; multi-head parallel attention merging uses a multi-head attention merging strategy to optimize QKV data layout, reducing additional data transposition and splicing operations of intermediate results. The present application solves the problems of memory explosion during reasoning of an existing triangular attention operator, irregular memory access in a triangular update operator, and redundant memory access.

Description

TECHNICAL FIELD

[0001] The application belongs to the technical field of molecular docking, and particularly relates to an efficient molecular docking method based on a geometric perception attention model. BACKGROUND

[0002] Molecular docking is a widely used computational method for predicting protein-ligand interactions, which is used to study the fusion of two or more molecular structures (e.g., proteins and ligands) together, and has become one of the most commonly used methods in modern drug discovery. With the increase in the size of the compound library, the computational complexity of traditional docking models has also increased, for example, docking tools such as GOLD, AutoDock, etc. In large-scale virtual screening tasks, they will show high computational and time costs, which further makes the computational method a new research paradigm for predicting protein-ligand interactions because of its excellent high-throughput screening capability and low computational cost. Therefore, it is crucial to study a high-performance molecular docking method.

[0003] In recent years, with the breakthrough progress of AlpahFold2 in the field of protein structure prediction, a series of deep learning-based methods have emerged to solve the molecular docking task. For example, graph neural network-based, Transformer-based, and diffusion model-based molecular docking methods. The advantages of such methods are: on the one hand, researchers can use the powerful learning ability of deep learning to improve the accuracy of molecular docking, and on the other hand, by using the powerful parallel processing capability of GPU, an end-to-end molecular docking model can be easily trained, and its inference process can be directly applied to drug screening, greatly improving the prediction speed of protein-ligand binding conformations. In addition, in order to overcome the problems of molecular docking models in terms of lack of universality and non-end-to-end architecture, when designing such models, they are usually based on existing general 3D molecular representation learning frameworks to ensure docking accuracy and screening speed.

[0004] Inspired by the Evoformer in AlphaFold2, GAABind combines triangle attention and triangle update into the general 3D molecular representation learning framework Uni-Mol, achieving 82.8% success rate in binding pose prediction, reaching the state-of-the-art performance. At the same time, the Pearson correlation coefficient correlation between predicted and achieved binding affinities is as high as 0.803, further proving the effectiveness of the MolTransformer architecture in the molecular docking task. Unlike the standard self-attention, triangle attention extends the traditional one-dimensional input information to two dimensions, which means that the dimension of the intermediate tensor can reach five when performing self-attention calculation. This design of updating molecular representation using the triangle inequality relationship ensures the physically reasonable geometric relationship between ligands and proteins. In addition, triangle update combines two triangle multiplication update modules in AlphaFold2 to realize the message passing between atom pairs. This process of incorporating geometric information into the attention mechanism effectively promotes the comprehensive encoding of intramolecular interactions of the input molecule, greatly enhancing the molecular docking model's ability in molecular representation learning.

[0005] Although the MolTransformer has achieved excellent performance in the molecular docking task, it still faces challenges such as low throughput and insufficient device memory resources when used in actual inference tasks, which cannot fully exploit the performance potential of hardware resources. The main reasons are as follows: First, the memory explosion problem faced by the triangle attention operator during execution, the instantiation of attention scores, the broadcast addition of bias terms and mask terms during the execution of multi-head attention calculation of triangle attention, which puts a huge pressure on device memory resources, seriously hindering the increase of data batch size in the molecular docking process. Second, the irregular cross-step memory access in triangle update seriously limits the improvement of device throughput. Third, the data redundancy memory access problem, which requires multiple Q, K, V tensors to be transposed before triangle attention calculation, causing frequent access of data in global memory.

[0006] Problem 1: Memory explosion problem of triangular attention operator reasoning. In actual operation, the triangular attention operator usually shows excessively high peak memory usage, which significantly affects the batch size that the model can support and the memory utilization efficiency. Problem 2: Irregular memory access problem in triangular update operator. When performing the transpose operation on a high-dimensional tensor, each thread needs to read the data to be transposed from the device memory first, and then write the read data back to the device memory according to the specified memory layout. The combined memory access requires the thread bundle to access the continuous elements in one instruction, which causes the thread bundle to be unable to simultaneously satisfy the combined memory access of reading and writing back. Problem 3: Redundant memory access problem. When writing the kernel that fuses the entire triangular attention, we need to first complete the transpose and splicing operations of Q, K and V. This additional transpose and splicing operation causes frequent access to the device memory.

[0007] In summary, predicting the docking between proteins and ligands plays a crucial role in drug discovery and virtual screening. As a new deep learning technology, the geometry-aware attention network successfully integrates the geometric information of molecules with the attention mechanism and has made remarkable achievements in molecular docking. However, the implementation of some key operations in this molecular docking model cannot fully utilize the hardware performance, and shows high latency and high memory occupation in actual drug virtual screening. SUMMARY

[0008] The present application proposes an efficient molecular docking method based on a geometry-aware attention model, which solves the existing memory explosion problem of triangular attention operator reasoning, irregular memory access problem in triangular update operator and redundant memory access problem.

[0009] The technical scheme of the present application is implemented as follows:

[0010] The efficient molecular docking method based on the geometry-aware attention model comprises the following parts:

[0011] S1, dynamic memory reduction strategy: the iterative evaluation and block technology are used to realize efficient parallel calculation of the attention matrix, and the shared memory is used to complete the communication of the intermediate results between different thread bundles;

[0012] S2, hierarchical fusion strategy: two high-performance fusion kernels are designed for the triangular attention, which are the fusion transpose kernel and the fusion triangular attention kernel, so that the data generated by each thread block is all saved in the register and the shared memory;

[0013] S3, tensor tiling data movement: for the triangular update operator, the mapping relationship between the thread and the data is reorganized by cross-dimension data moving, realizing the global combined memory access mode;

[0014] S4, multi-head parallel attention merging: using a multi-head attention merging strategy to optimize the QKV data layout, reducing the additional data transpose and splicing operations of the intermediate results.

[0015] Through the above technical solutions, for the memory explosion problem of triangular attention operator reasoning, S1 parallelizes the triangular attention operator, and S2 realizes the fusion of the entire triangular attention operator, successfully reducing the peak memory during model reasoning, the maximum data batch is expanded to 3.9 times of the original, achieving 2.83 times speedup; for irregular memory access problems, S3 realizes efficient data movement of multi-dimensional tensor data, and uses the merged memory access mode to improve the memory bandwidth of the GPU memory, achieving 1.75 times speedup; for redundant memory access problems, S4 fuses multiple tensor data transpose and splicing operations into one kernel, solving the redundancy problem of repeatedly accessing memory caused by multiple transpose operations, achieving 3.05 times speedup; finally, the molecular docking model optimized by S1, S2, S3 and S4 realizes 1.54 times speedup in the end-to-end molecular docking task.

[0016] Optionally, in step S1, the specific content of the dynamic memory reduction strategy is as follows:

[0017] S11, tile thread blocks along the sequence dimension direction, and allocate shared memory buffers according to the size of the hidden layer dimension. Each thread block loads different sizes of Q, K and V from global memory to shared memory, where Q represents the "question" that needs to pay attention to other problems, K represents the "keyword" of other positions, which is used to calculate the similarity with Q, and V represents the "information" actually carried by each position. Attention weight is finally weighted on V to generate the final attention output. The block sizes of Q, K and V are (Br, 2D), (Bc, 2D) and (Bc, D) respectively;

[0018] S12, iteratively calculate the thread blocks along the sequence direction, and constantly update the attention output matrix with the help of the local maximum value vector and the local summation term to realize parallel calculation of the softmax normalization operation. The specific formula is as follows:

[0019]

[0020] Where b represents the block size of the tiled piece, N represents the length of the input sequence, x i represents the i-th vector [(i-1)b:ib] of Q[k]K T , m i represents the local maximum value of x i , d′ i represents the local summation term of x i , V represents the V tensor participating in attention calculation, with a dimension of [B, N, N, D], and o′i represents a partial output tensor calculated by the i-th iteration;

[0021] S13, re-optimizing the layout of the thread bundles in the thread block, each thread bundle is responsible for loading different data blocks and performing a reduction operation on all output results in shared memory, and finally writing the output results back to the global memory.

[0022] Through the above technical solution, by applying the dynamic memory reduction strategy to the triangular attention operator, the storage of intermediate results in the attention calculation process is avoided, the peak memory of the attention operation is significantly reduced, and the entire triangular attention operator is accelerated by 2.83 times.

[0023] Optionally, in step S2, the specific content of the hierarchical fusion strategy is as follows:

[0024] S21, by adjusting the execution order of each operator in the triangular attention, the transpose operations of Q and K are placed before the batch matrix multiplication operation, and the two bias terms are combined into one bias term by addition, and the two mask terms are also combined into one mask term by addition, thereby realizing the combination of multiple calculation operations such as tensor transpose, matrix multiplication and addition, and bias and mask term addition;

[0025] S22, the fusion transpose kernel splices multiple Qs or Ks along the hidden layer dimension D, and each thread is responsible for carrying the data at the same head position in the two Qs or Ks to the global memory, so that the two data blocks originally stored in non-continuous storage positions are stored in a memory continuous area. The hidden layer dimension is expanded from D to 2D and is directly used for multi-head attention calculation;

[0026] S23, the fusion triangular attention kernel integrates all operators such as matrix multiplication, bias addition, mask addition, and Softmax normalization into one kernel, and the intermediate results generated by each thread block are all saved in registers and shared memory.

[0027] Through the above technical solution, by applying the hierarchical fusion strategy to the triangular attention operator, the transpose operations of Q and K are fused into one kernel, and the remaining matrix multiplication, bias addition, mask addition, and Softmax normalization are all fused into another kernel. The data generated by each thread block is all saved in registers and shared memory, avoiding the moving overhead of data between global memories and greatly reducing the memory usage burden of the device. In the end-to-end molecular docking task, an acceleration ratio of 1.54 times is achieved, and the maximum data batch supported by the model is expanded to 3.9 times of the original.

[0028] Optionally, in step S3, the specific content of the tensor tiling data movement is as follows:

[0029] S31. Determine the dimension of the data stored contiguously in the physical memory of the tensor, and load tensor slice data according to different dimensions. Each thread block loads a tensor slice of size [tilex, tiley] along the D dimension and N2 dimension.

[0030] S32. Tensor slice transpose operation is implemented by using a shared memory buffer. Each thread is responsible for writing four consecutive FP16 type data in memory into shared memory row by row. The entire thread block writes a tensor slice of size [64,64] into shared memory. Then each thread accesses the four FP16 data in shared memory column by column and writes them into global memory row by row, thereby realizing the tensor slice transpose operation.

[0031] S33. Each thread block loads the transposed data from shared memory and stores it along the dimensions of N2 and D to ensure that the merged access mode is followed when writing back to global memory.

[0032] By applying the tensor tiling data movement strategy to the triangle update operator, when the tensor shape is transformed from [N1, N2, C] to [C, N1, N2], a merge memory access mode is implemented during the transformation process. The thread block selects different movement strategies based on the innermost dimension after transpose, ensuring that the merged tensor is stored in the correct position, achieving a speedup of 1.70 times.

[0033] Optionally, in step S4, the specific details of multi-head parallel attention merging are as follows:

[0034] S41. Each thread block reads a data block with contiguous addresses in physical memory from multiple attention heads;

[0035] S42. Write the vector stored in the thread register back to the position with the same attention head, completing the entire head-dimension concatenation operation;

[0036] S43. By using vectorized memory access, threads in a thread bundle can access multiple consecutive elements at once, ensuring the maximum bandwidth utilization of global memory.

[0037] By applying the multi-head parallel attention merging strategy to the triangular attention operator, the above technical solution provides the input conditions for the subsequent attention calculation. The entire data movement process follows a merged memory access pattern, thereby improving memory bandwidth utilization and achieving a speedup of 3.05 times.

[0038] Optionally, in step S31, each thread is configured to read C consecutive elements along the input dimension, the size of C depending on the element type and cache line size, and the elements are stored in registers before being written back to SRAM and HBM.

[0039] Optionally, in step S31, each thread block additionally moves a K tensor with a size of [B, N2, N1, D] in addition to moving the Q tensor with a size of [B, N1, N2, D].

[0040] After the above technical solutions are adopted, the present application has the following advantages:

[0041] The present application solves the problem that the existing attention operator acceleration technology does not support triangular attention operator optimization. The existing acceleration technology, such as online attention technology, does not support such variant attention in a molecular docking model. Therefore, the present application aims to customize optimization strategies for such attention variants, and after the S1 and S2 technologies are adopted, the model peak memory can be effectively reduced, and the maximum data batch during model inference can be greatly improved.

[0042] The present application solves the problem of inefficient data movement of high-dimensional tensors. The existing tensor shape change operator cannot fully utilize the memory bandwidth of the GPU when performing high-dimensional tensor transposition, which seriously reduces the memory access efficiency during tensor transposition. By adopting the S3 technology in the present application, the memory access performance can be greatly improved by merging the memory resources of the device during high-dimensional tensor transposition.

[0043] The present application solves the problem of redundant memory access caused by the splicing operation of QKV tensors in the attention operator. By using the S4 technology in the present application, multiple groups of QKV tensors are successfully spliced along the hidden layer dimension, adapting the input form for subsequent attention calculation. BRIEF DESCRIPTION OF DRAWINGS

[0044] In order to more clearly illustrate the technical solutions in the embodiments of the present application or the prior art, the following will briefly introduce the drawings needed to be used in the embodiments or prior art description. Obviously, the drawings in the following description are only some embodiments of the present application, and for those skilled in the art, other drawings can also be obtained without creative labor.

[0045] Figure 1 The overall architecture of the GAABind model is shown in the figure.

[0046] Figure 2 The actual running time of each module in the encoding module when the input pair sequence length is 320x320 is shown in the figure.

[0047] Figure 3 The schematic diagram of the dynamic memory reduction strategy is shown in the figure.

[0048] Figure 4 The schematic diagram of the kernel fusion of MolTransformer is shown in the figure.

[0049] Figure 5 Tensor slice based data movement strategy

[0050] Figure 6 QK tensor transformation process with multi-head attention merge strategy

[0051] Figure 7 Performance of GAABind and FastBind with different configurations

[0052] Figure 8 Performance comparison of triangle attention and accelerated triangle attention

[0053] Figure 9 Performance comparison of tensor slice based data movement optimization

[0054] Figure 10 Performance comparison of multi-head parallel attention merge strategy DETAILED DESCRIPTION

[0055] The technical solutions in the embodiments of the present application will be clearly and completely described below with reference to the drawings in the embodiments of the present application. Obviously, the described embodiments are only part of the embodiments of the present application, rather than all the embodiments. Based on the embodiments in the present application, all other embodiments obtained by those skilled in the art without creative work fall within the scope of protection of the present application.

[0056] The application discloses an efficient molecular docking method based on a geometric perception attention model.

[0057] Embodiments

[0058] According to Figures 1 to 10 As shown in the figure, the efficient molecular docking method based on the geometric perception attention model includes the following contents.

[0059] 1. Related background

[0060] 1.1 GAABind model

[0061] GAABind is a Transformer-based model that takes a protein binding pocket with a 3D structure and a ligand in an unbound conformation as input, and outputs the protein pocket-ligand binding pose and binding affinity. The encoding blocks corresponding to the binding pocket and the ligand have the same structure but different weights, which is a common and effective design method in the field of molecular docking. This use of the same network structure for protein and ligand conformational learning helps to capture common features and interaction patterns between the two, thereby improving the accuracy of molecular docking.

[0062] The overall architecture of GAABind is as shown in Figure 1As shown, it includes four parts: the embedding layer, MolTransformer, the interaction module, and the prediction module. First, the embedding layer part initializes the extracted features of the input molecules to obtain the atom and pair embeddings of each molecule. Then, MolTransformer is used to update the obtained atom and pair embeddings. Finally, the interaction module iteratively updates the molecular interactions between the complex pair and the ligand pair embeddings, and the prediction module is used to predict the complex pair distance and the ligand pair distance, and directly generate the 3D coordinates of the ligand.

[0063] Figure 1 The GAABind overall architecture. The protein and ligand obtain pair embeddings through the embedding layer. Then, they enter the MolTransformer to continuously learn and update the pair embeddings. Finally, they pass through the interaction and prediction modules in turn to generate the protein-ligand binding affinity and the binding pose with the affinity. Among them, Figure 1 (a) Encoding module: MolTransformer is the main component of the model. (b) Triangle attention: used to enhance the mutual modeling between different atoms, and the hidden layer dimension of the representation is set to (D = 128). (c) Triangle update: realizes the update of the relationship between molecular information in high-dimensional space. (d) Conversion: further processes the pair representation to improve the accuracy of the model in predicting the ligand pose.

[0064] 1.2 Triangle inequality block

[0065] AlphaFold2 first realized the interaction of protein amino acid pair information using the triangle inequality relationship, which requires that the length of each edge {i,j} of d{i,j} be less than the sum of the lengths of edges {i,k} and {k,j}, i.e. d{i,j} < max(d{i,k}+d{k,j}). AlphaFold2 defines two triangle inequality blocks: triangle self-attention and triangle multiplication update, which are widely used in modern molecular docking models. GAABind optimizes the above two triangle inequality blocks and proposes triangle attention and triangle multiplication update, which perform well in the end-to-end molecular docking task.

[0066] Transformer is a sequence-to-sequence NLP model, and its core is self-attention, which uses attention mechanisms to describe the global dependency between input and output. This technology is widely used in drug discovery, computer vision, image processing, and recommendation systems, among other fields. For example, Figure 1As shown in (c), triangular attention differs from standard self-attention in three ways: First, in terms of input format, triangular attention uses a pair representation of the form [N, N] as input, instead of the common one-dimensional sequence of length N. Second, in addition to the three linear layers used to generate the Q, K, and V tensors, the pair representation needs to pass through an additional linear layer to obtain a bias term for calculating the attention score. Third, triangular attention combines AlphaFold2's "starting node attention" and "ending node attention," meaning that after calculating the dot product of Q and K, the two attention scores are added together before performing softmax normalization and multiplication with V. This mechanism of updating the pair representation using the triangular relationship between different sites in the pair enhances the modeling of interatomic interactions while ensuring a physically reasonable geometric relationship between the protein and ligand.

[0067] Initially, using triangular multiplication updates with outgoing / incoming edges was designed as a cost-effective alternative to triangular self-attention blocks. When combined with these two triangular inequality blocks, the model was able to produce protein structures with high accuracy. Figure 1 As shown in (b), the triangle update module combines the "triangle multiplication update with incoming edges" and "triangle multiplication update with outgoing edges" from AlphFold2. Specifically, it selects two rows in the representation, embeds each pair of edges with the same index in the two rows, maps them to a low-dimensional space, and applies a gating mechanism. Subsequently, it calculates the Haldman product for the generated row and column vectors respectively, and applies gating operations again, finally completing the update of the target edge representation. After the representation is processed by the triangle update module, the update of 3D molecular information relationships in high-dimensional space is realized, making the flow of information richer.

[0068] 2. In-depth analysis of MolTransformer

[0069] This invention conducts a detailed performance bottleneck analysis of MolTransformer, revealing that the key factor affecting model performance lies in the pair encoding module. To distinguish the sequence dimensions involved in the pair representation, this invention represents the input pair sequence as (Nr, Nc), where Nr = Nc = N, denotes the hidden layer dimension of the pair representation as Pm, and the number of attention heads as Head. Experimental results show that when the sequence dimension is N = 320, the pair encoding module accounts for 89.2% of the entire GAABind inference model, severely impacting model performance.

[0070] 2.1 Performance Analysis

[0071] To further analyze the performance characteristics of MolTransformer, this invention categorizes its operators into four types based on computational and memory access characteristics: 1) General matrix multiplication: This category includes matrix multiplication, batch matrix multiplication, and other intensive matrix computations. 2) Batch reduction: This category includes computationally less intensive operations such as Softmax and LayerNorm. 3) Element-wise operations: This category includes element-wise addition, activation, and broadcast addition. 4) Pure data movement operations: This category includes tensor shape transformation and copying.

[0072] This invention tested the actual average latency of each module in the encoding module when the input pair sequence was 320×320. The results are as follows: Figure 2 As shown, during the forward computation of the encoding module, 15.6% of the time is spent on matrix multiplication operations, 15.5% on batch reduction, 48.0% on element-wise operations, and 20.9% on data movement. First, the high time spent on element-wise operations is due to memory-intensive operations such as adding bias terms, mask terms, and summing attention scores in the triangular attention mechanism. Second, irregular step accesses in some element-wise operators and additional data merging operations cause significant time spent on data movement. These data indicate that further optimization of element-wise and data movement operations is needed to achieve better inference performance.

[0073] 2.2 Memory Consumption

[0074] This invention also reveals that MolTransformer faces a problem of excessively high peak memory usage during actual inference, severely wasting valuable memory resources on computing devices. GAABind has a total parameter count of only 17.95M, but due to the presence of the attention operator, the activation amount during model forward propagation becomes abnormally large, resulting in GAABind requiring a significant amount of memory during training and inference. Specifically, when actually executing the triangular attention operator, this operator requires at least B*N memory on a GPU. 3 The memory requirement is *Head*(Float16), where B is the batch size, N is the sequence length, and Head is the number of attention heads. This high memory requirement severely hinders the increase of batch size and model layer depth.

[0075] However, existing memory optimization techniques do not support this type of attention variant. Even with the introduction of activation checkpointing mechanisms to recalculate intermediate activation results during forward propagation as needed, the memory explosion problem faced by the model cannot be solved. Therefore, a memory-friendly parallel strategy needs to be designed to reduce the peak memory usage during the execution of the triangular attention operator, thereby reducing GPU memory consumption and computational overhead, and fully leveraging the performance potential of the computing hardware.

[0076] 3. FastBind model

[0077] Four techniques used by FastBind will be introduced below. In order to solve the performance inefficiency problem caused by the key operations in triangular attention and triangular update, such as memory explosion, irregular stride access and additional data movement. This invention proposes dynamic memory reduction, hierarchical fusion, tensor slice-based data movement strategy and multi-head parallel attention merging techniques, which can significantly reduce the peak memory and prediction delay of the GAABind molecular docking model when applied to the inference process of the GAABind molecular docking model, making it more efficient in practical applications.

[0078] 3.1 Memory reduction optimization

[0079] 3.1.1 High memory consumption

[0080] Efficient use of device memory is a long-term challenge for Transformer models. For this reason, various long sequence techniques have been proposed to alleviate the high memory consumption problem of Transformer models. However, as the sequence length and network depth continue to grow, GPU device memory is often quickly exhausted, severely limiting the training and inference efficiency of the model. Triangular attention also faces similar memory explosion problems.

[0081] In actual operation, the model will often show excessively high peak memory usage, which significantly affects the supported batch size and memory utilization efficiency. Specifically, there are two reasons for its low memory efficiency: on the one hand, for a four-dimensional sequence with an input shape of [N1, N2, H, D], after performing attention multiplication calculations, an intermediate result with a shape of [N, H, N, N] is generated and stored in the device memory. Since the dimension of D is much smaller than N, the storage of the intermediate result occupies too much memory, causing the model to have abnormal peak memory usage during inference. On the other hand, before performing Softmax normalization, triangular attention needs to perform two addition operations, respectively, the addition of the bias term and the addition of the mask term. The tensors involved in these operations need to be matched to the shape of the attention score through broadcasting, and this broadcasting operation further exacerbates the memory bottleneck. However, most techniques for reducing memory occupation do not support such attention operators, especially when additional bias terms and mask terms are included in the attention, the performance of their kernels is poor.

[0082] To address these challenges, the present application designs a memory-friendly high-performance kernel for triangular attention, which greatly reduces its peak memory requirement when executed on GPU. Specifically, each thread block loads the required patches from Q and K for dot product, then places the generated patches in registers, and implements the addition of bias term and mask term in registers, and finally performs softmax normalization and V multiplication operation. Next, two techniques adopted are introduced in detail.

[0083] 3.1.2 Dynamic memory reduction

[0084] Figure 3 The dynamic memory reduction strategy is adopted. The figure shows the hardware location of the data involved in the attention calculation process, and the attention score tensor responsible for each thread, where stage 1 represents the data stored in global memory, stage 2 represents the data stored in shared memory, and stage 3 represents the data stored in registers.

[0085] If the entire attention operator is not considered to be fused together, the time-consuming bias term addition and mask term addition operations in triangular attention can be easily fused into the softmax calculation. Since such addition follows the merged memory access pattern, it can basically eliminate the memory access overhead caused by broadcasting. However, the strategy of simply fusing the bias term and mask term with softmax does not solve the memory bottleneck caused by preserving attention scores, and the peak memory of the entire attention part is still not relieved. Therefore, still taking the fusion of the entire triangular attention as the premise, the dynamic memory reduction operator is designed.

[0086] Figure 3 The attention calculation process after adopting dynamic memory reduction optimization is shown. First, the thread block is tiled along the sequence dimension, and each thread block is responsible for loading different block sizes of Q, K and V from global memory to shared memory, where Q represents the "question" that needs to focus on other issues, K represents the "keyword" of other positions, which is used to calculate the similarity with Q, and V represents the "information" actually carried by each position, and the attention weight is finally weighted to V to generate the final attention output. The block size of Q, K and V is (Br, 2D), (Bc, 2D) and (Bc, D) respectively. Since Q and K have undergone additional splicing operation before loading, their hidden layer dimension size becomes 2D, which means that a larger shared memory buffer needs to be set to ensure that each thread block can carry all the data at one time. In order to improve the memory access efficiency, the patch size is adjusted to (64, 64), which also avoids register overflow. Then, the thread block will iteratively calculate along the sequence direction, and update the final attention output matrix by means of local maximum vector m and local summation term d to realize the parallel calculation of softmax normalization operation, the specific formula is as follows;

[0087]

[0088] Where b represents the size of the tile, N represents the length of the input sequence, and x i Indicates storage Q[k]K T The i-th vector [(i-1)b:ib], m i x represents i The local maximum value, d′ i x represents i The local summation term, where V represents the V tensor involved in the attention calculation, with dimensions [B, N, N, D], o′ i This represents a portion of the output tensor computed in the i-th iteration. By iteratively updating the output terms in this way, all computations of the triangular attention can be fused into a single kernel. This avoids writing intermediate results back to global memory, significantly reducing the peak memory usage of this operator.

[0089] In addition, considering the limited on-chip memory resources of the GPU, the four thread bundles in the thread block are arranged according to... Figure 4 The layout is as shown. Each thread bundle is only responsible for computing a portion of the attention map, avoiding loading too much input data in a single iteration and thus reducing the burden of shared memory utilization. Meanwhile, using a thread bundle-level operation primitive such as butterfly communication, the reduction operation of vector m and the summation term d is implemented in the thread register of a single thread bundle, which reduces the synchronization overhead of synchronization operations to some extent and facilitates communication between different thread bundles using shared memory. Finally, the reduction operation of all output results is implemented in shared memory, and the reduced results are merged and written back to global memory.

[0090] Compared to reducing the computational complexity of the triangular attention operator, optimizing its memory performance can actually bring more speed-up benefits. This is determined by the performance bottlenecks faced by the triangular attention operator when it is actually executed on specific hardware. After adopting dynamic memory reduction technology, the intermediate results generated by attention calculation only move between shared memory and registers. Only the final calculated attention output matrix is ​​saved to global memory, avoiding memory shortages caused by retaining intermediate activations.

[0091] 3.1.3 Layered Integration Strategy

[0092] Figure 4 The kernel fusion of MolTransformer; where (a) is the molecular transformer code block before fusion; and (b) is the molecular transformer code block after fusion of transpose and attention.

[0093] Operator fusion improves the computing efficiency by reducing the number of memory accesses, reducing the kernel startup overhead and enhancing the cache locality, and is a common technique to improve the computing and memory access performance of deep learning models. Specifically, if the output generated by one thread block is used by another thread block, a global memory synchronization needs to be performed to call a new cuda kernel, which means that the thread needs to write the intermediate result back to the global memory first so that it can be correctly called by the next kernel. When the amount of memory occupied by the intermediate result is small, the device memory burden is less affected, and when the amount of memory occupied is large, this operation will bring huge memory consumption and thread synchronization overhead, greatly affecting the performance of the GPU hardware. Figure 4 As shown in (a) of FIG. 1, the triangular attention contains data transpose, matrix multiplication, addition and reduction operators, and the data dependency relationship between the thread blocks of these operators makes the triangular attention have certain challenges in fusion, among which the data transpose, matrix multiplication and addition operations are relatively easy to fuse, while the fusion of the softmax reduction operation is more difficult, which requires us to develop a memory-efficient fusion strategy.

[0094] The present application designs two fusion kernels for triangular attention, which are fusion transpose kernel and fusion triangular attention kernel. The former will use the multi-head parallel attention merging strategy mentioned in 3.2.3, splice Q or V along the hidden layer dimension D, and the spliced hidden layer dimension changes from D to 2D, and the output generated can be directly used for the next multi-head attention calculation, avoiding unnecessary data layout conversion. The latter will fuse matrix multiplication, bias addition, mask addition and softmax normalization operations into one kernel, and save all the data generated by each thread block in registers and shared memory, avoiding the moving overhead of data between global memories, greatly reducing the memory usage burden of the device.

[0095] 3.2 Data movement optimization

[0096] 3.2.1 Irregular step memory access

[0097] Efficient use of the GPU's compute and memory hierarchy requires coordinating multiple optimizations. For example, when using triangle updates and triangle attention to build a model, a specific change in tensor shape needs to be completed before performing a batched matrix multiplication operation, and then a cublas or cutlass operator library is called to implement efficient multiplication of matrices. However, this cross-dimension data movement operation is often accompanied by a large amount of non-coalesced memory access, resulting in inefficient GPU memory access. Specifically, there are two main reasons for its inefficient memory access: first, when performing the transpose operation on a high-dimensional tensor, each thread needs to read the data to be transposed from the device memory first, and then write the read data back to the device memory according to the specified memory layout. The coalesced memory access requires the thread bundle to access consecutive elements in one instruction, which causes the thread bundle to fail to meet the coalesced memory access of reading and writing back. Therefore, the tensor shape change kernel needs to be redesigned to meet the coalesced access mode of data reading and writing. Second, when writing the kernel that fuses the entire triangle attention, the transpose and splicing operations of Q, K and V need to be completed first. The transpose operation refers to changing the shape of the Q, K and V tensors from [B, N1, N2, H, D] to [B, N1, H, N2, D], that is, moving the head dimension forward. The splicing operation refers to splicing multiple Q, K after transposition along the hidden layer dimension, that is, splicing two tensors with shapes [B, N, H, N, D] into [B, N, H, N, 2D]. This additional splicing operation causes frequent access to device memory. Therefore, an efficient shape change kernel needs to be designed for Q, K tensors to eliminate the additional data movement operation.

[0098] To address these challenges, the present application proposes two data movement techniques: 1) an efficient data movement strategy based on tensor slices. By storing data along the continuous dimension when writing back to physical memory, the goal of coalesced memory access for both reading and writing is met. 2) Multi-head asynchronous attention coalescing. This optimization aims to eliminate the additional data movement generated by the coalescing of multiple Q, K, while preparing the input for the later fusion of the entire attention.

[0099] 3.2.2 Tensor tiled data movement

[0100] The present application designs a data movement kernel based on tensor slices to achieve maximum memory bandwidth utilization for reading and writing on the GPU. Figure 5 The flow of data loaded by the thread in different storage locations on the GPU is shown, including high-bandwidth memory with large memory capacity but access speed only better than PCle, and on-chip memory resources with small capacity but fast access speed. Figure 5 The data movement strategy is based on tensor slices. Each thread block uses shared memory to perform high-dimensional device and splicing operations on data from different storage areas, and different colored matrix blocks represent c consecutive stored elements.

[0101] When the tensor shape is converted from [N1, N2, C] to [C, N1, N2], the whole moving process is as follows: first, it is necessary to judge which dimension of the transposed tensor stores physically continuous data. The purpose of this step is to determine the data that each thread block needs to process. Specifically, each thread block loads a tensor tile of size [tilex, tiley] along the D dimension and the N2 dimension. Such thread mapping conforms to the global memory coalesced access pattern, thereby effectively improving data access efficiency. Subsequently, the transposition operation of the tensor tile is realized by means of a shared memory buffer, and the mapping relationship between threads and data is reorganized. Finally, each thread block loads the transposed data from the shared memory and stores it along the N2 and D dimensions to ensure that the write-back to the global memory also follows the coalesced access pattern. In addition, considering that each thread in the thread bundle cannot completely eliminate the entire cache line when performing coalesced memory access on a single element, each thread is set to read C consecutive elements along the input dimension (see Figure 5 ), the size of C depends on the element type and the cache line size, and the elements are stored in registers before being written back to the static random access memory and high bandwidth memory.

[0102] In order to further adapt the operator fusion process of triangular update, the tensor tile moving strategy is extended. Specifically, each thread block needs to move not only the Q tensor of dimension [B, N1, N2, D], but also the K tensor of dimension [B, N2, N1, D], which requires the allocation of shared memory and registers that are twice the original size. In addition, different moving strategies are selected according to the innermost dimension after transposition to ensure that the coalesced tensor is stored in the correct position.

[0103] 3.2.3 Multi-head parallel attention coalescing

[0104] In order to ensure the smooth execution of multi-head attention matrix multiplication, the standard self-attention needs to be advanced before calculating the attention score. This means that the shapes of Q, K and V need to be converted from [B, N1, N2, H, D] to [B, N1, H, N2, D] before batch matrix multiplication can be performed. Unlike self-attention, triangular attention contains multiple Q, K and V, which requires the data layout of these tensors to be re-optimized to prepare data for subsequent attention fusion operations.

[0105] Figure 6The process of multi-head attention merging is demonstrated, and data blocks with the same attention head position are spliced together. Specifically, each thread block first reads a block of data blocks with consecutive addresses in physical memory from multiple attention heads. Since the head size dimension is not involved in transposition, the transposition operation of Q and K can be completed without the help of shared memory. Subsequently, the vectors stored in the thread registers are written back to the positions with the same attention head, thereby completing the entire head-dimension splicing operation. The entire data movement process follows the merging memory access mode, thereby improving the memory bandwidth utilization. In order to achieve the maximum bandwidth utilization of the global memory, the threads in the thread bundle move multiple elements at a time, that is, data movement is performed through vectorized access. In addition, another purpose of splicing multiple Q and K is to facilitate the subsequent calculation of attention scores. By changing the arrangement of Q and K in memory, conditions are provided for the fusion of the triangular attention operator in 3.1.3.

[0106] All the above optimizations are successfully integrated into FastBind, and the optimized model is directly used for molecular docking tasks. FastBind can accelerate molecular docking while still maintaining high-precision prediction results.

[0107] 4. Evaluation

[0108] 4.1 Experimental configuration

[0109] The proposed FastBind optimization strategy is evaluated on a platform equipped with Inter(R) Core(TM) i7-12700F and GeForce RTX 4090, and the software configuration is Pytorch 1.12.0, CUDA 11.5 and GCC 11.4.0. The public GAABind implementation is used as the benchmark, the model includes 10 encoding modules (number of attention heads Head = 4, embedding layer dimension D = 128), with a total of 17.59M parameters, and FastBind performance is tested under various scale configurations.

[0110] The data sets used in the experiment are consistent with GAABind, including PDBbind v2020 and CASF2016, which are widely used for training and testing of molecular docking models. The data set follows the data set filtering and splitting strategy proposed by Uni-Mol, which involves excluding all complexes with protein sequence similarity higher than 40% and ligand fingerprint similarity higher than 80% from the test data set. The training set contains 16563 complexes, and the validation set contains 1841 complexes. In addition, the test data set CASF2016 is composed of 285 protein complexes with high-quality crystal structures and reliable binding affinity measurements.

[0111] 4.2 End-to-end molecular docking performance

[0112] In this experiment, the performance of the proposed FastBind molecular docking model was evaluated, and the results are as follows: Figure 7 As shown, GAABind represents the baseline model without optimization strategy, FastBind represents the model accelerated by optimization strategy, and speedup ratio represents the speedup ratio achieved by the implementation with optimization strategy relative to the baseline implementation. Figure 7 The performance of GAABind and FastBind with different configurations is shown in the left graph, which shows the performance as the pair length changes when the batch size is 4, and the right graph shows the performance as the batch size changes when the pair length is 256.

[0113] from Figure 7 It can be seen that the execution time of molecular docking increases proportionally with the increase in pair length or input batch size. When the batch size is fixed (4), the speedup effect of FastBind becomes more significant with the increase in pair length. This trend is because the performance of the GAABind model is very sensitive to changes in pair length. When the pair length is large (320×320), the triangular attention module in the model takes approximately 35.6ms, accounting for 78% of the total inference time, significantly affecting the overall end-to-end molecular docking speed. Therefore, FastBind's speedup effect on this module is particularly significant for improving end-to-end performance. It can be seen that when the pair length is 384×384, FastBind achieves approximately 1.54 times the speedup in the entire end-to-end molecular docking task. When the pair length is fixed (256×256), FastBind achieves a relatively stable speedup. With a data batch size of 16, end-to-end molecular docking achieves approximately 1.41 times the speedup. Specifically, when the data batch size reaches 32, GAABind will experience insufficient memory, causing it to be unable to perform molecular docking tasks normally, while FastBind can still run normally and support larger data batches.

[0114] To further verify FastBind's characteristics as a memory-friendly model and the performance improvement brought by dynamic memory reduction on specific hardware, this invention tested the maximum data batch size supported by the model before and after optimization, and the results are shown in Table 1. When the pair length is small (128×128), the maximum data batch size increased from 145 to 255. When the pair length is large (384×384), the maximum data batch size supported by the model increased from 7 to 27. This result further demonstrates that FastBind can significantly reduce the peak memory usage of the model, thereby allowing the use of larger data batches in molecular docking tasks and greatly improving molecular docking efficiency.

[0115] Table 1 Maximum data batch allowed to execute on RTX 4090 by GAABind and FastBind with different pair length configurations

[0116]

[0117] 4.3 Kernel optimization performance

[0118] 4.3.1 Dynamic memory reduction performance

[0119] In this experiment, the performance of the proposed dynamic memory reduction strategy for accelerating triangular attention is verified, and the results are shown in Figure 8 , where Mol represents the triangular attention built using only PyTorch without any optimization, and SpMol is the attention operator after adopting the proposed dynamic memory reduction strategy.

[0120] Figure 8 The left graph shows the performance of SpMol with dynamic memory reduction under different pair lengths when the data batch is fixed at 16. When the pair length is small (less than 256x256), the speedup effect of dynamic memory reduction improves significantly as the pair length increases. When the pair length exceeds 256x256, the speedup effect stabilizes, which is due to the increase in input size, which requires more thread blocks to be started during kernel execution, resulting in increased reduction overhead between thread blocks. When the pair length is 384x384, SpMol can achieve an acceleration of about 2.60 times.

[0121] Figure 8 The right graph shows the performance of the dynamic memory reduction strategy under different batch sizes when the pair length is fixed at 256. When the data batch is less than 16, SpMol can consistently achieve an acceleration of about 2.83 times, further proving that the dynamic memory reduction strategy designed in this application can maintain stable speedup effects under different data batch sizes. It is worth noting that when the data batch reaches 32, the model using triangular attention will encounter a memory shortage problem. However, accelerating triangular attention avoids storing intermediate results during attention calculation, significantly reducing the peak memory requirement of attention operations. This optimization allows this application to use larger batches in inference tasks, thereby improving overall performance.

[0122] 4.3.2 Tensor tiling data movement optimization

[0123] This experiment verifies the performance of the data movement strategy based on tensor slices proposed in 3.2.2, and the test results are as follows Figure 9Pytorch DM (Pytorch Data Move) represents the Pytorch version using the default data move method, and Tiled DM represents the data move operation based on the tensor tile.

[0124] Figure 9 In the left graph, the performance of the triangular update module using the tensor tile data move optimization under different pair lengths is shown when the data batch is fixed (4). When the pair length is less than 320x320, the Tiled DM using the tensor tile data move can achieve an acceleration of about 1.60 to 1.75 times. When the pair length reaches 384x384, the acceleration of Tiled DM decreases, because as the pair length increases, the kernel starts too many thread blocks, resulting in a gradual decrease in the parallel efficiency of data move.

[0125] Figure 9 In the right graph, the speedup of the proposed optimization strategy under different data batches is shown when the pair length is fixed (384x384). Under different data batches, Tiled DM can stably achieve an acceleration of about 1.57 times. When the data batch is 4, Tiled DM can achieve an acceleration of about 1.70 times.

[0126] 4.3.3 Multi-head parallel attention merging

[0127] In this experiment, the performance of the multi-head parallel attention merging optimization proposed in section 3.2.3 is verified, and the results are shown in Figure 10 PyTorch DM (PyTorch Data Move) represents the benchmark data transpose operation, and M-Head DM represents the multi-head asynchronous data move operation for QKV, including QKV data transpose and splicing operations.

[0128] Figure 10 In the left graph, the speedup of the multi-head asynchronous data move strategy designed by the application is shown when the data batch is fixed (4) and different pair lengths are faced. It can be seen that when the pair length is small (128x128) or large (384x384), the optimization strategy proposed by the application can achieve an acceleration of about 3.0 times and 3.2 times, respectively.

[0129] Figure 10In the middle, the right figure shows the performance of the multi-head asynchronous data movement strategy with different data batch sizes when the length is fixed (256x256). When the data batch size is less than 16, the acceleration effect is basically stable. This is because when the shape of QKV is changed and multi-head splicing is performed, each thread block only carries data in the N, H and D dimensions. When the hardware resource occupancy rate is not saturated, the data carrying efficiency is less affected by the data batch size. Therefore, the multi-head parallel data movement strategy can achieve stable acceleration. When the data batch size is 4, an acceleration of about 2.56 times can be achieved.

[0130] 4.4 Verification

[0131] From the theoretical analysis point of view, any kernel optimization and parallel strategy will not change the calculation result. However, since FastBind introduces a custom CUDA kernel, different calculation methods may cause errors in some index accuracy. Therefore, the present application compares the numerical correctness of GAABind and FastBind on the CASF2016 dataset. The present application uses GAABind and FastBind to predict the same protein-ligand complex, and evaluates the performance of the two models in the binding pose prediction and binding affinity prediction tasks to verify the effectiveness of FastBind.

[0132] The experimental results are shown in Table 2. Among them, the present application uses the ligand root mean square deviation (RMSD, the smaller the better) of atomic coordinates to compare the predicted binding source with the ground truth, and the predefined RMSD threshold and the predicted RMSD percentile are used as the predicted percentage of evaluation indicators. In order to evaluate the protein-ligand binding affinity prediction, the present application calculates MAE, RMSE and Pearson correlation coefficient. The experimental results show that FastBind only has a slight decrease in accuracy in the binding pose prediction task, and still maintains high prediction ability in the binding affinity prediction task, which shows that FastBind can predict the same quality of protein-ligand complex as GAABind.

[0133] Table 2 Performance comparison of GAABind / FastBind in predicting binding pose and binding affinity on CASF2016 dataset

[0134]

[0135] The application accelerates the GAABind molecular docking model based on the triangle attention mechanism, analyzes the performance bottleneck existing in the benchmark reality, proposes a series of optimization and acceleration strategies, and builds an efficient FastBind model for executing molecular docking tasks in real-world scenarios. FastBind successfully solves the high peak memory problem of triangle attention by using a dynamic memory reduction strategy, so that larger data batches can be set on different hardware. When the batch size is 4 and the length is 384*384, the entire FastBind molecular docking model can achieve about 1.54 times acceleration effect. In addition, FastBind can be applied to any other molecular docking model with triangle attention, and the molecular docking speed is improved.

[0136] The above is only a preferred embodiment of the present application and is not intended to limit the present application. Any modification, equivalent replacement, improvement, etc. within the technical solution of the present application shall be included in the protection scope of the present application.

Claims

1. An efficient molecular docking method based on a geometric perception attention model, characterized in that, It includes the following parts: S1. Dynamic memory reduction strategy: It uses iterative evaluation and block technology to achieve efficient parallel computation of the attention matrix, and uses shared memory to complete the communication of intermediate results between different thread bundles; S2, Layered Fusion Strategy: Two high-performance fusion kernels are designed for triangular attention, namely the fusion transpose kernel and the fusion triangular attention kernel, so that all the data generated by each thread block is stored in registers and shared memory; In step S2, the specific details of the layered fusion strategy are as follows: S21. By adjusting the execution order of each operator in the triangular attention, the transpose operation of Q and K is placed before the batch matrix multiplication operation. Two bias terms are merged into one bias term by addition, and two mask terms are also merged into one mask term by addition, thereby realizing the union of multiple computational operations. The union of computational operations refers to tensor transpose, matrix multiplication and addition, and the addition of bias mask terms. S22. The fusion transpose kernel concatenates multiple Q or K pairs along the hidden layer dimension D of the thread block. Each thread is responsible for moving the data in the same head position of two Q or K pairs to global memory, so that the two data areas that were originally in non-contiguous storage locations are now stored in a contiguous memory region. The hidden layer dimension after concatenation is expanded from D to 2D and is directly used for multi-head attention calculation. Here, Q represents the "question" that needs to be paid attention to other issues, and K represents the "keywords" in other positions, which are used to calculate similarity with Q. S23. The fusion triangle attention kernel integrates matrix multiplication, bias term addition, mask term addition, and Softmax normalization operator into one kernel. The intermediate results generated by each thread block are all stored in registers and shared memory. S3, Tensor Tiling Data Movement: For the triangle update operator, the mapping relationship between threads and data is reorganized through cross-dimensional data movement to achieve a global merge memory access mode; In step S3, the specific details of the tensor tiling data movement are as follows: S31. Determine the dimension of the data stored contiguously in the physical memory of the tensor, and load tensor slice data according to different dimensions. Each thread block loads a tensor slice of size [tilex, tiley] along the D dimension and N2 dimension. S32. Tensor slice transpose operation is implemented by using a shared memory buffer. Each thread is responsible for writing four consecutive FP16 type data in memory into shared memory row by row. The entire thread block writes a tensor slice of size [64,64] into shared memory. Then each thread accesses the four FP16 data in shared memory column by column and writes them into global memory row by row, thereby realizing the tensor slice transpose operation. S33. Each thread block loads the transposed data from shared memory and stores it along the dimensions of N2 and D to ensure that the merged access mode is followed when writing back to global memory. S4. Multi-head parallel attention merging: Optimize QKV data layout using multi-head attention merging strategy to reduce additional data transposition and splicing operations in intermediate results. Here, V represents the "information" actually carried at each position, and the attention weights are finally weighted onto V to generate the final attention output.

2. The efficient molecular docking method based on a geometrical awareness attention model according to claim 1, characterized in that, In step S1, the specific details of the dynamic memory reduction strategy are as follows: S11. Lay out thread blocks along the sequence dimension and allocate shared memory buffers according to the size of the hidden layer dimension. Each thread block loads Q, K, and V of different sizes from global memory into shared memory, where the block sizes of Q, K, and V are (Br, 2D), (Bc, 2D), and (Bc, D), respectively. S12. Iterate the calculation of the thread block along the sequence direction, and continuously update the attention output matrix with the help of the local maximum vector and the local summation term to realize the parallel calculation of the softmax normalization operation. The specific formula is as follows. in, Indicates the size of the tiled sheet. Indicates the length of the input sequence. Indicates storage The i-th vector , express Local maximum value, express The local summation term, Let V represent the V tensor involved in the attention calculation, with dimensions [B, N, N, D]. This represents a portion of the output tensor calculated in the i-th iteration; S13. Optimize the layout of thread bundles in the thread block. Each thread bundle is responsible for loading different data blocks and performing reduction operations on all output results in shared memory. Finally, merge the output results and write them back to global memory.

3. The efficient molecular docking method based on a geometrical awareness attention model according to claim 1, characterized in that, In step S4, the specific details of multi-head parallel attention merging are as follows: S41. Each thread block reads a data block with contiguous addresses in physical memory from multiple attention heads; S42. Write the vector stored in the thread register back to the position with the same attention head, completing the entire head-dimension concatenation operation; S43. By using vectorized memory access, threads in a thread bundle can access multiple consecutive elements at once, ensuring the maximum bandwidth utilization of global memory.

4. The efficient molecular docking method based on a geometrical awareness attention model according to claim 1, characterized in that, In step S31, each thread is configured to read C consecutive elements along the input dimension, the size of C depending on the element type and cache line size, and the elements are stored in registers before being written back to SRAM and HBM.

5. The efficient molecular docking method based on a geometrical awareness attention model according to claim 1, characterized in that, In step S31, in addition to moving the Q tensor of size [B, N1, N2, D], each thread block also needs to move an additional K tensor of size [B, N2, N1, D].

Images