arXiv:2605.04215cs.LGcs.AI2026-05中稿 · publication in IJC…

根据输入预测生成长度,避免扩散模型浪费计算资源。

Predict-then-Diffuse: Adaptive Response Length for Compute-Budgeted Inference in Diffusion LLMs

论文配图:Predict-then-Diffuse: Adaptive Response Length for Compute-Budgeted Inference in Diffusion LLMs
图 1 · 摘自论文原文
  • 先预测响应长度,再按需运行扩散模型推理。
  • 实验显示计算量(FLOP)显著降低,且对数据分布偏斜鲁棒。
  • 适合需要高效推理的生产环境,尤其关注算力成本场景。

基于扩散的大型语言模型(D-LLMs)代表了生成式AI的前沿方向,其完全并行的词元生成方式可带来显著的吞吐量优势和更优的GPU利用率,但受限于生成前需固定输出长度。这一架构限制导致严重权衡:响应过长会浪费算力在语义无意义的填充词元上,过短则引发截断,需代价高昂的重新计算,引入不可预测的延迟峰值。为此,我们提出Predict-then-Diffuse框架,一种简单且模型无关的方法,通过先估计响应长度,再以该长度进行D-LLM推理,实现每条输入查询的算力预算化推理。核心是自适应响应长度预测器(AdaRLP),可根据输入预测最优响应长度。为防止低估导致重推,引入基于小幅度增加预测长度的数据驱动安全机制。整体框架既避免了填充词元的计算浪费,又保持输出质量。多数据集实验表明,相比默认的D-LLM推理机制,Predict-then-Diffuse显著降低计算成本(FLOP),且对数据分布偏斜具有鲁棒性。

原文摘要 · Abstract (English)

Diffusion-based Large Language Models (D-LLMs) represent a promising frontier in generative AI, offering fully parallel token generation that can lead to significant throughput advantages and superior GPU utilization over the traditional autoregressive paradigm. However, this parallelism is constrained by the requirement of a fixed-size response length prior to generation. This architectural limitation imposes a severe trade-off: oversized response length results in computational waste on semantically meaningless padding tokens, while undersized response length causes output truncation requiring costly re-computations that introduce unpredictable latency spikes. To tackle this issue, we propose Predict-then-Diffuse, a simple and model-agnostic framework that enables compute-budgeted inference per input query by first estimating the response length and then using it to run inference with D-LLM. At its core lies an Adaptive Response Length Predictor (AdaRLP), which estimates the optimal response length given an input query. As a measure against under-estimating the response length and re-running inference with a higher value, we introduce a data-driven safety mechanism based on a small increase of the predicted length. As a whole, our framework avoids wasting computation on padding tokens, at the same time preserving output quality. Experimental validation on multiple datasets demonstrates that Predict-then-Diffuse significantly reduces computational costs (FLOP) compared to the default D-LLM inference mechanism, while being robust to skewed data distributions.

扩散模型推理优化算力预算

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。