Tsinghua University · Electronic Engineering

Jiawei Zhang 张家炜

Ph.D. Student · Generative Modeling & Diffusion Models

I am a Ph.D. student in Information and Communication Engineering at Tsinghua University, advised by Prof. Yuantao Gu. I received my B.Eng. in Electronic Engineering from Tsinghua University in 2023.

My research interests include generative modeling, diffusion models, and model distillation. I also work on diffusion-based inference for Bayesian inverse problems.

jiawei-z23@mails.tsinghua.edu.cn
Google Scholar CV (PDF)
Watercolor portrait of Jiawei Zhang wearing sunglasses outdoors

Research Experience

Ant Research — Research Intern

Dec. 2025 – Present
  • Worked on one-step image generation by distilling Diffusion Transformers in the latent space of Representation Autoencoders (RAEs).
  • Developed a one-step distillation method based on Drifting Models, with an FID of 1.48 on ImageNet.
  • Conducted distributed training with PyTorch and FSDP.

Tencent Hunyuan — Research Intern

Jun. 2025 – Aug. 2025
  • Investigated diffusion distillation for efficient image editing based on FLUX Kontext.
  • Explored Consistency Models followed by DMD2 for few-step distillation.

Tsinghua University — Ph.D. Research

Aug. 2023 – Present
  • Developed diffusion-based inference methods for Bayesian inverse problems.
  • Proposed learnable extrapolation for few-step diffusion-based inverse problem solvers and a stage-wise framework for navigating the distortion–perception trade-off.
  • Three first-author papers at NeurIPS 2024, NeurIPS 2025, and ICML 2026.

Selected Publications

  1. Distilling Drifting Transformers with Representation Autoencoders

    Jiawei Zhang, Mengfei Xia, Gen Li, and Yuantao Gu.

    arXiv:2606.15553 [Paper]

  2. Stage-wise Distortion–Perception Traversal in Zero-shot Inverse Problems with Diffusion Models

    Jiawei Zhang, Ziyuan Liu, Leon Yan, Zhenyu Xiao, and Yuantao Gu.

    ICML 2026 [Paper] [Code]

  3. Improving Diffusion-based Inverse Algorithms under Few-Step Constraint via Learnable Linear Extrapolation

    Jiawei Zhang, Ziyuan Liu, Leon Yan, Gen Li, and Yuantao Gu.

    NeurIPS 2025 [Paper] [Code]

  4. Unleashing the Denoising Capability of Diffusion Prior for Solving Inverse Problems

    Jiawei Zhang, Jiaxin Zhuang, Cheng Jin, Gen Li, and Yuantao Gu.

    NeurIPS 2024 [Paper] [Code]

  5. Adaptive Polyak Step-Size for Momentum Accelerated Stochastic Gradient Descent with General Convergence Guarantee

    Jiawei Zhang, Cheng Jin, and Yuantao Gu.

    IEEE Transactions on Signal Processing [Paper] [Code]

  6. EPA: Neural Collapse Inspired Robust Out-of-Distribution Detector

    Jiawei Zhang, Yufan Chen, Cheng Jin, Lei Zhu, and Yuantao Gu.

    ICASSP 2024 [Paper]

Education

Tsinghua University

Aug. 2023 – Present

Ph.D. Student in Information and Communication Engineering
Department of Electronic Engineering · Advisor: Prof. Yuantao Gu

Tsinghua University

Aug. 2019 – Jun. 2023

B.Eng. in Electronic Engineering

Selected Honors & Awards

  • National Scholarship for Graduate Students2025
  • Tsinghua Future Scholar Fellowship2023–2028
  • Outstanding Graduate of Beijing Municipality & Tsinghua University2023
  • National Scholarship for Undergraduate Students2022