CoAtGIN: Marrying Convolution and Attention for Graph-based Molecule Property Prediction

Xuan Zhang, Cheng Chen, Zhaoxu Meng, Zhenghe Yang, Haitao Jiang, Xuefeng Cui

Abstract: Molecule property prediction based on computational strategies plays a key role in the process of drug discovery and design processes, such as DFT. However, these traditional methods are time-consuming, labor-intensive, and cannot satisfy the need for biomedicine. Owning to the development of deep learning, there are many variants of Graph Neural Networks (GNN) for molecular representation learning. However, the existing well-performing graph-based methods that have a number of parameters or light models cannot achieve good grades on various tasks. To manage the trade-off between efficiency and performance, we propose a novel model architecture, CoAtGIN, using both Convolution and Attention. At the local level, k-hop convolution is designed to capture long-range neighbor information. At the global level, in addition to using the virtual node to pass identical messages, we utilize linear attention to the aggregate global graph representation according to the importance of each node and edge. In the recent Open Graph Benchmark (OGB) Large-Scale Benchmark, CoAtGIN achieves the 0.0901 Mean Absolute Error (MAE) on the large-scale dataset PCQM4Mv2 with only 6.4 M model parameters. Moreover, using the linear attention block improves the performance, which helps to capture the global representation.