보상 기반 연결을 통한 디퓨전 언어 모델을 활용한 추론 시간 스케일링

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

대규모 언어 모델을 활용한 추론은 종종 여러 단계의 사고 과정을 생성함으로써 이점을 얻을 수 있지만, 기존의 집계 전략은 일반적으로 전체 경로 수준(예: 최적의 경로 선택 또는 최종 답변에 대한 투표)으로 이루어지며, 부분적이거나 "거의 정확한" 시도에서 발생하는 유용한 중간 결과를 버립니다. 본 연구에서는 "Stitching Noisy Diffusion Thoughts"라는 자기 일관성 프레임워크를 제안합니다. 이 프레임워크는 저렴한 디퓨전 샘플링을 통해 생성된 추론 과정을 재사용 가능한 단계 수준의 후보 풀로 변환합니다. 주어진 문제에 대해, (i) 마스크된 디퓨전 언어 모델을 사용하여 다양한 저비용 추론 경로를 여러 개 샘플링하고, (ii) 오프라인 프로세스 보상 모델(PRM)을 사용하여 모든 중간 단계를 평가하고, (iii) 가장 높은 품질의 단계를 여러 경로에서 연결하여 통합된 추론 과정을 생성합니다. 이 추론 과정은 이후 오토리그레시브(AR) 모델(솔버)에 입력되어 최종 답변을 재계산하도록 합니다. 이러한 모듈화된 파이프라인은 탐색(디퓨전)과 평가 및 솔루션 합성을 분리하여, 단일화된 통합 모델의 한계를 극복하면서도 광범위한 탐색을 유지합니다. 수학 추론 벤치마크에서, 단계 수준의 재조합은 더 어려운 문제에서 가장 큰 효과를 발휘하며, 추가적인 실험을 통해 최종 AR 솔버가 연결되었지만 불완전한 추론 과정을 정확한 답변으로 변환하는 데 중요한 역할을 한다는 것을 확인했습니다. 본 연구에서는 낮은 신뢰도의 디퓨전 샘플링과 병렬, 독립적인 실행을 사용하여, 훈련 없이 평균 정확도를 최대 23.8% 향상시켰습니다. 동시에, 기존의 디퓨전 모델(예: Dream, LLaDA) 및 통합 아키텍처(예: TiDAR)에 비해 최대 1.8배의 지연 시간 감소를 달성했습니다. 관련 코드는 https://github.com/roymiles/diffusion-stitching 에서 확인할 수 있습니다.

Reasoning with large language models often benefits from generating multiple chains-of-thought, but existing aggregation strategies are typically trajectory-level (e.g., selecting the best trace or voting on the final answer), discarding useful intermediate work from partial or "nearly correct" attempts. We propose Stitching Noisy Diffusion Thoughts, a self-consistency framework that turns cheap diffusion-sampled reasoning into a reusable pool of step-level candidates. Given a problem, we (i) sample many diverse, low-cost reasoning trajectories using a masked diffusion language model, (ii) score every intermediate step with an off-the-shelf process reward model (PRM), and (iii) stitch these highest-quality steps across trajectories into a composite rationale. This rationale then conditions an autoregressive (AR) model (solver) to recompute only the final answer. This modular pipeline separates exploration (diffusion) from evaluation and solution synthesis, avoiding monolithic unified hybrids while preserving broad search. Across math reasoning benchmarks, we find that step-level recombination is most beneficial on harder problems, and ablations highlight the importance of the final AR solver in converting stitched but imperfect rationales into accurate answers. Using low-confidence diffusion sampling with parallel, independent rollouts, our training-free framework improves average accuracy by up to 23.8% across six math and coding tasks. At the same time, it achieves up to a 1.8x latency reduction relative to both traditional diffusion models (e.g., Dream, LLaDA) and unified architectures (e.g., TiDAR). Code is available at https://github.com/roymiles/diffusion-stitching.