Generalization beyond the training distribution
Understanding how learning systems extrapolate across length, hardness, and composition—and what their successes reveal about learned algorithms.
I’m a PhD student in computer science at Princeton, advised by Sanjeev Arora.
My research focuses on how machine learning systems generalize beyond their training data. I’m particularly interested in self-improving models, synthetic data, and automated research.
Research
I combine controlled synthetic problems, empirical scaling, and theory to study what models learn and what they can discover next.
Self-improving learning systems
Studying synthetic data, weak-to-strong scaling, and feedback loops that let models create increasingly useful supervision.
Automating the research loop
Finding research problems where hypothesis generation becomes structured search, then designing agents that can explore that space in parallel.
Publications
Full CV →Fork-and-Flush: Escaping Idea Basins in Autoresearch Agents
Ziyang Cai, Christos Ziakas, Vasilis Kontonis, Tim Pearce, Siddhartha Sen, Akshay Krishnamurthy, Shivam Garg, Dimitris Papailiopoulos
arXivFull-bandwidth Transformer
Xi Wang, Ziyang Cai, Zheng Zhan, Harry Dong, Ying Fan, Gustavo de Rosa, Tim Pearce, John Langford
arXivT^2MLR: Transformer with Temporal Middle-Layer Recurrence
Ziyang Cai, Xingyu Zhu, Yihe Dong, Yinghui He, Sanjeev Arora
arXivGoedel-Architect: Streamlining Formal Theorem Proving with Blueprint Generation and Refinement
Jui-Hui Chung, Ziyang Cai, Zihao Li, Qishuo Yin, et al., Sanjeev Arora
arXivML-AutoResearch: Training Machine Learning Research Agents with Automatically Generated Environments
Ziyang Cai, Amir Saeidi, Harkirat Behl
arXivh1: Bootstrapping LLMs to Reason over Longer Horizons via Reinforcement Learning
Sumeet Ramesh Motwani, Alesia Ivanova, Ziyang Cai, Philip Torr, Riashat Islam, Shital Shah, Christian Schroeder de Witt, Charles London
arXivExtrapolation by Association: Length Generalization Transfer in Transformers
Ziyang Cai, Nayoung Lee, Avi Schwarzschild, Samet Oymak, Dimitris Papailiopoulos
arXivSelf-Improving Models Overcome Length and Hardness Generalization via Weak-to-Strong Scaling
Ziyang Cai*, Nayoung Lee*, Avi Schwarzschild, Kangwook Lee, Dimitris Papailiopoulos
arXivEverything Everywhere All at Once: LLMs Can In-Context Learn Multiple Tasks in Superposition
Zheyang Xiong, Ziyang Cai, John Cooper, et al.
arXivR&B: Domain Regrouping and Data Mixture Balancing for Efficient Foundation Model Training
Albert Ge, Tzu-Heng Huang, John Cooper, Avi Trost, Ziyi Chu, et al.
arXivDelving into Out-of-Distribution Detection with Vision-Language Representations
Yifei Ming, Ziyang Cai, Jiuxiang Gu, Yiyou Sun, Wei Li, Yixuan Li
arXivBackground
Before Princeton, I studied at UW–Madison and built automated synthetic-data pipelines for coding agents at Microsoft Research. I was previously advised by Dimitris Papailiopoulos, Junjie Hu, and Sharon Li.