생각의 순서 스케줄링: 확산 언어 모델 (Diffusion Language Models)에서의 사고 순서 학습

Masked diffusion language models (마스크 확산 언어 모델)은 토큰을 반복적으로 언마스킹 (unmasking) 함으로써 디코딩하며, 이때 언마스킹 순서는 생성 품질에 강력한 영향을 미치는 "사고의 순서 (order of thought)"를 정의하지만 일반적으로 휴리스틱 (heuristic)하게 선택됩니다. 우리는 Kullback-Leibler divergence (KL 발산)로 측정되고 모델의 경로별 로그 가능도 (pathwise log-likelihood)로 표현되는 순차적 디코딩 불일치 (sequential decoding mismatch)에 대한 다루기 쉬운 상한 (upper bound)을 도출하며, 이는 충분한 모델 표현력 하에서 엄밀함 (tightness)을 가집니다. 이 상한은 정렬된 궤적 (ordered trajectories)에 대해 조밀한 자기 인식 보상 (dense self-aware reward)을 유도하며, 순서 선택을 고정된 디노이저 (frozen denoiser)를 사용하는 원칙적인 정책 최적화 (policy optimization) 문제로 변환합니다. 우리는 이 아이디어를 Self-Aware Scheduling (SAS)으로 구현하였으며, 이는 Group Relative Policy Optimization (GRPO)을 사용하여 경량화된 순서 정책을 학습하고, any-order 및 semi-autoregressive (준-자기회귀) 디코딩 모두에 원활하게 적용됩니다. 1B MDM을 이용한 Sudoku 작업에서 SAS는 퍼즐 정확도를 최적의 휴리스틱 스케줄인 82.0%에서 91.8%로 향상시켰으며, 학습된 궤적을 따른 2단계 미세 조정 (fine-tuning)을 통해 97.5%에 도달했습니다. LLaDA-8B를 이용한 수학적 추론 작업에서 SAS는 GSM8K의 pass@1을 64%에서 76%로, MBPP를 39.5%에서 41%로 향상시켰으며, 생성 길이와 블록 크기에 관계없이 일관되게 휴리스틱 스케줄과 일치하거나 이를 능가했습니다. 프로젝트 페이지: https://jimmyxu123.github.io/SAS