I'm Wanzhou Lei, a third-year Ph.D. student in Applied Mathematics at the University of California, Berkeley, advised by Prof. Per-Olof Persson. Before coming to Berkeley, I received my M.S. in Data Science from Harvard University, where I was advised by Prof. Pavlos Protopapas, and my B.S. in Mathematics and Computer Science from New York University, where I was advised by Prof. Yunan Yang.
My research focuses on generative models for vision. I work on the representation spaces in which images are generated and restored, on better training of diffusion models, and on text conditioning in image generation. I am also interested in the theory of how attention learns, and in machine learning for scientific computing.
In my free time, I am a crazy butterfly enthusiast. I take ecological photos of butterflies all over the world, with an emphasis on Chinese and Neotropical species, and I am working on a book, Ecological Photo Guide to Butterflies in Tropical Andean Region. Some of my photos are in my photo gallery.
Feel free to send me an email to discuss anything.
RAESR restores real-world low-resolution images inside the fused latent of a frozen DINOv3-L encoder, with a single decoder pass and no diffusion sampling loop. It attains the best fidelity–perception trade-off among state-of-the-art restorers at 37 ms per 512×512 image.
SWIFT moves weak-to-strong guidance from sampling into training: each regression target is extrapolated away from a frozen weak denoiser, so the guide is dropped at inference. One fixed recipe lowers unguided ImageNet-256 FID in all 18 DiffusionBench configurations, and on RAEv2 it reaches 1.15 FID without guidance and 0.99 with it.
TET swaps the text encoder of a frozen text-to-image generator for an arbitrary new one, without retraining the backbone or using any image data. Distilling the model's behavior rather than the native encoder's features keeps near-native alignment on FLUX.2-klein with the encoder compressed by up to 85%, and a multilingual encoder makes the same frozen generator produce images in 12 languages zero-shot.
We prove that gradient-flow training of one-layer cross-attention escapes the 5/6 accuracy ceiling of flat attention on a synthetic multimodal task. The escape is a delayed feature-learning transition: values first fit three of the four classes, then the matched query–key score grows and flips the fourth, after which every fresh sample is classified correctly.
Led RAESR, one-pass real-world super-resolution in a frozen vision-foundation latent (first-authored paper, under review), and extended it to video with a 3.5M-parameter space-time transformer that aligns neighboring frames with frozen optical flow and corrects flicker and inconsistent detail across frames. The architecture is deployed in production in HappyHorse, Alibaba's flagship video generation model. Co-developed SWIFT, which moves weak-to-strong guidance from sampling into diffusion training.
Led the replacement of the non-smooth lithography simulation with a residual U-Net, making source-mask optimization end-to-end differentiable, and integrated the network into ASML's C++ optimization framework for production use. Delivered substantial reductions in edge placement error and faster convergence of gradient-based yield optimization.