마스크 확산 언어 모델을 위한 상태 기반 정책 최적화

초록

DiSPO는 마스크 확산 언어 모델(MDLM)에서 중간 마스크 상태를 고정하고, 해당 상태에서 캐시된 로짓을 재샘플링해 여러 후보 완성을 만든 뒤 동일한 최종 보상으로 평가하여, 새로 채워진 토큰에만 정책 그래디언트를 적용하는 플러그인 레이어이다. 이를 통해 기존의 최종 보상만 이용하는 정책 최적화보다 더 정교한 크레딧 할당이 가능해져 수학·계획 베엔치마크에서 일관된 성능 향상을 보인다.

상세 분석

본 논문은 마스크 확산 언어 모델(MDLM)이 여러 디노이징 단계에서 마스크된 토큰을 순차적으로 채워 나가는 특성을 활용해, 강화학습에서 흔히 겪는 ‘최종 보상만을 전체 토큰에 균등히 전파하는’ 문제를 근본적으로 해결하고자 한다. 기존의 terminal‑feedback 정책 최적화(예: diffu‑GRPO)는 최종 완성에 대한 스칼라 보상 R(q, o)만을 사용해 전체 시퀀스에 대한 likelihood‑ratio를 계산한다. 이는 중간 단계에서 어떤 토큰이 성공에 기여했는지 구분하기 어려워, 특히 복잡한 추론이나 계획 문제에서 학습 효율을 저하시킨다.

DiSPO는 이러한 한계를 ‘상태‑행동’ 관점을 도입해 극복한다. 각 디노이징 단계 t에서 현재 마스크된 위치 집합 Mₖ,ₜ를 상태 sₖ,ₜ=(q, xₖ,ₜ)로 정의하고, 해당 마스크 위치들을 한 번에 채우는 행동 aₖ,ₜ를 정책 πθ(·|sₖ,ₜ)로 모델링한다. 중요한 점은 행동을 ‘동시 채우기’로 묶음으로써, 한 상태에서 여러 대안(aₖ,ₜ,₁…aₖ,ₜ,𝑍)을 샘플링하고 각각을 완성 oₖ,ₜ,𝑧=FILL(xₖ,ₜ, aₖ,ₜ,𝑧)로 변환한다는 것이다. 여기서 FILL은 마스크된 위치만 교체하고 나머지는 그대로 유지하는 결정적 함수이다.

이때 보상은 여전히 최종 보상 R(q, oₖ,ₜ,𝑧)만을 사용한다. 동일한 보상 함수를 여러 대안에 적용함으로써, ‘같은 중간 상태에서 어떤 채우기가 더 좋은가’를 직접 비교한다. 보상의 평균을 베이스라인으로 삼아 Aₖ,ₜ,𝑧=Rₖ,ₜ,𝑧−\bar{R}ₖ,ₜ를 구하고, 이를 정책 그래디언트의 가중치로 활용한다. 이렇게 하면 중간 단계에서의 크레딧 할당이 훨씬 세밀해진다.

또한 DiSPO는 기존 rollout 과정에서 이미 캐시된 로짓을 재사용한다. 즉, 새로운 디노이징 스텝을 수행하지 않아도 동일한 상태에서 여러 후보를 생성할 수 있다. 이는 계산 비용을 크게 늘리지 않으면서도, ‘same‑state branching’이라는 새로운 학습 신호를 제공한다.

수식적으로는 상태별 마스크 토큰 서러게이트 ˜πθ를 정의하고, 이를 이용해 로그 확률 log ˜πθ(aₖ,ₜ,𝑧|sₖ,ₜ)를 계산한다. 이 서러게이트는 기존 diffu‑GRPO에서 사용한 단일 마스크 토큰 서러게이트(Equation 1)를 확장한 형태이며, 현재 마스크된 위치에만 집중한다. 따라서 정책 그래디언트는

∇θ L_step = - (1/Z) ∑₍𝑧₎ ρₖ,ₜ,𝑧(θ) Aₖ,ₜ,𝑧

와 같이 표현되며, 여기서 ρ는 현재와 이전 파라미터 사이의 likelihood‑ratio이다. 논문은 이 식이 기대값에서 고전적인 정책 그래디언트 ∇θ Jₜ(θ)와 일치함을 정리 4.1을 통해 증명한다.

또한 terminal‑feedback 손실 L_term과 step‑wise 손실 L_step을 가중합(α_term, α_step)으로 결합하면, 전체 기대 그래디언트가 단일 목표 J_seq(θ)+∑ₜ Jₜ(θ)의 그래디언트와 동일함을 정리 4.2에서 보인다. 즉, 두 손실을 동시에 최적화해도 이론적으로 모순이 없으며, 실제로는 중간 단계에서의 정밀한 신호가 최종 성능을 끌어올린다.

분산 감소 측면에서도 두 가지 기법을 제시한다. 첫째, 업데이트를 현재 단계에서 새로 채워진 토큰에만 제한함으로써 불필요한 파라미터 변동을 억제한다. 둘째, 같은 상태에서 여러 후보를 평균화(‘same‑state averaging’)함으로써 베이스라인에 대한 변동성을 감소시킨다(정리 4.3, 4.4).

실험에서는 LLaDA‑8B‑Instruct 모델에 DiSPO를 적용해, GSM8K·Math500 같은 수학 추론 벤치마크와 Sudoku·Countdown 같은 기획 문제에서 기존 diffu‑GRPO 대비 일관된 성능 향상을 기록했다. 특히 rollout 수와 optimizer 업데이트 횟수를 동일하게 맞춘 ‘matched compute’ 설정에서도 개선이 관찰돼, 추가 연산이 아니라 더 정교한 크레딧 할당이 효과임을 입증한다.

요약하면, DiSPO는 마스크 확산 언어 모델의 중간 상태를 정책 결정점으로 활용하고, 캐시된 로짓을 재활용해 비용 효율적으로 다중 후보를 생성·평가함으로써, 기존의 최종 보상만 이용하는 정책 최적화의 한계를 극복한다. 이 접근법은 향후 복잡한 추론·계획 작업에 MDLM을 적용할 때 보다 세밀한 학습 신호를 제공할 수 있는 유망한 방향을 제시한다.