Parallel Diffusion Solver via Residual Dirichlet Policy Optimization

TL;DR

EPD-Solver introduces an Ensemble Parallel Direction solver that reduces truncation error in diffusion sampling by evaluating multiple parallel gradients within each integration step and combining them with simplex weights. A two-stage optimization—distillation-based initialization followed by Residual Dirichlet Policy Optimization (RDPO) within a low-dimensional solver space—yields a parameter-efficient, RL-guided solver that aligns with human perceptual rewards while keeping the backbone fixed. The method includes a flexible EPD-Plugin to enhance existing samplers and achieves state-of-the-art quality under tight latency budgets across CIFAR-10, FFHQ, ImageNet, LSUN Bedroom, and large-scale T2I models like Stable Diffusion v1.5 and SD3-Medium, e.g., FID improvements at 5 NFE and HPSv2.1 = 0.2482 with ImageReward = 3.1121 at 20 NFEs. This work reinforces that leveraging low-dimensional geometric structure and solver-space RL can reconcile efficiency and fidelity in diffusion-based generation, with practical plug-in applicability for real-world deployment.

Abstract

Diffusion models (DMs) have achieved state-of-the-art generative performance but suffer from high sampling latency due to their sequential denoising nature. Existing solver-based acceleration methods often face significant image quality degradation under a low-latency budget, primarily due to accumulated truncation errors arising from the inability to capture high-curvature trajectory segments. In this paper, we propose the Ensemble Parallel Direction solver (dubbed as EPD-Solver), a novel ODE solver that mitigates these errors by incorporating multiple parallel gradient evaluations in each step. Motivated by the geometric insight that sampling trajectories are largely confined to a low-dimensional manifold, EPD-Solver leverages the Mean Value Theorem for vector-valued functions to approximate the integral solution more accurately. Importantly, since the additional gradient computations are independent, they can be fully parallelized, preserving low-latency sampling nature. We introduce a two-stage optimization framework. Initially, EPD-Solver optimizes a small set of learnable parameters via a distillation-based approach. We further propose a parameter-efficient Reinforcement Learning (RL) fine-tuning scheme that reformulates the solver as a stochastic Dirichlet policy. Unlike traditional methods that fine-tune the massive backbone, our RL approach operates strictly within the low-dimensional solver space, effectively mitigating reward hacking while enhancing performance in complex text-to-image (T2I) generation tasks. In addition, our method is flexible and can serve as a plugin (EPD-Plugin) to improve existing ODE samplers.

Figures (16)

• Figure 1: Comparison of various solvers on diffusion models. We compare the FID versus latency (ms) across different NFE settings on a NVIDIA 4090. Our proposed $\texttt{EPD-Solver}$ shows superior image quality without increasing latency.
• Figure 2: Computation graphs of various ODE solvers. (a) DDIM solver songdenoising (Euler's method) adopts the rectangle rule that uses the gradient at the start point: ${\mathbf{d}}_{t_{n+1}}=\bm{\epsilon}_\theta({\mathbf{x}}_{t_{n+1}},t_{n+1})$. disclose EDM solver karras2022elucidating (Heun's method) uses the trapezoidal rule that averages the gradients of both the start and the end timesteps, i.e., ${\mathbf{d}}_{t_{n+1}}=\bm{\epsilon}_\theta({\mathbf{x}}_{t_{n+1}},t_{n+1})$ and ${\mathbf{d}}_{t_{n}}'=\bm{\epsilon}_\theta({\mathbf{x}}_{t_{n}}',t_{n})$, where ${\mathbf{x}}_{t_n}'$ is the additional evaluation given by Euler's method. (c) AMED solver zhou2024fast optimizes a small network $g_\phi(\cdot)$ to output an intermediate timestep $s_n \in (t_n,t_{n+1})$ to compute the gradient: ${\mathbf{d}}_{s_n}=\bm{\epsilon}_\theta({\mathbf{x}}_{s_{n}},s_{n})$. Since AMED introduces a network in sequential computation, its latency is slightly higher than that of other solvers, as shown in \ref{['fig:fidvslatency']}. (d) Our $\texttt{EPD-Solver}$ leverage $K$ parallel gradients to achieve more accurate integral approximation. We optimize $K$ intermediate timesteps $\tau_n^1, \dots, \tau_n^K$, compute their gradients ${\mathbf{d}}_{\tau_{n}^1}, \dots, {\mathbf{d}}_{\tau_{n}^K}$, and combine them via a simplex-weighted sum.
• Figure 3: Cumulative explained variance ratio of sampling trajectories using DMs from EDM2 Karras2024edm2. We analyze the trajectory's orthogonal complement, i.e., the residuals after removing the linear component connecting $\mathbf{x}_{t_T}$ and $\mathbf{x}_{t_0}$. The rapid saturation at the two principle components (capturing $>97\%$ of the residual variance) indicates that the trajectory occurs almost within a single 2D plane.
• Figure 4: Stage 1: Distillation-Based Parameter Optimization (Top). We optimize the learnable solver parameters $\Theta_n$ by minimizing the trajectory reconstruction error against a high-precision teacher solver (e.g., DPM-Solver-2), providing a robust initialization for stage 2. Stage 2: Residual Dirichlet Policy Optimization (Bottom). To align generation with human preferences, we reformulate the solver as a stochastic policy parameterized by Dirichlet distributions (defined by $\bm{\alpha}^\mathsf{pos}_n$ and $\bm{\alpha}^\mathsf{mix}_n$). By sampling multiple parallel trajectories in the low-dimensional solver space and evaluating them with a reward model (e.g., HPSv2.1), we optimize the policy using PPO with a Reward-Leave-One-Out (RLOO) baseline.
• Figure 5: Qualitative comparison of T2I generation results using Stable Diffusion v1.5. We compare our EPD-Solver (20 NFE) against SoTA baselines including DDIM, DPM-Solver-2, EDM, and iPNDM (50 NFE). Our method achieves comparable or superior visual fidelity with significantly reduced inference steps. Qualitative results for SD3 are in \ref{['fig:sd3_512', 'fig:sd3_1024']}.

Theorems & Definitions (1)

• Theorem 1