Generative Diffusion Models Explained (Part 19): GAN as a Diffusion ODE

生成扩散模型漫谈（十九）：作为扩散ODE的GAN 苏剑林 · 2023-06-24 · Machine Learning optimization GAN diffusion

In the post Generative Diffusion Models Explained (Part 16): W-Distance ≤ Score Matching, we derived an inequality relating the Wasserstein distance to the score-matching loss of diffusion models, showing that the optimization objective of diffusion models bears some resemblance to that of WGANs. In this post, we'll look at the findings from MonoFlow: Rethinking Divergence GANs via the Perspective of Wasserstein Gradient Flows, which pushes this connection between GANs and diffusion models even further: a GAN can in fact be viewed as a diffusion ODE along a different time dimension!

These findings suggest that, although GANs and diffusion models appear on the surface to be two entirely different types of generative model, they actually share a great deal in common, and there is much they can learn from and borrow from each other.

Overview of the idea

As we know, the generator trained by a GAN is a direct, deterministic mapping $\boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z})$ from noise $\boldsymbol{z}$ to real samples, whereas the defining feature of diffusion models is "progressive generation": their generation process corresponds to sampling from a sequence of gradually changing distributions $p_0(\boldsymbol{x}_0),p_1(\boldsymbol{x}_1),\cdots,p_T(\boldsymbol{x}_T)$ (note: in the previous dozen or so posts, $\boldsymbol{x}_T$ denoted noise and $\boldsymbol{x}_0$ the target sample, with the sampling process going as $\boldsymbol{x}_T\to \boldsymbol{x}_0$; but for convenience of exposition here we reverse this and instead write $\boldsymbol{x}_0\to \boldsymbol{x}_T$). At first glance there really doesn't seem to be much in common — so how can we connect the two? more

Clearly, if we want to understand GANs from the perspective of diffusion models, we need to find a way to construct a sequence of gradually changing distributions. The generator $\boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z})$ itself is a one-shot transformation with no gradual variation involved — but we do know that the optimization of the model itself proceeds gradually. Could we use the historical trajectory $\boldsymbol{\theta}_t$ of the parameters $\boldsymbol{\theta}$ to build this sequence of gradually changing distributions? Specifically, suppose the generator is initialized as $\boldsymbol{\theta}_0$, and after $T$ steps of adversarial training reaches its optimal parameters $\boldsymbol{\theta}_T$, with intermediate parameters $\boldsymbol{\theta}_1,\boldsymbol{\theta}_2,\cdots,\boldsymbol{\theta}_{T-1}$ along the way. If we then define $\boldsymbol{x}_t = \boldsymbol{g}_{\boldsymbol{\theta}_t}(\boldsymbol{z})$, doesn't that give us a sequence of gradually changing $\boldsymbol{x}_0,\boldsymbol{x}_1,\cdots,\boldsymbol{x}_T$, and hence a sequence of gradually changing distributions $p_0(\boldsymbol{x}_0),p_1(\boldsymbol{x}_1),\cdots,p_T(\boldsymbol{x}_T)$?

If this idea works, then a GAN can be interpreted as a diffusion model along the (virtual) time dimension of gradient descent! Let's explore this line of thought below.

Flow of the gradient

First, let's recall the result on Wasserstein gradient flow from the previous post Gradient Flow: Exploring the Path to the Minimum: it showed that the equation

\begin{equation}\frac{\partial q_t(\boldsymbol{x})}{\partial t} = - \nabla_{\boldsymbol{x}}\cdot\big(q_t(\boldsymbol{x})\nabla_{\boldsymbol{x}}\log r_t(\boldsymbol{x})\big)\label{eq:w-flow}\end{equation}

minimizes the KL divergence between $p(\boldsymbol{x})$ and $q_t(\boldsymbol{x})$, i.e., $\lim\limits_{t\to\infty} q_t(\boldsymbol{x}) = p(\boldsymbol{x})$, where $r_t(\boldsymbol{x})=\frac{p(\boldsymbol{x})}{q_t(\boldsymbol{x})}$. If $p(\boldsymbol{x})$ represents the distribution of real samples, then, provided we can sample from $q_t(\boldsymbol{x})$, gradually pushing forward to $t\to\infty$ would let us sample from $p(\boldsymbol{x})$. According to Deriving the Continuity Equation and the Fokker-Planck Equation via Test Functions, sampling from $q_t(\boldsymbol{x})$ can be achieved via the following ODE:

\begin{equation}\frac{d\boldsymbol{x}}{dt} = \nabla_{\boldsymbol{x}}\log r_t(\boldsymbol{x})\label{eq:ode-core}\end{equation}

However, $r_t(\boldsymbol{x})$ in the equation above is unknown, so we cannot yet use this equation for sampling — we first need to find a way to estimate $r_t(\boldsymbol{x})$.

Estimation via discrimination

This is where the GAN discriminator comes in. Taking the original Vanilla GAN as an example, its training objective is

\begin{equation}\max_D\, \mathbb{E}_{\boldsymbol{x}\sim p(\boldsymbol{x})}[\log \sigma(D(\boldsymbol{x}))] + \mathbb{E}_{\boldsymbol{x}\sim q(\boldsymbol{x})}[\log (1 - \sigma(D(\boldsymbol{x})))]\label{eq:gan-d}\end{equation}

Here $D$ is the discriminator, $\sigma(t)=1/(1+e^{-t})$ is the sigmoid function, $p(\boldsymbol{x})$ is the distribution of real samples, and $q(\boldsymbol{x})$ is the distribution of fake samples. One can show (readers unfamiliar with this can refer to the "Supplementary Proof" section of RSGAN: The "Turing Test" Idea in Adversarial Models) that the theoretical optimum of the discriminator $D$ in the equation above is

\begin{equation}D(\boldsymbol{x}) = \log \frac{p(\boldsymbol{x})}{q(\boldsymbol{x})}\end{equation}

The more general f-GAN framework (see A Brief Introduction to f-GAN: A Production Line for GAN Models, Designing GANs: Yet Another GAN Production Line) gives a slightly different result, but it can likewise be shown that in each case the theoretical optimum of the discriminator is a function of $\frac{p(\boldsymbol{x})}{q(\boldsymbol{x})}$. In other words, as long as we can sample from $p(\boldsymbol{x})$ and $q_t(\boldsymbol{x})$, training a GAN discriminator via $\eqref{eq:gan-d}$ lets us estimate $r_t(\boldsymbol{x})=\frac{p(\boldsymbol{x})}{q_t(\boldsymbol{x})}$.

A step forward

At this point some readers might object: isn't this a chicken-and-egg argument? Didn't we introduce the estimate of $r_t(\boldsymbol{x})$ precisely so we could use equation $\eqref{eq:ode-core}$ to sample from $q_t(\boldsymbol{x})$? And now you're assuming we can already sample from $q_t(\boldsymbol{x})$ in order to estimate $r_t(\boldsymbol{x})$? Don't worry — here comes the crucial insight.

Suppose we have a generator $\boldsymbol{g}_{\boldsymbol{\theta}_t}(\boldsymbol{z})$ whose generated samples are equivalent to samples from $q_t(\boldsymbol{x})$, that is,

\begin{equation}\big\{\boldsymbol{g}_{\boldsymbol{\theta}_t}(\boldsymbol{z})\big|\,\boldsymbol{z}\sim \mathcal{N}(\boldsymbol{0},\boldsymbol{I})\big\}\quad = \quad\big\{\boldsymbol{x}_t\big|\,\boldsymbol{x}_t\sim q_t(\boldsymbol{x})\big\}\end{equation}

Then we can use this generator together with equation $\eqref{eq:gan-d}$ to estimate $r_t(\boldsymbol{x})$. Note that this is only the value of $r_t(\boldsymbol{x})$ at time $t$ — we don't know $r_t(\boldsymbol{x})$ at other times, so we can't directly use equation $\eqref{eq:ode-core}$ to complete the full sampling process. But we can push forward by one small step:

\begin{equation}\boldsymbol{x}_{t+1} = \boldsymbol{x}_t + \epsilon \nabla_{\boldsymbol{x}_t}\log r_t(\boldsymbol{x}_t) = \boldsymbol{x}_t + \epsilon \nabla_{\boldsymbol{x}_t} D(\boldsymbol{x}_t)\label{eq:forward}\end{equation}

Here $\epsilon$ is a small positive number representing the step size. Now we have the result of sampling at the next step, and we would like it to remain equivalent to sampling from the generator at the next step, i.e.,

\begin{equation}\begin{aligned} \big\{\boldsymbol{g}_{\boldsymbol{\theta}_{t+1}}(\boldsymbol{z})\big|\,\boldsymbol{z}\sim \mathcal{N}(\boldsymbol{0},\boldsymbol{I})\big\}\quad =& \quad\big\{\boldsymbol{x}_{t+1}\big|\,\boldsymbol{x}_{t+1}\sim q_{t+1}(\boldsymbol{x})\big\} \\[5pt] \quad =& \quad\big\{\boldsymbol{x}_{t+1}\big|\,\boldsymbol{x}_t + \epsilon \nabla_{\boldsymbol{x}_t} D(\boldsymbol{x}_t),\boldsymbol{x}_t\sim q_t(\boldsymbol{x})\big\} \end{aligned}\end{equation}

In other words, what we want to do is turn the motion of samples under the diffusion model into motion of the generator's parameters! To achieve this, we solve for $\boldsymbol{\theta}_{t+1}$ using the following loss:

\begin{equation}\boldsymbol{\theta}_{t+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}_t}(\boldsymbol{z}) - \epsilon \nabla_{\boldsymbol{g}}D(\boldsymbol{g}_{\boldsymbol{\theta}_t}(\boldsymbol{z}))\big\Vert^2\Big]\label{eq:gan-g0}\end{equation}

That is, we take $\boldsymbol{x}_t = \boldsymbol{g}_{\boldsymbol{\theta}_t}(\boldsymbol{z})$, advance it one step to get $\boldsymbol{x}_{t+1}$, and then try to make the new $\boldsymbol{g}_{\boldsymbol{\theta}_{t+1}}(\boldsymbol{z})$ match $\boldsymbol{x}_{t+1}$ as closely as possible. Once this round is done, we replace the original $\boldsymbol{\theta}_t$ with $\boldsymbol{\theta}_{t+1}$ and start a new round of iteration — that is, equations $\eqref{eq:gan-d}$ and $\eqref{eq:gan-g0}$ are executed alternately. Doesn't this already start to feel like a GAN?

The finishing touch

If that's not enough, we can refine things further to make it match GANs even more closely. Notice that the gradient of the function inside the expectation in equation $\eqref{eq:gan-g0}$ is:

\begin{equation}\begin{aligned} &\,\nabla_{\boldsymbol{\theta}}\Vert \boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z}) - \boldsymbol{g}_{\boldsymbol{\theta}_t}(\boldsymbol{z}) - \epsilon \nabla_{\boldsymbol{g}}D(\boldsymbol{g}_{\boldsymbol{\theta}_t}(\boldsymbol{z}))\Vert^2 \\ =&\,2\big\langle\boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z}) - \boldsymbol{g}_{\boldsymbol{\theta}_t}(\boldsymbol{z}) - \epsilon \nabla_{\boldsymbol{g}}D(\boldsymbol{g}_{\boldsymbol{\theta}_t}(\boldsymbol{z})), \nabla_{\boldsymbol{\theta}}\boldsymbol{g}_{\boldsymbol{\theta}}(\boldsymbol{z}) \big\rangle \\ \end{aligned}\end{equation}

Substituting in the current value $\boldsymbol{\theta}=\boldsymbol{\theta}_t$, we get

\begin{equation}-2\epsilon\big\langle \nabla_{\boldsymbol{g}}D(\boldsymbol{g}_{\boldsymbol{\theta}_t}(\boldsymbol{z})), \nabla_{\boldsymbol{\theta}_t}\boldsymbol{g}_{\boldsymbol{\theta}_t}(\boldsymbol{z}) \big\rangle = -2\epsilon\nabla_{\boldsymbol{\theta}_t}D(\boldsymbol{g}_{\boldsymbol{\theta}_t}(\boldsymbol{z}))\end{equation}

That is to say, if we only take a single gradient-based optimization step, then using equation $\eqref{eq:gan-g0}$ as the loss function is equivalent to using the following as the loss function (since their gradients differ only by a constant factor):

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

This is exactly one of the common forms of the generator loss. Alternately training with equations $\eqref{eq:gan-d}$ and $\eqref{eq:gan-g}$ gives us a familiar GAN variant.

In particular, the original paper also shows that the generator's loss function can be generalized to

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

where $h(\cdot)$ is an arbitrary monotonically increasing function, which corresponds to replacing $\log r_t(\boldsymbol{x})$ with $h(\log r_t(\boldsymbol{x}))$ in the Wasserstein gradient flow equation $\eqref{eq:w-flow}$. This is presumably where the name MonoFlow comes from (Monotonically increasing function + Wasserstein flow). We won't go through this proof here — interested readers can consult the original paper.

Reflections on the significance

In summary, the line of reasoning that lets us understand GANs as diffusion models is

$$\require{AMScd}\begin{CD} \cdots @>\quad>> \boldsymbol{g}_{\boldsymbol{\theta}_t}(\boldsymbol{z}) @> 式\eqref{eq:gan-d} >> r_t(\boldsymbol{x}) @> 式\eqref{eq:forward}>> \boldsymbol{x}_{t+1} @> 式\eqref{eq:gan-g0}>> \boldsymbol{g}_{\boldsymbol{\theta}_{t+1}}(\boldsymbol{z})@>\quad>>\cdots \end{CD}$$

where the core equation is $\eqref{eq:forward}$, which derives from equations $\eqref{eq:w-flow}$ and $\eqref{eq:ode-core}$ of the Wasserstein gradient flow — something we discussed in the previous post, Gradient Flow: Exploring the Path to the Minimum.

Some readers might ask: this perspective doesn't seem to give us anything more than what we already had with GANs — so why go to all this trouble to re-understand GANs? First, in the author's view, understanding GANs from the perspective of diffusion models — or unifying diffusion models and GANs — is in itself an interesting and enjoyable exercise, without necessarily needing to serve some practical purpose. Being interesting and fun is already its greatest value.

Second, as the authors of the original paper point out, the existing derivation of GANs is inconsistent with how GANs are actually trained in practice, whereas the diffusion perspective discussed here is consistent with the training process. In other words, if we take the actual training procedure as the standard, the existing derivation of GANs is, strictly speaking, wrong, while the diffusion perspective presented here is the correct one. How should we understand this? Take the GAN mentioned earlier as an example: the objectives of the discriminator and generator are, respectively,

\begin{gather}\max_D\, \mathbb{E}_{\boldsymbol{x}\sim p(\boldsymbol{x})}[\log \sigma(D(\boldsymbol{x}))] + \mathbb{E}_{\boldsymbol{x}\sim q(\boldsymbol{x})}[\log (1 - \sigma(D(\boldsymbol{x})))] \\ \min_q\mathbb{E}_{\boldsymbol{x}\sim q(\boldsymbol{x})}[-D(\boldsymbol{x})] \end{gather}

The usual proof proceeds by showing that the optimal solution of $D$ is $\log\frac{p(\boldsymbol{x})}{q(\boldsymbol{x})}$, then substituting this into the generator's loss function to find that it minimizes the KL divergence between $q(\boldsymbol{x}),p(\boldsymbol{x})$, so the optimal solution is $q(\boldsymbol{x})=p(\boldsymbol{x})$. But the training procedure corresponding to such a proof should really be: first, for an arbitrary $q(\boldsymbol{x})$, fully solve step $\max\limits_D$ (the resulting $D$ should be a function of $q(x)$, or in other words a function of the generator's parameters $\boldsymbol{\theta}$), and only then carry out step $\min\limits_q$ — rather than the alternating training that is actually used in practice. Under the diffusion-model-based understanding, however, the alternation is built into the design from the start, so it is much more consistent with the actual training process.

In short, understanding GANs from the perspective of diffusion models isn't just a new way of thinking about GANs — it's a perspective that is more faithful to the actual training process. For instance, it lets us explain why a GAN's generator shouldn't be trained for too many steps at a time: equations $\eqref{eq:gan-g}$ and $\eqref{eq:gan-g0}$ are only equivalent when a single optimization step is taken; if the GAN needs multiple optimization steps for the generator, then equation $\eqref{eq:gan-g0}$ should be used as the loss function instead. In fact, equation $\eqref{eq:gan-g0}$ is exactly the $KL\left(q(x)\Vert q^{o}(x)\right)$ term that the author previously proposed in Understanding Generative Models via Variational Inference: A Unified View of VAE, GAN, AAE, ALI — a term that ensures the generator maintains "continuity" rather than merely "novelty."

Summary

This post introduced MonoFlow, which shows that a GAN can be understood as a diffusion ODE along a different time dimension, thereby establishing a new perspective on GANs grounded in diffusion models. In particular, this perspective is more faithful to the actual training process than the conventional derivation of GANs.