Technical Summary: Synergistic Benefits of Joint Molecule Generation and Property Prediction

Problem Formulation

The paper addresses a long-standing challenge in machine learning: developing models that simultaneously excel in generative and predictive tasks. In the domain of drug design, success relies on balancing the generation of novel molecules from unexplored chemical spaces with robust property prediction that can extrapolate to these new molecules. Currently, molecule generation and property prediction are predominantly approached in separation. While transformer-based models are state-of-the-art for both tasks individually, joint training poses significant architectural and optimization challenges. Combining generative and predictive parts into a single model often leads to over-regularization or gradient interference between the generative and predictive objectives, causing models to forgo the potential synergistic benefits of joint learning.

Methodology: Hyformer

To address these challenges, the authors propose Hyformer, a transformer-based joint model that unifies generative and predictive functionalities within a single set of shared parameters.

Architecture

Hyformer combines an autoregressive transformer decoder with a bidirectional transformer encoder. The core innovation lies in an alternating attention mechanism controlled by a task token $[TASK] \in \{[LM], [PRED], [MLM]\}$:

• Generative Mode ($[LM]$): The model acts as a decoder using a causal self-attention mask ($M_\rightarrow$), restricting each token to attend only to its left context.
• Predictive/Pre-training Mode ($[PRED], [MLM]$): The model acts as an encoder using a bidirectional self-attention mask ($M_\leftrightarrow$), allowing full context access.

This alternating scheme alleviates gradient interference typical in joint models. The Jacobian of the softmax function under a causal mask alters the gradients of the attention module differently than under a bidirectional mask, preventing the objectives from conflicting during updates.

Training Strategy

Hyformer employs a two-stage training process:

1. Joint Pre-training: The model is pre-trained on unlabeled data using a combined loss function:
   $$\ell_{Hyformer} = \ell_{LM} + \mu\ell_{MLM} + \eta\ell_{PRED}$$
   Here, $\ell_{LM}$ is the language modeling loss, $\ell_{MLM}$ is the masked language modeling loss, and $\ell_{PRED}$ is a predictive loss based on analytically computable molecular descriptors (e.g., molecular weight, hydrophobicity). The balance between tasks is controlled by a vector of task probabilities $p_{[TASK]}$.
2. Fine-tuning: The model is fine-tuned on labeled downstream tasks using a joint loss:
   $$\ell_{Hyformer} = \ell_{LM} + \lambda\ell_{pred}$$
   The model alternates between generative and predictive tasks during fine-tuning based on task probabilities, omitting the masked language modeling loss to focus on the downstream task while retaining generative capabilities.

Sampling

Hyformer supports both unconditional and conditional sampling:

• Unconditional: Samples molecules $x \sim p_\theta(x)$ using the autoregressive part.
• Conditional: To generate molecules satisfying a property condition $Y$, the model samples $K$ candidates from $p_\theta(x)$ and accepts those where the jointly trained predictor $p_\theta(y|x)$ classifies them as having property $Y$. This is framed as a variant of best-of-$K$ sampling, leveraging the unified representation space for tighter alignment between generation and control.

Key Contributions

1. Novel Joint Model: Introduction of Hyformer, which unifies generative and predictive tasks in a single parameter set using an alternating attention mechanism.
2. Demonstration of Synergistic Benefits: Empirical evidence showing Hyformer outperforms baselines in:
   • Conditional molecule generation.
   • Out-of-distribution (OOD) property prediction.
   • Molecular representation learning (via probing).
3. Performance Parity: Showing that Hyformer rivals the performance of state-of-the-art purely generative and purely predictive models, proving that joint training does not necessitate a sacrifice in individual task performance.
4. Real-World Application: A use case in discovering novel antimicrobial peptides (AMPs), demonstrating the utility of joint modeling in drug design.

Experimental Results

Conditional Molecule Generation

Evaluated on the GuacaMol dataset, the jointly fine-tuned Hyformer achieved the lowest Mean Absolute Deviation (MAD) and Standard Deviation (SD) across target properties (QED, SA, logP) while maintaining high validity. It significantly outperformed non-joint variants and baselines like MolGPT and GraphGPT, confirming that joint fine-tuning enhances controllability.

Out-of-Distribution (OOD) Property Prediction

On the Hit Identification task from the Lo-Hi benchmark, Hyformer achieved the highest mean AUPRC across all datasets (DRD2, HIV, KDR, Sol), outperforming fingerprint-based baselines (k-NN, GB, SVM, MLP) and deep learning models (Chemformer, Graphformer). This demonstrates the model's ability to generalize to molecular scaffolds unseen during training.

Representation Learning

Using a novel probing protocol on the MoleculeNet benchmark, representations extracted from jointly pre-trained Hyformer were the most predictive across both linear and KNN probing. Hyformer achieved the best performance on 4 out of 10 datasets for linear probing and 5 out of 10 for KNN probing, outperforming state-of-the-art predictive models like Uni-Mol and other joint models.

Unconditional Generation and AMP Design

• Unconditional Generation: On GuacaMol and MOSES benchmarks, Hyformer consistently lay on the Pareto front, balancing distributional fidelity (FCD, KL divergence), validity, and uniqueness, rivaling state-of-the-art SMILES-based and graph-based models.
• Antimicrobial Peptide Design: In a real-world drug design task, Hyformer generated peptides with high fitness and AMP probability, significantly outperforming baselines (PepCVAE, AMPGAN, HydrAMP, AMP-Diffusion) while maintaining competitive perplexity and diversity. Attention analysis revealed the model prioritized biologically relevant amino acids (Arginine, Lysine, Tryptophan, Proline) associated with antimicrobial activity.

Significance and Claims

The paper claims that Hyformer successfully unifies molecular generation and property prediction for SMILES-based representations, opening avenues for integration into real-world drug discovery pipelines. The authors emphasize that the synergistic benefits of joint learning—specifically improved control over the generative process, robustness to OOD data, and learning of high-level predictive representations—are achievable without sacrificing the performance of individual tasks.

The authors acknowledge inherent trade-offs, noting that shared parameters may limit task-specific specialization. They suggest future work could explore dynamic or modular attention architectures to allocate capacity more flexibly. Furthermore, they note that the extent to which these synergistic benefits apply to other modalities, such as 3D structures or transcriptomics, remains an open question. The work is presented as a step toward tools that facilitate the design of new potential medications, rather than a solution for malicious applications.