Призма: эффективное масштабирование во время тестирования с помощью иерархического поиска и самопроверки для дискретных диффузионных языковых моделей

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

Аннотация

Вычислительные ресурсы на этапе логического вывода вновь стали практичным способом улучшения рассуждений больших языковых моделей (LLM). Большинство алгоритмов масштабирования на этапе тестирования (TTS) полагаются на авторегрессионное декодирование, которое плохо подходит для дискретных диффузионных языковых моделей (dLLM) из-за их параллельного декодирования всей последовательности. В результате разработка эффективных и производительных методов TTS для раскрытия полного генеративного потенциала dLLM остается слабо изученной проблемой. Для её решения мы предлагаем Prism (Pruning, Remasking, and Integrated Self-verification Method) — эффективный фреймворк TTS для dLLM, который (i) выполняет иерархический поиск по траекториям (HTS), динамически отсекая и перераспределяя вычисления в раннем-среднем окне денойзинга, (ii) вводит локальное ветвление с частичным перемаскированием для исследования разнообразных реализаций при сохранении токенов с высокой уверенностью и (iii) заменяет внешние верификаторы на самопроверяемую обратную связь (SVF), получаемую с помощью промптов для самооценки промежуточных завершений. На четырех тестах по математическим рассуждениям и генерации кода для трех dLLM, включая LLaDA 8B Instruct, Dream 7B Instruct и LLaDA 2.0-mini, наш Prism демонстрирует благоприятный баланс между производительностью и эффективностью, достигая качества лучшего из 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.