棱镜框架:基于分层搜索与自验证的离散扩散语言模型高效测试时扩展方法
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
摘要
推理时计算已重新成为提升大语言模型推理能力的实用方法。多数测试时缩放算法依赖自回归解码,这种机制与离散扩散语言模型的全序列并行解码特性不相容。因此,开发高效测试时缩放方法以释放离散扩散语言模型的全部生成潜力,仍是探索不足的挑战。为此,我们提出Prism(剪枝、重掩码与集成自验证方法),这是一个针对离散扩散语言模型的高效测试时缩放框架,其具备三大核心机制:(一)执行分层轨迹搜索,在去噪过程的前中期动态剪枝并重新分配计算资源;(二)引入局部分支与部分重掩码技术,在保留高置信度标记的同时探索多样化实现路径;(三)通过中间生成结果的自评估提示获取自验证反馈,替代外部验证器。在LLaDA 8B Instruct、Dream 7B Instruct和LLaDA 2.0-mini三种离散扩散语言模型上进行的数学推理与代码生成四项基准测试表明,Prism实现了优越的性能-效率平衡,能以显著更少的函数评估次数达到最优N选性能。代码已发布于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.