Uniform Diffusion Models Revisited: Leave-One-Out Denoiser and Absorbing State Reformulation
AuthorsSamson Gourevitch, Yazid Janati, Dario Shariatian, Umut Simsekli, Eric Moulines, Eric P. Xing, Alain Durmus
This paper shows that a subtle mismatch in how uniform diffusion models are trained and sampled can be fixed with a leave-one-out view, leading to better language generation and a cleaner alternative formulation.
Key results
On a representative OpenWebText run, top-p sampling at matched entropy improved generation perplexity for the leave-one-out model relative to the denoiser baseline.
AUDM adds a small auxiliary module for the latent absorbing variable U on top of the shared 170M-parameter DiT backbone.
The additional 7M-parameter module corresponds to about 4.1% parameter overhead in the absorbing-state uniform diffusion setup.
The language models on OpenWebText and One Billion Word are trained from scratch for 1M steps.
What the paper found
This paper shows that in Uniform Diffusion Models (UDM), the standard plug-in bridge parameterization is not optimized by the usual denoising posterior but by a leave-one-out posterior that predicts token ℓ from all noisy tokens except x_t^ℓ. The authors derive exact conversions among the denoiser, the leave-one-out predictor, and the concrete score, and use them to train a UDM leave-one-out model directly with standard cross-entropy; on OpenWebText (OWT) and One Billion Word (LM1B), this consistently improves validation perplexity and the generation frontier, with top-p sampling at matched entropy reducing Gen-PPL from 54.8 to 42.2 on a representative OWT run. They also show that applying top-p or temperature after converting a denoiser into leave-one-out form is better than applying sampling heuristics to the denoiser itself, and they build an informed Gibbs predictor-corrector sampler without any auxiliary network, which Pareto-dominates ancestral sampling at the same runtime. Separately, the paper reframes UDM as an absorbing-state process with a random latent absorbing token U, yielding Absorbing-State UDM (AUDM) and a resampled variant, ReAUDM, that preserve the UDM joint law while recovering masked-diffusion-like carry-over and remasking operations. On language modeling, AUDM adds only 7M parameters, about 4.1 percent overhead, and improves likelihood over UDM and all zero-shot baselines; on Sudoku, leave-one-out UDM reaches nearly perfect solve rate, while AUDM and ReAUDM match or exceed masked diffusion, supporting the claim that the UDM–MDM gap is driven more by parameterization and sampling design than by the corruption marginals themselves.
Original abstract
Discrete diffusion models are often trained through clean-data prediction, but the prediction can be used in different ways to define the reverse dynamics. In Masked Diffusion Models (MDM) these choices largely coincide, whereas in Uniform Diffusion Models (UDM) they do not. We show that the standard plug-in bridge parameterization for UDM is not optimized by the denoising posterior, but by a leave-one-out posterior that predicts each clean token without using its own noisy observation. This identifies a mismatch between the plug-in ELBO and the usual cross-entropy denoising objective. We characterize the leave-one-out target and derive exact conversions between the denoiser, the leave-one-out posterior, and the score. These conversions allow us to disentangle parameterization and training objective. Our results also lead to inference improvements without any additional training through an informed predictor-corrector sampler and improved temperature sampling based on the leave-one-out predictor. We further introduce an absorbing-state reformulation of uniform diffusion that preserves the UDM joint law while decomposing it into masked-diffusion-like sampling operations, with simpler denoising posteriors, carry-over unmasking, and a natural remasking mechanism. On language modeling, leave-one-out parameterizations consistently improve UDM generation, while the absorbing construction matches or surpasses masked diffusion. These results suggest that the empirical gap between masked and uniform diffusion is driven less by the choice of marginals themselves than by parameterization and sampling design. The code and models can be found at https://github.com/samsongourevitch/rev_udm.
Read the original paperMore in Generative Models
Browse all 63 papers →RULER: Instance-aware Rubric Rewards for SVG Generation
Hangyu Ran, Yuhao Zheng, Yingying Zhang, Kevin Qinghong Lin, Han Peng
RULER uses instruction-specific visual rubrics as reinforcement-learning rewards to make SVG generation more faithful, stylish, and resistant to reward hacking.
Think Before You Score: Thinking Reward Model for Visual Generation
Xuehai Bai, Zhenchen Tang, Yang Shi, Dianyi Wang, Tengfei Liu, Wanshun Su, Xuanyu Zhu, Ruohui Wang, Haiwen Diao, Haotian Wang, Xiaoling Gu, Yuanxing Zhang
A visual reward model that first decides what matters in each image-generation case, then scores outputs with detailed rubrics to provide better training signals.
WanPE: Towards Cinematic Prompt Enhancement for Modern Text-to-Video Generation
Yubo Zhu, Yawen Shao, Ziyun Dai, Zixun Fang, Kai Zhu, Siyang Sun, Haolan Xue, Chuxin Wang, Tingyu Weng, Jingming Luo, Chen Shi, Lianghua Huang, Yufeng Ai, Yuzheng Wang, Wenyuan Zhang, Yu Shang, Yuxiang Bao, Zoubin Bi, Jie Xiao, Jinbo Xing, Jiaxing Zhao, Chongyang Zhong, Hengjian Chen, Chenwei Xie, Akide Liu, Zhehan Kan, Yu Liu, Wei Zhai, Sheng Zhong, Wei Tong
WanPE turns ordinary text prompts into director-level cinematic plans, substantially improving the quality and consistency of long-form AI-generated videos.