Qitian Wu
Research Scientist, Google Research
New York, NY·qitianwu228 [AT] gmail.com
👋 About Me
🔬 I am a Research Scientist at Google Research, working on LLM post-training and agents. My recent projects:
- Gemini post-training for Text-to-SQL generation: post-trained Gemini to the top of the BIRD Text-to-SQL leaderboard. (Synthetic data generation, sandbox environments with tool use, reward models)
- Agents for systems optimization: built an agentic system that iteratively optimizes production code in evolving environments, deployed in compilers. (Agentic harness, long-horizon reasoning, endpoint serving)
- Gemini reasoning for network-state observability: partnered with infrastructure teams to scope large-scale cluster telemetry into an LLM reasoning task. (Data curation and filtering, evaluation pipelines)
🧬 Previously, as an Eric and Wendy Schmidt Center Postdoctoral Fellow at the Broad Institute of MIT and Harvard (with Caroline Uhler), I led the pretraining of a biomedical foundation model from scratch, from data pipelines to distributed training on GPU clusters, and applied it to discovery of therapeutic targets in biomedicine (paper).
⚡ During my PhD, I worked as a part-time research intern at Amazon Web Services for three years, designing linear-complexity attention for 1M-token contexts and developing scalable Transformer architectures (NodeFormer, DIFFormer and SGFormer) as widely used open-source models, with 1,300+ citations and 900+ GitHub stars in total. Earlier, I was an applied research intern at Tencent WeChat, building deep learning models for industry-scale recommender systems.
🎓 I received my Ph.D., M.E. and B.E. degrees from Shanghai Jiao Tong University, advised by Junchi Yan, and collaborated with David Wipf, Philip S. Yu, Hongyuan Zha and Michael Bronstein during my PhD. My research has been recognized with the Microsoft Research PhD Fellowship (11 recipients in Asia), the Baidu PhD Fellowship (10 recipients worldwide) and the Elsevier/Stanford World's Top 2% Scientists list in Artificial Intelligence. I serve as an Area Chair for NeurIPS and ICLR.
📝 Selected Publications
My research has been extensively published in NeurIPS, ICML, ICLR, JMLR and TPAMI with 4,700+ citations in total. Full list on Google Scholar.
LLMs, Agents & Foundation Models
Resolving Action Bottleneck: Agentic Reinforcement Learning Informed by Token-Level Energy
L. He, J. Zhu, Y. Zhou, Z. Gu, J. Liu, W.-C. Huang, H. P. Zou, D. Wipf, P. S. Yu, Qitian Wu
Qitian Wu, Q. Gong, L. Yuan, Z. Li, O. Ashenberg, F. Chen, R. Xavier, C. Uhler
SLMRec: Distilling Large Language Models into Small for Sequential Recommendation
W. Xu, Qitian Wu, Z. Liang, J. Han, X. Ning, Y. Shi, W. Lin, Y. Zhang
Lock Before Reveal: Valid Causal Inference When LLM Agents Search Relational Data
L. He*, Qitian Wu*, D. Wipf (* Equal contribution)
Scalable Transformer Architectures
NodeFormer: A Scalable Graph Structure Learning Transformer for Node Classification
Qitian Wu, W. Zhao, Z. Li, D. Wipf, J. Yan
SGFormer: Simplifying and Empowering Transformers for Large-Graph Representations
Qitian Wu, W. Zhao, C. Yang, H. Zhang, F. Nie, H. Jiang, Y. Bian, J. Yan
📦 Integrated into PyTorch Geometric as a built-in model and attention layer
DIFFormer: Scalable (Graph) Transformers Induced by Energy Constrained Diffusion
Qitian Wu, C. Yang, W. Zhao, Y. He, D. Wipf, J. Yan
🔥 Rated top 8 of 4,966 ICLR 2023 submissions
PAINET: A Principled Efficient Transformer for 3D Dynamics Modeling
K. Yang, Y. Huang, J. Tao, W. Wang, Qitian Wu
Supercharging Graph Transformers with Advective Diffusion
Qitian Wu, C. Yang, K. Zeng, M. Bronstein
Theory of Machine Learning & Generative Models
Transformers from Diffusion: A Unified Framework for Neural Message Passing
Qitian Wu, D. Wipf, J. Yan
DiffPuter: Empowering Diffusion Models for Missing Data Imputation
H. Zhang, L. Fang, Qitian Wu, P. S. Yu
Bridging Explicit and Implicit Deep Generative Models via Neural Stein Estimators
Qitian Wu, R. Gao, H. Zha
Generalization under Distribution Shift
Handling Distribution Shifts on Graphs: An Invariance Perspective
Qitian Wu, H. Zhang, J. Yan, D. Wipf
Graph Out-of-Distribution Generalization via Causal Intervention
Qitian Wu, F. Nie, C. Yang, T. Bao, J. Yan
🏆 Selected Awards
🤝 Leadership & Service
| Conferences | Area Chair, NeurIPS and ICLR; Senior Program Committee, AAAI and IJCAI |
| Journals | Action Editor, TMLR; Editorial Board Reviewer, JMLR |