Complete molecular generation based on attention-equivariant geometric diffusion models

This paper proposes an attention-equivariant diffusion model (comprising AEGNN and RBNN) for molecular generation tasks. AEGNN leverages multi-head self-attention to jointly update edge features, atomic coordinates, and node features within the molecular graph. Under strict rotational and translational equivariance, it progressively reconstructs atom types, 3D structures, and chemical bonds through a reverse diffusion process. RBNN further enhances generative performance by co-training two models, both constructed by stacking identical basic modules (AEM in this paper). The Knowledge Generator (KG) employs a deeper stack to capture complex nonlinear relationships, generating high-precision initial results that fit the ground truth. In contrast, the Residual Refiner (RR) uses a shallower structure, focusing on fitting the residual between KG's output and the true result. This design reduces computational overhead while strengthening the model's correction capability. The two models are connected in a cascaded manner: the output of KG serves as input to RR, and the final prediction is obtained by adding RR's residual correction to the initial output of KG. Experiments on the QM9 dataset demonstrate significant improvements in the uniqueness and novelty of generated molecules, indicating that AEGDM can effectively explore the chemical space and facilitate the discovery of candidate molecules with innovative structures. Furthermore, the RBNN mechanism accelerates experimental iteration and enhances model performance, providing crucial technical support for iterative optimization of molecular generation models.