스케줄링 사고: 확산 언어 모델에서 사고 순서 학습

Scheduling Thoughts: Learning the Order of Thought in Diffusion Language Models

마스크된 확산 언어 모델은 토큰을 반복적으로 마스크 해제하면서 디코딩하는데, 이때 마스크 해제 순서는 생성 품질에 큰 영향을 미치는 "사고 순서"를 정의하지만, 일반적으로 휴리스틱하게 선택됩니다. 본 연구에서는 모델의 경로별 로그 우도를 사용하여 측정되며 Kullback-Leibler 발산으로 표현되는 순차적 디코딩 불일치에 대한 추정 가능한 상한을 유도하며, 충분한 모델 표현력을 갖는 경우 이 상한은 정확해집니다. 이러한 상한은 정렬된 시퀀스에 대한 밀집형 자기 인식 보상을 유발하여, 고정된 노이즈 제거기를 가진 순서 선택 문제를 체계적인 정책 최적화 문제로 정의합니다. 본 연구에서는 Self-Aware Scheduling (SAS)이라는 방법을 제안하며, SAS는 Group Relative Policy Optimization을 사용하여 경량의 순서 정책을 학습하고, 모든 순서 및 준-자기 회귀 디코딩에 원활하게 적용됩니다. 10억 개의 매개변수를 가진 MDM 모델로 Sudoku 문제를 해결할 때, SAS는 휴리스틱 스케줄링(82.0%) 대비 퍼즐 정확도를 91.8%까지 향상시키며, 학습된 경로를 활용한 추가 미세 조정 시 97.5%에 도달합니다. 또한, LLaDA-8B 모델로 수학적 추론 문제를 해결할 때, SAS는 GSM8K 데이터셋에서 pass@1 정확도를 64%에서 76%로, MBPP 데이터셋에서 39.5%에서 41%로 향상시키며, 생성 길이 및 블록 크기에 관계없이 휴리스틱 스케줄링과 동등하거나 그 이상의 성능을 보입니다. 프로젝트 페이지: https://jimmyxu123.github.io/SAS

Masked diffusion language models decode by iteratively unmasking tokens, where the unmasking order defines an "order of thought" that strongly influences generation quality yet is typically chosen heuristically. We derive a tractable upper bound on the sequential decoding mismatch, measured by the Kullback-Leibler divergence and expressed in terms of the model's pathwise log-likelihood, with tightness under sufficient model expressivity. This bound induces a dense self-aware reward over ordered trajectories, casting order selection as a principled policy optimization problem with a frozen denoiser. We instantiate this idea as Self-Aware Scheduling (SAS), which learns a lightweight order policy using Group Relative Policy Optimization and applies seamlessly to both any-order and semi-autoregressive decoding. On Sudoku with 1B MDM, SAS improves puzzle accuracy from 82.0% (best heuristic schedule) to 91.8%, and reaches 97.5% with second-stage fine-tuning along learned trajectories. On mathematical reasoning with LLaDA-8B, SAS improves pass@1 on GSM8K from 64% to 76% and on MBPP from 39.5% to 41%, consistently matching or exceeding heuristic schedules across generation lengths and block sizes. Project page: https://jimmyxu123.github.io/SAS