DiffGap

A diffusion model for pocket-conditioned molecule generation is trained on inputs it will never see again: at step t it gets a noised version of the real ligand and learns to recover the clean one. At sampling time no real ligand exists — each step's input is whatever the model produced at the step before, and across a thousand steps, small coordinate and atom-type errors compound into implausible geometry. The distribution the network trained on and the one it meets have diverged, the exposure bias familiar from autoregressive sequence models arriving here through the reverse diffusion chain.

DiffGap closes that gap by making the model rehearse its own mistakes during training. At each step it re-predicts the clean molecule from the current noisy state, re-noises that prediction back to step t through the forward process, and then chooses between two conditions: with probability p_T it trains on the true noisy sample as usual, and otherwise it trains on this "pseudo-molecule", a state constructed to look like what inference will hand it. A temperature annealing schedule drives p_T down as training proceeds, so the network starts on clean supervision and shifts onto its own predictions as they improve.

Because the change lives in the training loop rather than the architecture, it is applied unchanged to two existing pocket-conditioned diffusion models, TargetDiff and BindDM, and released as retrained checkpoints for both. The work comes from Peidong Liu and Xianggen Liu's group at Sichuan University, with co-author affiliations at Tsinghua University and Xidian University. It appeared on arXiv in November 2024 under the name GapDiff and was published at ACML 2025 as DiffGap.

#Key Features

• Pseudo-molecule estimation: The network re-predicts the clean molecule from the current state and the pocket, and that estimate is re-noised through the forward posterior, giving a training condition that mimics inference without a full sampling rollout.
• Probabilistic branch selection: Each step trains on the ground-truth noisy state with probability p_T and on the pseudo-molecule otherwise, so supervision never collapses onto the model's own output.
• Temperature annealing: An arc-shaped decay curve, chosen by ablation over linear and scheduled-sampling curves, with a floor of 0.5 that keeps half the steps on ground truth and prevents the self-reinforcing drift seen when p_T fell too far.
• Architecture-agnostic: The same procedure retrains two different pocket-conditioned diffusion backbones, and both retrained checkpoints improve on their unmodified baselines.
• No sampling overhead: The extra prediction runs only in training, costing roughly 12% more training time; inference is unchanged.

#Technical Details

The backbone is an E(n)-equivariant graph neural network of nine equivariant layers running a 1,000-step DDPM schedule over ligand atom coordinates and one-hot atom types, conditioned on the protein pocket. Training uses CrossDocked2020 filtered at 1 Å pose RMSD, leaving 100,000 protein-ligand pairs for training and 100 pockets for testing, with Adam at learning rate 1e-4, batch size 4, and a 200,000-step budget.

The headline DiffGap results are the BindDM-based checkpoint. Mean Vina Score improves to -6.28 (median -6.90) from -5.92 (-6.81) for BindDM, mean Vina Dock reaches -8.43, and the High Affinity rate — the share of generated molecules docking better than the reference ligand — rises to 68.9% from 64.8% in the mean. Applied instead to TargetDiff, mean Vina Score moves from -5.47 to -6.51 and mean Vina Dock from -7.80 to -8.54. Geometry improves alongside affinity: Jensen-Shannon divergence between generated and reference all-atom pair-distance distributions falls to 0.065 from 0.101 for the reproduced baseline, with the lowest divergence on five of eight carbon bond types. Drug-likeness and synthesizability are held rather than improved, at mean QED 0.51 and SA 0.59 with diversity 0.75.

#Applications

The released checkpoints sample drug-like 3D ligands for any protein pocket supplied as a structure, so a team with a solved or modeled binding site can generate candidate chemotypes without a reference ligand and without retraining. Because the contribution is a training procedure rather than a network, groups training their own structure-based generators can adopt it without changing their architecture or sampling code.

#Impact

DiffGap makes a concrete case that part of the quality gap in 3D molecule diffusion is a train-inference mismatch rather than a capacity limit, recoverable by changing what the model is conditioned on during training. Evaluation is entirely computational, resting on AutoDock Vina scores and cheminformatics proxies rather than measured binding, and confined to CrossDocked2020 and PDBbind. The repository vendors the TargetDiff and BindDM codebases wholesale, so the specific contribution is the modified training loop and the two retrained checkpoints rather than the whole tree; those checkpoints sit on Hugging Face behind a one-paragraph README with no license declaration. The authors note that the mechanism is specific to diffusion models and that the principles behind its gains are not yet understood well enough to optimize further.