Метод зоряной маскированной дифузии

Метод зоряной маскированной дифузии

В masked diffusion LLM есть довольно тупая проблема: если токен уже размаскирован, всё, он заморожен. Ошиблась модель на раннем шаге — дальше она уже не исправляет ошибку, а пытается построить вокруг неё текст. Такие несогласованности особенно заметны при маленьком количестве шагов. Существующие методы перемаскирования дают модели шанс передумать, но требуют сотен шагов, что убивает главное преимущество диффузии — скорость.

Авторы из T-Tech, ВШЭ и Constructor University представили на ICML метод, в котором все зашумлённые состояния зависят только от чистого текста, но не друг от друга. На каждом шаге модель предсказывает полную чистую последовательность, а следующее состояние сэмплируется заново из неё, ничего не наследуя от текущего: любой токен можно вернуть в маску и переписать, причём дообучения это не требует. С нуля так генерировать нельзя — перемаскирование на ранних шагах разрушает несложившийся контекст, поэтому первые ~70% шагов черновик пишется обычным способом, и лишь потом включается режим правки. Куда именно править, указывает отдельный лёгкий предсказатель ошибок, обученный на промахах модели при генерации. Главный минус такого подхода — отсутствие возможности вставлять и удалять токены.

В результате, за 128 шагов метод обгоняет по качеству и разнообразию текста лучший из прежних подходов с бюджетом в 512 шагов, что позволяет генерировать почти втрое быстрее (3.4 против 9.2 секунды на H200). Предсказатель ошибок, обученный на веб-текстах, переносится на код, математику и новости без дообучения.