기술 요약: 마스킹 확산 언어 모델을 위한 확산 상태 정책 최적화

문제 정의

마스킹 확산 언어 모델 (MDLMs) 은 여러 번의 탈노이즈 단계를 거쳐 마스킹된 위치를 반복적으로 채워 텍스트를 생성합니다. 정책 최적화 (PO) 방법은 언어 모델의 추론 및 정렬 능력을 성공적으로 향상시켜 왔으나, 대부분의 기존 접근법은 **최종 결과에만 기반한 보상 (terminal-only rewards)**에 의존합니다. 이러한 방법에서는 최종 완성본에 스칼라 보상을 할당한 후, 이를 모든 중간 생성 결정에 역전파합니다.

이러한 접근 방식은 **거친 신용 할당 (coarse credit assignment)**을 초래합니다. 이는 모든 로컬 채우기 결정의 전체 시퀀스를 단일 단위로 취급하여, 롤아웃 과정에서 이미 생성된 풍부한 중간 상태 - 행동 정보를 활용하지 못합니다. 구체적으로, 어떤 중간 탈노이즈 단계에서든 모델은 현재 마스킹된 모든 위치에 대해 토큰 수준의 분포를 생성하지만, 표준 최종 피드백 PO 는 이러한 세분성을 무시하고 최종 결과에 대한 특정 기여도와 관계없이 모든 결정에 동일한 결과 신호를 할당합니다.

방법론: 확산 상태 정책 최적화 (DISPO)

저자들은 추가적인 다단계 확산 롤아웃 없이 중간 채우기 결정을 직접 최적화하도록 설계된 플러그인 신용 할당 계층인 **확산 상태 정책 최적화 (DISPO)**를 제안합니다.

핵심 개념

DISPO 는 MDLM 의 생성 과정을 상태 - 행동 프레임워크로 재해석합니다:

• 상태 ($s_{k,t}$): 탈노이즈 단계 $t$에서의 중간 부분 마스킹 시퀀스.
• 행동 ($a_{k,t}$): 해당 상태에서 현재 마스킹된 위치들의 결합된 채우기.

알고리즘적 메커니즘

1. 롤아웃 및 캐싱: 고정된 행동 파라미터 ($\theta_{old}$) 를 사용한 표준 롤아웃 동안, 모델은 부분 마스킹 시퀀스의 궤적을 생성합니다. 결정적으로, 방문한 모든 상태에서의 모든 마스킹 위치에 대한 로짓이 캐시됩니다.
2. 상태 조건부 분기: 선택된 중간 상태들에서 DISPO 는 롤아웃에서 캐시된 로짓을 사용하여 마스킹 위치들에 대한 대체 채우기를 재샘플링함으로써 "분기"합니다. 이는 동일한 중간 상태에서 여러 개의 "반사실적" 완성본 ($Z$개의 분기) 을 생성합니다.
3. 점수 매기기 및 이점: 각 분기된 완성본은 표준 PO 에서 사용하는 동일한 최종 보상 함수 ($R$) 를 사용하여 점수가 매겨집니다. 해당 특정 상태에서의 $Z$개 완성본들의 보상에 기반하여 각 분기에 대한 그룹 상대 이점이 계산됩니다.
4. 토큰 로컬 업데이트: 정책 기울기는 현재 상태에서의 새롭게 채워진 토큰 (행동) 만을 업데이트하도록 계산되며, 이전에 채워진 토큰들은 고정된 상태의 일부로 간주됩니다.
5. 대리 가능도: MDLM 에서 정확한 시퀀스 가능도의 계산 불가능성을 처리하기 위해, DISPO 는 계산 가능한 **한 단계 마스킹 토큰 대리 모델 ($\tilde{\pi}_\theta$)**을 활용합니다. 이를 통해 업데이트 중인 마스킹 위치에 대해 구체적으로 가능도 비율을 계산할 수 있습니다.

훈련 목적 함수

DISPO 는 고정 상태 기대 반환 목적 함수를 정의합니다. 이는 가중 합을 통해 기본 최종 피드백 최적화기 (예: diffu-GRPO 또는 SPG) 와 결합될 수 있습니다:
$$\mathcal{L}(\theta) = \alpha_{base} \mathcal{L}_{base}(\theta) + \alpha_{step} \mathcal{L}_{step}(\theta)$$
여기서 $\mathcal{L}_{step}$은 분기된 완성본들로부터 유도된 상태별 손실입니다. 이론적 분석은 이 결합된 목적 함수가 최종 목표와 시간 단계 평균 고정 상태 목표로 구성된 혼합 목적 함수에 대한 유효한 정책 기울기 신호를 제공함을 확인시켜 줍니다.

주요 기여

1. 상태 조건부 신용 할당: 본 논문은 중간 마스킹 시퀀스를 상태로, 마스킹 채우기를 행동으로 취급하는 방법을 도입하여, 표준 최종 피드백 PO 가 결여된 세분화된 신용 할당을 가능하게 합니다.
2. 캐시된 로짓을 통한 효율적 분기: 주요 혁신은 동일한 상태에서 대체 완성본을 생성하기 위해 롤아웃 캐시된 로짓을 재사용하는 것입니다. 이를 통해 추가적인 다단계 확산 순방향 통과 없이 동일한 상태 간 비교가 가능해지며, 지배적인 훈련 비용 (롤아웃 및 최적화기 단계) 을 베이스라인과 일치시킵니다.
3. 이론적 형식화: 저자들은 DISPO 를 고정 상태 목적 함수로 형식화하고 정책 기울기 추정기를 유도합니다. 그들은 단계별 손실이 고정 상태 목적 함수에 대한 편향되지 않은 (스케일된) 추정치임을 증명하며, 이를 최종 손실과 결합하면 잘 정의된 혼합 목적 함수를 최적화함을 보여줍니다.
4. 분산 감소 전략: 이 방법은 분산 분석에 영감을 받은 두 가지 설계 선택을 채택합니다:
   • 토큰 로컬 업데이트: 모든 토큰을 업데이트하는 것보다 실행 가능한 (마스킹된) 토큰으로 기울기를 제한함으로써 분산을 감소시킵니다.
   • 동일 상태 평균화: $Z$개의 분기에 대해 평균을 내어 분산을 $O(1/Z)$만큼 감소시킵니다.

실험 결과

저자들은 LLaDA-8B-Instruct에서 DISPO 를 평가했으며, 이를 diffu-GRPO와 SPG라는 두 가지 기본 최적화기와 통합했습니다.

• 벤치마크: 실험은 수학 추론 (GSM8K, MATH500) 및 기호 계획 (Sudoku, Countdown) 에서 수행되었습니다.
• 성능: 일치된 계산 예산 (동일한 수의 롤아웃, 그룹 크기, 최적화기 단계) 하에서 DISPO 는 모든 벤치마크에서 두 기본 최적화기를 일관되게 향상시켰습니다.
  • 계획 작업에서 향상폭이 특히 두드러졌습니다 (예: diffu-GRPO 를 사용한 경우 Sudoku 정확도가 약 14% 에서 약 31% 로 증가).
  • 향상은 다양한 생성 예산 ($N_{gen} \in \{128, 256, 512\}$) 에서도 지속되었습니다.
• 애블레이션 연구:
  • 상태별 vs 최종: 상태별 목적 함수만으로도 최종 결과만 있는 베이스라인보다 성능이 우수했으며, 이들의 결합이 최상의 결과를 낳았습니다.
  • 타이밍: 완료에 가까운 상태에 초점을 맞춘 후기 편향 시간 단계 샘플링이 균일하거나 초기 편향 샘플링보다 더 좋은 성능을 보였으며, 이는 최종 완료에 가까운 중간 상태들이 더 신뢰할 수 있는 신용 신호를 제공함을 시사합니다.
  • 기울기 분산: 실험은 실행 가능한 토큰만 업데이트하고 분기 크기 $Z$를 증가시키는 것이 기울기 분산을 현저히 감소시켰음을 확인했습니다.
• 오류 분석: Sudoku 에 대한 "조기 약속 (premature commitments)"에 대한 분석은 DISPO 가 제약 위반의 첫 발생을 베이스라인에 비해 더 늦은 탈노이즈 단계로 이동시켰음을 보여주었으며, 이는 중간 행동을 최적화함으로써 초기에 발생하여 수정하기 어려운 오류를 피하는 데 도움이 됨을 나타냅니다.

중요성 및 주장

본 논문은 DISPO 가 마스킹 확산 정책 최적화를 위한 범용 플러그인으로 작용한다고 주장합니다. 그 주요 중요성은 다음과 같습니다:

• 중간 정보 활용: MDLM 생성 과정에 내재된 중간 상태 - 행동 정보가 완전히 새로운 궤적을 생성하는 계산 오버헤드 없이 세분화된 신용 할당에 활용될 수 있음을 보여줍니다.
• 상호 보완성: DISPO 는 보상 설계 (예: SAPO) 나 궤적 수준 목적 함수와 같은 기존 방법과 상호 보완적으로 제시됩니다. SAPO 가 중간 상태를 기반으로 보상을 설계하는 반면, DISPO 는 고정 상태 하에서 채우기 행동을 직접 최적화합니다.
• 효율성: 캐시된 로짓을 재사용함으로써 DISPO 는 지배적인 훈련 비용 (다단계 롤아웃) 을 증가시키지 않고 이러한 개선을 달성하여 기존 MDLM 훈련 파이프라인에 실용적인 추가 요소가 됩니다.

저자들은 현재 구현이 반복적인 대리 모델 및 보상 평가로 인해 일부 월클럭 오버헤드를 수반하지만, 알고리즘적 비용은 베이스라인과 일치하며 성능 향상은 견고하다고 언급합니다. 그들은 DISPO 를 확산 기반 언어 모델에서 더 효율적이고 효과적인 추론을 향한 한 걸음으로 위치시킵니다.