Modèles de diffusion uniformes revisités : débruitage leave-one-out et reformulation de l'état absorbant

Résumé

Les modèles de diffusion discrets sont souvent entraînés par prédiction des données propres, mais cette prédiction peut être utilisée de différentes manières pour définir la dynamique inverse. Dans les modèles de diffusion masquée (MDM), ces choix coïncident largement, contrairement aux modèles de diffusion uniforme (UDM) où ce n'est pas le cas. Nous montrons que le paramétrage par pont d'insertion standard pour les UDM n'est pas optimisé par le postérieur de débruitage, mais par un postérieur leave-one-out qui prédit chaque jeton propre sans utiliser sa propre observation bruitée. Cela identifie un décalage entre l'ELBO par insertion et l'objectif de débruitage usuel par entropie croisée. Nous caractérisons la cible leave-one-out et dérivons des conversions exactes entre le débruitage, le postérieur leave-one-out et le score. Ces conversions nous permettent de dissocier le paramétrage et l'objectif d'entraînement. Nos résultats conduisent également à des améliorations d'inférence sans aucun entraînement supplémentaire grâce à un échantillonneur prédicteur-correcteur informé et à un échantillonnage par température amélioré basé sur le prédicteur leave-one-out. Nous introduisons en outre une reformulation par état absorbant de la diffusion uniforme qui préserve la loi conjointe des UDM tout en la décomposant en opérations d'échantillonnage de type diffusion masquée, avec des postérieurs de débruitage plus simples, un démasquage par report et un mécanisme de remasquage naturel. Sur la modélisation du langage, les paramétrages leave-one-out améliorent systématiquement la génération des UDM, tandis que la construction par état absorbant égale ou dépasse la diffusion masquée. Ces résultats suggèrent que l'écart empirique entre la diffusion masquée et uniforme est moins dû au choix des marginales elles-mêmes qu'à la conception du paramétrage et de l'échantillonnage. Le code et les modèles sont disponibles à l'adresse 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.