Pi-Flow Policy-Based Few-Step Generation via Imitation Distillation
Hozy Summary
Setup & Model
Setups) Flow Matching Model
- \(p(x_0)\) : the (latent) data probability density for a datapoint \(x_0\in\mathbb{R}^D\)
- \(x_t = \alpha_t x_0 + \sigma_t \epsilon\) : the diffused noisy data
- where
- \(t\in(0,1],\; \epsilon\sim\mathcal{N}(\mathbf{0,I}),\; \alpha_t = 1-t,\; \sigma_t = t\).
- cf.)
- Marginal Density : \(\displaystyle p(x_t) = \int_{\mathbb{R}^D}\mathcal{N}\left(x_t;\;\alpha_t x_0, \sigma_t^2\mathbf{I}\right)p(x_0)\text{d}x_0\)
- Optimal transport map via probability flow ODE
- \(\displaystyle\frac{\text{d}x_t}{\text{d}t} = \dot{x}_t = \frac{x_t - \mathbb{E}_{x_0\sim p(x_0\mid x_t)}[x_0]}{t} = \frac{x_t - \int_{\mathbb{R}^D}x_0p(x_0\mid x_t)\text{d}x_0}{t}\).
- with the denoising posterior \(p(x_0\mid x_t) := \displaystyle\frac{\mathcal{N}\left(x_t;\;\alpha_t x_0, \sigma_t^2\mathbf{I}\right) p(x_0)}{p(x_t)}\)
- where
- Objective
- \(\displaystyle\mathcal{L}_\theta = \mathbb{E}_{t, x_0, x_t}\left[\frac{1}{2}\Vert u - G_\theta(x_t,t) \Vert^2\right]\) with sample velocity of \(\displaystyle u :=\frac{x_t-x_0}{t}\)
- for a neural network \(G_\theta(x_t, t)\) with learnable parameters
- \(\displaystyle\mathcal{L}_\theta = \mathbb{E}_{t, x_0, x_t}\left[\frac{1}{2}\Vert u - G_\theta(x_t,t) \Vert^2\right]\) with sample velocity of \(\displaystyle u :=\frac{x_t-x_0}{t}\)
Model) pi-Flow: Policy-Based Few-Step Generation
- Def.)
- \(\pi:\mathbb{R}^D\times\mathbb{R}\rightarrow\mathbb{R}^D\) : a network-free policy function that maps a state \((x_t, t)\) to a flow velocity
- Requirement for being network-free
- \(\pi\) needs to describe a single ODE trajectory fully determined by its initial state \((x_{t_{\text{src}}}, t_{\text{src}})\)
- Requirement for being network-free
- \(\pi:\mathbb{R}^D\times\mathbb{R}\rightarrow\mathbb{R}^D\) : a network-free policy function that maps a state \((x_t, t)\) to a flow velocity
- Training)
- Adapt a flow model to output the entire dynamic policy that governs the full trajectory
- i.e.) For the policy function space \(\mathcal{F}:=\left\{\pi:\mathbb{R}^D\times\mathbb{R}\rightarrow\mathbb{R}^D\right\}\) train a policy generator network \(G_\phi:\mathbb{R}^D\times\mathbb{R}\rightarrow\mathcal{F}\)
- Adapt a flow model to output the entire dynamic policy that governs the full trajectory
- Two-step Sampling)
- Single Policy Generation Step
- Given the initial state \((x_{t_{\text{src}}}, t_{\text{src}})\), sample a policy \(\pi\leftarrow G_\phi(x_{t_{\text{src}}}, t_{\text{src}})\)
- Multiple Policy Integration Substeps
- Integrate the ODE by querying policy velocity over multiple steps
- \(\displaystyle x_{t_{\text{dst}}} \leftarrow x_{t_{\text{src}}} + \int_{t_{\text{src}}}^{t_{\text{dst}}} \pi(x_t, t) \text{d}t\).
- Integrate the ODE by querying policy velocity over multiple steps
- Single Policy Generation Step
- Props.)
- Decouples network evaluation steps from the ODE integration substeps.
- Advantages)
- Few network evaluation for generation
- Capable of following well-established imitation learning (IL) approaches to directly match the policy velocity \(\pi(x_t, t)\) to the teacher velocity \(G_\theta(x_t, t)\).
- How?) It can execute dense integration substeps like standard flow matching teacher.
- Advantages)
- Decouples network evaluation steps from the ODE integration substeps.
3. Policy Generation
3.1 Dynamic-x_0 policy (DX)
- Def.)
- \(\pi(x_t, t) := \displaystyle\frac{x_t - \hat{x}_{0}^{(t)}}{t}\) where \(\hat{x}_{0}^{(t)}\) approximates the posterior moment \(\mathbb{E}_{x_0\sim p(x_0\mid x_t)[x_0]}\)
- Model)
- For evenly spaced times \(t_1,\ldots,t_N\in[t_\text{dst}, t_{src}]\),
- set a grid of \(\hat{x}_0^{(t_i)}\;(i=1,\ldots,N)\)
- Expand the output of \(G_\phi(x_{t_{src}}, t_{src}) = u\) to \(x_0\in\mathbb{R}^N\) using the \(u\text{-to-}x_0\) reparameterization.
- Predict \(\hat{x}_0^{(t_i)}\;(i=1,\ldots,N)\) by a single evaluation of the student network \(G_\phi(x_{t_{src}}, t_{src})\).
- Obtain \(\hat{x}_0^{(t)}\) for \(t\in[t_{\text{dst}}, t_{\text{src}}]\) using the linear interpolation.
- For evenly spaced times \(t_1,\ldots,t_N\in[t_\text{dst}, t_{src}]\),
- Props.)
- The posterior moment depends only on \(t\).
- Fast and expressive (we may choose any \(N\))
- Limited robustness because \(\hat{x}_{0}^{(t)}\) is not adaptive to perturbations in \(x_t\).
3.2 GMFlow Policy
- Motivation)
- Stronger robustness!
- Idea)
- GM Flow
- Expand the network output channels to predict a factorized Gaussian mixture (GM) velocity distribution \(q(u\mid x_{t_{\text{src}}}) = \displaystyle\prod_{i=1}^L\sum_{k=1}^K A_{ik}\mathcal{N}\left( u_i; \mu_{ik}, s^2\mathbf{I} \right)\)
- for
- hyperparameters
- \(L\) : sequence length
- \(C\) : channel size
- \(K\) : a hyperparameter specifying the number of mixture components
- GM parameters predicted by network \(G_\phi\)
- \(A_{ik}\in\mathbb{R}_+\),
- \(\mu_{ik}\in\mathbb{R}^C\),
- \(s\in\mathbb{R}_+\),
- hyperparameters
- i.e.) The student network \(G_\phi\) maps the initial state \(x_{t_{\text{src}}}\) to multiple denoising modes that parameterize the GMFlow policy.
- for
- This GMFlow policy enables a closed-form velocity expression at future state \((x_t, t),\;\forall 0\lt t\lt t_{\text{src}}\) described by
- \(q(x_0\mid x_{t_{\text{src}}}) = \displaystyle\sum_{k=1}^K A_k'\mathcal{N}\left(x_0;\; \mu_k', s'^2\mathbf{I}\right)\).
- Expand the network output channels to predict a factorized Gaussian mixture (GM) velocity distribution \(q(u\mid x_{t_{\text{src}}}) = \displaystyle\prod_{i=1}^L\sum_{k=1}^K A_{ik}\mathcal{N}\left( u_i; \mu_{ik}, s^2\mathbf{I} \right)\)
- GM Flow
Derivation from Appendix F
- Model)
- Policy function \(\pi:\mathbb{R}^C\times\mathbb{R}\rightarrow\mathbb{R}^C\) is given by
\(\begin{aligned} \pi(x_t, t) &= \frac{x_t-\mathbb{E}_{x_0\sim q(x_0\mid x_t)}[x_0]}{t} \\ &= \frac{x_t-\sum_{k=1}^K A_k' \mu_k'}{t} \end{aligned}\)
- Policy function \(\pi:\mathbb{R}^C\times\mathbb{R}\rightarrow\mathbb{R}^C\) is given by
4. Policy-based Image Distillation (π-ID)
- Idea)
- DAgger-style on-policy imitation
- Settings)
- \(\text{NFE}\) : the number of evaluation
- Then the \(\text{NFE}\)-number of intervals will be set as \([0, \frac{1}{\text{NFE}}], [\frac{1}{\text{NFE}}, \frac{2}{\text{NFE}}], \ldots, [\frac{\text{NFE}-1}{\text{NFE}}, 1]\)
- \(G_\theta\) : pre-trained teacher
- \(G_\phi\) : student
- \(c\) : condition
- \(\text{NFE}\) : the number of evaluation
- Training)
- Sample \(t_\text{src}\sim\left\{ \frac{1}{\text{NFE}}, \frac{2}{\text{NFE}}, \ldots, 1 \right\}\)
- i.e.) the interval will be \(\left[t_\text{src}-\frac{1}{\text{NFE}}, t_\text{src}\right]\) with \(t_{\text{dst}}=t_\text{src}-\frac{1}{\text{NFE}}\)
- Initialize \(x_{t_\text{src}}\)
- Data Dependent (real data)
- Reverse Denoising from random noise
- cf.) Both have roughly the same computational cost and performance.
- Get
- \(\pi\leftarrow G_\phi(x_{t_\text{src}}, t_\text{src}, c)\),
- \(\pi_D\leftarrow \text{stopgrad}(\pi)\).
- For \(t\sim\mathcal{U}(t_\text{src}-\frac{1}{\text{NFE}}, t_\text{src})\)
- \(x_t\leftarrow x_{t_\text{src}} + \displaystyle\int_{t_\text{src}}^t \pi_D(x_s, s)\text{d}s\) : student roll out
- cf.) Authors mentioned small substep size of 1/128.
- \(\mathcal{L}_\phi\leftarrow\mathcal{L}_\phi + \displaystyle\frac{1}{2}\Vert G_\theta(x_t, t, c) - \pi(x_t ,t)\Vert^2\) : \(\ell_2\)-loss with the teacher query and the student output
- \(x_t\leftarrow x_{t_\text{src}} + \displaystyle\int_{t_\text{src}}^t \pi_D(x_s, s)\text{d}s\) : student roll out
- Optimize \(\phi\leftarrow\text{Adam}(\phi, \nabla_\phi\mathcal{L}_\phi)\).
- Sample \(t_\text{src}\sim\left\{ \frac{1}{\text{NFE}}, \frac{2}{\text{NFE}}, \ldots, 1 \right\}\)
4.1 Data-dependent on-policy π-ID
4.2 Data-free on-policy π-ID
Enjoy Reading This Article?
Here are some more articles you might like to read next: