DiffBench Meets DiffAgent: End-to-End LLM-Driven Diffusion Acceleration Code Generation

AI summary

Overview

• Research area: Computer vision and generative modeling (diffusion models), combined with large language model agents for automated code generation and performance optimization.
• Technical level: Advanced. The paper assumes familiarity with diffusion sampling pipelines, model acceleration techniques, GPU software stacks, and LLM-based agent architectures.
• One-sentence scope: The paper introduces DiffBench, a 604-task benchmark for LLM-generated diffusion acceleration code, and DiffAgent, a multi-agent LLM framework that plans, writes, debugs, and genetically optimizes that code, reporting pass-rate gains over direct LLM code generation.

What This Paper Is About

Diffusion models must run many iterative denoising steps, which makes inference slow and expensive, and choosing how to combine acceleration techniques (fast samplers, feature caching, token merging, gated activation, mixed precision) normally requires deep expert knowledge for each model architecture and deployment target. The paper asks whether large language models can be made to autonomously produce correct, efficient diffusion acceleration code that satisfies a user's speedup and quality requirements. To answer this, the authors build DiffBench to measure how well LLMs do this today, and DiffAgent to improve their performance through planning, debugging, domain knowledge, and genetic-algorithm-guided tuning.

Key Contributions

1. Problem formalization: The authors formalize diffusion acceleration code generation as a task, defining inputs (user prompt plus reference code from open-source diffusion libraries) and outputs (a new diffusion inference implementation meeting specified requirements), with the model needing to both select effective acceleration methods and tune their parameters.
2. DiffBench: A benchmark of 604 tasks with prompts and ground-truth implementations, organized into five difficulty levels (41, 116, 261, 93, and 93 tasks respectively) and covering U-Net backbones (SD1.5, SD2.1, SDXL) and transformer variants (DiT, PixArt-α, PixArt-Σ), text-to-image, class-to-image and image-to-image conditioning, resolutions from 256×256 to 1024×1024, three samplers (DDIM, DPM-Solver, UniPC), and four acceleration techniques (ToMe, DeepCache, T-Gate, FP16).
3. DiffAgent: An agentic framework with a planning agent, coding agent, debugging agent (using the Reflexion architecture), and a genetic-algorithm-based selector that scores candidates by a weighted sum of relative quality loss and speedup and feeds results back into the next planning round.
4. Systematic evaluation: Experiments on four LLMs (GPT-4.1, Claude Sonnet 4, Gemini 2.5 Flash, o3-mini) on DiffBench, including per-level pass rates, a five-category error-mode analysis, and leave-one-out ablations of the agent's modules.

Main Findings

• Baseline LLMs struggle on DiffBench: Without DiffAgent, average pass rates S_p are 54.30% for Claude Sonnet 4, 14.24% for GPT-4.1, 12.09% for Gemini 2.5 Flash, and 11.92% for o3-mini. The paper states these models achieve pass rates below 35% on the acceleration challenges despite strong performance on standard coding tasks.
• DiffAgent improves every base model: Average pass rate rises from 11.92% to 56.46% for o3-mini (reported as +44.54) and from 54.30% to 81.59% for Claude Sonnet 4 (reported as +27.29). The abstract describes improvements of 54.30% to 81.59% across all benchmarks.
• Strongest results on compositional tasks: With DiffAgent, Claude Sonnet 4 reaches 90.24% on Level 1, 91.38% on Level 2, 99.23% on Level 3, 33.33% on Level 4, and 63.44% on Level 5.
• Hard, constrained levels remain the bottleneck: Level 4 (explicit speedup target) and Level 5 (latency constraint) are the weakest categories for every configuration tested.
• Common code errors are largely fixed: At Level 5, the Compile Error rate for Claude Sonnet 4 falls from 31.18% to 2.15% with DiffAgent; Key Attributes Error also drops, and Low Quality Error is reported as eliminated (0%) in several cases.
• All three modules matter (ablation on Claude Sonnet 4): Removing any module lowers average S_p from 81.59% into the 64.90%–67.88% range and lowers hard-task achievement rate S_a from 68.27% to between 8.16% and 62.02%.
• The genetic algorithm is decisive for hard tasks: Without it, Level 4 and Level 5 pass rates fall to 4.30% and S_a to 8.16%.
• The knowledge base gives the broadest benefit: Its removal causes the largest absolute decline in overall pass rate, a 16.69% reduction.
• The debugging agent mainly improves robustness: Without it, the Level 5 pass rate declines by approximately 30% and S_a by 6.25%.
• Performance plateaus at P = 7 and T_sel = 4: On Levels 4 and 5, going from P=4 to P=7 to P=10 and T_sel=2 to 4 to 6 yields diminishing gains beyond P=7 and T_sel=4, so the authors set those values.

Methodology in Plain English

The authors built a simulated developer environment. Each DiffBench task gives an LLM a natural-language request (which model, how many sampling steps, which acceleration methods, what speedup on what hardware) plus reference code, and asks for a working accelerated implementation. Generated code is checked in three stages: first a static comparison of key parameters against ground truth (pipeline class, model ID, scheduler class, inference steps, acceleration methods, resolution, preprocessors); second, an absolute quality check using CLIP-Score on 10 few-shot examples drawn from COCO, with samples below a threshold δ failing; third, a relative comparison against a reconstructed baseline, using quality loss L (smaller is better) and speedup U (larger is better), with raw latency τ used for Level 5. Only code passing all three stages counts as a pass. To create Level 4 and Level 5 targets, the authors ran a 50-iteration search on a 36-sample validation set under a CLIP-Score degradation bound σ, labeled qualifying configurations as "medium," then scaled the maximum speedup by Δ1 and Δ2 to make "easy" and "hard" variants (hard variants may have no valid solution), with σ=5%, Δ1=0.8, and Δ2=1.2 in the experiments.

DiffAgent then wraps the LLM. A planning agent writes a baseline plan plus acceleration plans that specify which optimization methods to use, in what order, with what parameters. A coding agent turns each plan into code, guided by reference templates from established diffusion libraries. A debugging agent finds errors and proposes fixes (up to T_debug=3 iterations before restarting generation, up to T_code=5 restart cycles before backtracking to replanning; worst case, the LLM is invoked T_code × T_debug times per prompt). Finally, a genetic-algorithm selector scores each runnable implementation by a weighted combination of quality loss and speedup, keeps the best P candidates, terminates early if requirements are met, and otherwise samples M promising "offspring" and sends their plans plus a feedback report back to the planner for the next generation, up to T_sel iterations. Defaults are P=7 and M=4, with P and M set to zero when no optimization is needed; the paper notes a default T_sel of 5 in the framework description and uses T_sel=4 in the hyperparameter experiments.

Why This Matters

• Research impact: DiffBench provides the first benchmark, per the paper, for LLM-generated diffusion pipelines and acceleration code, giving the community a standardized way to measure agent performance on a task that requires both domain expertise and environment feedback. It also provides a template for evaluating LLM agents on optimization problems with hard numerical constraints, not just functional correctness.
• Real-world applications:
  • Deploying text-to-image, class-to-image, and image-to-image pipelines at lower inference cost by automatically combining fast samplers, feature reuse, token merging, gated activation, and FP16.
  • Meeting latency budgets on specific hardware platforms, which is the explicit objective of Level 5 tasks.
  • Reducing the manual engineering required to port and accelerate diffusion models across SD1.5, SD2.1, SDXL, DiT, PixArt-α, and PixArt-Σ.
  • Automatically tuning acceleration parameters to hit a target speedup while keeping CLIP-Score degradation within a bound.
• Industry relevance: The authors are affiliated with AMD, Peking University, and Tsinghua University, and the framework targets hardware-specific deployment constraints, making it directly relevant to accelerator vendors, cloud inference providers, and teams serving generative models under cost and latency pressure.

Future Directions

• Closing the gap on constrained tasks: Even the best configuration (Claude Sonnet 4 with DiffAgent) reaches only 33.33% pass rate and 56.61% achievement rate at Level 4, leaving substantial room for improvement on explicit speedup targets.
• Broadening architectures and modalities: The benchmark covers image models; the abstract motivates the work with image and video generation, but video diffusion is not part of the 604 tasks, and the paper lists no results for it.
• Expanding acceleration and hardware coverage: The benchmark includes four acceleration techniques (ToMe, DeepCache, T-Gate, FP16) and three samplers; extending to more methods and more deployment platforms is a natural next step.
• Controlling LLM invocation cost: The framework can call the LLM up to T_code × T_debug times per prompt per plan in the worst case, and the paper reports no explicit token, latency, or cost accounting for the agent workflow itself.

Target Audience

• Researchers working on diffusion model acceleration, efficient generative inference, and sampling schedulers.
• Researchers and engineers building LLM agents for code generation, program optimization, and evolutionary or self-reflective search.
• ML systems and compiler engineers who need to automate performance tuning across model architectures and hardware targets.
• Technical leads and deployment engineers at hardware vendors or inference-serving platforms who must deliver diffusion models under strict latency and cost constraints.

Authors’ abstract

Diffusion models have achieved remarkable success in image and video generation. However, their inherently multiple step inference process imposes substantial computational overhead, hindering real-world deployment. Accelerating diffusion models is therefore essential, yet determining how to combine multiple model acceleration techniques remains a significant challenge. To address this issue, we introduce a framework driven by large language models (LLMs) for automated acceleration code generation and evaluation. First, we present DiffBench, a comprehensive benchmark that implements a three stage automated evaluation pipeline across diverse diffusion architectures, optimization combinations and deployment scenarios. Second, we propose DiffAgent, an agent that generates optimal acceleration strategies and codes for arbitrary diffusion models. DiffAgent employs a closed-loop workflow in which a planning component and a debugging component iteratively refine the output of a code generation component, while a genetic algorithm extracts performance feedback from the execution environment to guide subsequent code refinements. We provide a detailed explanation of the DiffBench construction and the design principles underlying DiffAgent. Extensive experiments show that DiffBench offers a thorough evaluation of generated codes and that DiffAgent significantly outperforms existing LLMs in producing effective diffusion acceleration strategies.