Jiawei Zhang

Jiawei Zhang

Assistant professor in CS @ UW-Madison

University of Wisconsin–Madison

Biography

Jiawei Zhang is an assistant professor in the Department of Computer Sciences at University of Wisconsin–Madison. Previously, he was a postdoctoral fellow supported by The MIT Postdoctoral Fellowship For Engineering Excellence in the Laboratory for Information & Decision Systems (LIDS) at MIT, working with Prof. Asuman Ozdaglar and Prof. Saurabh Amin. He obtained the Ph.D. degree in Computer and Information Engineering from the Chinese University of Hong Kong, Shenzhen, with an advisory of Prof. Zhi-Quan (Tom) Luo, and was honored with the Presidential Award for Outstanding Doctoral Students. He obtained the B.Sc. in Mathematics (Hua Loo-Keng Talent Program) from the University of Science and Technology of China.

Interests
  • Efficient, scalable, generalizable, and robust optimization algorithms for large-scale optimization problems
  • Fundamental theories and algorithms for generative models, including LLMs and diffusion models
  • Optimization and decision-making for engineering applications, including energy systems and communication systems
Education

  • Ph.D. in Computer and Information Engineering

    The Chinese University of Hong Kong, Shenzhen (CUHK-Shenzhen)

  • B.Sc. in Mathematics (Hua Loo-Keng Talent Program)

    University of Science and Technology of China (USTC)

News

Slides on our recent works on optimal first-order algorithms for stochastic constrained optimization, theoretical analysis of GRPO from the adaptive gradient descent perspective, and lossless, efficient inference for diffusion language models, have been uploaded.
Congratulations to my first undergraduate intern, Ruichuan, for getting a Spotlight at ICML with his very first submission! Many thanks to my collaborator Ahmet Alacaoglu for co-supervising him with me.
Our paper "Addressing misspecification in contextual optimization", arXiv preprint arXiv:2409.10479, has been accepted by ICML 2025.
Our paper "Stochastic Smoothed Primal-Dual Algorithms for Nonconvex Optimization with Linear Inequality Constraints", arXiv preprint arXiv:2504.07607, has been accepted by ICML 2025.
I will join the Department of Computer Sciences at the University of Wisconsin–Madison as a tenure-track professor.

Research Topics

  • Efficient, scalable, generalizable, and robust optimization algorithms for large-scale optimization problems, including LLM training, constrained optimization, and bilevel optimization.

  • Fundamental theories and algorithms for generative models, including LLMs and diffusion models, focusing on efficient and reliable training and inference.

  • Optimization and decision-making for engineering applications, including energy systems and communication systems.

Prospective Students

I am looking for students interested in the research areas mentioned above or related topics about optimization, machine learning and their applications to science and engineering starting in 2025. They should have either a strong background in math and good coding skills, or a strong coding ability and adept theoretical insights. Students with backgrounds in mathematics, statistics, computer science, electronic engineering, and other related fields are welcome to apply.

Recent Publications

Cheng Ge, Caitlyn Heqi Yin, Hao Liang, Jiawei Zhang (2025). Why GRPO Needs Normalization: A Local-Curvature Perspective on Adaptive Gradients. Annual Conference on Neural Information Processing Systems (NeurIPS 2025) ER Workshop.

Cite

Shutong Wu, Jiawei Zhang (2025). Free Draft-and-Verification: Toward Lossless Parallel Decoding for Diffusion Large Language Models. arXiv preprint arXiv:2510.00294.

Cite

Kihyun Kim, Jiawei Zhang, Asuman Ozdaglar, Pablo A Parrilo (2025). Population-Proportional Preference Learning from Human Feedback: An Axiomatic Approach. arXiv preprint arXiv:2506.05619.

Cite

Wei Shen, Zhang Yaxiang, Minhui Huang, Mengfan Xu, Jiawei Zhang, Cong Shen (2025). MLorc: Momentum Low-rank Compression for Large Language Model Adaptation . arXiv preprint arXiv:2506.01897.

Cite

Wei Shen, Ruichuan Huang, Minhui Huang, Cong Shen, Jiawei Zhang (2025). On the convergence analysis of muon. arXiv preprint arXiv:2505.23737.

Cite

See all publications

Preprint Papers

Shutong Wu, Jiawei Zhang (2025). Free Draft-and-Verification: Toward Lossless Parallel Decoding for Diffusion Large Language Models. arXiv preprint arXiv:2510.00294.

Cite

Kihyun Kim, Jiawei Zhang, Asuman Ozdaglar, Pablo A Parrilo (2025). Population-Proportional Preference Learning from Human Feedback: An Axiomatic Approach. arXiv preprint arXiv:2506.05619.

Cite

Wei Shen, Zhang Yaxiang, Minhui Huang, Mengfan Xu, Jiawei Zhang, Cong Shen (2025). MLorc: Momentum Low-rank Compression for Large Language Model Adaptation . arXiv preprint arXiv:2506.01897.

Cite

Wei Shen, Ruichuan Huang, Minhui Huang, Cong Shen, Jiawei Zhang (2025). On the convergence analysis of muon. arXiv preprint arXiv:2505.23737.

Cite

Jiawei Zhang, Wenqiang Pu, Zhi-Quan Luo (2022). On the Iteration Complexity of Smoothed Proximal ALM for Nonconvex Optimization Problem with Convex Constraints. arXiv preprint arXiv:2207.06304.

Cite

Teaching Experience

Award

Outstanding Teaching Assistant granted by School of Science and Engineering, The Chinese University of Hong Kong, Shenzhen, China

Awards

The MIT Postdoctoral Fellowship For Engineering Excellence, Massachusetts Institute of Technology, 2023
Presidential Award for Outstanding Doctoral Students, The Chinese University of Hong Kong, Shenzhen, 2021
PhD Fellowship, Shenzhen Research Institute of Big Data, 2020-2021

Services

  • Action Editor, Transactions on Machine Learning Research (TMLR).

  • Area Chair, International Conference on Machine Learning (ICML 2026).

Contact