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. The AEGNN leverages multi-head self-attention to jointly update edge features, atomic coordinates, and node features within a molecular graph. Subject to strict rotational and translational equivariance, it progressively reconstructs atom types, three-dimensional structures, and chemical bonds via a reverse diffusion process. The RBNN further improves generative performance through co-training two submodels, both built by stacking identical basic modules (denoted as AEM in this paper). The Knowledge Generator (KG) adopts a deeper layer stack to capture complex nonlinear relationships, producing high-precision initial predictions that align with the ground truth. By contrast, the Residual Refiner (RR) adopts a shallower architecture, focusing on fitting the residual between the KG’s output and the true target values. This design cuts computational overhead while boosting the model’s error correction capacity. The two modules are connected in a cascade structure: the output of the KG acts as the input to the RR, and the final prediction is obtained by adding the residual correction from the RR to the initial output of the KG. Experiments conducted on the QM9 dataset show substantial improvements in the uniqueness and novelty of generated molecules, demonstrating that the AEGDM can efficiently explore chemical space and aid the discovery of structurally innovative candidate molecules. Moreover, the RBNN mechanism speeds up experimental iteration and elevates overall model performance, offering critical technical support for the iterative optimization of molecular generation frameworks.