Co-Diffusion: An Affinity-Aware Two-Stage Latent Diffusion Framework for Generalizable Drug-Target Affinity Prediction

Abstract

Predicting drug-target affinity is fundamental to virtual screening and lead optimization. However, existing deep models often suffer from representation collapse in stringent cold-start regimes, where the scarcity of labels and domain shifts prevent the learning of transferable pharmacophores and binding motifs. In this paper, we propose Co-Diffusion, a novel affinity-aware framework that redefines DTA prediction as a constrained latent denoising process to enhance generalization. Co-Diffusion employs a two-stage paradigm: Stage I establishes an affinity-steered latent manifold by aligning drug and target embeddings under an explicit supervised objective, ensuring that the latent space reflects the intrinsic binding landscape. Stage II introduces modality-specific latent diffusion as a stochastic perturb-and-denoise regularizer, forcing the model to recover consistent affinity semantics from noisy structural representations. This approach effectively mitigates the reconstruction-regression conflict common in generative DTA models. Theoretically, we show that Co-Diffusion maximizes a variational lower bound on the joint likelihood of drug structures, protein sequences, and binding strength. Extensive experiments across multiple benchmarks demonstrate that Co-Diffusion significantly outperforms state-of-the-art baselines, particularly yielding superior zero-shot generalization on unseen molecular scaffolds and novel protein families-paving a robust path for in silico drug prioritization in unexplored chemical spaces.

Figures (7)

• Figure 1: Overview of the drug discovery process. The proposed Co-Diffusion framework operates at the predictive modeling stage, providing robust affinity estimates to guide the selection of candidates before in vitro testing.
• Figure 2: Latent Diffusion Process in Training and Inference. (a) In training, the VAE encoder produces latent $\mathbf{z}_0$, which is perturbed to $\mathbf{z}_t$ by adding noise. A UNet is trained to predict the noise, following the Noise-Prediction Training Objective. (b) In inference, the reverse denoising process transforms sampled noise $\mathbf{z}_T$ back to $\mathbf{z}_0$, which is decoded to generate the target sequence, as described in Generation.
• Figure 3: The graphical structure for Co-Diffusion model.
• Figure 4: Overview of Co-Diffusion.
• Figure 5: Performance comparison on Davis and KIBA under random split.