프리즘: 이산 확산 언어 모델을 위한 계층적 탐색 및 자가 검증 기반 효율적 테스트 시점 확장 기법

Prism: Efficient Test-Time Scaling via Hierarchical Search and Self-Verification for Discrete Diffusion Language Models

February 2, 2026

저자: Jinbin Bai, Yixuan Li, Yuchen Zhu, Yi Xin, Qingyu Shi, Aosong Feng, Xiaohong Liu, Molei Tao, Jianru Xue, Xiangtai Li, Ming-Hsuan Yang

cs.AI

초록

추론 시 계산(compute)은 LLM 추론 성능을 향상시키는 실용적인 방법으로 다시 부각되고 있습니다. 대부분의 테스트 타임 스케일링(TTS) 알고리즘은 자기회귀 디코딩에 의존하는데, 이는 전체 시퀀스에 대해 병렬 디코딩을 수행하는 이산 확산 언어 모델(dLLM)에는 적합하지 않습니다. 그 결과 dLLM의 전체 생성 잠재력을 끌어내기 위한 효과적이고 효율적인 TTS 방법 개발은 아직 충분히 탐구되지 않은 과제로 남아 있습니다. 이를 해결하기 위해 우리는 dLLM을 위한 효율적인 TTS 프레임워크인 Prism(Pruning, Remasking, and Integrated Self-verification Method)을 제안합니다. Prism은 (i) 초기부터 중간 디노이징 창에서 계산을 동적으로 가지치기하고 재배치하는 계층적 궤적 탐색(HTS)을 수행하며, (ii) 높은 신뢰도를 가진 토큰을 보존하면서 다양한 구현을 탐색하기 위해 부분 재마스킹과 함께 지역 분기(Local branching)를 도입하고, (iii) 외부 검증기를 중간 완성본에 대한 자기 평가 프롬프트를 통해 얻은 자체 검증 피드백(SVF)으로 대체합니다. LLaDA 8B Instruct, Dream 7B Instruct, LLaDA 2.0-mini를 포함한 세 가지 dLLM에 대한 네 가지 수학적 추론 및 코드 생성 벤치마크에서 우리의 Prism은 유리한 성능-효율성 트레이드오프를 달성하며, 최고의 N-샘플링(Best-of-N) 성능을 훨씬 더 적은 함수 평가 횟수(NFE)로 일치시켰습니다. 코드는 https://github.com/viiika/Prism에서 공개되었습니다.

English

Inference-time compute has re-emerged as a practical way to improve LLM reasoning. Most test-time scaling (TTS) algorithms rely on autoregressive decoding, which is ill-suited to discrete diffusion language models (dLLMs) due to their parallel decoding over the entire sequence. As a result, developing effective and efficient TTS methods to unlock dLLMs' full generative potential remains an underexplored challenge. To address this, we propose Prism (Pruning, Remasking, and Integrated Self-verification Method), an efficient TTS framework for dLLMs that (i) performs Hierarchical Trajectory Search (HTS) which dynamically prunes and reallocates compute in an early-to-mid denoising window, (ii) introduces Local branching with partial remasking to explore diverse implementations while preserving high-confidence tokens, and (iii) replaces external verifiers with Self-Verified Feedback (SVF) obtained via self-evaluation prompts on intermediate completions. Across four mathematical reasoning and code generation benchmarks on three dLLMs, including LLaDA 8B Instruct, Dream 7B Instruct, and LLaDA 2.0-mini, our Prism achieves a favorable performance-efficiency trade-off, matching best-of-N performance with substantially fewer function evaluations (NFE). The code is released at https://github.com/viiika/Prism.