De Novo Molecule Generation

De novo molecular design is the task of generating new molecules with desired properties from scratch, rather than modifying known compounds. It represents a fundamental shift from traditional medicinal chemistry, which searches a pre-enumerated chemical library, to a generative approach that can explore regions of chemical space that were never previously synthesized or conceived.

Generative models for molecules learn a distribution over chemical structures from training data (typically drug-like molecules from ZINC or ChEMBL) and can then sample new molecules from this distribution, ideally biased toward desired property profiles such as high potency, drug-likeness, and synthetic accessibility.

What De Novo Generation Is and Why It Matters

The appeal of de novo generation is that it decouples drug candidate exploration from the limitations of physical compound libraries. A medicinal chemist exploring analogs of a hit compound is constrained to a small neighborhood of chemical space. A generative model trained on millions of drug-like molecules can, in principle, propose structurally novel scaffolds that no one has synthesized yet but that satisfy all design constraints.

In practice, the challenge is ensuring that generated molecules are:

• Valid: Chemically valid structures (correct valences, no impossible bonds)
• Novel: Not simply memorized training compounds
• Unique: Diverse, not repetitive
• Drug-like: Satisfying criteria like QED, TPSA, molecular weight
• Synthesizable: Actually makeable in a chemistry laboratory
• Active: Having the desired biological property (the hardest constraint)

Molecular Representations for Generation

For introductory implementation, SMILES strings are the simplest entry point because standard sequence models apply directly. The ZINC dataset provides ~250 million drug-like SMILES for training.

SMILES-Based VAE Architecture

The Variational Autoencoder (VAE) for SMILES, introduced by Gomez-Bombarelli et al. (2018) in "Automatic Chemical Design Using a Data-Driven Continuous Representation of Molecules," encodes SMILES strings into a continuous latent space and decodes points in latent space back to SMILES strings.

Architecture overview:

• Encoder: 1D convolutional layers or an LSTM that reads a character-level SMILES sequence and outputs the mean (mu) and log-variance (log_var) of the latent distribution.
• Sampling: A latent vector z is sampled from N(mu, sigma) using the reparameterization trick: z = mu + sigma * epsilon where epsilon is from N(0, I).
• Decoder: A GRU or LSTM that generates SMILES characters one by one from z, acting as a language model conditioned on the latent code.
• Loss: Reconstruction loss (cross-entropy over character sequences) + KL divergence term (KL(q(z|x) || p(z))) with a beta-weighting schedule (beta-VAE) to balance reconstruction quality against latent space regularity.

The key property of the VAE latent space is that it enables smooth interpolation and local search: moving slightly in latent space should produce structurally similar but not identical molecules. This enables Bayesian optimization in latent space to find molecules with high predicted property scores.

Key Validity Metrics: Validity, Uniqueness, Novelty, QED, SA Score

Both QED and SA Score are computed with RDKit. The Frechet ChemNet Distance (FCD), analogous to FID in image generation, measures how similar the distribution of generated molecules is to a reference set using ChemNet embeddings.

Code Example: SMILES VAE with PyTorch

import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import Dataset, DataLoader
import numpy as np
from rdkit import Chem
from rdkit.Chem import QED, Descriptors
import re

# -------------------------------------------------------
# Vocabulary and tokenization for SMILES
# -------------------------------------------------------
# Simple character-level tokenization
SMILES_CHARS = (
    ['PAD', 'SOS', 'EOS'] +
    list('CNOSPFClBrI') +    # atoms
    list('()[]=#-+@./\\1234567890%') +  # bonds and rings
    list('chons')            # aromatic atoms
)
CHAR2IDX = {c: i for i, c in enumerate(SMILES_CHARS)}
IDX2CHAR = {i: c for c, i in CHAR2IDX.items()}
VOCAB_SIZE = len(SMILES_CHARS)
PAD_IDX, SOS_IDX, EOS_IDX = 0, 1, 2

MAX_LEN = 120  # Maximum SMILES length

def tokenize_smiles(smiles: str) -> list:
    """Tokenize SMILES into characters, handling Cl and Br as single tokens."""
    pattern = r'Cl|Br|[A-Za-z#()[\]=@+\-./\\%0-9]'
    return re.findall(pattern, smiles)

def smiles_to_tensor(smiles: str) -> torch.Tensor:
    """Convert SMILES string to padded index tensor of length MAX_LEN+2."""
    tokens = tokenize_smiles(smiles)
    indices = [SOS_IDX] + [CHAR2IDX.get(t, PAD_IDX) for t in tokens] + [EOS_IDX]
    indices = indices[:MAX_LEN + 2]
    indices += [PAD_IDX] * (MAX_LEN + 2 - len(indices))
    return torch.tensor(indices, dtype=torch.long)

def tensor_to_smiles(tensor: torch.Tensor) -> str:
    """Convert index tensor back to SMILES string."""
    chars = []
    for idx in tensor.tolist():
        if idx == EOS_IDX:
            break
        if idx in (PAD_IDX, SOS_IDX):
            continue
        chars.append(IDX2CHAR.get(idx, ''))
    return ''.join(chars)

# -------------------------------------------------------
# VAE Model
# -------------------------------------------------------
class SmilesVAE(nn.Module):
    def __init__(self, vocab_size: int, embed_dim: int = 128,
                 latent_dim: int = 256, hidden_dim: int = 512):
        super().__init__()
        self.latent_dim = latent_dim

        # Encoder: embedding + GRU
        self.embed = nn.Embedding(vocab_size, embed_dim, padding_idx=PAD_IDX)
        self.encoder_gru = nn.GRU(embed_dim, hidden_dim, num_layers=2,
                                   batch_first=True, bidirectional=True)
        self.fc_mu     = nn.Linear(hidden_dim * 2, latent_dim)
        self.fc_logvar = nn.Linear(hidden_dim * 2, latent_dim)

        # Decoder: latent -> initial hidden state + GRU
        self.latent_to_hidden = nn.Linear(latent_dim, hidden_dim)
        self.decoder_gru = nn.GRU(embed_dim + latent_dim, hidden_dim,
                                   num_layers=2, batch_first=True)
        self.output_proj = nn.Linear(hidden_dim, vocab_size)

    def encode(self, x: torch.Tensor):
        """Encode token sequence to mu and log_var."""
        emb = self.embed(x)
        _, hidden = self.encoder_gru(emb)
        # Concatenate forward and backward last hidden states
        h = torch.cat([hidden[-2], hidden[-1]], dim=1)
        return self.fc_mu(h), self.fc_logvar(h)

    def reparameterize(self, mu: torch.Tensor, log_var: torch.Tensor) -> torch.Tensor:
        """Sample z using reparameterization trick."""
        std = (0.5 * log_var).exp()
        eps = torch.randn_like(std)
        return mu + eps * std

    def decode(self, z: torch.Tensor, target: torch.Tensor = None,
               max_len: int = MAX_LEN + 2) -> torch.Tensor:
        """
        Decode latent vector z to logits over vocabulary.
        If target is provided, use teacher forcing.
        """
        batch_size = z.size(0)
        device = z.device

        # Initialize decoder hidden state from latent vector
        h0 = torch.tanh(self.latent_to_hidden(z)).unsqueeze(0).repeat(2, 1, 1)

        if target is not None:
            # Teacher forcing during training
            emb = self.embed(target[:, :-1])  # Input: all except last token
            # Concatenate latent vector at each timestep
            z_expanded = z.unsqueeze(1).expand(-1, emb.size(1), -1)
            dec_input = torch.cat([emb, z_expanded], dim=-1)
            output, _ = self.decoder_gru(dec_input, h0)
            logits = self.output_proj(output)
        else:
            # Autoregressive generation (inference)
            token = torch.full((batch_size,), SOS_IDX, dtype=torch.long, device=device)
            hidden = h0
            logits_list = []
            for _ in range(max_len):
                emb = self.embed(token).unsqueeze(1)
                z_exp = z.unsqueeze(1)
                dec_in = torch.cat([emb, z_exp], dim=-1)
                out, hidden = self.decoder_gru(dec_in, hidden)
                step_logits = self.output_proj(out.squeeze(1))
                logits_list.append(step_logits)
                token = step_logits.argmax(dim=-1)
            logits = torch.stack(logits_list, dim=1)
        return logits

    def forward(self, x: torch.Tensor):
        mu, log_var = self.encode(x)
        z = self.reparameterize(mu, log_var)
        logits = self.decode(z, target=x)
        return logits, mu, log_var

# -------------------------------------------------------
# VAE Loss
# -------------------------------------------------------
def vae_loss(logits, targets, mu, log_var, beta=1.0):
    """
    ELBO loss for SMILES VAE.
    logits: (batch, seq_len, vocab_size)
    targets: (batch, seq_len) - shifted right (target[:, 1:])
    """
    # Reconstruction loss: cross-entropy over sequence
    batch_size, seq_len, vocab_size = logits.shape
    recon_loss = nn.functional.cross_entropy(
        logits.reshape(-1, vocab_size),
        targets[:, 1:].reshape(-1),  # Skip SOS in target
        ignore_index=PAD_IDX,
        reduction='mean'
    )
    # KL divergence: -0.5 * sum(1 + log_var - mu^2 - var)
    kl_loss = -0.5 * torch.mean(1 + log_var - mu.pow(2) - log_var.exp())
    return recon_loss + beta * kl_loss, recon_loss, kl_loss

# -------------------------------------------------------
# Training loop (simplified)
# -------------------------------------------------------
def train_smiles_vae(smiles_list, n_epochs=50, batch_size=256, lr=1e-3,
                     latent_dim=256, device='cpu'):
    """Train a SMILES VAE on a list of canonical SMILES strings."""
    # Prepare dataset
    tensors = [smiles_to_tensor(smi) for smi in smiles_list
               if Chem.MolFromSmiles(smi) is not None]
    dataset = torch.stack(tensors)
    loader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

    model = SmilesVAE(VOCAB_SIZE, latent_dim=latent_dim).to(device)
    optimizer = optim.Adam(model.parameters(), lr=lr)

    # KL annealing: ramp beta from 0 to 1 over first 10 epochs
    for epoch in range(n_epochs):
        beta = min(1.0, epoch / 10.0)
        model.train()
        total_loss = 0.0
        for batch in loader:
            batch = batch.to(device)
            logits, mu, log_var = model(batch)
            loss, recon, kl = vae_loss(logits, batch, mu, log_var, beta=beta)
            optimizer.zero_grad()
            loss.backward()
            nn.utils.clip_grad_norm_(model.parameters(), max_norm=5.0)
            optimizer.step()
            total_loss += loss.item()
        if (epoch + 1) % 5 == 0:
            print(f"Epoch {epoch+1}: loss={total_loss/len(loader):.4f}, beta={beta:.2f}")
    return model

# -------------------------------------------------------
# Generation and evaluation
# -------------------------------------------------------
def generate_molecules(model, n_samples=1000, latent_dim=256, device='cpu', temp=1.0):
    """Sample molecules from the VAE latent space."""
    model.eval()
    with torch.no_grad():
        z = torch.randn(n_samples, latent_dim, device=device) * temp
        logits = model.decode(z)
        token_ids = logits.argmax(dim=-1)  # (n_samples, max_len)
    return [tensor_to_smiles(ids) for ids in token_ids]

def evaluate_generation(smiles_list: list, train_smiles_set: set) -> dict:
    """Compute validity, uniqueness, novelty, and mean QED/SA."""
    from rdkit.Chem import RDConfig
    import os, sys
    sys.path.append(os.path.join(RDConfig.RDContribDir, 'SA_Score'))
    import sascorer

    valid = [smi for smi in smiles_list if Chem.MolFromSmiles(smi) is not None]
    unique = list(set(valid))
    novel = [smi for smi in unique if smi not in train_smiles_set]

    qed_scores, sa_scores = [], []
    for smi in novel:
        mol = Chem.MolFromSmiles(smi)
        qed_scores.append(QED.qed(mol))
        sa_scores.append(sascorer.calculateScore(mol))

    return {
        'validity':  len(valid)  / len(smiles_list),
        'uniqueness': len(unique) / max(len(valid), 1),
        'novelty':   len(novel)  / max(len(unique), 1),
        'mean_qed':  np.mean(qed_scores) if qed_scores else 0.0,
        'mean_sa':   np.mean(sa_scores)  if sa_scores  else 10.0,
        'n_generated': len(smiles_list),
        'n_valid':    len(valid),
        'n_novel':    len(novel),
    }

Property-Constrained Generation

A plain VAE generates from a prior distribution without explicit property constraints. To generate molecules with desired properties, several approaches exist:

Bayesian Optimization in Latent Space

Train a surrogate model (Gaussian Process or neural network) to predict a property score from the latent vector. Use Bayesian optimization (BO) to find latent vectors that maximize the predicted score. Decode these latent vectors to candidate molecules. This is the approach from the Gomez-Bombarelli et al. JCIM 2018 paper, which showed 10x penalized LogP improvement over the training set.

Conditional VAE (cVAE)

Condition both the encoder and decoder on the target property. The encoder receives both the SMILES and a property label; the decoder is conditioned on z plus the desired property value. This enables direct sampling at a target property value without post-hoc optimization.

Reinforcement Learning (REINVENT)

REINVENT by AstraZeneca fine-tunes a pre-trained SMILES language model with reinforcement learning, using a reward function that scores generated molecules on desired properties (predicted activity, QED, SA score, etc.). REINVENT is currently one of the most widely used generative drug design tools in industry due to its practical effectiveness and flexible reward compositing.

Other Architectures

Common Pitfalls

• Mode collapse: The VAE generates the same few molecules repeatedly. Caused by KL weight too high or insufficient decoder capacity. Use KL annealing to warm up the KL term gradually.
• Low chemical diversity: Generated molecules cluster around common scaffolds from the training set. Increase sampling temperature (sigma multiplier greater than 1.0) or use a diversity-promoting objective.
• Invalid SMILES rate: A SMILES VAE typically produces 10-20% invalid strings. Using SELFIES encoding eliminates this but requires a different tokenizer. Alternatively, use constrained decoding with beam search.
• Posterior collapse: The encoder ignores the input and the decoder generates from the prior alone. Diagnosed by mu and log_var being approximately constant regardless of input. Addressed with free-bits or beta-VAE strategies.
• High SA Score: Generated molecules often have high synthetic accessibility scores because the model does not learn synthetic feasibility constraints. Add SA score as a reward term in RL-based fine-tuning.

Further Reading

• Gomez-Bombarelli et al. (2018): Automatic Chemical Design with VAE
• REINVENT4 by AstraZeneca (GitHub)
• Jin et al. (2018): Junction Tree VAE
• Hoogeboom et al. (2022): Equivariant Diffusion for Molecule Generation
• SELFIES Documentation
• MOSES: Molecular Sets Benchmark

Try it live

Explore a related model hands-on: open AI Drug Design: Molecular Generation, Docking & ADMET Pipeline and change its parameters while it runs. Everything stays in your browser.