PRISM框架：离散扩散语言模型推理时成本优化新路径

成本概述

本文聚焦于离散扩散语言模型（dLLMs）在推理时（Test-Time Scaling）的成本问题。随着大模型能力提升的焦点从“训练时扩展”转向“推理时扩展”，dLLMs作为一种非自回归生成模型，因其并行、全局双向上下文的特性，在复杂推理任务中展现出独特优势。然而，传统面向自回归模型设计的推理方法，如Best-of-N，在dLLMs上应用时面临成本高昂、效率低下的问题。本文旨在分析dLLMs推理时的成本构成，介绍PRISM框架如何通过优化推理搜索策略降低成本，同时保持或提升模型性能。

典型场景

dLLMs的推理时扩展常见于需要复杂推理的场景，如数学问题求解、逻辑推理、代码生成等。在这些场景中，模型需要生成多个候选答案，并通过搜索与验证框架选择最优解。然而，传统方法往往默认模型为自回归生成，导致在dLLMs上应用时效率低下，成本高昂。

成本构成

dLLMs推理时的成本主要包括计算成本和存储成本。计算成本主要来源于模型推理过程中的函数调用次数，即采样轨迹数量与每条轨迹去噪步数的乘积。存储成本则涉及模型推理过程中产生的中间状态和最终答案的存储。在传统Best-of-N方法中，所有候选答案都要完整跑完，即便中途质量不佳的轨迹也会消耗完整预算，导致计算成本高昂。

影响因素

dLLMs推理时的成本受多种因素影响，包括：

• 采样轨迹数量（N）：采样轨迹越多，计算成本越高。
• 每条轨迹去噪步数（T）：去噪步数越多，计算成本越高。
• 模型复杂度：模型越复杂，每步去噪的计算量越大。
• 数据质量：数据质量影响模型输出的稳定性，进而影响推理过程中的裁剪策略。
• 推理预算：推理预算限制了可采样的轨迹数量和去噪步数，影响最终答案的质量。

成本评估方法

评估dLLMs推理时的成本，需明确业务目标、拆解资源模型、建立用量口径，并区分固定成本与弹性成本。具体步骤如下：

1. 明确业务目标：确定推理任务类型、答案质量要求、推理时间限制等。
2. 拆解资源模型：将推理过程拆分为采样、去噪、裁剪、验证等阶段，明确各阶段的资源需求。
3. 建立用量口径：定义采样轨迹数量、每条轨迹去噪步数、裁剪比例等关键指标。
4. 区分固定成本与弹性成本：固定成本包括模型加载、初始化等一次性开销；弹性成本随采样轨迹数量和去噪步数变化。
5. 评估峰值与平均值：考虑推理任务中的峰值场景，如突发高并发推理请求。
6. 设计预算阈值：为关键资源设置预算线、预警线和异常增长监控，避免成本超支。

成本优化路径

针对dLLMs推理时成本高昂的问题，PRISM框架通过优化推理搜索策略降低成本。具体优化路径如下：

1. 资源规格优化：根据实际负载调整模型复杂度，避免过度配置。例如，在简单推理任务中，可使用较小规模的dLLMs。
2. 弹性伸缩：根据推理任务的峰谷动态调整采样轨迹数量和去噪步数。在高峰时段增加采样轨迹数量，提升答案质量；在低谷时段减少采样轨迹数量，降低成本。
3. 推理搜索策略优化：PRISM框架将推理过程拆分为早期随机探索、中期渐进裁剪和后期精修三个阶段。在高噪声阶段保持较宽的候选集合以保留多样性；在早中期去噪窗口使用自验证信号裁剪低质量轨迹，并把计算资源重新分配给更有前景的候选；最终只保留较小数量的轨迹继续完成精修。这种设计使得PRISM的实际复杂度接近O(N + KT)，其中K是最终保留的较小候选宽度，相比传统Best-of-N的O(NT)显著降低成本。
4. 存储生命周期管理：对推理过程中产生的中间状态和最终答案进行分层存储，控制长期存储和备份成本。例如，将中间状态存储在内存中，将最终答案存储在对象存储中，并设置合理的保留周期。
5. 日志与监控优化：控制日志采集范围、保留周期和索引粒度，避免日志成本失控。同时，建立推理任务监控体系，实时跟踪推理进度、成本消耗和答案质量，及时发现并处理异常情况。

成本与性能平衡

在优化dLLMs推理时成本的过程中，需兼顾性能、稳定性和可用性。例如，在减少采样轨迹数量和去噪步数时，需确保最终答案的质量不受影响；在裁剪低质量轨迹时，需避免误裁剪有潜力的轨迹。此外，还需考虑推理任务的实时性要求，避免因过度优化成本而导致推理时间过长。

常见成本浪费

在dLLMs推理过程中，常见的成本浪费包括：

• 闲置资源：未充分利用的采样轨迹和去噪步数导致计算资源浪费。
• 过度配置：使用过大规模的dLLMs或过高的去噪步数导致成本高昂。
• 无效日志：采集过多无关日志导致存储成本上升。
• 重复存储：对中间状态和最终答案进行重复存储导致存储成本浪费。

风险与注意事项

在优化dLLMs推理时成本的过程中，需注意以下风险：

• 稳定性风险：过度裁剪低质量轨迹可能导致答案质量下降或遗漏最优解。
• 安全性风险：降低推理预算可能影响模型的安全性验证和审计能力。
• 容量不足风险：在高峰时段减少采样轨迹数量可能导致推理任务积压或超时。
• 恢复能力下降风险：减少备份和存储可能导致数据丢失或恢复困难。

为避免上述风险，需在优化成本的同时，建立完善的监控和预警体系，及时发现并处理异常情况。同时，需定期评估优化效果，根据业务需求和模型性能调整优化策略。

总结

本文分析了dLLMs推理时的成本构成、影响因素和评估方法，介绍了PRISM框架如何通过优化推理搜索策略降低成本。在优化成本的过程中，需兼顾性能、稳定性和可用性，避免常见成本浪费和风险。通过合理规划资源、优化推理搜索策略、控制存储和日志成本等措施，可实现dLLMs推理时的高效、低成本运行。