FastGen-PDD: Parallel Decoding Distillation
for Image and Video Generation

Parallel Decoding Distillation (PDD): Intuition

Sampling diffusion and flow matching models using classical ordinary differential equation (ODE) solvers requires many small denoising steps, resulting in hundreds of network evaluations. The common strategy of trajectory-based distillation methods for training a few-step generator is to train a student model to predict a large-step update, essentially merging multiple small teacher updates. In Parallel Decoding Distillation (PDD), we take a different approach: instead of merging the teacher updates into one large update, we predict multiple teacher updates in a single network evaluation. Specifically, we discretize time \(t\in[0,1]\) into a fixed sequence of \(N\in\mathbb{N}\) intervals \([t_{n},t_{n+1}]\) for \(n=0,1,\ldots,N-1\) and group them into blocks of size \(L\leq N\). During training, the student model learns to predict the teacher's ODE-solver step for all intervals within a block using a single forward pass. Then, at generation, each student evaluation advances \(L\) intervals, as illustrated in Figure 2.

(a) Teacher

(b) Student training

(c) Student Generation

Figure 3(b) illustrates our student architecture, alongside the teacher in Figure 3(a). We start from the teacher architecture and simply repeat the final linear layer \(N\) times, where \(N\) is the grid size. That is, we use a distinct final linear layer for each interval \([t_n,t_{n+1}]\) in the discretization. Predicting \(N\) grid-level outputs, rather than only \(L\) block-level outputs, enables sampling with varying NFEs at inference time. We achieve this by training with variable block sizes and then adjusting the block size at generation according to the desired NFE. Moreover, since a weighted sum of linear transformations is itself a linear transformation, we can fuse the relevant linear layers at generation. As a result, for each block, we only need to evaluate and store a single fused linear layer; see Figure 3(c). Lastly, this architecture allows the student weights to be initialized seamlessly from the teacher weights.

(a) Student and teacher during training

(b) Parallel decoding loss sketch

To train our parallel decoder (student), we employ an on-policy training algorithm. The goal is to align each student output with the corresponding teacher ODE-solver step. That is, we align the \(n\)-th student output with the teacher ODE-solver step over the interval \([t_n, t_{n+1}]\). We illustrate a single optimization step with block size \(L=4\) in Figure 4. First, we sample a noisy state \(X_n\) at time step \(n\). Then, starting from \(X_n\), we roll out the student over the next block of intervals:

\[ X_n=\bar{X}_n,\ \bar{X}_{n+1},\ldots,\bar{X}_{n+L-1}. \]

Next, we sample a single state \(\bar{X}_k\) from the student trajectory, where \(k\in\{n,n+1,\ldots,n+L-1\}\). Finally, we match the \(k\)-th student output to the teacher's ODE-solver step evaluated at \(\bar{X}_k\). Figure 5 shows a simulation of PDD training on a toy problem. For each iteration, we plot the largest-error sample from the batch. Notably, the student trajectory becomes aligned with the teacher trajectory (labeled “Optimal trajectory”) as training progresses.