기술 요약: 스케줄링에 관한 고찰 (SAS)

1. 문제 정의

마스크드 확산 모델(Masked Diffusion Models, MDMs)은 토큰을 반복적으로 언마스킹(unmasking)하여 이산 시퀀스를 생성하며, 이는 자기회귀(autoregressive) 모델의 고정된 좌에서 우로 향하는 인과적 구조(left-to-right factorization)에 대한 유연한 대안을 제공한다. 그러나 이러한 유연성은 결정적인 자유도를 도입한다: 바로 사고의 순서(unmasking schedule)이다.

기존의 접근 방식들은 휴리스틱 스케줄(예: 로컬 신뢰도, 마진 또는 엔트로피를 기반으로 토큰 선택)에 의존한다. 저자들은 이러한 휴리스틱이 "근시안적"이라고 주장한다: 즉, 전역적인 디코딩 궤적을 최적화하기보다는 즉각적인 다음 단계만을 최적화한다는 것이다. 결과적으로, 이들은 국소적으로는 확실하지만 미래 단계에 대한 제약 조건을 효과적으로 전파하지 못하는 "쉬운" 토큰들을 선택하는 경우가 많다. 또한, 전문가가 설계한 논리적 순서(예: 인간의 스도쿠 풀이 전략)라 할지라도 특정 사전 학습된 확산 모델에게는 최적이 아닐 수 있다. 왜냐로, 모델의 조건부 예측은 인간의 논리와는 다른 경로를 따라 가장 신뢰할 수 있는 경로를 따를 수도 있기 때문이다.

핵심 문제는 $n!$개의 가능한 순서를 탐색하는 것이 불가능하다는 점을 고려하면서, 사전 학습된 확산 모델을 재학습시키지 않고도 생성 품질을 극대화하는 최적의 언마스킹 순서를 학습하는 것이다.

2. 방법론: 자기 인식 스케줄링 (Self-Aware Scheduling, SAS)

본 논문은 베이스 확산 모델($p_\theta$)을 동결(frozen) 상태로 유지하면서 경량화된 순서 정책($v_\phi$)을 학습하는 플러그 앤 플레이(plug-and-play) 프레임워크인 **자기 인식 스케줄링(SAS)**을 제안한다.

이론적 기초

저자들은 확산 디코딩을 잠재적 순서 생성 과정(latent-order generative process)으로 프레임화한다. 이들의 목표는 실제 데이터 분포 $\pi(x)$와 디코딩 절차에 의해 유도된 분포 $P^{seq}_{\theta,\phi}(x)$ 사이의 쿨백-라이블러(KL) 발산을 최소화하는 것이다.

• 불가능성(Intractability): 모든 $n!$개의 디코딩 순서에 대해 평균을 내야 하므로, 이 KL 발산을 직접 최소화하는 것은 불가능하다.
• 다룰 수 있는 상한(Tractable Bound): 저자들은 이 미스매치에 대한 다룰 수 있는 상한을 도출한다. 그들은 KL 발다를 최소화하는 것이 데이터-정책 결합 분포와 모델-정책 결합 분포 사이의 **결합 KL 미스매치(joint KL mismatch)**를 최소화하는 것과 (상수 차이를 제외하고) 동일함을 보여준다.
• 자기 인식 보상(Self-Aware Reward): 이 유도는 고정된 디노이저(denoiser) 하에서의 타겟 시퀀스에 대한 **경로별 로그 가능도(pathwise log-likelihood)**로 정의되는 조밀하고 모델 스스로를 인식하는 보상 함수를 산출한다:
  $$R_\theta(x, \sigma) = \log p_\theta(x | \sigma) = \sum_{t=1}^n \log p_{i_t}^\theta(x_{i_t} | \tilde{x}^{(t-1)}(\sigma))$$
  이 보상은 특정 순서로 토큰을 드러내도록 강제되었을 때 모델이 데이터를 얼마나 잘 설명할 수 있는지에 따라 각 공개 결정에 크레딧을 부여한다.

알고리즘

• 정책 파라미터화: 순서 정책 $v_\phi$는 현재 마스킹된 위치에 대한 범주형 분포를 출력하는 경량 신경망(예: MLP 또는 Transformer)이다.
• 최적화: 정책은 **그룹 상대 정책 최적화(Group Relative Policy Optimization, GRPO)**를 사용하여 학습된다. 과정은 유한 호라이즌 마르코프 결정 과정(MDP)으로 정식화된다:
  • 상태(State): 부분적으로 드러난 시퀀스.
  • 행동(Action): 다음으로 언마스킹할 인덱스 선택.
  • 보상(Reward): 현재 컨텍스트 하에서의 정답 토큰에 대한 로그 가능도 (Teacher-forced).
• 2단계 프레임워크:
  1. 순서 학습: 기존 모델을 위한 최적의 궤적을 찾기 위해 $\theta$를 동결한 채 $\phi$를 최적화한다.
  2. 2단계 미세 조정: 학습된 순서 정책을 고정하고, 학습된 궤적을 따라 확산 모델 $\theta$를 미세 조정하여 오류를 더욱 줄인다(선택 사항).

3. 주요 기여

1. 이론: 본 논문은 디코딩 미스매치(KL 발산)를 디노이저의 누적 경로별 로그 가능도와 연결하는 다룰 수 있는 상한을 도출한다. 이는 모델 자신의 가능도를 순서 선택을 위한 조밀한 보상으로 사용하는 것에 대한 이론적 정당성을 확립한다.
2. 알고리즘 (SAS): GRPO를 통해 언마스킹 순서 정책을 학습하는 플러그 앤 플레이 프레임워크를 도입한다. 이는 **임의 순서 순차 디코딩(any-order sequential decoding)**과 준-자기회귀(블록) 디코딩(semi-autoregressive (block) decoding) 모두에 적용 가능하다.
3. 실증적 검증: 다양한 추론 작업(스도쿠, 수학, 코드)과 모델 규모(1B ~ 8B 파라미터)에 걸쳐, 최적화된 사고 순서가 휴리스틱 및 전문가 스케줄보다 실질적이고 일관된 개선을 가져온다는 것을 입증한다.

4. 실험 결과

스도쿠 (1B 모델)

• 성능: SAS는 퍼즐 정확도를 82.0%(최상의 휴리스틱: Confidence)에서 **91.8%**로 향상시켰다.
• 비교: 학습된 정책은 인간 전문가의 순서(76.6%)를 크게 앞질렀으며, 이는 불완전한 확산 모델의 경우 최적의 경로가 반드시 인간의 논리적 경로와 일치하지 않을 수 있음을 시사한다.
• 분석: 시각화를 통해 정책이 "쉬운 것에서 어려운 것으로" 향하는 커리큘럼을 채택하여, 높은 확신을 가진 토큰(naked singles)을 먼저 해결함으로써 제약 조건을 전파하지만, 이득이 되는 지점에서는 인간의 논리에서 벗어남을 확인했다.
• 2단계: 학습된 궤적을 따라 모델을 미세 조정하면 정확도가 **97.5%**까지 추가로 상승한다.

수학적 추론 (LLaDA-8B)

• GSM8K: Pass@1이 **64%**에서 **76%**로 향상되었다.
• MBPP (코드): Pass@1이 **39.5%**에서 **41%**로 향상되었다.
• 일반화: 학습된 정책은 다양한 생성 길이(128 ~ 512)와 블록 크기(준-자기회귀 디코딩)에 대해 견고함을 유지하며, 좌에서 우로 향하는 휴리스틱을 포함한 모든 베이스라인을 일관되게 능가했다.

교차 도메인 전이

• 수학 작업에서 학습된 정책은 코드 작업으로 효과적으로 전이되었으며(그 반대도 마찬가지), 테스트된 모든 설정에서 가장 강력한 휴리스틱 베이스라인을 능가했다. 이는 정책이 작업 특정적인 패턴을 암기하기보다는 전이 가능한 스케줄링 원칙(정보 밀도 관련)을 학습함을 나타낸다.

보상 비교

• 희소한(sparse) 0/1 터미널 보상(정답/오답)과 비교했을 때, 조밀한(dense) 자기 인식 보상이 더 안정적이고 종종 더 우수한 성능을 제공했으며, 이는 종말 피드백(terminal feedback)보다 모델 일관적인 조밀한 신호의 이점을 강조한다.

5. 의의 및 주장

본 논문은 디코딩 순서를 제어 가능하고 최적화 가능한 "추론 프리미티브(reasoning primitive)"로 규정한다. 저자들은 무엇을 생성할지를 배우는 것을 넘어, 각 정보에 언제 확언(commit)할지를 배우는 것이 가능하다고 주장한다.

• 원칙적인 최적화: 로컬 불확실성을 최적화하는 휴리스틱 스케줄과 달리, SAS는 전역 궤적 가능도를 최적화하여 디코딩 경로를 모델의 예측 강점과 일치시킨다.
• 효율성: 이 접근법은 전체 확산 백본에 대한 비용이 많이 드는 강화 학습 업데이트 대신 경량화된 정책만을 학습하므로 계산적으로 효율적이다.
• 자기 개선: 이 프레임워크는 반복적인 스케줄링과 모델 정제를 통해 자기 개선형 확산 LLM으로 나아가는 실질적인 경로를 제시하며, 여기서 모델은 자신의 가장 신뢰할 수 있는 추론 경로를 스스로 학습한다.

저자들은 교사 강제(teacher-forced) 훈련과 자유 실행(free-running) 추론 사이의 잠재적 간극, 그리고 현재 정책 학습을 위한 작업별 학습 의존성(교차 도메인 전이가 유망하지만 아직 완전히 규명되지 않음)을 한계점으로 인정한다.