[Paper Review] PrefixMol: Target- and Chemistry-aware Molecule Design via Prefix Embedding

Zhangyang Gao, Yuqi Hu+2 more|arXiv (Cornell University)|Feb 14, 2023

Computational Drug Discovery Methods4 citations

TL;DR

PrefixMol proposes a unified generative model for molecule design that uses learnable prefix embeddings to jointly control target binding affinity (via 3D protein pockets) and multiple chemical properties (e.g., QED, LogP, SA, Lipinski). By conditioning attention mechanisms on these prefix embeddings, the model achieves strong controllability in both single- and multi-condition generation, outperforming prior methods and revealing interdependencies between molecular properties through attention analysis.

ABSTRACT

Is there a unified model for generating molecules considering different conditions, such as binding pockets and chemical properties? Although target-aware generative models have made significant advances in drug design, they do not consider chemistry conditions and cannot guarantee the desired chemical properties. Unfortunately, merging the target-aware and chemical-aware models into a unified model to meet customized requirements may lead to the problem of negative transfer. Inspired by the success of multi-task learning in the NLP area, we use prefix embeddings to provide a novel generative model that considers both the targeted pocket's circumstances and a variety of chemical properties. All conditional information is represented as learnable features, which the generative model subsequently employs as a contextual prompt. Experiments show that our model exhibits good controllability in both single and multi-conditional molecular generation. The controllability enables us to outperform previous structure-based drug design methods. More interestingly, we open up the attention mechanism and reveal coupling relationships between conditions, providing guidance for multi-conditional molecule generation.

Motivation & Objective

• To address the challenge of generating molecules that satisfy both target-specific binding affinity and desired chemical properties in a unified framework.
• To overcome negative transfer in multi-task learning by using prefix embeddings as conditional prompts instead of joint training of separate models.
• To enable precise, controllable generation of molecules under multiple constraints, including protein-ligand binding affinity and key drug-like properties.
• To uncover hidden coupling relationships between molecular conditions through attention mechanism analysis.

Proposed method

• The model employs learnable prefix embeddings that are injected into the query and key matrices of the self-attention mechanism to condition generation on target and chemical conditions.
• Each condition—such as protein pocket (VINA score), QED, LogP, SA, and Lipinski rule compliance—is encoded as a separate prefix embedding via auxiliary networks.
• The prefix embeddings act as contextual prompts, influencing the attention distribution and guiding the autoregressive generation of SMILES sequences.
• The attention mechanism is analyzed via gradient-based perturbation to compute partial derivatives, revealing how changes in one condition affect attention weights across others.
• A relation matrix R is constructed by summing absolute partial derivatives across conditions, visualizing interdependencies between properties.
• The model is trained on an extended CrossDocked dataset enriched with molecular property labels, enabling joint optimization of multiple objectives.

Experimental results

Research questions

• RQ1 Can a unified deep generative model effectively control both target binding affinity and multiple chemical properties in molecule generation?
• RQ2 How can negative transfer in multi-condition molecular generation be mitigated when combining diverse constraints?
• RQ3 What are the underlying coupling relationships between different molecular properties during conditional generation?
• RQ4 To what extent does the attention mechanism reveal interdependencies between conditions in multi-conditional generation?

Key findings

• PrefixMol outperforms previous structure-based drug design methods in generating molecules with desired binding affinities (VINA scores) and chemical properties.
• The model achieves strong controllability in both single- and multi-condition generation, with property values closely matching user-specified targets.
• When jointly controlling QED and SA, the model achieves a Lipinski score of 5.0, indicating optimal compliance with drug-likeness rules.
• VINA binding affinity is most sensitive to changes in QED, LogP, and SA, not to its own value, explaining why multi-condition control outperforms single-condition optimization.
• The attention mechanism reveals that LogP and QED are the most correlated properties, with strong cross-controllability observed in the relation matrix.
• Visualization of generated molecules shows structural novelty compared to reference compounds, indicating the model generates truly new, high-affinity ligands rather than simple modifications.