Parallel Decoding Distillation for Image and Video Generation (PDD)
Hozy Summary
- Idea)
- Learn a parallel decoder that predicts multiple denoising steps in a single network evaluation
2 Generative Flow Matching
2.1 Teacher Flow Matching Model
- Model)
- \((X_t)_{0\le t\le 1}\) : a flow process s.t.
- \(\mathcal{X}\) : the data space
- \(v:\mathcal{X}\times[0,1]\rightarrow\mathcal{X}\) : a velocity field
- \(p_0:\mathcal{X}\rightarrow\mathbb{R}_{\ge 0}\) : a source distribution
- \(\displaystyle \frac{\text{d}}{\text{d}t}X_t = v_t(X_t),\quad X_0\sim p_0\) : the marginal process
- \(p:\mathcal{X}\times[0,1]\rightarrow\mathbb{R}_{\ge0}\) : the probability path
- where \(X_t\sim p_t,\quad\forall t\in[0,1]\)
- \(X_t = (1-t) X_0 + t X_1\) : the linear timestep scheduler
- where \(X_0\sim p_0, \quad X_1\sim p_1, \text{ and } t\in[0,1]\)
- \((X_t)_{0\le t\le 1}\) : a flow process s.t.
- Sampling)
- For a discretized time interval \(0=t_0 \lt t_1 \lt \cdots \lt t_N = 1\)
- \(X_{n+1} = X_n + (t_{n+1} - t_n) u_n (X_n)\) : the ODE solution where
- \(\displaystyle u_n(X_n) = \frac{1}{t_{n+1}-t_n}\int_{t_n}^{t_{n+1}} v_t(X_t) \text{d}t\) : the mean velocity of the \(n\)-th interval \([t_{n}, t_{n+1}]\)
- \(X_{n+1} = X_n + (t_{n+1} - t_n) u_n (X_n)\) : the ODE solution where
- Solvers)
- Euler Solver : \(u_n(X_n) \approx v_{t_n}(X_n)\)
- Runge-Kutta methods (Family of higher order solvers)
- Midpoint Method
- Def.)
\(\begin{cases} X_{\text{mid}} = X_n + \frac{t_{n+1}-t_n}{2} v_{t_n}(X_n) \\ t_{\text{mid}} = \frac{t_{n+1}+t_n}{2} \end{cases}\). - Then, \(u_n(X_n)\approx v_{t_{\text{mid}}}(X_{\text{mid}})\)
- Def.)
- Midpoint Method
- For a discretized time interval \(0=t_0 \lt t_1 \lt \cdots \lt t_N = 1\)
3. Parallel Decoding Distillation
Parallel Decoder
- Settings)
- \(v_t\) : a pre-trained (teacher) flow model with a flow process \((X_t)_{0\le t\le 1}\)
- i.e.) Outputs a single velocity given a single state as \(v_{t_n}(X_{t_n})\in\mathcal{X}\)
- Architecture)
- \(v_t(x) = W H_t(x)\) where
- \(H_t\) : the backbone that outputs the final hidden state
- \(W\) : the final linear layer
- \(v_t(x) = W H_t(x)\) where
- \(n\in\{1,\ldots,N\}\) : the start index of a block from the time discretization \(0=t_0 \lt t_1 \lt \cdots \lt t_N = 1\)
- \(L\) : the block size
- Thus, each block consists \(L\) indices of \(\{n,\ldots,n+L-1\}\)
- \(v_t\) : a pre-trained (teacher) flow model with a flow process \((X_t)_{0\le t\le 1}\)
- Model)
- Parallel Decoder)
- \(\bar{u}_n^\theta(\cdot\mid X_n)\in\mathcal{X}^L\).
- i.e.) Outputs \(L\) (blocks size) mean velocities for the target block.
- However, in implementation the model outputs \(N\) values.
- Props.)
- Same backbone as the teacher flow model.
- Architecture)
- \(\bar{u}_n^\theta(k\mid x_n) = W_k^\theta H_{t_n}^\theta (x_n)\).
- It learns \(N\) linear layers of \(W_0^\theta, \ldots, W_{N-1}^\theta\).
- i.e.) Outputs \(L\) (blocks size) mean velocities for the target block.
- \(\bar{u}_n^\theta(\cdot\mid X_n)\in\mathcal{X}^L\).
- Mean Velocity Prediction)
- \(\underbrace{\bar{u}_n^\theta(k\mid X_n)}_{\text{model pred.}} \approx \underbrace{u_k(X_k)}_{\text{teacher mean velocity}},\quad k=n,\ldots,n+L-1\).
- i.e.) Predict the teacher mean velocity over the subinterval \([t_k, t_{k+1}]\) with a model conditioned only on the initial state \(X_n\).
- How?)
- Intra-block Process (Parallelized Process)
- Given
- the initial value of \(\bar{X}_n = X_n\)
- we may get
- \(\bar{X}_{k+1} = \bar{X}_k + (t_{k+1} - t_{k})\bar{u}_n^\theta(k\mid X_n)\) for \(k\sim\mathcal{U}\{n,\ldots,n+L-1\}\)
- cf.) Single evaluation available.
- Why?) \(\bar{u}_n^\theta(\cdot\mid X_n)\) depends only on \(X_n\).
- Thus, the \(L\) jump within an interval from \(t=t_{n}\) to \(t=t_{n+L}\) is possible by
- \(\bar{X}_{n+L} = \bar{X}_{n} + (t_{n+1}-t_n)\bar{u}_n^\theta(n\mid X_n) + \cdots + (t_{n+L}-t_{n+L-1})\bar{u}_n^\theta(n+L-1\mid X_n)\).
- i.e.) the block step rule below.
- \(\bar{X}_{n+L} = \bar{X}_{n} + (t_{n+1}-t_n)\bar{u}_n^\theta(n\mid X_n) + \cdots + (t_{n+L}-t_{n+L-1})\bar{u}_n^\theta(n+L-1\mid X_n)\).
- Given
- Intra-block Process (Parallelized Process)
- \(\underbrace{\bar{u}_n^\theta(k\mid X_n)}_{\text{model pred.}} \approx \underbrace{u_k(X_k)}_{\text{teacher mean velocity}},\quad k=n,\ldots,n+L-1\).
- Parallel Decoder)
- Sampling)
- Block Step Rule (Inter-block process)
- \(\bar{X}_{n+L} = X_n + \displaystyle \underbrace{\sum_{k=n}^{n+L-1}(t_{k+1}-t_k)\bar{u}_n^\theta(k\mid X_n)}_{\text{Parallelized Process (intra-block)}}\).
- Repeat the above block-step rule \(N/L\) times and get \(\bar{X}_N\)
- At inference, the block update can be exactly simplified using a fused linear layer:
- \(\bar{X}_{n+L} = \bar{X}_n + (t_{n+L} - t_n) W_{n:n+L}^\theta H_{t_n}^\theta(\bar{X}_n)\) where
- \(W_{n:n+L}^\theta = \displaystyle\sum_{k=n}^{n+L-1} \Delta_k W_k^\theta\) : the fused linear layer
- \(\Delta_k = \displaystyle\frac{t_{k+1}-t_k}{t_{n+L}-t_n}\).
- \(\bar{X}_{n+L} = \bar{X}_n + (t_{n+L} - t_n) W_{n:n+L}^\theta H_{t_n}^\theta(\bar{X}_n)\) where
- At inference, the block update can be exactly simplified using a fused linear layer:
- Block Step Rule (Inter-block process)
- Training)
- Objective)
- \(\mathcal{L}_{\text{PD}}(\theta) = \mathbb{E}\bigg[ \Big\Vert \bar{u}_n^\theta(k\mid X_n) - \overbrace{u_k(\text{stopgrad}(\underbrace{\bar{X}_k}_{\text{from student}}))}^{\text{from teacher}} \Big\Vert^2 \bigg]\).
- Objective)
Enjoy Reading This Article?
Here are some more articles you might like to read next: