Predict-Then-Diffuse: Enabling Adaptive Inference of Computational Budget for Diffusion Language Models

The research team at the University of Bergamo in Italy proposed the Predict-Then-Diffuse framework, addressing the core issue of diffusion language models (Diffusion LLMs) needing to pre-determine response lengths. By predicting response lengths to optimize inference efficiency, it significantly reduces computational costs while maintaining output quality. The framework adopts the "predict first, diffuse later" approach to solve the resource waste or output truncation problems caused by fixed-length strategies.

After the success of diffusion models in the image domain, they were applied to NLP. However, Diffusion LLMs need to determine a fixed response length before generation, unlike autoregressive models (e.g., GPT) which generate token by token and can stop naturally. This constraint leads to a trade-off dilemma: setting a length that's too long wastes computation on meaningless padding tokens; setting it too short results in output truncation requiring retries, causing latency spikes and resource waste. In real-world scenarios, query lengths are diverse, making the "one-size-fits-all" strategy difficult to adapt.

The framework consists of three steps: 1. Response Length Prediction: Use a model-agnostic Adaptive Response Length Predictor (AdaRLP) to estimate the optimal length; 2. Safety Margin Mechanism: Add a data-driven safety margin to the predicted value to balance efficiency and completeness; 3. Diffusion Generation: Perform diffusion generation with the adjusted length to avoid padding waste and truncation risks.

The project provides two core Jupyter Notebooks:

• Analytical Simulation Notebook (ptd_analytical_simulation.ipynb): Train the AdaRLP predictor, evaluate performance, simulate and verify theoretical boundaries, and output prediction data;
• Empirical Profiling Comparison Notebook (ptd_empirical_profiling_comparison.ipynb): Measure FLOPs, GPU time, and memory usage, comparing three strategies: baseline (original prediction), fallback (with safety margin), and fixed length. Project dependencies are managed via pyproject.toml and uv, supporting Python 3.13+ and NVIDIA GPUs.

Verification across multiple datasets shows:

• Significant reduction in computational cost: Reduced FLOPs consumption compared to the default mechanism, improving hardware utilization or lowering costs;
• Stable output quality: Accurate prediction and safety margin ensure content is accurate and complete;
• Strong robustness: Adapts to the long-tail distribution of real-world queries (most are short, a few are long).

This technology is of great significance for the deployment of diffusion language models:

• Cloud service optimization: Helps vendors optimize resource allocation, reduce operational costs, and provide predictable response times;
• Edge devices: Enables efficient model operation in resource-constrained environments;
• Real-time applications: Avoids latency fluctuations from truncation retries (e.g., dialogue systems);
• Green AI: Reduces computational energy consumption, aligning with sustainable development trends.

Current limitations: Length prediction requires historical data, and the accuracy of predicting completely new queries needs improvement; the safety margin depends on the distribution of training data, and recalibration is needed when scenarios change. Future directions: Online learning to allow continuous improvement of the predictor; multi-task adaptation for different tasks (code generation, Q&A, etc.); dynamic length adjustment during generation; combining technologies like speculative decoding to further improve efficiency.

The Predict-Then-Diffuse framework solves the fixed-length constraint problem through the "predict-execute" paradigm, which is a key progress in the practicalization of diffusion language models. It provides reference implementations and experimental data for researchers and engineers focusing on LLM inference efficiency, cost control, or edge deployment. As technology matures, such computational budget optimization techniques will become standard configurations for deployment.