Пересмотр равномерных диффузионных моделей: деноизер с исключением одного и переформулировка поглощающего состояния

Аннотация

Дискретные диффузионные модели часто обучаются с помощью предсказания чистых данных, но это предсказание может использоваться различными способами для определения обратной динамики. В моделях маскированной диффузии (MDM) эти варианты в значительной степени совпадают, тогда как в моделях равномерной диффузии (UDM) это не так. Мы показываем, что стандартная параметризация плагинного моста для UDM не оптимизируется шумоподавляющим апостериорным распределением, а скорее апостериорным распределением с исключением одного, которое предсказывает каждый чистый токен, не используя его собственное зашумленное наблюдение. Это выявляет несоответствие между плагинной вариационной нижней границей (ELBO) и обычной кросс-энтропийной целью шумоподавления. Мы характеризуем цель с исключением одного и выводим точные преобразования между шумоподавителем, апостериорным распределением с исключением одного и скором. Эти преобразования позволяют разделить параметризацию и цель обучения. Наши результаты также приводят к улучшению вывода без дополнительного обучения за счет информированного сэмплера предиктор-корректор и улучшенной температурной выборки на основе предиктора с исключением одного. Мы дополнительно вводим поглощающую переформулировку равномерной диффузии, которая сохраняет совместный закон UDM, разлагая его на операции выборки, подобные маскированной диффузии, с более простыми шумоподавляющими апостериорными распределениями, переносом снятия маски и естественным механизмом повторной маскировки. В задаче языкового моделирования параметризации с исключением одного последовательно улучшают генерацию UDM, в то время как поглощающая конструкция достигает результатов, сопоставимых с маскированной диффузией или превосходящих ее. Эти результаты позволяют предположить, что эмпирический разрыв между маскированной и равномерной диффузией объясняется скорее параметризацией и дизайном выборки, чем выбором самих маргинальных распределений. Код и модели доступны по адресу 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.