34.4 Conditional Generation and Evaluation

1. Classifier-Free Guidance

In practice, training a separate property classifier at every noise level to guide diffusion has proven fragile: the classifier must handle hundreds of noise scales, errors compound across denoising steps, and each new target property requires retraining from scratch. Classifier-free guidance eliminates this entire bottleneck.

Classifier-free guidance (Ho & Salimans, 2022) is the standard method for conditioning diffusion models on target properties without a separate classifier. The idea: train a single model that can operate both conditionally (given a property label \(c\)) and unconditionally (without any label). At sampling time, extrapolate away from the unconditional prediction toward the conditional prediction:

where \(w > 1\) is the guidance scale and \(\varnothing\) denotes the null condition (the unconditional model). When \(w = 1\), this reduces to standard conditional generation. When \(w > 1\), the guidance pushes the model further toward the conditional prediction. The resulting samples match the condition more strongly but lose diversity.

Classifier-free guidance avoids training a separate classifier or property predictor on noisy intermediate states. That per-noise-level approach is brittle and computationally expensive. The mechanism works by jointly training conditional and unconditional pathways in a single network (randomly dropping the condition label during training), then amplifying the difference between the two predictions at inference to steer generation. Use classifier-free guidance when your target property can be represented as a label or embedding at training time. For properties that require external black-box evaluation (docking scores, lab assays), REINFORCE (covered below) applies instead.

During training, the condition \(c\) is randomly dropped (replaced with \(\varnothing\)) with some probability \(p_{\text{uncond}}\) (typically 10 to 20%). This trains both the conditional and unconditional models simultaneously in a single network. In short: conditioning turns a generative model from a random molecule factory into a targeted design tool, and the guidance scale is the dial between creative exploration and focused optimization.

import torch
import torch.nn as nn
import torch.nn.functional as F


class GuidedDiffusionModel(nn.Module):
    """Diffusion model with classifier-free guidance for molecular generation.

    Trains a single noise predictor that handles both conditional and
    unconditional generation. At inference, extrapolates toward the condition.
    """
    def __init__(self, denoiser: nn.Module, num_properties: int,
                 cond_embed_dim: int = 128, p_uncond: float = 0.1):
        super().__init__()
        self.denoiser = denoiser
        self.p_uncond = p_uncond

        # Property conditioning embedding
        self.cond_embed = nn.Sequential(
            nn.Linear(num_properties, cond_embed_dim),
            nn.SiLU(),
            nn.Linear(cond_embed_dim, cond_embed_dim)
        )
        # Null embedding for unconditional generation
        self.null_embed = nn.Parameter(torch.randn(cond_embed_dim))

    def forward(self, x_t: torch.Tensor, t: torch.Tensor,
                condition: torch.Tensor = None, force_uncond: bool = False):
        """Predict noise, optionally conditioned on target properties."""
        batch_size = x_t.size(0)

        if condition is not None and not force_uncond:
            # During training: randomly drop condition
            if self.training:
                mask = torch.rand(batch_size, device=x_t.device) < self.p_uncond
                cond_emb = self.cond_embed(condition)
                null_emb = self.null_embed.unsqueeze(0).expand(batch_size, -1)
                # Replace some conditions with null embedding
                cond_emb = torch.where(mask.unsqueeze(-1), null_emb, cond_emb)
            else:
                cond_emb = self.cond_embed(condition)
        else:
            cond_emb = self.null_embed.unsqueeze(0).expand(batch_size, -1)

        return self.denoiser(x_t, t, cond_emb)

    @torch.no_grad()
    def guided_sample(self, x_t: torch.Tensor, t: torch.Tensor,
                      condition: torch.Tensor, guidance_scale: float = 3.0):
        """Apply classifier-free guidance during sampling."""
        # Conditional prediction
        noise_cond = self.forward(x_t, t, condition)
        # Unconditional prediction
        noise_uncond = self.forward(x_t, t, force_uncond=True)

        # Extrapolate: move further in the conditional direction
        noise_guided = noise_uncond + guidance_scale * (noise_cond - noise_uncond)
        return noise_guided

2. Inpainting and Motif Scaffolding

Classifier-free guidance steers the entire generation toward a target property, but sometimes you need finer control: keeping part of a structure fixed while generating the rest around it.

A particularly powerful form of conditional generation is inpainting: generating part of a structure while keeping another part fixed. For protein design, this enables motif scaffolding: given a functional motif (a binding site, a catalytic triad, an epitope), generate a complete protein backbone that positions the motif correctly and provides structural stability.

During the reverse diffusion process, each denoising step replaces values at fixed positions with the ground-truth motif noised to the current level. Scaffold regions generate freely while remaining structurally consistent with the motif. Figure 34.4.1 illustrates diffusion inpainting for motif scaffolding.

@torch.no_grad()
def inpainting_sample(model, motif_positions: torch.Tensor,
                      motif_mask: torch.Tensor, shape: tuple,
                      T: int, alpha_bars: torch.Tensor,
                      alphas: torch.Tensor, betas: torch.Tensor,
                      device: torch.device):
    """Generate structures with fixed motif regions via inpainting.

    Args:
        model: trained noise predictor
        motif_positions: ground truth positions for motif residues (N, 3)
        motif_mask: boolean mask, True for motif residues (N,)
        shape: output shape (N, 3)
        T: number of diffusion steps
    """
    # Start from noise
    x = torch.randn(shape, device=device)

    for t_val in reversed(range(T)):
        t = torch.full((1,), t_val, device=device, dtype=torch.long)
        alpha_bar_t = alpha_bars[t_val]

        # Standard reverse step for scaffold regions
        noise_pred = model(x, t)
        alpha_t = alphas[t_val]
        beta_t = betas[t_val]

        x_denoised = (1 / torch.sqrt(alpha_t)) * (
            x - (beta_t / torch.sqrt(1 - alpha_bar_t)) * noise_pred
        )

        if t_val > 0:
            noise = torch.randn_like(x)
            x_denoised = x_denoised + torch.sqrt(beta_t) * noise

        # Inpainting: replace motif positions with correctly noised ground truth
        if t_val > 0:
            motif_noise = torch.randn_like(motif_positions)
            alpha_bar_prev = alpha_bars[t_val - 1]
            motif_noisy = (torch.sqrt(alpha_bar_prev) * motif_positions +
                          torch.sqrt(1 - alpha_bar_prev) * motif_noise)
        else:
            motif_noisy = motif_positions  # final step: use clean motif

        # Blend: motif positions from ground truth, scaffold from generation
        mask = motif_mask.unsqueeze(-1).expand_as(x)
        x = torch.where(mask, motif_noisy, x_denoised)

    return x

3. REINFORCE for Non-Differentiable Rewards

Both classifier-free guidance and inpainting rely on information available during training: property labels or known structural motifs. Yet many of the properties that matter most for discovery can only be evaluated after a molecule is fully generated.

Classifier-free guidance works when the target property can be encoded as a conditioning signal during training. But many discovery-relevant properties cannot: docking scores require running a physics-based simulation, synthetic accessibility scores involve graph-based heuristics, and experimental binding affinity is only available after wet-lab testing. For these non-differentiable reward functions, policy gradient methods from reinforcement learning provide the solution.

The REINFORCE algorithm (Williams, 1992) treats the generative model as a policy \(\pi_\theta(x)\) that produces molecular samples \(x\), and the property evaluator as a reward function \(R(x)\). The objective is to maximize the expected reward:

The gradient of this objective is given by the log-derivative trick (rewriting the gradient of an expectation as an expectation of the score function weighted by the reward, so that we never need to differentiate through \(R\) itself):

This gradient does not require differentiating through the reward function \(R(x)\); it only requires evaluating \(R(x)\) and computing the log-probability \(\log \pi_\theta(x)\) of the generated sample. For autoregressive molecular generators, \(\log \pi_\theta(x)\) is the sum of log-probabilities at each generation step. For diffusion models, it can be approximated via the probability flow ordinary differential equation (ODE), which recasts the stochastic diffusion process as a deterministic trajectory whose log-likelihood can be computed exactly (Section 34.2).

import torch
import torch.nn.functional as F
from typing import Callable, List
from dataclasses import dataclass


@dataclass
class RewardConfig:
    """Configuration for multi-property reward function."""
    weights: dict          # property_name -> weight
    thresholds: dict       # property_name -> minimum acceptable value
    baseline_ema: float    # exponential moving average decay for baseline


class REINFORCETrainer:
    """REINFORCE optimizer for molecular generators with non-differentiable rewards.

    Fine-tunes a pretrained autoregressive molecular generator
    to maximize a multi-property reward function.
    """
    def __init__(self, generator: torch.nn.Module,
                 reward_fns: dict,  # name -> callable(smiles) -> float
                 config: RewardConfig,
                 lr: float = 1e-5):
        self.generator = generator
        self.reward_fns = reward_fns
        self.config = config
        self.optimizer = torch.optim.Adam(generator.parameters(), lr=lr)
        self.baseline = 0.0  # running baseline for variance reduction

    def compute_reward(self, smiles_list: List[str]) -> torch.Tensor:
        """Compute multi-property reward for a batch of molecules."""
        rewards = torch.zeros(len(smiles_list))
        property_values = {}

        for name, fn in self.reward_fns.items():
            values = torch.tensor([fn(s) for s in smiles_list])
            property_values[name] = values

            # Apply threshold: zero reward if below minimum
            threshold = self.config.thresholds.get(name, float('-inf'))
            mask = values >= threshold

            # Weighted contribution
            weight = self.config.weights.get(name, 1.0)
            rewards += weight * values * mask.float()

        return rewards, property_values

    def train_step(self, batch_size: int = 64) -> dict:
        """One REINFORCE training step."""
        self.generator.train()

        # Generate molecules and collect log-probabilities
        samples, log_probs = self.generator.sample_with_log_prob(batch_size)

        # Decode to SMILES (generator-specific)
        smiles_list = self.generator.decode_to_smiles(samples)

        # Compute rewards (non-differentiable)
        rewards, properties = self.compute_reward(smiles_list)
        rewards = rewards.to(log_probs.device)

        # Update baseline (variance reduction): subtracting a near-mean
        # value from each reward centers the advantages around zero,
        # which sharply reduces gradient variance without introducing bias.
        self.baseline = (self.config.baseline_ema * self.baseline +
                        (1 - self.config.baseline_ema) * rewards.mean().item())

        # REINFORCE gradient: (R - baseline) * grad log pi
        advantage = rewards - self.baseline
        loss = -(advantage * log_probs).mean()

        # Optional: add entropy bonus to prevent mode collapse
        entropy = -(torch.exp(log_probs) * log_probs).mean()
        loss = loss - 0.01 * entropy

        self.optimizer.zero_grad()
        loss.backward()

        # Gradient clipping for stability
        torch.nn.utils.clip_grad_norm_(self.generator.parameters(), 1.0)
        self.optimizer.step()

        return {
            'loss': loss.item(),
            'mean_reward': rewards.mean().item(),
            'baseline': self.baseline,
            'entropy': entropy.item(),
            'valid_frac': sum(1 for s in smiles_list if s is not None) / len(smiles_list),
            'properties': {k: v.mean().item() for k, v in properties.items()}
        }

4. Evaluation Metrics for Molecular Generators

How do you know if your molecular generator is any good? The evaluation landscape for generative models in chemistry and biology has matured into a set of standard metrics, organized into four categories.

Validity measures whether generated molecules satisfy basic chemical rules: correct valences, no impossible bonds, parseable SMILES strings. For SMILES-based generators, validity rates in published benchmarks range from 30% (early variational autoencoders, or VAEs) to 99%+ (graph-based generators with action masking, where the model's output space is restricted at each step to only chemically valid bond additions), a roughly threefold improvement driven primarily by switching from string representations to graph representations that enforce valence rules by construction. For 3D structure generators, validity includes bond length and angle distributions matching reference data.

Uniqueness measures the fraction of non-duplicate molecules among valid outputs. A generator that produces the same molecule 1000 times has 100% validity but 0.1% uniqueness. Mode collapse, where the generator's output distribution concentrates on a small number of nearly identical samples, in generative adversarial networks (GANs) manifests as low uniqueness.

Novelty measures the fraction of generated molecules not present in the training set. A generator that memorizes and replays training data has 100% validity and uniqueness but 0% novelty. High novelty is essential for discovery.

Distribution similarity measures how closely the generated distribution matches a reference distribution. The Frechet ChemNet Distance (FCD; Preuer et al., 2018) computes the Frechet distance (a measure of how far apart two multivariate Gaussian fits to the activation vectors are) between the activations of a pretrained ChemNet model (a neural network trained on molecular activity data to learn general-purpose chemical representations) on generated and reference molecules, analogous to the Frechet Inception Distance (FID) score for images. Lower FCD indicates more similar distributions.

Internal diversity measures how different the generated molecules are from one another, complementing the four categories above. A generator may achieve high validity, uniqueness, and novelty yet still produce molecules clustered around a single scaffold. Internal diversity, computed as the average pairwise Tanimoto distance between molecular fingerprints, detects this clustering: values near 1.0 indicate highly diverse outputs, while values near 0.0 indicate near-identical structures differing only in minor substituents.

from rdkit import Chem
from rdkit.Chem import Descriptors, QED, AllChem
from collections import Counter
import numpy as np
from typing import List, Dict


class MolecularGeneratorEvaluator:
    """Evaluate molecular generators on standard metrics.

    Computes validity, uniqueness, novelty, property distributions,
    and internal diversity for a set of generated SMILES strings.
    """
    def __init__(self, training_smiles: set):
        self.training_smiles = training_smiles

    def evaluate(self, generated_smiles: List[str]) -> Dict[str, float]:
        """Compute all standard evaluation metrics."""
        n_total = len(generated_smiles)

        # Parse molecules, filtering invalid SMILES
        valid_mols = []
        valid_smiles = []
        for smi in generated_smiles:
            mol = Chem.MolFromSmiles(smi)
            if mol is not None:
                # Canonicalize for uniqueness checking
                canon = Chem.MolToSmiles(mol)
                valid_mols.append(mol)
                valid_smiles.append(canon)

        n_valid = len(valid_mols)
        validity = n_valid / n_total if n_total > 0 else 0.0

        # Uniqueness: fraction of distinct valid molecules
        unique_smiles = set(valid_smiles)
        uniqueness = len(unique_smiles) / n_valid if n_valid > 0 else 0.0

        # Novelty: fraction not in training set
        novel = unique_smiles - self.training_smiles
        novelty = len(novel) / len(unique_smiles) if unique_smiles else 0.0

        # Property distributions for valid molecules
        properties = self._compute_properties(valid_mols)

        # Internal diversity: average pairwise Tanimoto distance
        diversity = self._internal_diversity(valid_mols[:1000])  # cap for speed

        return {
            'validity': validity,
            'uniqueness': uniqueness,
            'novelty': novelty,
            'diversity': diversity,
            'n_valid': n_valid,
            'n_unique': len(unique_smiles),
            'n_novel': len(novel),
            **{f'mean_{k}': v for k, v in properties.items()}
        }

    def _compute_properties(self, mols: list) -> dict:
        """Compute molecular property statistics."""
        qed_scores = [QED.qed(m) for m in mols]
        mw = [Descriptors.MolWt(m) for m in mols]
        logp = [Descriptors.MolLogP(m) for m in mols]

        return {
            'QED': np.mean(qed_scores) if qed_scores else 0.0,
            'MW': np.mean(mw) if mw else 0.0,
            'LogP': np.mean(logp) if logp else 0.0,
        }

    def _internal_diversity(self, mols: list) -> float:
        """Average pairwise Tanimoto distance (higher = more diverse)."""
        if len(mols) < 2:
            return 0.0
        fps = [AllChem.GetMorganFingerprintAsBitVect(m, 2, 2048) for m in mols]
        from rdkit.DataStructs import BulkTanimotoSimilarity
        total_dist = 0.0
        n_pairs = 0
        for i in range(min(len(fps), 500)):  # cap pairwise comparisons
            sims = BulkTanimotoSimilarity(fps[i], fps[i+1:])
            total_dist += sum(1.0 - s for s in sims)
            n_pairs += len(sims)
        return total_dist / n_pairs if n_pairs > 0 else 0.0


# Usage example
evaluator = MolecularGeneratorEvaluator(training_smiles=set())
# results = evaluator.evaluate(generated_smiles_list)
# print(f"Validity: {results['validity']:.1%}")
# print(f"Uniqueness: {results['uniqueness']:.1%}")
# print(f"Novelty: {results['novelty']:.1%}")
# print(f"Diversity: {results['diversity']:.3f}")

5. Benchmarks: MOSES, GuacaMol, and TDC

Standardized benchmarks enable fair comparison of molecular generators. Three dominate the field:

MOSES (Polykovskiy et al., 2020) provides a curated subset of ZINC, a freely available database of commercially available compounds for virtual screening, with standardized train/test splits and a comprehensive metric suite (validity, uniqueness, novelty, FCD, fragment similarity, scaffold similarity). It is the most widely used benchmark for unconditional molecular generation (as of 2024, MOSES is widely regarded as largely saturated for state-of-the-art models, with top methods achieving near-perfect scores on most metrics; the PMO benchmark described below and task-specific TDC leaderboards have become the preferred evaluation targets for goal-directed generation).

GuacaMol (Brown et al., 2019) focuses on goal-directed generation: given a target molecular property profile (e.g., high QED, LogP, the octanol-water partition coefficient measuring how readily a molecule crosses cell membranes, in range, specific scaffold), how well can the generator produce molecules matching the specification? This is more relevant for discovery than unconditional generation.

Therapeutic Data Commons (TDC) (Huang et al., 2021) provides real-world drug discovery tasks including absorption, distribution, metabolism, excretion, and toxicity (ADMET) prediction, drug-target interaction, and molecular generation benchmarks with clinically relevant endpoints. TDC is the closest to real discovery workflows and connects to the domain-specific applications in Chapter 49.

from tdc import Evaluator

# Load the standard molecular generation evaluator
evaluator = Evaluator(name='Diversity')
diversity = evaluator(generated_smiles)

evaluator = Evaluator(name='Validity')
validity = evaluator(generated_smiles)

evaluator = Evaluator(name='FCD')
fcd_score = evaluator(generated_smiles, reference_smiles)

# Or run the full MOSES benchmark suite
from tdc import BenchmarkGroup
group = BenchmarkGroup(name='MolGen')
results = group.evaluate(generated_smiles)

6. Putting It Together: The Conditional Generation Pipeline

A complete conditional generation pipeline for molecular discovery combines the components from this section into a coherent workflow. Figure 34.4.2 illustrates the five stages and the feedback loop that connects evaluation back to conditioning.

Stage 1: Pretrain. Train an unconditional or weakly conditional generative model on a large molecular dataset (ZINC, ChEMBL, a curated database of bioactive compounds with measured activity against biological targets, PubChem). This gives the model broad chemical knowledge: valid bond patterns, common functional groups, realistic size distributions.

Stage 2: Condition. Add conditioning signals via classifier-free guidance (differentiable properties) or prepare for REINFORCE fine-tuning (non-differentiable properties). For structure-based design, condition on the target protein pocket geometry.

Stage 3: Generate and filter. Generate a large pool of candidates (10,000+), filter by validity and basic property thresholds, and deduplicate.

Stage 4: Evaluate. Score filtered candidates with the full property prediction pipeline (docking, ADMET, synthesizability). Compute the standard metrics to verify the generator is working correctly.

Stage 5: Iterate. Feed top-scoring candidates back as conditioning signals or training data for the next round. This iterative refinement typically converges toward the target property region while maintaining diversity.

The next section implements this complete pipeline.

Exercises

1. (Conceptual) Explain why the guidance scale \(w\) in classifier-free guidance can be interpreted as controlling a temperature parameter on the conditional distribution. What happens to the entropy of the generated distribution as \(w \to \infty\)? How does this relate to the exploitation side of the exploration-exploitation trade-off from Chapter 1?
2. (Coding) Implement the MolecularGeneratorEvaluator from this section and evaluate two molecular generation strategies: (a) random sampling from a pretrained SMILES VAE, and (b) latent space optimization of the same VAE using a QED predictor as the objective. Compare validity, uniqueness, novelty, and mean QED for 1000 molecules from each strategy. Does the optimized generator sacrifice any metric to improve QED?
3. (Analysis) The REINFORCE estimator has notoriously high variance. Implement REINFORCE with and without the baseline variance reduction technique from this section. Generate 100 batches of 64 molecules each and plot the gradient variance (estimated as the variance of the per-sample gradient contributions) across batches. How much does the baseline reduce variance? Experiment with the baseline EMA decay parameter: what value gives the best variance-reward trade-off?