고정하지 말고 수정하라: 자기 수정형 마스크 확산 언어 모델을 위한 샘플러 매칭 학습 (Sampler-Matched Training)

마스크 확산 언어 모델 (Masked Diffusion Language Models, MDLMs)은 각 디노이징 (denoising) 단계마다 모든 위치를 재예측하지만, 표준 샘플러 (standard samplers)는 토큰이 한 번 공개되면 이를 확정 지어 버리며, 이로 인해 수정 능력을 활용하지 못하게 됩니다. 기존 방식들은 확정된 토큰을 수정하기 위해 휴리스틱 (heuristic) 또는 학습된 메커니즘을 추가하거나, 재예측하기 전에 다시 [MASK]로 재마스킹 (remask)하는 방식을 취합니다. 보조 모듈 없이 가시적인 토큰을 직접 수정하는 원칙적인 샘플러에 대한 연구는 아직 미흡한 상태입니다. 우리는 추가적인 모듈이나 보조 패스 (auxiliary passes) 없이 가시적 토큰에서 가시적 토큰으로의 직접적인 수정을 허용하는, 교정 스타일의 역방향 업데이트 (reverse update)로부터 유도된 파라미터 프리 (parameter-free) 샘플러인 D3IM을 소개합니다. 또한 D3IM은 우리가 보존 편향 (preservation bias)이라고 명명한 모델 측면의 장애물을 밝혀냅니다. 즉, 모델이 잘못 확정된 자신의 토큰을 수정하기보다 그대로 재현하려는 경향이 있다는 것입니다. 우리는 D3IM의 샘플링 과정을 시뮬레이션하는 경량화된 사후 학습 (post-training) 절차인 SCOPE (Self-Conditioned On Prediction Errors)를 통해 이 문제를 해결합니다. 64단계의 디노이징 (denoising)을 수행한 LLaDA-8B 모델에서, SCOPE+D3IM은 표준 언마스킹 (unmasking)을 사용한 기존 LLaDA-8B 대비 GSM8K에서 +13.0 (68.3%), MATH-500에서 +4.8 (23.6%), HumanEval에서 +15.3 (29.3%), MBPP에서 +10.4 (30.8%)의 성능 향상을 보였으며, 수학 및 HumanEval 작업에서는 디노이징 (denoising) 단계가 많아질수록 성능 향상 폭이 더욱 커졌습니다.