Generative Diffusion Models Explained (XX): From ReFlow to WGAN-GP

生成扩散模型漫谈（二十）：从ReFlow到WGAN-GP 苏剑林 · 2023-06-28 · Machine Learning optimization GAN gradient diffusion

In the previous post Generative Diffusion Models Explained (XIX): GAN as a Diffusion ODE], we explained how to understand GAN as a diffusion ODE unfolding along a different time dimension. In short, GAN essentially converts the motion of samples in a diffusion model into the motion of the generator's parameters! However, that derivation relied on relatively complex and somewhat self-contained material such as Wasserstein gradient flows, which made it hard to connect smoothly with the earlier posts in the diffusion series — technically there was a bit of a "fault line."

In my view, ReFlow, introduced in Generative Diffusion Models Explained (XVII): General Steps for Constructing an ODE (Part 2)], is the most intuitive way to understand diffusion ODEs. Since we can understand GAN from the perspective of a diffusion ODE, there must also exist a way of understanding GAN starting from ReFlow. After some experimentation, I managed to derive a result resembling WGAN-GP directly from ReFlow.

Theory Recap

The reason I call "ReFlow the most intuitive way to understand diffusion ODEs" is that it is itself extremely flexible and stays very close to actual experimental code — it can establish, via an ODE, a mapping from an arbitrary noise distribution to the target data distribution, and the training objective is very direct, requiring no convoluted detours to line up with the code you'd actually implement.

Specifically, suppose $\boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0)$ is random noise sampled from the prior distribution, and $\boldsymbol{x}_1\sim p_1(\boldsymbol{x}_1)$ is a real sample drawn from the target distribution (note: in earlier posts we usually had $\boldsymbol{x}_T$ as the noise and $\boldsymbol{x}_0$ as the target sample; here, for convenience, we swap the roles). ReFlow allows us to specify an arbitrary trajectory of motion from $\boldsymbol{x}_0$ to $\boldsymbol{x}_1$. For simplicity, ReFlow chooses a straight line, i.e.,

\begin{equation}\boldsymbol{x}_t = (1-t)\boldsymbol{x}_0 + t \boldsymbol{x}_1\label{eq:line}\end{equation}

Now let's derive the ODE it satisfies:

\begin{equation}\frac{d\boldsymbol{x}_t}{dt} = \boldsymbol{x}_1 - \boldsymbol{x}_0\end{equation}

This ODE is simple, but not practical, because what we actually want is to generate $\boldsymbol{x}_1$ from $\boldsymbol{x}_0$ via the ODE, whereas the ODE above places the very thing we want to generate on the right-hand side of the equation — which is putting the cart before the horse. To fix this, ReFlow's idea is simple: learn a function of $\boldsymbol{x}_t$ to approximate $\boldsymbol{x}_1 - \boldsymbol{x}_0$, and once learned, use it to replace $\boldsymbol{x}_1 - \boldsymbol{x}_0$, i.e.,

\begin{equation}\boldsymbol{\varphi}^* = \mathop{\text{argmin}}_{\boldsymbol{\varphi}} \mathbb{E}_{\boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0),\boldsymbol{x}_1\sim p_1(\boldsymbol{x}_1)}\left[\frac{1}{2}\Vert\boldsymbol{v}_{\boldsymbol{\varphi}}(\boldsymbol{x}_t, t) - (\boldsymbol{x}_1 - \boldsymbol{x}_0)\Vert^2\right]\label{eq:s-loss}\end{equation}

together with

\begin{equation}\frac{d\boldsymbol{x}_t}{dt} = \boldsymbol{x}_1 - \boldsymbol{x}_0\quad\Rightarrow\quad\frac{d\boldsymbol{x}_t}{dt} = \boldsymbol{v}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t, t)\label{eq:ode-core}\end{equation}

We have already shown previously that, under the assumption that $\boldsymbol{v}_{\boldsymbol{\varphi}}(\boldsymbol{x}_t, t)$ has unlimited fitting capacity, this new ODE indeed achieves the transformation of samples from distribution $p_0(\boldsymbol{x}_0)$ to distribution $p_1(\boldsymbol{x}_1)$.

Relative Motion

One important property of ReFlow is that it places no restriction on the form of the prior distribution $p_0(\boldsymbol{x}_0)$, meaning we can replace the prior with any distribution we like — for instance, a distribution produced by transforming through a generator $\boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z})$:

\begin{equation}\boldsymbol{x}_0\sim p_0(\boldsymbol{x}_0)\quad\Leftrightarrow\quad \boldsymbol{x}_0 = \boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z}),\,\boldsymbol{z}\sim \mathcal{N}(\boldsymbol{0},\boldsymbol{I})\end{equation}

After substituting this into equation $\eqref{eq:s-loss}$ and completing training, we can use equation $\eqref{eq:ode-core}$ to transform any $\boldsymbol{x}_0 = \boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z})$ into a real-looking sample $\boldsymbol{x}_1$.

However, we're not content to stop here. As mentioned earlier, GAN turns the motion of samples in a diffusion model into motion of the generator's parameters, and the same trick can be applied within this ReFlow framework: suppose the generator's current parameters are $\boldsymbol{\theta}_{\tau}$; we want the change in $\boldsymbol{\theta}_{\tau}\to \boldsymbol{\theta}_{\tau+1}$ to mimic the effect of taking one small forward step of equation $\eqref{eq:ode-core}$:

\begin{equation}\boldsymbol{\theta}_{\tau+1} = \mathop{\text{argmin}}_{\boldsymbol{\theta}}\mathbb{E}_{\boldsymbol{z}\sim \mathcal{N}(\boldsymbol{0},\boldsymbol{I})}\Big[\big\Vert \boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z}) - \boldsymbol{g}_{\boldsymbol{\theta}_{\tau}}(\boldsymbol{z}) - \epsilon\,\boldsymbol{v}_{\boldsymbol{\varphi}^*}(\boldsymbol{g}_{\boldsymbol{\theta}_{\tau}}(\boldsymbol{z}), 0)\big\Vert^2\Big]\label{eq:g-loss}\end{equation}

Note that the $t$ appearing in equations $\eqref{eq:s-loss}$ and $\eqref{eq:ode-core}$ carries a different meaning from the $\tau$ in the parameter $\boldsymbol{\theta}_{\tau}$ — the former is the ODE's time parameter, the latter is the training progress, hence the different notation. Also, $\boldsymbol{g}_{\boldsymbol{\theta}_{\tau}}(\boldsymbol{z})$ appears as the $\boldsymbol{x}_0$ of the ODE, so stepping forward a small amount gives $\boldsymbol{x}_{\epsilon}$, and the time $t$ to be substituted into $\boldsymbol{v}_{\boldsymbol{\varphi}^*}(\boldsymbol{x}_t, t)$ is $0$.

Now we have a new $\boldsymbol{g}_{\boldsymbol{\theta}_{\tau+1}}(\boldsymbol{z})$, which in theory produces a distribution closer to the real distribution (since we've taken one small step forward). We then treat it as the new $\boldsymbol{x}_0$, substitute it into equation $\eqref{eq:s-loss}$ to train, and once training is complete, substitute it into equation $\eqref{eq:g-loss}$ to optimize the generator, and so on — this gives an alternating training procedure much like GAN's.

WGAN-GP

So, can we tie this process quantitatively to existing GAN formulations? Yes! It turns out to connect to WGAN-GP] with its gradient penalty.

Let's first look at the loss function $\eqref{eq:s-loss}$. Expanding the expectation term gives

\begin{equation}\frac{1}{2}\Vert\boldsymbol{v}_{\boldsymbol{\varphi}}(\boldsymbol{x}_t, t)\Vert^2 - \langle\boldsymbol{v}_{\boldsymbol{\varphi}}(\boldsymbol{x}_t, t),\boldsymbol{x}_1 - \boldsymbol{x}_0\rangle + \frac{1}{2}\Vert\boldsymbol{x}_1 - \boldsymbol{x}_0\Vert^2\end{equation}

The third term does not depend on the parameter $\boldsymbol{\varphi}$, so dropping it doesn't affect the result. Now suppose $\boldsymbol{v}_{\boldsymbol{\varphi}}$ has strong enough fitting capacity that we don't need to explicitly feed in $t$; then, as a loss function, the expression above is equivalent to

\begin{equation}\frac{1}{2}\Vert\boldsymbol{v}_{\boldsymbol{\varphi}}(\boldsymbol{x}_t)\Vert^2 - \langle\boldsymbol{v}_{\boldsymbol{\varphi}}(\boldsymbol{x}_t),\boldsymbol{x}_1 - \boldsymbol{x}_0\rangle = \frac{1}{2}\Vert\boldsymbol{v}_{\boldsymbol{\varphi}}(\boldsymbol{x}_t)\Vert^2 - \left\langle\boldsymbol{v}_{\boldsymbol{\varphi}}(\boldsymbol{x}_t),\frac{d\boldsymbol{x}_t}{dt}\right\rangle\end{equation}

$\boldsymbol{v}_{\boldsymbol{\varphi}}(\boldsymbol{x}_t)$ is a vector-valued function with matching input and output dimensions; let's further assume it is the gradient of some scalar function $D_{\boldsymbol{\varphi}}(\boldsymbol{x}_t)$, i.e., $\boldsymbol{v}_{\boldsymbol{\varphi}}(\boldsymbol{x}_t)=\nabla_{\boldsymbol{x}_t} D_{\boldsymbol{\varphi}}(\boldsymbol{x}_t)$. Then the expression above becomes

\begin{equation}\frac{1}{2}\Vert\nabla_{\boldsymbol{x}_t} D_{\boldsymbol{\varphi}}(\boldsymbol{x}_t)\Vert^2 - \left\langle\nabla_{\boldsymbol{x}_t} D_{\boldsymbol{\varphi}}(\boldsymbol{x}_t),\frac{d\boldsymbol{x}_t}{dt}\right\rangle = \frac{1}{2}\Vert\nabla_{\boldsymbol{x}_t} D_{\boldsymbol{\varphi}}(\boldsymbol{x}_t)\Vert^2 - \frac{d D_{\boldsymbol{\varphi}}(\boldsymbol{x}_t)}{dt}\end{equation}

Assuming $D_{\boldsymbol{\varphi}}(\boldsymbol{x}_t)$ varies fairly smoothly, $\frac{d D_{\boldsymbol{\varphi}}(\boldsymbol{x}_t)}{dt}$ should be close to its finite difference $D_{\boldsymbol{\varphi}}(\boldsymbol{x}_1)-D_{\boldsymbol{\varphi}}(\boldsymbol{x}_0)$ evaluated at the two points $t=0,t=1$, so the loss function above is approximately

\begin{equation}\frac{1}{2}\Vert\nabla_{\boldsymbol{x}_t} D_{\boldsymbol{\varphi}}(\boldsymbol{x}_t)\Vert^2 - D_{\boldsymbol{\varphi}}(\boldsymbol{x}_1) + D_{\boldsymbol{\varphi}}(\boldsymbol{x}_0)\end{equation}

Readers familiar with GANs will probably find this quite recognizable — it's exactly the discriminator loss of WGAN with a gradient penalty! Even the construction $\eqref{eq:line}$ of the point $\boldsymbol{x}_t$ used in the gradient penalty term is identical (a linear interpolation between real and fake samples)! The only difference is that the original WGAN-GP's gradient penalty is centered at 1, whereas here it is centered at zero — but in fact, articles such as WGAN-div: An Unsung Hero That Filled In WGAN's Gaps] and Optimization Algorithms from a Dynamical Systems Perspective (IV): The Third Stage of GAN] have already shown that a gradient penalty centered at zero typically performs better.

In other words, under this particular parameterization and set of assumptions, the loss function $\eqref{eq:s-loss}$ is in fact equivalent to the discriminator loss of WGAN-GP. As for the generator loss, in the previous post Generative Diffusion Models Explained (XIX): GAN as a Diffusion ODE] we already showed that when $\boldsymbol{v}_{\boldsymbol{\varphi}}(\boldsymbol{x}_t)=\nabla_{\boldsymbol{x}_t} D_{\boldsymbol{\varphi}}(\boldsymbol{x}_t)$, the gradient of a single-step optimization of equation $\eqref{eq:g-loss}$ is equivalent to the gradient of

\begin{equation}\boldsymbol{\theta}_{\tau+1} = \mathop{\text{argmin}}_{\boldsymbol{\theta}}\mathbb{E}_{\boldsymbol{z}\sim \mathcal{N}(\boldsymbol{0},\boldsymbol{I})}[-D(\boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z}))]\end{equation}

which is exactly the generator loss of WGAN-GP as well.

Summary

In this post, I attempted to derive the connection between WGAN-GP and diffusion ODEs starting from ReFlow. This angle turns out to be relatively simpler and more intuitive, and it avoids relatively complex notions such as Wasserstein gradient flows.