[논문 리뷰] Self-Rewarding Sequential Monte Carlo for Masked Diffusion Language Models

Ziwei Luo, Ziqi Jin외 3명|arXiv (Cornell University)|2026. 02. 02.

Topic Modeling인용 수 0

한 줄 요약

Self-Rewarding Sequential Monte Carlo(SR-SMC)를 도입하여 마스킹 확산 언어 모델의 추론 시 확장성을 위한 궤적 수준의 신뢰도 기반 샘플링을 가능하게 하여, 샘플 품질과 다양성을 추가 학습 없이 개선합니다.

ABSTRACT

This work presents self-rewarding sequential Monte Carlo (SMC), an inference-time scaling algorithm enabling effective sampling of masked diffusion language models (MDLMs). Our algorithm stems from the observation that most existing MDLMs rely on a confidence-based sampling strategy, where only tokens with the highest prediction confidence are preserved at each step. This restricts the generation to a noise-sensitive, greedy decoding paradigm, resulting in an inevitable collapse in the diversity of possible paths. We address this problem by launching multiple interacting diffusion processes in parallel, referred to as particles, for trajectory exploration. Importantly, we introduce the trajectory-level confidence as a self-rewarding signal for assigning particle importance weights. During sampling, particles are iteratively weighted and resampled to systematically steer generation towards globally confident, high-quality samples. Our self-rewarding SMC is verified on various masked diffusion language models and benchmarks, achieving significant improvement without extra training or reward guidance, while effectively converting parallel inference capacity into improved sampling quality. Our code is available at https://github.com/Algolzw/self-rewarding-smc.

연구 동기 및 목표

• Greedy MDLM 샘플링의 다양성 한계를 동기부여하고 해결합니다.
• 궤적 수준의 신뢰도를 자기 보상 신호로 사용하는 일반적인 SR-SMC 프레임워크를 제안합니다.
• 추가 보상 모델 없이 SR-SMC가 MDLM 및 dLLM 전반의 샘플 품질을 향상시킨다는 것을 보여줍니다.
• 다양한 모델과 벤치마크에서 광범위한 실험을 통해 확장성을 demonstrated합니다.

제안 방법

• N개 상호작용하는 확산 프로세스(입자)를 유지하여 여러 궤적을 병렬로 탐색합니다.
• 업데이트된 토큰들에 대한 곱으로 궤적 수준의 신뢰도(식 13)를 정의하여 입자를 가중합니다.
• 적응적 리마스킹 및 unmask할 토큰을 선택하는 정책을 갖춘 역확산 커널을 사용합니다(식 9–10 및 정책 식 7–8).
• ESS를 기반으로 적응적 재샘샘플링을 포함한 표준 SMC 단계(재샘플링, 전파, 재가중)를 적용합니다(식 14).
• 온도 제어가 가능한 이산 토큰 샘플링을 위해 Gumbel-Max 트릭을 사용합니다(식 15).
• 궤적 수준의 신뢰도가 bootstrap SMC 설정에서 자연스러운 자기 보상임을 이론적으로 정당화합니다(보정 3.1).

실험 결과

연구 질문

• RQ1 SR-SMC가 추가 학습이나 외부 보상 없이 MDLM 및 dLLM의 샘플링 품질을 개선할 수 있는가?
• RQ2 궤적 수준 가중치가 MDLM에서 토큰 수준의 신뢰도에 비해 더 나은 탐색과 다양성을 가능하게 하는가?
• RQ3 SR-SMC가 표준 벤치마크에서 서로 다른 MDLM 및 확산 기반 LLM에 대해 어떻게 수행하는가?
• RQ4 입자 수와 온도가 SR-SMC의 성능과 안정성에 미치는 영향은 무엇인가?

주요 결과

• SR-SMC는 MDLM(MDLM, BD3-LMs) 및 dLLMs(LLaDA-1.5, Dream-7B) 전반의 생성 당혹도(Gen. PPL)와 샘플 품질을 일관되게 향상시킵니다.
• 입자 수(N)를 늘리면 점진적인 이득이 발생하며, 특히 N=3 또는 4에서 두드러진 개선이 나타납니다.
• SR-SMC가 적용된 블록 디코딩 형태는 특정 BD3-LMs 구성에서 Gen. PPL을 20 미만으로 낮추며 자동회귀 기반과의 차이를 좁힙니다.
• SR-SMC는 GSM8K, MATH, HumanEval, MBPP 벤치마크에서 성능을 향상시키며 모델 및 길이에 따라 평균 약 2–4포인트의 이득을 보입니다.
• SR-SMC는 샘플링 온도에 대해 강건하며 저온에서 그리디 디코딩에 비해 반복성을 줄입니다.
• 약한 입자 수와 제로샷 설정에서도 SR-SMC가 상당한 이득을 보이는 일종의 제거 분석이 있습니다.
• 분석에 따르면 입자들 간에 추월이 벌어지는 블록의 비율이 무시하지 못할 만큼 존재하여 SR-SMC가 비그리드 궤적을 탐색한다는 것을 확인합니다.