Test-Time Scaling with Diffusion Language Models via Reward-Guided Stitching

1. 문제 정의 (Problem)

대규모 언어 모델 (LLM) 을 활용한 추론은 종종 여러 사고의 사슬 (Chain-of-Thought, CoT) 을 생성하는 것이 유리합니다. 그러나 기존 집계 전략은 대부분 **궤적 수준 (trajectory-level)**에서 작동합니다. 즉, 가장 좋은 전체 경로를 선택하거나 최종 답변에 대한 다수결 투표 (voting) 를 수행하는 방식입니다.

• 한계점: 이러한 방식은 부분적으로 맞거나 거의 정답에 도달했으나 중간에 실패한 시도의 유용한 중간 단계 (intermediate work) 를 폐기하게 됩니다.
• 확산 언어 모델 (Diffusion LM) 의 기회와 과제: 확산 모델은 병렬 디노이징을 통해 다양한 저비용 추론 경로를 생성할 수 있지만, 단일 긴 경로를 처음부터 다시 생성하는 데 드는 비용은 여전히 높거나, 생성된 경로가 노이즈가 많아 신뢰도가 낮을 수 있습니다. 또한, 기존 하이브리드 모델 (확산 + 자기회귀) 은 탐색, 평가, 생성 단계를 강하게 결합하여 유연성이 떨어지고 추론 정확도가 자기회귀 (AR) 모델에 비해 낮아지는 경향이 있습니다.

2. 방법론 (Methodology)

저자들은 **"Stitching Noisy Diffusion Thoughts"**라는 새로운 프레임워크를 제안합니다. 이는 확산 샘플링을 완전한 후보 솔루션 집합이 아닌, **재사용 가능한 중간 추론 단계의 풀 (pool)**로 변환하는 자기 일관성 (self-consistency) 기반 접근법입니다.

프로세스는 크게 세 단계로 구성됩니다:

1. 탐색 (Explore):

   • 입력 문제를 기반으로 **마스크 확산 언어 모델 (Masked Diffusion LM)**을 사용하여 $N$개의 다양한 저비용 추론 궤적 (CoT) 을 샘플링합니다.
   • 확산 모델은 초기에 완전히 마스킹된 템플릿에서 시작하여 병렬적으로 토큰을 채워나가며, 낮은 신뢰도 토큰은 재마스킹하는 방식을 사용합니다.

2. 평가 (Evaluate):

   • 생성된 모든 궤적의 **각 중간 단계 (step-level)**를 오프 - 더 - 쉘 (off-the-shelf) **프로세스 보상 모델 (Process Reward Model, PRM)**로 점수화합니다.
   • 이는 전체 궤적이 실패하더라도 유용한 중간 유도 단계를 유지할 수 있게 해주는 핵심 단계입니다.

3. 스티칭 및 재계산 (Stitch + Recompute):

   • 스티칭: 궤적 전체를 선택하는 대신, 모든 궤적에서 가장 높은 품질 (높은 PRM 점수) 을 가진 단계들을 수집하여 **복합적인 근거 (composite rationale)**로 연결합니다. 이 과정은 시간 순서를 유지하면서 다른 궤적의 보완적인 하부 유도 단계를 통합합니다.
   • 재계산: 이렇게 만들어진 스티칭된 근거 (evidence) 와 원래 문제를 조건으로 가벼운 자기회귀 (AR) 솔버를 실행하여 최종 답변을 다시 계산합니다. AR 솔버는 스티칭된 근거의 불일치를 해결하고 누락된 연결 고리를 채워 최종 정답을 도출합니다.

이 아키텍처는 **탐색 (확산 모델)**과 **평가/해결 (PRM 및 AR 솔버)**을 모듈식으로 분리하여, 확장성과 유연성을 동시에 확보합니다.

3. 주요 기여 (Key Contributions)

1. 단계 수준 (Step-level) 자기 일관성 프레임워크: 궤적 수준의 선택을 대체하여, 확산으로 샘플링된 CoT 에서 PRM 점수가 매겨진 단계를 선택하고 스티칭하는 새로운 방식을 도입했습니다.
2. 모듈러 파이프라인 설계: 확산 모델 (저비용 탐색), PRM (세분화된 단계 평가), 경량 AR 솔버 (최종 답변 생성) 를 분리하여 각 구성 요소의 역할을 명확히 하고, 이를 통해 테스트 시간 확장 (Test-time scaling) 의 효율성을 입증했습니다.
3. 성능 및 효율성 입증: 수학 및 코딩 벤치마크에서 기존 AR 모델 및 확산/하이브리드 모델 대비 정확도 - 지연 시간 (Accuracy-Latency) 트레이드오프를 크게 개선했습니다.

4. 실험 결과 (Results)

저자들은 수학 (GSM8K, MATH500) 및 코딩 (HumanEval, MBPP) 태스크에서 제로샷 (zero-shot) 설정으로 실험을 수행했습니다.

• 정확도 향상:
  • 기존 확산 모델 (LLaDA, Dream 등) 대비 평균 **23.8%~30.6%**의 정확도 향상을 기록했습니다.
  • 강력한 자기회귀 (AR) 베이스라인 (예: Qwen3-8B) 과 비교했을 때, 평균 4.3%p 더 높은 정확도를 보이거나 유사한 수준을 유지했습니다.
  • 특히 어려운 문제 (MATH500 등) 일수록 단계 수준의 재조합의 이점이 큽니다.
• 지연 시간 및 효율성:
  • 기존 확산 모델 및 통합 아키텍처 (TiDAR) 대비 최대 1.8 배의 지연 시간 감소를 달성했습니다.
  • 순차적 포워드 패스 (forward passes) 는 최대 9.85 배 감소시켰습니다 (확산 모델의 병렬 탐색과 AR 솔버의 단일 실행 덕분).
  • 낮은 신뢰도 샘플링: 확산 모델을 낮은 신뢰도 (low-confidence) 로 실행하여 노이즈가 많은 경로를 생성하더라도, PRM 과 스티칭을 통해 높은 정확도를 유지하며 계산 비용을 크게 절감할 수 있음을 보였습니다.
• 비교 분석:
  • 단순 다수결 투표나 전체 CoT 연결보다 신뢰도 기반 스티칭이 가장 효과적이었습니다.
  • AR 솔버가 스티칭된 불완전한 근거를 정답으로 변환하는 데 결정적인 역할을 함을 어블레이션 (ablation) 연구를 통해 확인했습니다.

5. 의의 및 결론 (Significance & Conclusion)

이 논문은 학습이 필요 없는 (training-free) 테스트 시간 확장 전략으로서, 확산 언어 모델의 잠재력을 극대화하는 새로운 패러다임을 제시합니다.

• 핵심 통찰: 추론 과정에서 "전체 궤적"을 하나의 단위로 보는 것이 아니라, "유용한 중간 단계"를 재사용하고 결합하는 것이 더 효율적이고 강력합니다.
• 실용성: 고비용의 긴 자기회귀 추론 없이도, 확산 모델의 병렬 탐색 능력과 AR 모델의 정밀한 생성 능력을 결합하여 저지연 (low-latency) 환경에서도 고품질 추론이 가능함을 입증했습니다.
• 미래 방향: 이 접근법은 계산 자원이 제한된 환경에서 복잡한 추론 태스크를 수행할 때, 기존 단일 모델 아키텍처의 한계를 극복하는 효과적인 대안이 될 수 있습니다.

요약하자면, 이 연구는 "노이즈가 많은 확산 추론 경로들을 조각내어 가장 좋은 부분들만 이어붙이고 (Stitching), 이를 바탕으로 최종 답을 다시 계산하는" 방식을 통해, 정확도와 속도라는 상충되는 목표를 동시에 달성하는 혁신적인 방법을 제시했습니다.