Le masquage réflexif multi-tour suscite le raisonnement dans les modèles de diffusion par masques

Résumé

Alors que le raisonnement sur les modèles autorégressifs (AR) s'effectue souvent par un raisonnement en chaîne de pensée et une réflexion, leur raffinement des sorties précédentes repose encore sur une génération entièrement séquentielle, même lorsque seules des modifications locales sont nécessaires. En revanche, le mécanisme de masquage dans les modèles de diffusion à masque (MDM) supporte naturellement des modifications locales explicites sur les sorties antérieures, permettant un raffinement sélectif sans avoir à écarter les réponses précédentes et à en générer une nouvelle à partir de zéro. Bien que cette propriété s'aligne plus étroitement sur la manière dont les humains corrigent leurs erreurs par un raffinement local itératif, les MDM existants ne prennent pas en charge le masquage et le débruitage multi-tours. Nous proposons le Masquage Réflexif (RM), qui suscite une telle capacité de raisonnement intrinsèque dans les MDM via un post-entraînement léger. Le RM offre une mise à l'échelle native au moment du test, où un MDM revisite et révise itérativement ses sorties antérieures en fonction du contexte évolutif. Pour exploiter les enseignements des tours précédents, à l'instar du raisonnement AR, nous introduisons la Référence d'Historique, un mécanisme sans paramètre qui tire parti des états de débruitage intermédiaires lors de la révision. Notre approche ne nécessite aucune modification architecturale et s'applique facilement aux MDM existants. Sur diverses tâches et modalités, notamment la génération de texte, le Sudoku et l'édition d'images, le Masquage Réflexif surpasse systématiquement les méthodes de base standard basées sur le masquage et démontre une grande généralité, positionnant le RM comme un élément fondamental pour le raisonnement sur les 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.