Modelos de Difusión Uniforme Revisitados: Denoizador de Exclusión de Uno y Reformulación del Estado Absorbente

Resumen

Los modelos de difusión discretos a menudo se entrenan mediante la predicción de datos limpios, pero esta predicción puede utilizarse de diferentes maneras para definir la dinámica inversa. En los Modelos de Difusión Enmascarada (MDM), estas elecciones coinciden en gran medida, mientras que en los Modelos de Difusión Uniforme (UDM) no es así. Demostramos que la parametrización de puente plug-in estándar para UDM no se optimiza mediante la posterior de denoising, sino mediante una posterior de dejar uno fuera que predice cada token limpio sin utilizar su propia observación ruidosa. Esto identifica una discrepancia entre el ELBO plug-in y el objetivo habitual de denoising por entropía cruzada. Caracterizamos el objetivo de dejar uno fuera y derivamos conversiones exactas entre el denoiser, la posterior de dejar uno fuera y el score. Estas conversiones nos permiten desvincular la parametrización del objetivo de entrenamiento. Nuestros resultados también conducen a mejoras en la inferencia sin necesidad de entrenamiento adicional, mediante un muestreador predictor-corrector informado y un muestreo de temperatura mejorado basado en el predictor de dejar uno fuera. Además, introducimos una reformulación de estado absorbente de la difusión uniforme que preserva la ley conjunta de UDM mientras la descompone en operaciones de muestreo similares a la difusión enmascarada, con posteriores de denoising más simples, desenmascaramiento por arrastre y un mecanismo natural de reenmascaramiento. En el modelado del lenguaje, las parametrizaciones de dejar uno fuera mejoran consistentemente la generación de UDM, mientras que la construcción absorbente iguala o supera a la difusión enmascarada. Estos resultados sugieren que la brecha empírica entre la difusión enmascarada y la uniforme está impulsada menos por la elección de las marginales en sí mismas que por el diseño de la parametrización y el muestreo. El código y los modelos se pueden encontrar en https://github.com/samsongourevitch/rev_udm.

English

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.