Многошаговое рефлексивное маскирование индуцирует рассуждение в маскированных диффузионных моделях.

Аннотация

В то время как рассуждение на основе авторегрессионных (AR) моделей часто осуществляется посредством рассуждения по цепочке мыслей и рефлексии, их доработка предыдущих результатов по-прежнему опирается на полностью последовательную генерацию, даже когда требуются лишь локальные правки. Напротив, механизм маскирования в моделях маскированной диффузии (MDM) естественным образом поддерживает явные локальные правки предыдущих результатов, позволяя проводить выборочную доработку без отбрасывания предыдущих ответов и генерации новых с нуля. Хотя это свойство ближе к тому, как люди исправляют ошибки путем итеративной локальной доработки, существующие MDM не поддерживают многократное маскирование и шумоподавление. Мы предлагаем рефлексивное маскирование (RM), которое активирует такую внутреннюю способность к рассуждению в MDM с помощью легковесного пост-обучения. RM обеспечивает нативное масштабирование во время тестирования, при котором MDM итеративно пересматривает и корректирует свои предыдущие результаты на основе развивающегося контекста. Для использования insight’ов из предыдущих итераций, подобно AR-рассуждению, мы дополнительно вводим историческую ссылку — механизм без параметров, который использует промежуточные состояния шумоподавления в процессе корректировки. Наш подход не требует изменения архитектуры и легко применим к существующим MDM. На разнообразных задачах и модальностях, включая генерацию текста, судоку и редактирование изображений, рефлексивное маскирование последовательно превосходит стандартные методы на основе маскирования и демонстрирует сильную общность, позиционируя RM как фундаментальный примитив для рассуждения на основе MDM.

English

While reasoning on autoregressive (AR) models is often performed by chain-of-thought reasoning and reflection, their refinement of previous outputs still relies on fully sequential generation, even when only local edits are needed. In contrast, the masking mechanism in Mask Diffusion Models (MDMs) naturally supports explicit local edits on previous outputs, allowing selective refinement without discarding previous answers and generating another from scratch. While this property more closely aligns with how humans correct mistakes by iterative local refinement, existing MDMs do not support multi-turn masking and denoising. We propose Reflective Masking (RM), which elicits such an intrinsic reasoning capability in MDMs via lightweight post-training. RM provides a native test-time scaling, where an MDM iteratively revisits and revises its prior outputs based on evolving context. To exploit insights from previous turns like AR reasoning, we further introduce History Reference, a parameter-free mechanism that leverages intermediate denoising states during revision. Our approach requires no architectural changes and is easily applicable to existing MDMs. Across diverse tasks and modalities, including text generation, Sudoku, and image editing, Reflective Masking consistently outperforms standard masking-based baselines and demonstrates strong generality, positioning RM as a fundamental primitive for reasoning on MDMs.