About

I am a Ph.D. student in the Machine Learning Department at Carnegie Mellon University, advised by Andrej Risteski and Aditi Raghunathan. Before CMU I received my B.S. in Mathematics from the School of Mathematical Sciences, Peking University, where I was advised by Liwei Wang.

My current research is aimed at building reliable foundation models. I work on this from two directions — the architecture of the model itself, and the algorithms used to train it — with the goal of making training more stable while making the model more creative, rather than trading one for the other.

Previously I worked on graph representation learning and on the theoretical foundations of reinforcement learning. My work on quantifying the expressive power of graph neural networks received an ICLR 2024 Best Paper Honorable Mention as first author, and was presented as an Oral.

As an undergraduate I was also fortunate to collaborate with Prof. Chi Jin at Princeton on the theory of partially observable and multi-agent reinforcement learning, and with Prof. Yuejie Chi at CMU on robust multi-agent RL.

Currently interested in:

• LLM reasoning & post-training

  RL for language models, entropy control, and mitigating sharpening and premature confidence. 6 papers

  • Demystifying Entropy Control in LLM RL TrainingICML 2026 Spotlight
  • Differential Smoothing Mitigates Sharpening and Improves LLM ReasoningICML 2026
  • Understanding and Mitigating Premature Confidence for Better LLM ReasoningCOLM 2026
  • Pretraining Data Scale Reverses the AdamW–Muon Fine-Tuning GapCOLM 2026 Workshop Oral
  • Lossless Anti-Distillation SamplingPreprint
  • Momentum Streams for Optimizer-Inspired TransformersPreprint

• AI agents for mathematical research

  Using LLMs to attack open problems in mathematics — most recently the Gilbert–Pollak conjecture on the Steiner ratio. 1 paper

  • Towards Solving the Gilbert–Pollak Conjecture via Large Language ModelsICML 2026

• RL & imitation learning

  Foundations of reinforcement and imitation learning, and robust multi-agent RL. 4 papers

  • When Does Online Imitation Learning Help in LLM Post-Training?Preprint
  • Breaking the Curse of Multiagency in Robust Multi-Agent RLICML 2025
  • Taming the Curses of Multiagency in Robust Markov GamesPreprint
  • Demystifying Entropy Control in LLM RL TrainingICML 2026 Spotlight

• Graph representation learning

  GNN expressiveness, together with algebraic and structural graph theory. 2 papers

  • Beyond Weisfeiler-Lehman: A Quantitative Framework for GNN ExpressivenessICLR 2024 Best Paper HM, Oral
  • Homomorphism Expressivity of Spectral Invariant GNNsICLR 2025 Oral

I am always happy to talk about any of the above — feel free to reach out at jgai@andrew.cmu.edu.

Education

• 2025present

  Carnegie Mellon University

  Ph.D. in Machine Learning · Machine Learning Department

  Advised by Andrej Risteski and Aditi Raghunathan
  Pittsburgh, PA, USA
• 20212025

  Peking University

  B.S. in Mathematics · School of Mathematical Sciences

  Advised by Liwei Wang
  GPA 3.957 / 4.00 — ranked 1 / 217
  Beijing, China

Publications

* denotes equal contribution; † denotes corresponding author.

Research Experience

• CMU

  Graduate Research Assistant

  Advised by Aditi Raghunathan and Andrej Risteski
• CMU

  Research Intern

  Advised by Yuejie Chi and Laixi Shi
  Robust multi-agent reinforcement learning
• Princeton

  Research Intern

  Advised by Chi Jin
  Partially observable and multi-agent reinforcement learning theory
• PKU

  Undergraduate Researcher

  Advised by Liwei Wang
  Graph representation learning and GNN expressiveness

Talks

Selected Honors

• 2024ICLR 2024 Best Paper Honorable Mention — as first author
• 2025Outstanding Graduate of Peking University
• 2024SenseTime Scholarship — 25 undergraduates nationwide in AI
• 2024National Scholarship of China — highest undergraduate honor in China
• 2023National Scholarship of China
• 2024Merit Student of Peking University
• 2023Merit Student of Peking University
• 2023First Prize, 14th Chinese College Students' Mathematical Contest
• 2023Winner Prize, 14th S.-T. Yau College Student Mathematics Contest
• 2020First Prize, 37th Chinese Physics Olympiad