Multi-Turn Reflective Masking이 Mask Diffusion Models에서 추론을 이끈다

AR 기반 모델의 다-turn 추론은 초기 결정의 오류가 누적되며 전체 시퀀스를 재생성해야 한다. Reflective Masking(RM)은 출력의 일부를 선택적으로 재마스킹하고 수정하는 방식으로, 출력의 일부를 현 시점에서만 수정하고 재생성 비용을 줄인다. History Reference(HR)는 denoising 궤적의 과거 상태를 참조해 일관성과 수정 안정성을 높이며, 이는 이미지 편집, 구조적 추론(Sudoku) 및 자유 생성 텍스트 과제에서 일관된 성능 향상을 뒷받침한다.

단계별 요약: (1) 기존 MDM의 비마스크 토큰은 재생성 없이 유지되지만, RM은 각 위치에서 keep, re-mask, reveal 중 하나의 선택을 학습한다. (2) 입력 표준화: t번째 스텝에서 tilde x^(t)와 per-position 확률 p_theta를 이용해 각 위치의 다음 상태를 결정한다. (3) History Reference는 과거 상태를 회전-가중 합으로 요약한 history embedding을 현재 입력에 더해, 수정 판단의 주변 맥락 정보를 강화한다. (4) HR과 RM의 조합은 테스트 타임에서의 수정 효율과 정확성을 높이며, 텍스트/이미지/구조적 추론에서 성능이 향상된다.

1. Reflective Masking의 기본Inference Rule: t 시점에서 비마스크 위치 i의 상태 tilde x_i^(t)에서 p_theta(M|tilde x) > p_theta(x_i|tilde x) 이면 MASK 유지/재마스킹, 그렇지 않으면 원래 토큰 유지. tilde x_i^(t) == MASK인 경우 argmax p_theta(v|tilde x)로 v를 Reveal. 2) History Reference(HR): 각 위치 i에 대해 history embedding a_i^(t) = sum_{k=0}^t gamma^{t-k} R_{k-t} e_i^(k) 로 누적; 현재 입력에 더해 모델이 과거 상태 정보를 활용하도록 한다. 3) RM 학습 목표: z_i ∈ {w_i, MASK, x_i*}에 대해 τ(z_i, x_i*)를 oracle로 정의하고, per-position CE loss를 최소화한다. 4) 데이터 구성: x*에서 일부 위치를 corrupt하고, wrong-token과 MASK로 트레이젝토리를 시뮬레이션한다. 5) 구현 메모: per스텝 입력은 a^(t)와 c, HR로 확장된 입력이며, O(1) 시간 복잡도로 누적 임베딩 업데이트가 가능하도록 설계되었다. 6) 비제한성 및 이론 보장: Bayes-최적화에 근거한 Revision 정책 학습과 HR의 정보 증가가 위험을 감소시킨다는 이론적 보장을 제시한다.

메인 벤치마크: 이미지 편집에서 RM은 Edit Precision 99.73, Edit Coverage 73.02, MAE-RGB 3.613, PSNR 34.759, SSIM 0.9744, VQAScore 85.17, User Study 68.2로 Baseline 대비 우수. Sudoku 수정에서 RM+HR의 Exact Accuracy 93.4%, Valid Rate 93.6%, Replay Mistake 0.03, Conflict Cells 0.236으로 최적. Text reasoning에서 MATH/MBPP/ARC-Challenge에서 RM이 Vanilla SFT 대비 개선; 예: MBPP에서 39.4%로 8.8% 향상. Case 1-5에서 RM이 다중 과제에 걸쳐 일관된 성능 증가를 보임. 제한: 현재의 RM은 현재 Block Diffusion 모델의 한계를 넘지 못하며, 더 큰 규모의 학습에서 Transfer 가능성은 추가 연구 필요. 향후 Block Diffusion 모델에 대한 적용은 전역 컨텍스트 수정의 한계를 넘는 방향으로 필요.

아키텍처 변경 없이 RM을 적용하는 것이 목표다. 1) RM의 기본 인퍼런스 규칙은 비마스크 위치에서 MASK보다 MASK의 확률이 더 높으면 재마스킹하고, MASK 위치에서는 토큰을 Reveal한다. 2) HR은 history embedding과 rotary encoding으로 구성되며, 시차Δ에 따라 회전 행렬을 적용해 past 토큰 임베딩을 누적한다. 3) 학습 목표는 per-position oracle actionsτ(z_i, x_i*)를 예측하는 cross-entropy loss이며, 이를 통해 RM의 예측 분포가 Bayesian-optimal에 수렴하도록 한다. 4) Synthetic trajectory 샘플링으로 트레이닝 데이터를 구성하고, 각 위치의 수렴 행동을 학습한다. 5) 테스트 타임 확장성을 위한 일정한 업데이트 규칙이 보장되며, per-step 입력 a^(t)와 c만으로도 업데이트가 가능하다. 6) 제한점으로는 현재의 RM은 더 큰 시퀀스나 더 긴 의사결정에 대해 추가 연구가 필요하다는 점이다.

제한점으로, 본 연구는 비교적 간단한 long-horizon 문제 및 현재의 베이스 MDM의 한계를 가지는 작업들에 집중한다. 대규모 스케일링/대규모 데이터로의 전이성 평가가 제한적이다. 향후 stronger block diffusion 모델에의 적용에서 전체 컨텍스트 수정이 필요할 수 있으며, 전역 수정의 효율적 확장을 위한 추가 메커니즘이 요구된다.