Learning Multimodal Graph-to-Graph Translation for Molecular Optimization

Wengong Jin; Kevin Yang; Regina Barzilay; Tommi Jaakkola

Source-linked AI summary

Learning Multimodal Graph-to-Graph Translation for Molecular Optimization

Wengong Jin, Kevin Yang, Regina Barzilay, Tommi Jaakkola

arXiv:1812.01070v3Dec 3, 2018cs.LGcs.AIcs.NEstat.ML

TL;DR

Molecular optimization requires translating paired molecular graphs into better molecules despite multiple valid improvements for each input. The paper combines a junction tree encoder-decoder with latent codes and adversarial training, and reports better, more diverse molecules than baselines across optimization tasks.

Problem
Molecular optimization must learn diverse mappings from input molecules to improved target molecules because many valid modifications may exist for one input.
Method
The model uses a junction tree encoder-decoder, low-dimensional latent vectors for multimodal outputs, and adversarial training to align generated molecules with valid targets.
Results
The model outperforms previous state-of-the-art baselines across three molecular optimization tasks and translates inputs into diverse compounds.
Takeaways & Limitations
Combining variational junction-tree translation with adversarial training generates better and more diverse molecules than the evaluated baselines.
Takeaways & Limitations
Diversity calculations exclude molecules with no valid translation, separating translation accuracy from diversity measurement.

Abstract

from arXiv · show

We view molecular optimization as a graph-to-graph translation problem. The goal is to learn to map from one molecular graph to another with better properties based on an available corpus of paired molecules. Since molecules can be optimized in different ways, there are multiple viable translations for each input graph. A key challenge is therefore to model diverse translation outputs. Our primary contributions include a junction tree encoder-decoder for learning diverse graph translations along with a novel adversarial training method for aligning distributions of molecules. Diverse output distributions in our model are explicitly realized by low-dimensional latent vectors that modulate the translation process. We evaluate our model on multiple molecular optimization tasks and show that our model outperforms previous state-of-the-art baselines.

1 INTRODUCTION

The paper formulates molecular optimization as graph-to-graph translation from paired molecules, emphasizing diverse valid improvements. It introduces a junction tree encoder-decoder with latent codes and adversarial training, and reports superior performance across tasks.

Molecular optimization is framed as translating input molecular graphs into graphs with better chemical properties using paired molecular data.
Multiple strategies can improve a molecule, so the model must represent diverse translation outputs rather than a single transformation.
The proposed junction tree encoder-decoder uses stochastic latent codes to capture meaningful molecular variations during decoding.
Adversarial training aligns generated graph distributions with valid target molecules to reduce invalid translations.
The model is evaluated on three molecular optimization tasks and outperforms state-of-the-art graph-generation and MMPA baselines across different tasks.

2 RELATED WORK

Prior molecular optimization methods use generative modeling, reinforcement learning, graph representations, and rule-based matched molecular pair analysis. The paper positions neural graph translation and adversarial graph regularization as alternatives to explicit transformation rules.

Earlier molecular optimization approaches include generative modeling and reinforcement learning, with a shift from SMILES representations toward molecular graphs.
MMPA hard-codes matched molecular pairs into graph transformation rules, requiring potentially millions of rules to cover complex patterns.
The paper replaces explicit MMPA rules with neural networks that learn molecular transformations from data.
Its adversarial regularization over graphs is inspired by multimodal image translation and adversarial continuous-representation methods for text style transfer.

3 JUNCTION TREE ENCODER-DECODER

The junction tree encoder-decoder represents molecules with graph and scaffold structures, then decodes a junction tree before assembling a molecular graph. Attention, message passing, and candidate-attachment scoring support hierarchical and potentially diverse graph generation.

Tree and graph encoder: Molecules are represented as molecular graphs plus junction trees whose clusters are valid chemical substructures such as rings and bonds.
Tree and graph encoder: A unified graph message-passing encoder embeds both junction-tree and molecular-graph nodes into continuous vectors.
Junction tree decoder: The tree decoder expands the junction tree top-down, using recurrent message updates, attention over source encodings, and predictions for labels and topology.
Graph decoder: The graph decoder non-deterministically assembles neighboring clusters because one junction tree can correspond to multiple molecules.
Graph decoder: Candidate molecular attachments are scored with graph message passing, sum-pooled representations, and dot products against encoded source graph vectors.
Graph decoder: Training uses teacher forcing with ground-truth junction trees, while testing assembles one neighborhood at a time in decoded-tree order.

4 MULTIMODAL GRAPH-TO-GRAPH TRANSLATION

The model learns a multimodal mapping between molecule domains by combining latent-code variation with adversarial regularization of generated molecular structures. Its variational junction tree encoder-decoder reconstructs target molecules while adversarial training aligns generated outputs with the target domain and helps avoid invalid translations.

Problem: The task is many-to-many translation: a single source molecule can have multiple valid modifications in the target domain.The model must therefore generate a diverse set of outputs for each input molecule.
Latent variation: Low-dimensional latent vectors z explicitly encode multimodal output variation and are sampled from a standard Gaussian prior.The latent distribution is regularized toward the prior to support efficient sampling.
Latent variation: The latent code is derived from the ground-truth target molecule during training, and the decoder reconstructs that molecule from both its encoding and source molecule X.The posterior is inferred from paired molecules, with reconstruction and KL regularization forming the variational objective.
Adversarial scaffold regularization: Adversarial training makes generated molecular structures follow the target-domain distribution by training a discriminator against real target molecules.Because tree and graph decisions are discrete, adversarial regularization operates on continuous representations derived from decoder hidden states.
Adversarial scaffold regularization: Continuous decoder representations are constructed while incrementally expanding or backtracking through the decoded junction tree, using soft label distributions for translated samples.The discriminator is trained on representations of teacher-forced real trees and generated translated trees.

5 EXPERIMENTS

The experiments evaluate graph-to-graph molecular translation across three optimization tasks, measuring accuracy, diversity, and novelty. The proposed models outperform key baselines across tasks while producing diverse outputs, though novelty is lower than some baselines.

Experimental setup: The evaluation covers penalized logP, QED, and DRD2 optimization tasks using paired molecules subject to property-improvement and similarity constraints.Training and evaluation enforce molecular similarity to prevent arbitrary mappings that ignore the input molecule.
Evaluation metrics: Translation accuracy uses 20 latent-code decodings for penalized logP and success rates for QED and DRD2 target ranges.A QED or DRD2 translation succeeds when at least one candidate satisfies the similarity constraint and reaches the specified property range.
Results: Across all three tasks, the proposed models outperform MMPA by a large margin, while graph-to-graph translation performs better than VSeq2Seq.The authors attribute stronger performance against JT-VAE and GCPN to direct supervision from parallel data; adversarial training provides a slight improvement over VJTNN.
Scope: VJTNN+GAN is evaluated only on QED and DRD2 tasks with well-defined target domains explicitly constrained by property ranges.This scope boundary limits direct interpretation of the adversarially trained model on the penalized logP task.
Results: The models achieve higher diversity than MMPA on two tasks and than VSeq2Seq on three tasks, using average pairwise Tanimoto distance among valid translations.Diversity is bounded by the maximum distance allowed by the similarity constraint.
Results: On QED and DRD2, the models discover new compounds most of the time, but less frequently than MMPA and GCPN.Those baselines nevertheless have much lower translation success rates.

6 CONCLUSION

The paper concludes that combining a variational junction tree encoder-decoder with adversarial training generates better and more diverse molecules than the baselines. Diversity calculations exclude source molecules with no valid translation to separate diversity from translation accuracy.

6 CONCLUSION: Combining the variational junction tree encoder-decoder with adversarial training generates better and more diverse molecules than the baselines.This conclusion summarizes the evaluated graph-to-graph translation models for molecular optimization.
6 CONCLUSION: Diversity scores exclude failure cases with no valid translation so that diversity is not confounded with translation success rate.Otherwise, models with lower success rates would always have lower measured diversity.

A MODEL ARCHITECTURE

The model combines tree and graph encoders with attention-based decoding, augmenting graph message passing with tree messages to distinguish otherwise locally isomorphic attachments. Adversarial soft tree decoding enables discriminator gradients during training.

Tree and Graph Encoder: Tree encoding uses a tree GRU for stable iterative message passing, while graph encoding uses the same parameterization with separate parameters.Edge features are omitted in the tree encoder because they are always zero.
Tree Decoder Attention: Decoder attention uses a bilinear function between the decoder state and normalized source tree and graph vectors.
Graph Decoder: Graph decoding scores candidate attachments with a graph neural architecture while accounting for the subtree beneath each junction-tree neighbor.The architecture addresses cases where exchanging equally labeled neighboring clusters would otherwise produce locally isomorphic graphs.
Graph Decoder: Tree messages are incorporated into graph message passing to preserve junction-tree positions and avoid local isomorphism.The index αv marks atom positions in the junction tree and retrieves messages summarizing the subtree under i.
Adversarial Scaffold Regularization: Soft tree decoding replaces ground-truth inputs with predicted label distributions, allowing gradients to propagate from the discriminator.The recursive decoder predicts topology, creates children when pt ≥ 0.5, predicts node labels, and computes inter-node messages.

B EXPERIMENTAL DETAILS

The experiments specify model and adversarial-training hyperparameters, property-calculation tools, matching-pair baselines, and task-specific molecular-pair curation procedures.

Training Details: The model uses hidden state dimension 300, latent code dimension |z| = 8, six tree-encoder iterations, three graph-encoder iterations, and 3.9M parameters.The VSeq2Seq baseline uses one-layer bidirectional and unidirectional LSTMs for encoding and decoding.
Training Details: All models train with Adam for 20 epochs at learning rate 0.001, annealed by 0.9 each epoch.
Training Details: Adversarial training uses a three-layer discriminator with hidden dimension 300, LeakyReLU activations, N = 5 discriminator iterations, and gradient penalty weight β = 10.
Property Calculation: Penalized logP, QED, and DRD2 properties are computed using RDKit-based functionality, with the DRD2 predictor downloaded from the REINVENT repository.
MMPA Procedure: MMPA baselines use mmpdb transformation rules from ZINC, with DRD2 additionally using the Olivecrona dataset, and translate each test molecule K = 20 times.
Dataset Curation: Training molecular pairs are curated by Tanimoto similarity and task-specific property-improvement or source-target range criteria.The penalized logP task starts from 250K ZINC molecules; QED and DRD2 use sim(X, Y ) ≥ 0.4.