Modeling Equivariant Neural Networks for Hardware Acceleration, a Case Study on the Molecular Docking Tool DiffDock

The equivariant neural network (ENN) exploits the symmetry of problems by enforcing equivariance with the neural network. ENN has a significant impact on various real-world scenarios, including molecular modeling, protein biology, and 3D point cloud. However, the intricate computational flow of ENN hampers its further advancement for applications that require high runtime speed. For instance, virtual screening, which is an early stage of computer-aided drug discovery, demands efficient screening (i.e., molecular docking) of large molecular databases to identify potential drug candidates. Among the available molecular docking tools, an SE(3) ENN named Diffdock stands out for its exceptional docking accuracy and serves as a prominent tool in the drug discovery research. Although customized hardware such as FPGA could accelerate such applications, the huge design effort is unaffordable. SystemC is an efficient solution to model the hardware designs with higher level abstraction. In this study, we propose the first SystemC model for equivariant neural networks, specifically targeting Diffdock, with the aim of accelerating its performance using customized hardware such as the FPGA. Experiments demonstrate that our model on DiffDock can achieve speedups of up to 3.20× and an average of 1.84× when compared to the GPU implementation.