Zuobai Zhang

PerturbDiff: Functional Diffusion for Single-Cell Perturbation Modeling

Yashi Zhang

Hongyu Guo

Building Virtual Cells that can accurately simulate cellular responses to perturbations is a long-standing goal in systems biology. A fundam… (voir plus)ental challenge is that high-throughput single-cell sequencing is destructive: the same cell cannot be observed both before and after a perturbation. Thus, perturbation prediction requires mapping unpaired control and perturbed populations. Existing models address this by learning maps between distributions, but typically assume a single fixed response distribution when conditioned on observed cellular context (e.g., cell type) and the perturbation type. In reality, responses vary systematically due to unobservable latent factors such as microenvironmental fluctuations and complex batch effects, forming a manifold of possible distributions for the same observed conditions. To account for this variability, we introduce PerturbDiff, which shifts modeling from individual cells to entire distributions. By embedding distributions as points in a Hilbert space, we define a diffusion-based generative process operating directly over probability distributions. This allows PerturbDiff to capture population-level response shifts across hidden factors. Benchmarks on established datasets show that PerturbDiff achieves state-of-the-art performance in single-cell response prediction and generalizes substantially better to unseen perturbations. See our project page (https://katarinayuan.github.io/PerturbDiff-ProjectPage/), where code and data will be made publicly available (https://github.com/DeepGraphLearning/PerturbDiff).

2026-04-29

International Conference on Machine Learning (accepté)

doi.org

openreview.net

Fast Proteome-Scale Protein Interaction Retrieval via Residue-Level Factorization

Narendra Chaudhary

Qian Cong

Jian Zhou

Sanchit Misra

Jian Tang

Protein-protein interactions (PPIs) are mediated at the residue level. Most sequence-based PPI models consider residue-residue interactions … (voir plus)across two proteins, which can yield accurate interaction scores but are too slow to scale. At proteome scale, identifying candidate PPIs requires evaluating nearly *all possible protein pairs*. For

2025-12-31

International Conference on Learning Representations (Accept (Poster))

openreview.net

Scaling Atomistic Protein Binder Design with Generative Pretraining and Test-Time Compute

Kieran Didi

Zuobai Zhang

Guoqing Zhou

Danny Reidenbach

Zhonglin Cao

Sooyoung Cha

Tomas Geffner

Christian Dallago

Jian Tang

Michael Bronstein

Martin Steinegger

Emine Kucukbenli

Arash Vahdat

Karsten Kreis

Protein interaction modeling is central to protein design, which has been transformed by machine learning with broad applications in drug di… (voir plus)scovery and beyond. In this landscape, structure-based de novo binder design is most often cast as either conditional generative modeling or sequence optimization via structure predictors ("hallucination"). We argue that this is a false dichotomy and propose Complexa, a novel fully atomistic binder generation method unifying both paradigms. We extend recent flow-based latent protein generation architecture and leverage the domain-domain interactions of monomeric computationally predicted protein structures to construct Teddymer, a new large-scale dataset of synthetic binder-target pairs for pretraining. Combined with high-quality experimental multimers, this enables training a strong base model. We then perform inference-time optimization with this generative prior, unifying the strengths of previously distinct generative and hallucination methods. Complexa sets a new state of the art in computational binder design benchmarks: it delivers markedly higher in-silico success rates than existing generative approaches, and our novel test-time optimization strategies greatly outperform previous hallucination methods under normalized compute budgets. We further demonstrate explicit interface hydrogen bond optimization, fold class-guided binder generation, and extensions to small molecule targets and enzyme design tasks, again surpassing prior methods. Code, models and new data will be publicly released.

2025-12-31

International Conference on Learning Representations (Accept (Oral))

openreview.net

Towards All-Atom Foundation Models for Biomolecular Binding Affinity Prediction

Liang Shi

Zuobai Zhang

Huiyu Cai

Santiago Miret

Zhi Yang

Jian Tang

Biomolecular interactions play a critical role in biological processes. While recent breakthroughs like AlphaFold 3 have enabled accurate mo… (voir plus)deling of biomolecular complex structures, predicting binding affinity remains challenging mainly due to limited high-quality data. Recent methods are often specialized for specific types of biomolecular interactions, limiting their generalizability. In this work, we repurpose AlphaFold 3 for representation learning to predict binding affinity, a non-trivial task that requires shifting from generative structure prediction to encoding observed geometry, simplifying the heavily conditioned trunk module, and designing a framework to jointly capture sequence and structural information. To address these challenges, we introduce the **Atom-level Diffusion Transformer (ADiT)**, which takes sequence and structure as inputs, employs a unified tokenization scheme, integrates diffusion transformers, and removes dependencies on multiple sequence alignments and templates. We pre-train three ADiT variants on the PDB dataset with a denoising objective and evaluate them across protein-ligand, drug-target, protein-protein, and antibody-antigen interactions. The model achieves state-of-the-art or competitive performance across benchmarks, scales effectively with model size, and successfully identifies wet-lab validated affinity-enhancing antibody mutations, establishing a generalizable framework for biomolecular interactions. We plan to release the code upon acceptance.

2025-12-31

International Conference on Learning Representations (Accept (Poster))

openreview.net

Consistent Synthetic Sequences Unlock Structural Diversity in Fully Atomistic De Novo Protein Design

Danny Reidenbach

Zhonglin Cao

Zuobai Zhang

Kieran Didi

Tomas Geffner

Guoqing Zhou

Jian Tang

Christian Dallago

Arash Vahdat

Emine Kucukbenli

Karsten Kreis

High-quality training datasets are crucial for the development of effective protein design models, but existing synthetic datasets often inc… (voir plus)lude unfavorable sequence-structure pairs, impairing generative model performance. We leverage ProteinMPNN, whose sequences are experimentally favorable as well as amenable to folding, together with structure prediction models to align high-quality synthetic structures with recoverable synthetic sequences. In that way, we create a new dataset designed specifically for training expressive, fully atomistic protein generators. By retraining La-Proteína, which models discrete residue type and side chain structure in a continuous latent space, on this dataset, we achieve new state-of-the-art results, with improvements of +54% in structural diversity and +27% in co-designability. To validate the broad utility of our approach, we further introduce Proteína-Atomística, a unified flow-based framework that jointly learns the distribution of protein backbone structure, discrete sequences, and atomistic side chains without latent variables. We again find that training on our new sequence-structure data dramatically boosts benchmark performance, improving Proteína-Atomística’s structural diversity by +73% and co-designability by +5%. Our work highlights the critical importance of aligned sequence-structure data for training high-performance de novo protein design models. All data will be publicly released.

2025-09-23

NeurIPS.cc/2025/Workshop/AI4Science (poster)

openreview.net

Bering: joint cell segmentation and annotation for spatial transcriptomics with transferred graph embeddings

Kang Jin

Zuobai Zhang

Ke Zhang

Francesca Viggiani

Claire Callahan

Jian Tang

Bruce J. Aronow

Jian Shu

Single-cell spatial transcriptomics such as in-situ hybridization or sequencing technologies can provide subcellular resolution that enables… (voir plus) the identification of individual cell identities, locations, and a deep understanding of subcellular mechanisms. However, accurate segmentation and annotation that allows individual cell boundaries to be determined remains a major challenge that limits all the above and downstream insights. Current machine learning methods heavily rely on nuclei or cell body staining, resulting in the significant loss of both transcriptome depth and the limited ability to learn latent representations of spatial colocalization relationships. Here, we propose Bering, a graph deep learning model that leverages transcript colocalization relationships for joint noise-aware cell segmentation and molecular annotation in 2D and 3D spatial transcriptomics data. Graph embeddings for the cell annotation are transferred as a component of multi-modal input for cell segmentation, which is employed to enrich gene relationships throughout the process. To evaluate performance, we benchmarked Bering with state-of-the-art methods and observed significant improvement in cell segmentation accuracies and numbers of detected transcripts across various spatial technologies and tissues. To streamline segmentation processes, we constructed expansive pre-trained models, which yield high segmentation accuracy in new data through transfer learning and self-distillation, demonstrating the generalizability of Bering.

2025-07-17

Nature Communications (publié)

doi.org

Cell ontology guided transcriptome foundation model

Manqi Zhou

Boyu Han

Transcriptome foundation models (TFMs) hold great promises of deciphering the transcriptomic language that dictate diverse cell functions by… (voir plus) self-supervised learning on large-scale single-cell gene expression data, and ultimately unraveling the complex mechanisms of human diseases. However, current TFMs treat cells as independent samples and ignore the taxonomic relationships between cell types, which are available in cell ontology graphs. We argue that effectively leveraging this ontology information during the TFM pre-training can improve learning biologically meaningful gene co-expression patterns while preserving TFM as a general purpose foundation model for downstream zero-shot and fine-tuning tasks. To this end, we present **s**ingle **c**ell, **Cell**-**o**ntology guided TFM (scCello). We introduce cell-type coherence loss and ontology alignment loss, which are minimized along with the masked gene expression prediction loss during the pre-training. The novel loss component guide scCello to learn the cell-type-specific representation and the structural relation between cell types from the cell ontology graph, respectively. We pre-trained scCello on 22 million cells from CellxGene database leveraging their cell-type labels mapped to the cell ontology graph from Open Biological and Biomedical Ontology Foundry. Our TFM demonstrates competitive generalization and transferability performance over the existing TFMs on biologically important tasks including identifying novel cell types of unseen cells, prediction of cell-type-specific marker genes, and cancer drug responses. Source code and model weights are available at https://github.com/DeepGraphLearning/scCello.

2024-12-10

Neural Information Processing Systems (Accept (spotlight))

doi.org

openreview.net

Multi-Scale Representation Learning for Protein Fitness Prediction

Zuobai Zhang

Pascal Notin

Yining Huang

Aurelie Lozano

Vijil Chenthamarakshan

Debora Marks

Payel Das

Jian Tang

Designing novel functional proteins crucially depends on accurately modeling their fitness landscape. Given the limited availability of func… (voir plus)tional annotations from wet-lab experiments, previous methods have primarily relied on self-supervised models trained on vast, unlabeled protein sequence or structure datasets. While initial protein representation learning studies solely focused on either sequence or structural features, recent hybrid architectures have sought to merge these modalities to harness their respective strengths. However, these sequence-structure models have so far achieved only incremental improvements when compared to the leading sequence-only approaches, highlighting unresolved challenges effectively leveraging these modalities together. Moreover, the function of certain proteins is highly dependent on the granular aspects of their surface topology, which have been overlooked by prior models. To address these limitations, we introduce the Sequence-Structure-Surface Fitness (S3F) model - a novel multimodal representation learning framework that integrates protein features across several scales. Our approach combines sequence representations from a protein language model with Geometric Vector Perceptron networks encoding protein backbone and detailed surface topology. The proposed method achieves state-of-the-art fitness prediction on the ProteinGym benchmark encompassing 217 substitution deep mutational scanning assays, and provides insights into the determinants of protein function. Our code is at https://github.com/DeepGraphLearning/S3F.

2024-12-10

Conference on Neural Information Processing Systems (Accept (poster))

doi.org

openreview.net

Pretrainable geometric graph neural network for antibody affinity maturation

Huiyu Cai

Zuobai Zhang

Mingkai Wang

Bozitao Zhong

Yanling Wu

Tianlei Ying

Jian Tang

Quanxiao Li

Yuxuan Zhong

Increasing the binding affinity of an antibody to its target antigen is a crucial task in antibody therapeutics development. This paper pres… (voir plus)ents a pretrainable geometric graph neural network, GearBind, and explores its potential inin silicoaffinity maturation. Leveraging multi-relational graph construction, multi-level geometric message passing and contrastive pretraining on mass-scale, unlabeled protein structural data, GearBind outperforms previous state-of-the-art approaches on SKEMPI and an independent test set. A powerful ensemble model based on GearBind is then derived and used to successfully enhance the binding of two antibodies with distinct formats and target antigens. ELISA EC50values of the designed antibody mutants are decreased by up to 17 fold, andKDvalues by up to 6.1 fold. These promising results underscore the utility of geometric deep learning and effective pretraining in macromolecule interaction modeling tasks.

2024-09-05

Nature Communications (publié)

doi.org

Augmenting Evolutionary Models with Structure-based Retrieval

Yining Huang

Zuobai Zhang

Jian Tang

Debora Susan Marks

Pascal Notin

2024-06-16

ICML.cc/2024/Workshop/ML4LMS (poster)

openreview.net

Str2str: A Score-Based Framework for Zero-Shot Protein Conformation Sampling

Jiarui Lu

Bozitao Zhong

Zuobai Zhang

Jian Tang

The dynamic nature of proteins is crucial for determining their biological functions and properties, for which Monte Carlo (MC) and molecula… (voir plus)r dynamics (MD) simulations stand as predominant tools to study such phenomena. By utilizing empirically derived force fields, MC or MD simulations explore the conformational space through numerically evolving the system via Markov chain or Newtonian mechanics. However, the high-energy barrier of the force fields can hamper the exploration of both methods by the rare event, resulting in inadequately sampled ensemble without exhaustive running. Existing learning-based approaches perform direct sampling yet heavily rely on target-specific simulation data for training, which suffers from high data acquisition cost and poor generalizability. Inspired by simulated annealing, we propose Str2Str, a novel structure-to-structure translation framework capable of zero-shot conformation sampling with roto-translation equivariant property. Our method leverages an amortized denoising score matching objective trained on general crystal structures and has no reliance on simulation data during both training and inference. Experimental results across several benchmarking protein systems demonstrate that Str2Str outperforms previous state-of-the-art generative structure prediction models and can be orders of magnitude faster compared to long MD simulations. Our open-source implementation is available at https://github.com/lujiarui/Str2Str

2024-05-06

International Conference on Learning Representations (Accept (poster))

doi.org

openreview.net

Fusing Neural and Physical: Augment Protein Conformation Sampling with Tractable Simulations

Bozitao Zhong

The protein dynamics are common and important for their biological functions and properties, the study of which usually involves time-consum… (voir plus)ing molecular dynamics (MD) simulations *in silico*. Recently, generative models has been leveraged as a surrogate sampler to obtain conformation ensembles with orders of magnitude faster and without requiring any simulation data (a "zero-shot" inference). However, being agnostic of the underlying energy landscape, the accuracy of such generative model may still be limited. In this work, we explore the few-shot setting of such pre-trained generative sampler which incorporates MD simulations in a tractable manner. Specifically, given a target protein of interest, we first acquire some seeding conformations from the pre-trained sampler followed by a number of physical simulations in parallel starting from these seeding samples. Then we fine-tuned the generative model using the simulation trajectories above to become a target-specific sampler. Experimental results demonstrated the superior performance of such few-shot conformation sampler at a tractable computational cost.

2024-03-03

ICLR.cc/2024/Workshop/GEM (poster)

openreview.net

La plateforme Mila Ventures

Boussole des politiques en IA

Séances d'information virtuelles

Zuobai Zhang

Publications

La plateforme Mila Ventures

Boussole des politiques en IA

Séances d'information virtuelles

Mots-clés populaires:

Zuobai Zhang

Publications