Modelos de Difusão Uniformes Revisitados: Denoiser Leave-One-Out e Reformulação do Estado Absorvente

Resumo

Modelos de difusão discretos são frequentemente treinados por meio da predição de dados limpos, mas a predição pode ser usada de diferentes maneiras para definir a dinâmica reversa. Em Modelos de Difusão Mascarados (MDM), essas escolhas coincidem em grande parte, enquanto em Modelos de Difusão Uniformes (UDM) elas não coincidem. Mostramos que a parametrização de ponte plug-in padrão para UDM não é otimizada pelo posterior de remoção de ruído, mas sim por um posterior leave-one-out que prediz cada token limpo sem utilizar sua própria observação ruidosa. Isso identifica uma incompatibilidade entre o ELBO plug-in e o objetivo usual de entropia cruzada para remoção de ruído. Caracterizamos o alvo leave-one-out e derivamos conversões exatas entre o removedor de ruído, o posterior leave-one-out e o score. Essas conversões nos permitem separar a parametrização e o objetivo de treinamento. Nossos resultados também levam a melhorias na inferência sem qualquer treinamento adicional por meio de um amostrador preditor-corretor informado e uma amostragem com temperatura melhorada baseada no preditor leave-one-out. Introduzimos ainda uma reformulação de estado absorvente da difusão uniforme que preserva a lei conjunta da UDM enquanto a decompõe em operações de amostragem semelhantes à difusão mascarada, com posteriores de remoção de ruído mais simples, desmascaramento de arrasto e um mecanismo natural de remascaramento. Em modelagem de linguagem, parametrizações leave-one-out melhoram consistentemente a geração da UDM, enquanto a construção absorvente iguala ou supera a difusão mascarada. Esses resultados sugerem que a lacuna empírica entre difusão mascarada e uniforme é menos impulsionada pela escolha das marginais em si do que pela parametrização e pelo design de amostragem. O código e os modelos podem ser encontrados em 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.