I am Ziheng Cheng, a third-year PhD student in Department of IEOR, UC Berkeley and fortunately supervised by Xin Guo. Prior to that, I got my B.S. degree in School of Mathematical Sciences, Peking University, supervised by Cheng Zhang. I was also very fortunate to have worked with Song Mei, Kun Yuan, Tengyu Ma. My research interests span broadly in statistics, optimization and machine learning, including language models and diffusion models, reinforcement learning, distributed optimization, sampling and variational inference. If you are interested in my research, please feel free to contact me.

News

• Sep, 2026 A new paper during internship at Seed!
   Periodic Weak Spots: Phase Sensitivity from Chunked KV-Cache Compression
   We identify a key limitation of LLM architectures with KV-cache compression, such as Deepseek V4 series.
• Jul, 2026 A new paper on arxiv!
   Actor-Critic Learning for Extended Mean Field Control with Deterministic Policies
• Jul, 2026 One paper accepted at COLM 2026!
   Multi-Mask Diffusion Language Models for Few-Step Generation
• May, 2026 One paper accepted at ICML 2026!
   Multi-Objective Learning for Diffusion Models: A Statistical Theory under Semi-Supervised Learning
• Jan, 2026 Join Bytedance Seed at San Jose as an intern! Looking forward to working on LLM pretraining!
• Sep, 2025 A new paper on arxiv!
   Deterministic Policy Gradient for Reinforcement Learning with Continuous Time and State
   We develop a groundbreaking paradigm for continuous-time deep RL.
• Sep, 2025 Two papers accepted at NeurIPS 2025!
   OVERT: A Benchmark for Over-Refusal Evaluation on Text-to-Image Models
   Provable Sample-Efficient Transfer Learning Conditional Diffusion Models via Representation Learning
• Jan, 2025 One paper accepted at ICLR 2025!
   Convergence of Distributed Adaptive Optimization with Local Updates
   The first end-to-end convergence guarantee of distributed Adam with local updates! Extremely solid techniques!

Selected Publications

• (Preprint) Periodic Weak Spots: Phase Sensitivity from Chunked KV-Cache Compression
  Xingyu Zhu*, Yi (Luke) Pu*, Ziheng Cheng*, Ang Lv*, Jing Liu, Lexing Ying, Yiyuan Ma, Xin Dong
  [Arxiv]

• (ICML 2026) Multi-Objective Learning for Diffusion Models: A Statistical Theory under Semi-Supervised Learning
  Ziheng Cheng*, Yixiao Huang*, Hanlin Zhu, Haoran Geng, Somayeh Sojoudi, Jitendra Malik, Pieter Abbeel, Xin Guo
  [Arxiv]

• (Preprint) Deterministic Policy Gradient for Reinforcement Learning with Continuous Time and State
  Ziheng Cheng, Xin Guo, Yufei Zhang
  [Arxiv]

• (NeurIPS 2025) OVERT: A Benchmark for Over-Refusal Evaluation on Text-to-Image Models
  Ziheng Cheng*, Yixiao Huang*, Hui Xu, Somayeh Sojoudi, Xuandong Zhao, Dawn Song, Song Mei
  [Arxiv]

• (NeurIPS 2025) Provable Sample-Efficient Transfer Learning Conditional Diffusion Models via Representation Learning
  Ziheng Cheng, Tianyu Xie, Shiyue Zhang, Cheng Zhang
  [Arxiv]

• (ICLR 2025) Convergence of Distributed Adaptive Optimization with Local Updates
  Ziheng Cheng, Margalit Glasgow
  [Arxiv]

• (ICML 2024) Kernel Semi-Implicit Variational Inference
  Ziheng Cheng*, Longlin Yu*, Tianyu Xie, Shiyue Zhang, Cheng Zhang
  [Arxiv]

• (ICLR 2024) Momentum Benefits Non-IID Federated Learning Simply and Provably
  Ziheng Cheng*, Xinmeng Huang*, Pengfei Wu, Kun Yuan
  [Arxiv]

Industry Experiences

• Bytedance Seed
  Research Intern, working on LLM pretraining
  Jan. 2026 – Present

• Microsoft Research Asia
  Research Intern, working on data selection
  Oct. 2023 – May. 2024