Yangtian Zhang

01About

I'm a Ph.D. student in Computer Science at Yale, co-advised by Prof. Rex Ying and Prof. David van Dijk. Before Yale I worked with Prof. Jian Tang at Mila on geometric generative modeling, and I hold a B.Eng. in Computer Science from the ACM Honors Class at Shanghai Jiao Tong University.

My research is on generative modeling, mostly diffusion and flow models for discrete and structured data: discrete diffusion for language that conditions on its whole trajectory (CaDDi), diffusion over permutations via reflected soft ranks (Soft-Rank Diffusion), and, in earlier work, geometric models on the torus and in SE(3) (DiffPack, E3Bind).

I'm equally interested in steering and accelerating these models. I've used policy gradients to learn a model's own generation order (ICML'26 Spotlight), RL to post-train LLMs to refine sequences through verifiable edits (STRIDE), and at Meta I'm working on distilling diffusion language models for efficient generation.

02News

May 2026Joining Meta as a Research Scientist Intern this summer, working with Dr. Fei Sha in Menlo Park.
May 2026Three papers at ICML 2026. Variational Learning for Insertion-based Generation is a Spotlight (top 2.2%); Soft-Rank Diffusion and STRIDE were also accepted.
May 2026CoPersona accepted at KDD 2026.
Sep 2025CaDDi accepted at NeurIPS 2025. See you in San Diego!
Jun 2025Joining Google DeepMind as a Student Researcher in New York.
May 2025FlowCF accepted at KDD 2025.
Sep 2024Our survey on diffusion models for recommender systems is online.
Sep 2023DiffPack accepted at NeurIPS 2023.
Jan 2023E3Bind accepted at ICLR 2023.
Aug 2021We released TorchDrug, a machine learning platform for drug discovery.

03Publications

Variational Learning for Insertion-based Generation

Yangtian Zhang*, Zhe Wang*, Arthur Gretton, Rex Ying, David van Dijk, Michalis Titsias, Jiaxin Shi (*equal contribution)

ICML 2026Spotlight · top 2.2%

Treat the generation order as a latent variable. Tokens are inserted anywhere, the order is learned by variational inference, and left-to-right AR and discrete diffusion fall out as special cases.

arXiv PDF

We introduce a variational learning framework for insertion-based generation, where sequences are constructed via a learned-order autoregressive process that dynamically selects insertion positions. By formulating the generation order as a latent variable, our method enables flexible-length sequence construction and unifies left-to-right autoregression with discrete diffusion under a single principled framework.
```
@article{zhang2026insertion,
  title={Variational Learning for Insertion-based Generation},
  author={Zhang, Yangtian and Wang, Zhe and Gretton, Arthur and Ying, Rex and van Dijk, David and Titsias, Michalis and Shi, Jiaxin},
  journal={International Conference on Machine Learning (ICML), Spotlight (top 2.2\%)},
  year={2026}
}
```
Learning Permutation Distributions via Reflected Diffusion on Ranks

Sizhuang He*, Yangtian Zhang*, Shiyang Zhang, David van Dijk (*equal contribution)

ICML 2026

Soft-Rank Diffusion. Lift a permutation to continuous soft ranks, diffuse them with reflection at the boundaries, and denoise with contextual Plackett–Luce heads.

arXiv PDF

We propose Soft-Rank Diffusion, a discrete diffusion framework for learning probability distributions on the finite symmetric group. Rather than relying on shuffle-based random walks, we lift permutations to a continuous latent representation of order and define a structured soft-rank forward process. We further introduce contextualized generalized Plackett-Luce (cGPL) denoisers for the reverse process, enabling flexible and stable generative modeling of permutation distributions with applications to ranking and combinatorial structure learning.
```
@article{he2026permutation,
  title={Learning Permutation Distributions via Reflected Diffusion on Ranks},
  author={He, Sizhuang and Zhang, Yangtian and Zhang, Shiyang and van Dijk, David},
  journal={International Conference on Machine Learning (ICML)},
  year={2026}
}
```
STRIDE: Post-Training LLMs to Reason and Refine Bio-Sequences via Edit Trajectories

Daiheng Zhang, Shiyang Zhang, Sizhuang He, Yangtian Zhang, Syed Asad Rizvi, David van Dijk

ICML 2026

An LLM that writes its reasoning as an executable chain of INSERT / DELETE / REPLACE edits, trained with SFT on shortest edit paths and then RL.

arXiv PDF Code

STRIDE (Sequence Trajectory Refinement via Internalized Denoising Emulation) is a post-training framework that trains an LLM to emit executable trajectories of atomic edits (INSERT/DELETE/REPLACE) as a verifiable reasoning trace for variable-length refinement. STRIDE combines supervised fine-tuning on Levenshtein-aligned shortest edit demonstrations with group-based policy optimization to align edit trajectories with task rewards while preserving coherent editing behavior. STRIDE improves variable-length protein editing success from 42% to 89% while increasing novelty from 47% to 97%, and yields stronger validity and controllability compared to diverse baselines.
```
@article{zhang2026stride,
  title={STRIDE: Post-Training LLMs to Reason and Refine Bio-Sequences via Edit Trajectories},
  author={Zhang, Daiheng and Zhang, Shiyang and He, Sizhuang and Zhang, Yangtian and Rizvi, Syed Asad and van Dijk, David},
  journal={International Conference on Machine Learning (ICML)},
  year={2026}
}
```
CoPersona: Collaborative Persona Graphs for Robust LLM Personalization

Yangtian Zhang*, Leyao Wang*, Hiren Madhu, Ngoc Bui, Walter Roznyatovskiy, Rex Ying (*equal contribution)

KDD 2026

Sparse user histories leave facets of a persona unobserved. Split each history into facets, link peers facet-by-facet in a multiplex graph, and borrow what is missing.

arXiv PDF

Real-world LLM personalization is often constrained by sparse and skewed user histories. CoPersona is a graph-based collaborative personalization framework that completes sparse user profiles by borrowing signals from behaviorally similar peers. It decomposes interaction histories into facet-level representations and explicitly models peer-to-peer, facet-level alignment through a multiplex persona graph, then combines non-parametric peer retrieval with parametric graph reasoning at inference time. Experiments across multiple domains and model scales demonstrate consistent improvements over strong baselines.
```
@article{zhang2026copersona,
  title={CoPersona: Collaborative Persona Graphs for Robust LLM Personalization},
  author={Zhang, Yangtian and Wang, Leyao and Madhu, Hiren and Bui, Ngoc and Roznyatovskiy, Walter and Ying, Rex},
  journal={ACM SIGKDD Conference on Knowledge Discovery and Data Mining},
  year={2026}
}
```
Non-Markovian Discrete Diffusion with Causal Language Models

Yangtian Zhang*, Sizhuang He*, Daniel Levine, Lawrence Zhao, David Zhang, Syed A Rizvi, Emanuele Zappala, Rex Ying, David van Dijk (*equal contribution)

NeurIPS 2025

CaDDi. Each denoising step conditions on the entire trajectory, not just the last state, so a causal LM can run discrete diffusion and reuse pretrained weights unchanged.

arXiv PDF

Discrete diffusion models offer a flexible, controllable approach to structured sequence generation, yet they still lag behind causal language models in expressive power. A key limitation lies in their reliance on the Markovian assumption, which restricts each step to condition only on the current state, leading to potential uncorrectable error accumulation. We introduce CaDDi, a discrete diffusion model that conditions on the entire generative trajectory, thereby lifting the Markov constraint and allowing the model to revisit and improve past states. By unifying sequential (causal) and temporal (diffusion) reasoning in a single non-Markovian transformer, CaDDi also treats standard causal language models as a special case and permits the direct reuse of pretrained LLM weights with no architectural changes. Empirically, CaDDi outperforms state-of-the-art discrete diffusion baselines on natural-language benchmarks, substantially narrowing the remaining gap to large autoregressive transformers.
```
@article{zhang2025caddi,
  title={Non-Markovian Discrete Diffusion with Causal Language Models},
  author={Zhang, Yangtian and He, Sizhuang and Levine, Daniel and Zhao, Lawrence and Zhang, David and Rizvi, Syed A and Zappala, Emanuele and Ying, Rex and van Dijk, David},
  journal={Advances in Neural Information Processing Systems},
  year={2025}
}
```
Flow Matching for Collaborative Filtering

Chengkai Liu*, Yangtian Zhang*, Jianling Wang, Rex Ying, James Caverlee (*equal contribution)

KDD 2025

FlowCF. Flow from a behavior-guided prior to binary implicit feedback with a discrete flow, for accurate and fast generative recommendation.

arXiv PDF Code

Generative models have shown great promise in collaborative filtering by capturing the underlying distribution of user interests and preferences. However, existing approaches struggle with inaccurate posterior approximations and misalignment with the discrete nature of recommendation data. We propose FlowCF, a flow-based recommendation system leveraging flow matching for collaborative filtering, with (1) a behavior-guided prior that aligns with user behavior patterns to handle sparse and heterogeneous user-item interactions, and (2) a discrete flow framework that preserves the binary nature of implicit feedback while keeping the stable training and efficient inference of flow matching. FlowCF achieves state-of-the-art accuracy across datasets with the fastest inference speed.
```
@article{liu2025flow,
  title={Flow Matching for Collaborative Filtering},
  author={Liu, Chengkai and Zhang, Yangtian and Wang, Jianling and Ying, Rex and Caverlee, James},
  journal={ACM SIGKDD Conference on Knowledge Discovery and Data Mining},
  year={2025}
}
```
CaLMFlow: Volterra Flow Matching using Causal Language Models

Sizhuang He, Daniel Levine, Ivan Vrkic, Marco Francesco Bressana, David Zhang, Syed Asad Rizvi, Yangtian Zhang, Emanuele Zappala, David van Dijk

Preprint 2024

Flow matching recast as a Volterra integral equation and solved by a causal language model that attends over the whole discretized path.

arXiv PDF
```
@article{he2024calmflow,
  title={CaLMFlow: Volterra Flow Matching using Causal Language Models},
  author={He, Sizhuang and Levine, Daniel and Vrkic, Ivan and Bressana, Marco Francesco and Zhang, David and Rizvi, Syed Asad and Zhang, Yangtian and Zappala, Emanuele and van Dijk, David},
  journal={arXiv preprint arXiv:2410.05292},
  year={2024}
}
```
A Survey on Diffusion Models for Recommender Systems

Jianghao Lin, Jiaqi Liu, Jiachen Zhu, Yunjia Xi, Chengkai Liu, Yangtian Zhang, Yong Yu, Weinan Zhang

Preprint 2024

arXiv
DiffPack: A Torsional Diffusion Model for Autoregressive Protein Side-Chain Packing

Yangtian Zhang*, Zuobai Zhang*, Bozitao Zhong, Sanchit Misra, Jian Tang (*equal contribution)

NeurIPS 2023

Diffuse only what can move. Side-chain torsions are denoised on the torus, one angle at a time from χ₁ to χ₄, with a 60× smaller model than prior art.

arXiv PDF

Accurately predicting the conformation of protein side-chains given their backbones is important for structure prediction, design and protein-protein interactions. Existing machine learning methods treat the problem as regression and overlook the constraints of constant covalent bond lengths and angles. DiffPack is a torsional diffusion model that learns the joint distribution of side-chain torsional angles, the only degrees of freedom in side-chain packing, by diffusing and denoising on the torsional space. To avoid issues from perturbing all four angles at once, it generates them autoregressively from χ1 to χ4. DiffPack improves angle accuracy by 11.9% and 13.5% on CASP13 and CASP14 with 60× fewer parameters, and also improves AlphaFold2's side-chain predictions.
```
@article{zhan2023diffpack,
  title={DiffPack: A Torsional Diffusion Model for Autoregressive Protein Side-Chain Packing},
  author={Zhang, Yangtian and Zhang, Zuobai and Zhong, Bozitao and Misra, Sanchit and Tang, Jian},
  journal={Advances in Neural Information Processing Systems},
  year={2023}
}
```
E3Bind: An End-to-End Equivariant Network for Protein-Ligand Docking

Yangtian Zhang*, Huiyu Cai*, Chence Shi, Bozitao Zhong, Jian Tang (*equal contribution)

ICLR 2023

Dock like AlphaFold folds. An equivariant network that iteratively refines the ligand pose inside the pocket, end to end.

arXiv PDF

In silico prediction of the ligand binding pose to a given protein target is a crucial but challenging task in drug discovery. This work focuses on blind flexible self-docking. Inspired by AlphaFold2, we propose E3Bind, an end-to-end equivariant network that iteratively updates the ligand pose, modeling the protein-ligand interaction through the geometric constraints in docking and the local context of the binding site. Experiments on standard benchmarks demonstrate superior performance over traditional and recent deep learning methods.
```
@article{zhang2022e3bind,
  title={E3Bind: An End-to-End Equivariant Network for Protein-Ligand Docking},
  author={Zhang, Yangtian and Cai, Huiyu and Shi, Chence and Zhong, Bozitao and Tang, Jian},
  journal={Proceedings of the International Conference on Learning Representations (ICLR)},
  year={2023}
}
```

PEER: A Comprehensive and Multi-Task Benchmark for Protein Sequence Understanding

Minghao Xu, Zuobai Zhang, Jiarui Lu, Zhaocheng Zhu, Yangtian Zhang, Chang Ma, Runcheng Liu, Jian Tang

NeurIPS 2022 · Datasets & Benchmarks

Paper Website

@article{xu2022peer,
  title={PEER: A Comprehensive and Multi-task Benchmark for Protein Sequence Understanding},
  author={Xu, Minghao and Zhang, Zuobai and Lu, Jiarui and Zhu, Zhaocheng and Zhang, Yangtian and Chang, Ma and Liu, Runcheng and Tang, Jian},
  journal={Advances in Neural Information Processing Systems (Datasets and Benchmarks Track)},
  volume={35},
  pages={35156--35173},
  year={2022}
}

TorchDrug: A Powerful and Flexible Machine Learning Platform for Drug Discovery

Zhaocheng Zhu, Chence Shi, Zuobai Zhang, Shengchao Liu, Minghao Xu, Xinyu Yuan, Yangtian Zhang, Junkun Chen, Huiyu Cai, Jiarui Lu, Chang Ma, Runcheng Liu, Louis-Pascal Xhonneux, Meng Qu, Jian Tang

Preprint 2022

arXiv Website

@article{zhu2022torchdrug,
  title={TorchDrug: A Powerful and Flexible Machine Learning Platform for Drug Discovery},
  author={Zhu, Zhaocheng and Shi, Chence and Zhang, Zuobai and Liu, Shengchao and Xu, Minghao and Yuan, Xinyu and Zhang, Yangtian and Chen, Junkun and Cai, Huiyu and Lu, Jiarui and others},
  journal={Preprint},
  year={2022}
}

Imitation Learning via Multi-Step Occupancy Measure Matching

Minghuan Liu, Hangyu Wang, Yangtian Zhang, Minkai Xu, Zhengbang Zhu, Weinan Zhang

Preprint 2022

PDF

@article{liuautogail,
  title={Imitation Learning via Multi-Step Occupancy Measure Matching},
  author={Liu, Minghuan and Wang, Hangyu and Zhang, Yangtian and Xu, Minkai and Zhu, Zhengbang and Zhang, Weinan},
  journal={Preprint},
  year={2022}
}

04Experience

2026 –
Meta · Research Scientist Intern · with Fei Sha
Distilling diffusion language models for efficient generation.
Menlo Park
2025
Google DeepMind · Student Researcher · with Jiaxin Shi & Michalis Titsias
Discrete diffusion and learning-to-order insertion models (→ ICML'26 Spotlight).
New York
2024
Microsoft Research AI4Science · Research Intern · with Tao Qin
Multi-modal foundation models for scientific discovery.
Beijing
2021 – 24
Mila – Quebec AI Institute · Research Assistant · with Jian Tang
Geometric generative models for proteins and molecular docking; TorchDrug & TorchProtein.
Montréal
2020 – 22
APEX Lab, SJTU · Undergraduate Researcher · with Yong Yu & Weinan Zhang
Multi-step imitation learning.
Shanghai

Education

2024 –
Yale University
Ph.D. in Computer Science
2018 – 22
Shanghai Jiao Tong University
B.Eng. in Computer Science, ACM Honors Class · summa cum laude

Service

Reviewer for NeurIPS, ICML, ICLR, and AAAI.