JTreeformer: Graph-Transformer via Latent-Diffusion Model for Molecular Generation

TL;DR

JTreeformer addresses the challenge of generating chemically valid and diverse molecules by transforming graph generation into junction-tree sequence generation and integrating a diffusion process in latent space. The model fuses a GCN-augmented encoder with a DAGCN-based decoder and trains with a reconstruction-plus-auxiliary loss, while a DDIM latent-diffusion model improves sampling efficiency and quality. Empirical results on MOSES and QM9 show state-of-the-art performance, with notable improvements in validity, novelty, and molecular diversity, and ablations confirm the importance of DAGCN and feature design. The approach offers a scalable, controllable platform for molecular discovery with potential impact on drug design and materials science.

Abstract

The discovery of new molecules based on the original chemical molecule distributions is of great importance in medicine. The graph transformer, with its advantages of high performance and scalability compared to traditional graph networks, has been widely explored in recent research for applications of graph structures. However, current transformer-based graph decoders struggle to effectively utilize graph information, which limits their capacity to leverage only sequences of nodes rather than the complex topological structures of molecule graphs. This paper focuses on building a graph transformer-based framework for molecular generation, which we call \textbf{JTreeformer} as it transforms graph generation into junction tree generation. It combines GCN parallel with multi-head attention as the encoder. It integrates a directed acyclic GCN into a graph-based Transformer to serve as a decoder, which can iteratively synthesize the entire molecule by leveraging information from the partially constructed molecular structure at each step. In addition, a diffusion model is inserted in the latent space generated by the encoder, to enhance the efficiency and effectiveness of sampling further. The empirical results demonstrate that our novel framework outperforms existing molecule generation methods, thus offering a promising tool to advance drug discovery (https://anonymous.4open.science/r/JTreeformer-C74C).

Figures (7)

• Figure 1: Architecture of JTreeformer,which involves an encoder-decoder structure with attention mechanisms and graph convolutions (DAGCN and mask-attention for decoder). In the latent space of the molecule generated by the encoder, the latent diffusion process is utilized to sample and generate molecular structures iteratively.
• Figure 2: An example of prediction in junction node position. F1 denotes the father of the precursor, while F2(F) denotes the father of the current node. Case 0 indicates that the father and precursor of the current node are the same node. Case 1 indicates that precursor and current share the same father node, while Case 2 and 3 are two examples to illustrate the special relationship between current and its father F2.
• Figure 3: Sampling by pre trained DDIM model in Latent space generated by JTreeformer in different diffusion steps of 50, 100, 200, 1000 of removing noise from the original samplings.
• Figure 4: The leftmost one is the input molecule, followed by their potential spatial neighbors, which can be seen to have similar scaffolding and functional groups to the input molecule.
• Figure 5: Latent space generated by JTreeformer. Different colour is utilized to represent properties including the Weight, TPSA, logP of molecules. A synthesized figure is displayed first, which takes the value of Weight, TPSA, logP as RGB separately.