Prisma: Scalabilità Efficiente al Momento del Test tramite Ricerca Gerarchica e Auto-Verifica per Modelli Linguistici a Diffusione Discreta
Prism: Efficient Test-Time Scaling via Hierarchical Search and Self-Verification for Discrete Diffusion Language Models
February 2, 2026
Autori: 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
Abstract
Il calcolo al momento dell'inferenza è riemerso come un metodo pratico per migliorare il ragionamento dei LLM. La maggior parte degli algoritmi di scaling al momento del test (TTS) si basa su decodifica autoregressiva, che è poco adatta ai modelli linguistici a diffusione discreta (dLLM) a causa della loro decodifica parallela sull'intera sequenza. Di conseguenza, lo sviluppo di metodi TTS efficaci ed efficienti per sbloccare il pieno potenziale generativo dei dLLM rimane una sfida poco esplorata. Per affrontare ciò, proponiamo Prism (Pruning, Remasking, and Integrated Self-verification Method), un framework TTS efficiente per dLLM che (i) esegue una Ricerca Gerarchica delle Traiettorie (HTS) che pota dinamicamente e riassegna il calcolo in una finestra di denoising da iniziale a intermedia, (ii) introduce il branching locale con remasking parziale per esplorare implementazioni diverse preservando i token ad alta confidenza, e (iii) sostituisce i verificatori esterni con un Feedback Auto-Verificato (SVF) ottenuto tramite prompt di auto-valutazione su completamenti intermedi. Su quattro benchmark di ragionamento matematico e generazione di codice su tre dLLM, inclusi LLaDA 8B Instruct, Dream 7B Instruct e LLaDA 2.0-mini, il nostro Prism raggiunge un favorevole compromesso prestazioni-efficienza, eguagliando le prestazioni del best-of-N con un numero sostanzialmente inferiore di valutazioni di funzione (NFE). Il codice è rilasciato su 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.