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

arXiv: https://arxiv.org/abs/2603.11125 · HTML: https://mlsanity.com/paper/2603.11125/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Flow matching, time series and science, one of 14 directions in Diffusion models.
- Builds on 8 of its 33 references that are in the map.
- First posted on arXiv 11 Mar 2026.

## Placed in

- Diffusion models / [Flow matching, time series and science](https://mlsanity.com/diffusion/flow-matching-time-series-and-science/)

## Builds on (most cited first, 8 of 8)

- [Denoising Diffusion Probabilistic Models](https://mlsanity.com/paper/2006.11239/) (2020)
- [High-Resolution Image Synthesis with Latent Diffusion Models](https://mlsanity.com/paper/2112.10752/) (2021)
- [DeepDTA: Deep Drug-Target Binding Affinity Prediction](https://mlsanity.com/paper/1801.10193/) (2018)
- [MolTrans: Molecular Interaction Transformer for Drug Target Interaction Prediction](https://mlsanity.com/paper/2004.11424/) (2020)
- [Interpretable bilinear attention network with domain adaptation improves drug-target prediction](https://mlsanity.com/paper/2208.02194/) (2022)
- [Re-Dock: Towards Flexible and Realistic Molecular Docking with Diffusion Bridge](https://mlsanity.com/paper/2402.11459/) (2024)
- [Bayesian neural network with pretrained protein embedding enhances prediction accuracy of drug-protein interaction](https://mlsanity.com/paper/2012.08194/) (2020)
- [DTIAM: A unified framework for predicting drug-target interactions, binding affinities and activation/inhibition mechanisms](https://mlsanity.com/paper/2312.15252/) (2023)

## Abstract (arXiv)

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.