Accelerating Image Generation with
Sub-Path Linear Approximation Model

Based on the above analysis, in this paper, we introduce Sub-Path Linear (SL) ODEs to model approximated sub-paths in the original PF-ODE trajectories, which gives a progressive estimation for $dist_{\Delta}$. For a sampled sub-path $\{\bm{x}_{t^{\prime}}\}_{t\in[t-k,t]}$ on a solution trajectory dictated by Eq. 3, we interpolate a linear path from $(\bm{x}_{t-k},t-k)$ to $(\bm{x}_{t},t)$, guided by the vector direction of $dist_{\Delta}(\bm{x}_{t},\bm{x}_{t-k},t)$. To distinguish the impacts of $dist_{\Delta}$ and $dist_{0,{\bm{\theta}}}$, we account for the drift component in the linear approximated path, causing a shift on coefficient from $(\bm{x}_{t-k},t-k)$ to $(\frac{\alpha(t)}{\alpha(t-k)}\bm{x}_{t-k},t-k)$. The points on the approximated path $\{\bm{x}_{\gamma,t}\}_{\gamma\in[0,1]}$ are thus computed as:

for a sampled $(\bm{x}_{t-k},t-k)$ and $(\bm{x}_{t},t)$.

Since $\bm{x}_{t}$ and $\bm{x}_{t-k}$ conform to distributions governed by the PF-ODE, our linear transformation effectively defines a linear ODE from distribution $\frac{\alpha(t)}{\alpha(t-k)}\bm{x}_{t-k}\sim p_{t-k,k}(\bm{x}_{t-k})$ to $\bm{x}_{t}\sim p_{t}(\bm{x}_{t})$ over $\gamma$, where $p_{t,k}(\bm{x}_{t})$ has the property $p_{t,k}(\bm{x}_{t}|\bm{x}_{0})=\mathcal{N}(\alpha(t+k)\bm{x}_{0},\left[\frac{\alpha(t+k)\sigma(t)}{\alpha(t)}\right]^{2}I)$:

We denote it as Sub-Path Linear (SL) ODE. To apply the approximation strategy on the SL-ODE, the Denoiser and generation function replacing $\bm{x}_{t}$ with $\bm{x}_{\gamma,t}$ are given by:

Incorporating these into Eq. 11, we derive the sub-path learning objective for our SPLAM model as :

which gives a progressive estimation for the otherwise intractable $dist_{\Delta}$ objective. The value for $\sigma(\gamma,t)$ can be precisely estimated from the distribution $p_{t}(\bm{x}_{t})$ and $p_{t-k}(\bm{x}_{t-k})$ but has a complex expression. Empirically we utilize an approximate result as $\sigma(\gamma,t)=(1-\gamma)\frac{\alpha(t)}{\alpha(t-k)}\sigma(t-k)+\gamma*\sigma(t)$. Compared to consistency models which adopt Eq. 10 or Eq. 11, our $\mathcal{L}_{\text{}}$ maintains a progressive estimation for $dist_{\Delta}$ and a consistent estimation for $dist_{0,{\bm{\theta}}}$, which enables the learning for large skipping step size. The overall prediction error can still be assessed by the aggregate of approximated errors between the sub-path endpoints and the approximated error between these points is continuously optimized through the SL-ODEs. Consequently, the optimization for the approximated errors in our SPLAM could be significantly improved. Our approach could further benefit from the increased skipping step size, allowing our method to generate images of higher quality with reduced sampling steps in more efficient training.

In this paper, we adopt pre-trained Stable Diffusion (SD) models [34] to obtain the solution PF-ODE trajectories which we build our SL-ODEs upon, and we call the approach Sub-path Linear Approximation Distillation (SPLAD). To achieve conditional generation with the conditional input $c$, the noise prediction model is parameterized as $\bm{\epsilon}_{\bm{\theta}}(\bm{z}_{t},c,t)$ [46, 23]. We also introduce $\gamma$ into the prediction models for solving our SL-ODEs, and leverage the $\gamma-$conditioned training where $\gamma$ is converted to Fourier embeddings and fed into the models as an input. Specifically, to predict $\bm{z}_{0}$ in the latent space, the generation function for SPLAM is defined as:

where $\bm{f}_{\bm{\theta}}(\bm{z}_{\gamma,t},c,\gamma,t)$ mirrors Eq. 14 while replacing $\bm{\epsilon}_{\bm{\theta}}(\bm{z}_{\gamma,t},\gamma,t)$ with the conditional form $\bm{\epsilon}_{\bm{\theta}}(\bm{z}_{\gamma,t},c,\gamma,t)$. The functions $c_{\text{skip}}$ and $c_{\text{out}}$ are leveraged to ensure that $\bm{F}_{\bm{\theta}}(\bm{z}_{1,0},c,1,0)\equiv\bm{z}_{0}$ (we regard $F_{\bm{\theta}}$ as the same expression of $f_{\bm{\theta}}$ since $c_{\text{skip}}(t)\ll c_{\text{out}}(t)$ for most time steps). Integrating this with Eq. 9, our SPLAD approach minimizes the following objective:

where $\mathcal{U}$ denotes the uniform distribution, and $k$ is a pre-determined skipping step size. The $\alpha(t)$ in Eq. 9 is omitted due to its negligible effect on optimization in practice. The term $\hat{\bm{z}}^{\Phi}_{1,t-k}=\hat{\bm{z}}^{\Phi}_{t-k}$ is estimated using ODE solvers $\Phi(\cdot\cdot\cdot;\phi)$ derived from teacher models. In this paper DDIM [42] is employed as our choice from the advanced solvers of LDMs. Moreover, to improve the estimation of $\hat{\bm{z}}^{\Phi}_{t-k}$, we apply the Multiple Estimation which executes the solver $\Phi(\cdot\cdot\cdot,\phi)$ multiple times with a reduced skipping step size $k_{\phi}$. Denoting $t_{\phi,i}=t-i*k_{\phi}$ and initializing $\hat{\bm{z}}^{\Phi}_{t_{\phi,0}}=\bm{z}_{t}$, the multiple estimation is iteratively executed as:

for $i=0,1,2,...;i\leq\frac{k}{k_{\phi}}-1$, where $\emptyset$ denotes no conditional inputs and $w$ is a fixed guidance scale which controls the effect of conditional generation [10] from the conditional input $c$.

The pseudo-code for SPLAD is presented in Algorithm 1. SPLAD shares a similar training pipeline with consistency models[44, 26] but can be distinguished as it optimizes the sub-path learning objectives based on the SL-ODEs and utilizes the $\gamma$-conditioned training. For a pair of input noise and time step $(\bm{z}_{t},t)$, SPLAM gives the prediction of the denoised latent $\hat{\bm{z}}_{0}$ as:

for one-step generation, adhering strictly to the $\gamma=1$ condition. We also use the same iterative sample strategy as illustrated in [44] which could improve the quality of the generated images. In practice, we set the $\gamma\text{-}$embedding to $\mathbf{0}$ for $\gamma=1$, thereby allowing the weights associated with trained $\gamma$-embeddings to be discarded post-training. Thus our Sub-Path Linear Approximation Model (SPLAM) necessitates no additional parameters beyond the training phase and can be utilized the same as teacher models.

On text-to-image generation task, we train two models with pre-trained Stable Diffusion-V1.5 (SDv1.5) and Stable Diffusion-V2.1-base (SDv2.1-base) as teacher models respectively. Following the setting of [26], the training dataset is one subset of LAION-5B[39]: LAION-Aesthetics-6+. We choose DDIM-Solver as the ODE solver $\phi$ at skipping step $k_{\phi}=20$.

For evaluation, we adopt the commonly used FID and CLIP Score metrics. The results are reported on both SDv1.5 and SDv2.1-base backbones, thus verifying the generalizability of our method. For the experiment of distilling SDv2.1-base, we bench-mark our model on two test sets, LAION-Aesthetics-6+ as used in LCM [26] and MSCOCO2014-30k for zero-shot generalization. We also reproduce a SDv2.1-base LCM according to the training configuration outlined in [26] while replacing the $w$-condition with the fixed guidance scale, which has also improved the performance. We generally set the guidance scale for distilling SDv2.1-base to 8 and skipping step size to 20, which is consistent with  [26]. For the experiment of distilling SDv1.5, we compare our model with state-of-the-art generative models including foundation diffusion models, GANs, and accelerated diffusion models. The guidance scale is set to $3$ to obtain the optimal FID, and we adopt the huber [43] loss for our SPLAD metric. The skipping step size is set to 100 for SPLAM which has shown fast convergence. We examine our method on two commonly used benchmarks, MSCOCO2014-30k and MSCOCO2017-5k. More implementation details are provided in the supplementary materials.

The results for SDv2.1-base are presented in Tab. 1, we use DDIM [42], DPM [23], DPM++ [24] and LCM [26] as baselines. It reveals that our SPLAM surpasses baseline methods nearly across both test sets, at each step, and on both FID and CLIP Score metrics. We suppose that the close results on LAION are caused by overfitting, since the test set and train set are sourced from the same data collection. For SDv1.5 under the guidance scale $w=3$, the quantitative results are demonstrated in Tab. 2(a) and Tab. 2(c). Our model with 4 steps gets FID-30k of 10.06 and FID-5k of 20.77, which outperforms all other accelerated diffusion models, including flow-based method InstaFlow [21] and techniques that introduce GAN objectives such as UFOGen [49] and DMD [50]. Furthermore, SPLAM showcases commensurate results with state-of-the-art foundation generative models such as DALL-E2 [32]. Even in two steps, SPLAM has achieved a competitive performance of FID-30k 12.31 with parallel algorithms. In practical scenarios, a higher guidance scale $w$ is typically favored to enhance the resultant image quality. Accordingly, we trained our SPLAM with $w$ set to 8 and bench-mark it against a range of advanced diffusion methodologies, as delineated in Tab. 2(c). In this regime, SPLAM also demonstrates significant advantages, achieving state-of-the-art performance with a four-step FID-30k of 13.39 which exceeds other models by a large margin and is close to the teacher model. Notably, the FID-30k of our model with only two steps reaches 14.50, surpassing the four-step LCM and DMD. While DMD training consumes over one hundred A100 GPU days, which is more than 16 times our training duration.