Multi-Turn Reflective Masking Elicits Reasoning in Mask Diffusion Models

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.

한 줄 요약

Mask Diffusion Models에 반복적 자기 수정 능력을 부여하는 Reflective Masking과 History Reference 메커니즘을 제안한다.

핵심 기여도

Reflective Masking을 통해 MDM이 반복적 자기 수정을 수행하도록 유도, 기존 MDM의 단일 턴 마스킹 한계를 극복.
History Reference라는 파라미터 없는 메커니즘을 도입, 이전 수정 과정의 중간 상태를 활용해 반복 수정 일관성을 향상.
텍스트 생성, 수수도, 이미지 편집 등 다양한 태스크에서 기존 마스킹 기반 기법 대비 10~15% 성능 향상.
2개의 H100 80GB GPU에서 약 5시간만에 학습 가능, 기존 아키텍처 변경 방식 대비 훨씬 효율적.

핵심 아이디어

기존 MDM은 마스킹을 단일 턴으로 제한하여 반복 수정이 불가능했으나, Reflective Masking은 모델이 생성 중인 불확실성을 기반으로 마스킹을 반복적으로 수행하도록 유도한다. 이는 MDM이 생성된 토큰을 선택적으로 재수정하며, 전체 시퀀스 재생성을 피할 수 있게 한다.

또한, History Reference는 MDM이 이전 수정 과정에서 생성된 중간 denoising 상태를 유지하고 이를 반복 수정 시 활용하게 함으로써, 반복 중에 동일한 오류를 반복하는 것을 방지한다. 이는 AR 모델과 달리, MDM이 명시적으로 이전 예측을 참조할 수 있는 새로운 메커니즘을 제공한다.

기술적 접근법

Reflective Masking: 마스킹을 모델의 내부 결정 과정으로 구현, 불확실한 토큰을 선택적으로 재수정.
History Reference: 중간 denoising 상태를 저장하고 반복 수정 시 활용, 파라미터 없이 구현.
학습: 기존 아키텍처 변경 없이 가벼운 post-training 수행, 2개의 H100 GPU에서 약 5시간 학습.
데이터 생성: 모델의 본래 출력 분포와 일치하는 신뢰할 수 있는 신호를 생성하는 전략 사용.

주요 결과

텍스트 생성, 수수도, 이미지 편집 등 다양한 태스크에서 기존 마스킹 기반 기법 대비 10~15% 성능 향상.
Sudoku 태스크에서 Reflective Masking 적용 시 정확도 +12.3%, History Reference 추가 시 +7.8% 추가 개선.
이미지 편집에서 사용자 지시에 따른 수정 정확도 +14.1% 향상.
텍스트 생성에서 반복 오류 감소율 22.6% 기록.

의의 및 한계

Reflective Masking은 MDM이 반복적 자기 수정을 통해 오류를 줄이고 계산 효율성을 높이는 새로운 방식을 제시한다. 특히, AR 모델과 달리 MDM이 명시적인 반복 수정을 통해 더 정확하고 일관된 결과를 생성할 수 있다는 점에서 학술적 의의가 크다.

하지만, 반복 수정 과정에서 일부 태스크에서는 수렴 속도가 느려지는 문제가 발생할 수 있다. 또한, History Reference는 파라미터 없이 구현되지만, 복잡한 태스크에서는 추가적인 메모리 관리가 필요할 수 있다.

실용적 활용

Reflective Masking과 History Reference는 텍스트 생성, 이미지 편집, 코드 생성 등 다양한 생성 모델이 반복 수정을 통해 정확도를 높이는 데 활용 가능하다. 특히, 오류 수정이 중요한 의료, 법률, 소프트웨어 개발 분야에서 실용적 가치가 높다.