Portrait of Zuobai Zhang is unavailable

Zuobai Zhang

Alumni
Website
Google Scholar

Publications

PerturbDiff: Functional Diffusion for Single-Cell Perturbation Modeling
Xinyu Yuan
Xixian Liu
Yashi Zhang
Zuobai Zhang
Hongyu Guo
Jian Tang
Building Virtual Cells that can accurately simulate cellular responses to perturbations is a long-standing goal in systems biology. A fundam… (see more)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 (accepted)
doi.org
openreview.net
Fast Proteome-Scale Protein Interaction Retrieval via Residue-Level Factorization
Jianan Zhao
Zhihao Zhan
Narendra Chaudhary
Xinyu Yuan
Zuobai Zhang
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 … (see more)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
openreview.net
La-Proteina: Atomistic Protein Generation via Partially Latent Flow Matching
Tomas Geffner
Kieran Didi
Zhonglin Cao
Danny Reidenbach
Zuobai Zhang
Christian Dallago
Emine Kucukbenli
Karsten Kreis
Arash Vahdat
Recently, many generative models for de novo protein structure design have emerged. Yet, only few tackle the difficult task of directly gene… (see more)rating fully atomistic structures jointly with the underlying amino acid sequence. This is challenging, for instance, because the model must reason over side chains that change in length during generation. We introduce La-Proteina for atomistic protein design based on a novel partially latent protein representation: coarse backbone structure is modeled explicitly, while sequence and atomistic details are captured via per-residue latent variables of fixed dimensionality, thereby effectively side-stepping challenges of explicit side-chain representations. Flow matching in this partially latent space then models the joint distribution over sequences and full-atom structures. La-Proteina achieves state-of-the-art performance on multiple generation benchmarks, including all-atom co-designability, diversity, and structural validity, as confirmed through detailed structural analyses and evaluations. Notably, La-Proteina also surpasses previous models in atomistic motif scaffolding performance, unlocking critical atomistic structure-conditioned protein design tasks. Moreover, La-Proteina is able to generate co-designable proteins of up to 800 residues, a regime where most baselines collapse and fail to produce valid samples, demonstrating La-Proteina's scalability and robustness.
2025-12-31
International Conference on Learning Representations (Accept (Poster))
doi.org
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… (see more)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
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… (see more)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
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… (see more)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
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… (see more) 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 (published)
doi.org
Proteina: Scaling Flow-based Protein Structure Generative Models
Tomas Geffner
Kieran Didi
Zuobai Zhang
Danny Reidenbach
Zhonglin Cao
Jason Yim
Mario Geiger
Christian Dallago
Emine Kucukbenli
Arash Vahdat
Karsten Kreis
Recently, diffusion- and flow-based generative models of protein structures have emerged as a powerful tool for de novo protein design. Here… (see more), we develop *Proteina*, a new large-scale flow-based protein backbone generator that utilizes hierarchical fold class labels for conditioning and relies on a tailored scalable transformer architecture with up to
2025-04-22
International Conference on Learning Representations (Accept (Oral))
doi.org
openreview.net
Cell ontology guided transcriptome foundation model
Xinyu Yuan
Zhihao Zhan
Zuobai Zhang
Manqi Zhou
Jianan Zhao
Boyu Han
Yue Li
Jian Tang
Transcriptome foundation models (TFMs) hold great promises of deciphering the transcriptomic language that dictate diverse cell functions by… (see more) 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… (see more)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… (see more)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 (published)
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
openreview.net