GAN Models from an Energy Perspective (II): GAN = ``Analysis'' + ``Sampling''

Energy Distribution Model

First, we have a set of data x_1, x_2, \dots, x_n \sim p(x). We hope to fit it using a probabilistic model. The model we choose is: q_{\theta}(x) = \frac{e^{-U_{\theta}(x)}}{Z_{\theta}} where U_{\theta} is an undetermined function with parameters \theta, which we call the “energy function,” and Z_{\theta} is the normalization factor (partition function): Z_{\theta} = \int e^{-U_{\theta}(x)}dx \label{eq:z} Such a distribution can be called an “energy distribution,” also known in physics as the “Boltzmann distribution.”

As for why we choose such an energy distribution, there are many explanations. It can be said to be inspired by physics, or by the principle of maximum entropy, or you can simply think it is because this distribution is relatively easy to handle. Regardless, this distribution is very common and practical; the softmax activation we use so frequently is actually an assumption of this distribution.

The current difficulty lies in how to solve for the parameters \theta, and the source of the difficulty is that the partition function [eq:z] is usually difficult to calculate explicitly. Of course, despite the practical difficulties in calculation, it does not prevent us from continuing the derivation.

Confrontation of Positive and Negative Phases

To solve for the parameters \theta, we first define the log-likelihood function: \mathbb{E}_{x\sim p(x)} \big[\log q_{\theta}(x)\big] We want it to be as large as possible, which means we want: L_{\theta}=\mathbb{E}_{x\sim p(x)} \big[-\log q_{\theta}(x)\big] to be as small as possible. To this end, we use gradient descent on L_{\theta}. We have: \begin{aligned} \nabla_{\theta}\log q_{\theta}(x)=&\nabla_{\theta}\log e^{-U_{\theta}(x)}-\nabla_{\theta}\log Z_{\theta}\\ =&-\nabla_{\theta} U_{\theta}(x)-\frac{1}{Z_{\theta}}\nabla_{\theta} Z_{\theta}\\ =&-\nabla_{\theta} U_{\theta}(x)-\frac{1}{Z_{\theta}}\nabla_{\theta} \int e^{-U_{\theta}(x)}dx\\ =&-\nabla_{\theta} U_{\theta}(x)+\frac{1}{Z_{\theta}} \int e^{-U_{\theta}(x)}\nabla_{\theta} U_{\theta}(x) dx\\ =&-\nabla_{\theta} U_{\theta}(x)+\int \frac{e^{-U_{\theta}(x)}}{Z_{\theta}}\nabla_{\theta} U_{\theta}(x) dx\\ =&-\nabla_{\theta} U_{\theta}(x)+\mathbb{E}_{x\sim q_{\theta}(x)}\big[\nabla_{\theta} U_{\theta}(x)\big] \end{aligned} Therefore: \nabla_{\theta} L_{\theta} = \mathbb{E}_{x\sim p(x)}\big[\nabla_{\theta} U_{\theta}(x)\big] - \mathbb{E}_{x\sim q_{\theta}(x)}\big[\nabla_{\theta} U_{\theta}(x)\big] \label{eq:q-grad} This means the update formula for gradient descent is: \theta \leftarrow \theta - \varepsilon \Big(\mathbb{E}_{x\sim p(x)}\big[\nabla_{\theta} U_{\theta}(x)\big] - \mathbb{E}_{x\sim q_{\theta}(x)}\big[\nabla_{\theta} U_{\theta}(x)\big]\Big) \label{eq:q-grad-gd} Note the characteristics of equation [eq:q-grad]: it is the difference between the mean of \nabla_{\theta} U_{\theta}(x) under the true distribution and the fitted distribution. This is the famous decomposition of “positive phase” and “negative phase” in machine learning. Equation [eq:q-grad] reflects the confrontation between the positive and negative phases, which some people correspond to our dreaming process.

Theoretical Analysis vs. Experimental Sampling

In fact, equations [eq:q-grad] and [eq:q-grad-gd] show that the theoretical analysis of the energy distribution model we initially assumed is not difficult. However, when it comes to implementation in experiments, we find that we must complete sampling from q_{\theta}: \mathbb{E}_{x\sim q_{\theta}(x)}. That is, given a specific U_{\theta}(x), we need to find a way to sample a batch of x from q_{\theta}(x)=e^{-U_{\theta}(x)}/Z_{\theta}.

However, at present, we do not have much experience in sampling from q_{\theta}(x)=e^{-U_{\theta}(x)}/Z_{\theta}. For us, the convenient sampling process is as follows: z\sim q(z),\quad x = G_{\varphi}(z) where q(z) represents the standard normal distribution. That is, we can sample a z from the standard normal distribution and then transform it into the desired x through a fixed model G_{\varphi}. This means the theoretical expression for this distribution is: q_{\varphi}(x) = \int \delta\big(x - G_{\varphi}(z)\big)q(z)dz \label{eq:q-varphi} The problem is that if we use q_{\varphi}(x) to replace the original q_{\theta}(x), sampling becomes convenient, but similar theoretical derivations become difficult. In other words, we simply cannot derive a result similar to [eq:q-grad-gd].

The Birth of GAN

So, a whimsical idea is: can we combine the two and let each play to its strengths where it excels?

Since \mathbb{E}_{x\sim q_{\theta}(x)} in equation [eq:q-grad-gd] is difficult to implement, why not just replace this part with \mathbb{E}_{x\sim q_{\varphi}(x)}: \theta \leftarrow \theta - \varepsilon \Big(\mathbb{E}_{x\sim p(x)}\big[\nabla_{\theta} U_{\theta}(x)\big] - \mathbb{E}_{x\sim q_{\varphi}(x)}\big[\nabla_{\theta} U_{\theta}(x)\big]\Big) Which is: \theta \leftarrow \theta - \varepsilon \Big(\mathbb{E}_{x\sim p(x)}\big[\nabla_{\theta} U_{\theta}(x)\big] - \mathbb{E}_{x=G_{\varphi}(z),z\sim q(z)}\big[\nabla_{\theta} U_{\theta}(x)\big]\Big) \label{eq:q-grad-gd-new} Now sampling is convenient, but the premise is that q_{\varphi}(x) must be close enough to q_{\theta}(x) (because q_{\theta}(x) is the standard, correct one). Therefore, we use KL divergence to measure the difference between the two: \begin{aligned} KL\big(q_{\varphi}(x)\big\Vert q_{\theta}(x)\big)=&\int q_{\varphi}(x) \log \frac{q_{\varphi}(x)}{q_{\theta}(x)}dx \\ =& - H_{\varphi}(X) + \mathbb{E}_{x\sim q_{\varphi}(x)}\big[U_{\theta}(x)\big]+\log Z_{\theta} \end{aligned} The premise for equation [eq:q-grad-gd-new] to be effective is that q_{\varphi}(x) and q_{\theta}(x) are close enough, meaning the above expression is small enough. For a fixed q_{\theta}(x), Z_{\theta} is a constant, so the optimization objective for \varphi is: \varphi =\mathop{\text{argmin}}_{\varphi} - H_{\varphi}(X) + \mathbb{E}_{x\sim q_{\varphi}(x)}\big[U_{\theta}(x)\big] \label{eq:varphi-gd} Here H_{\varphi}(X) = - \int q_{\varphi}(x) \log q_{\varphi}(x) dx represents the entropy of q_{\varphi}(x). - H_{\varphi}(X) wants the entropy to be as large as possible, which implies diversity; \mathbb{E}_{x\sim q_{\varphi}(x)}[U_{\theta}(x)] wants the image potential energy to be as small as possible, which implies authenticity.

On the other hand, note that equation [eq:q-grad-gd-new] is actually the gradient descent formula for the objective: \theta =\mathop{\text{argmin}}_{\theta} \mathbb{E}_{x\sim p(x)}\big[U_{\theta}(x)\big] - \mathbb{E}_{x=G_{\varphi}(z),z\sim q(z)}\big[U_{\theta}(x)\big] \label{eq:theta-gd} So we find that the entire process is actually the alternating gradient descent of [eq:theta-gd] and [eq:varphi-gd]. As mentioned in the first article, this objective for \theta may bring numerical instability. Based on the reasons stated in the first article, real samples should be near local minima, so we can add a gradient penalty term to [eq:theta-gd]. The final process is: \begin{aligned} \theta =& \mathop{\text{argmin}}_{\theta} \mathbb{E}_{x\sim p(x)}\big[U_{\theta}(x)\big] - \mathbb{E}_{x=G_{\varphi}(z),z\sim q(z)}\big[U_{\theta}(x)\big] + \lambda \mathbb{E}_{x\sim p(x)}\big[\Vert \nabla_x U_{\theta}(x)\Vert^2\big]\\ \varphi =& \mathop{\text{argmin}}_{\varphi} - H_{\varphi}(X) + \mathbb{E}_{x=G_{\varphi}(z),z\sim q(z)}\big[U_{\theta}(x)\big] \end{aligned} \label{eq:gan-energy} This is the GAN model based on gradient penalty. We already “brainstormed” it in “GAN Models from an Energy Perspective (I),” and now we have derived it from the mathematical analysis of energy models.

Therefore, GAN is essentially the result of energy models and sampling models each playing to their strengths and avoiding their weaknesses.

Maximum Entropy and Mutual Information

First, we can utilize the fact that x=G_{\varphi}(z): x=G_{\varphi}(z) implies the conditional probability q_{\varphi}(x|z) = \delta\big(x - G(z)\big), which is a deterministic model. It can also be understood as a Gaussian distribution \mathcal{N}(x;G_{\varphi}(z),0) with mean G(z) and variance 0.

Then we consider the mutual information I(X,Z): \begin{aligned} I_{\varphi}(X,Z)=&\iint q_{\varphi}(x|z)q(z)\log \frac{q_{\varphi}(x|z)}{q_{\varphi}(x)}dxdz\\ =&\iint q_{\varphi}(x|z)q(z)\log q_{\varphi}(x|z) dxdz - \iint q_{\varphi}(x|z)q(z) \log q_{\varphi}(x)dxdz\\ =&\int q(z)\left(\int q_{\varphi}(x|z)\log q_{\varphi}(x|z) dx\right)dz + H(X) \end{aligned} Now we have found the relationship between I_{\varphi}(X,Z) and H_{\varphi}(X). Their difference is: \int q(z)\left(\int q_{\varphi}(x|z)\log q_{\varphi}(x|z) dx\right)dz \triangleq -H_{\varphi}(X|Z) In fact, H_{\varphi}(X|Z) is called the “conditional entropy.”

If we are dealing with discrete distributions, then because x=G_{\varphi}(z) is deterministic, q_{\varphi}(x|z)\equiv 1, so H_{\varphi}(X|Z) is 0, meaning I_{\varphi}(X,Z)=H_{\varphi}(X). If it is a continuous distribution, as mentioned before, it can be understood as a Gaussian distribution \mathcal{N}(x;G_{\varphi}(z),0) with variance 0. We can first consider the case of constant variance \mathcal{N}(x;G(z),\sigma^2), where we find H_{\varphi}(X|Z) \sim \log \sigma^2 is a constant, and then let \sigma \to 0. Although the result is infinity, infinity is in principle not calculable, but in fact, the variance does not need to be zero; it only needs to be small enough that it is indistinguishable to the naked eye.

Therefore, in general, we can determine that the mutual information I_{\varphi}(X,Z) and entropy H_{\varphi}(X) differ only by an insignificant constant. Thus, in equation [eq:gan-energy], H_{\varphi}(X) can be replaced by I_{\varphi}(X,Z): \begin{aligned} \theta =& \mathop{\text{argmin}}_{\theta} \mathbb{E}_{x\sim p(x)}\big[U_{\theta}(x)\big] - \mathbb{E}_{x=G_{\varphi}(z),z\sim q(z)}\big[U_{\theta}(x)\big] + \lambda \mathbb{E}_{x\sim p(x)}\big[\Vert \nabla_x U_{\theta}(x)\Vert^2\big]\\ \varphi =& \mathop{\text{argmin}}_{\varphi} - I_{\varphi}(X,Z) + \mathbb{E}_{x=G_{\varphi}(z),z\sim q(z)}\big[U_{\theta}(x)\big] \end{aligned} \label{eq:gan-energy-2} Now we want to minimize - I_{\varphi}(X,Z), which is to maximize the mutual information I_{\varphi}(X,Z). Intuitively, this is not hard to understand because this term is used to prevent mode collapse. If mode collapse occurs, then almost any z generates the same x, and the mutual information between X and Z will certainly not be large.

However, changing the objective from H_{\varphi}(X) to I_{\varphi}(X,Z) seems like just a formal transformation and doesn’t appear to solve the problem yet. Fortunately, we have already researched maximizing mutual information. The method is in the section “The Essence of Mutual Information” in “Mutual Information in Deep Learning: Unsupervised Feature Extraction.”. In other words, there is already a solution for directly estimating mutual information; readers can refer to that article directly.

Mutual Information and Information Lower Bound

If an exact estimation of mutual information is not required, one can use the idea from InfoGAN to obtain a lower bound of mutual information and then optimize this lower bound.

Starting from the definition of mutual information: I_{\varphi}(X,Z)=\iint q_{\varphi}(x|z)q(z)\log \frac{q_{\varphi}(x|z)q(z)}{q_{\varphi}(x)q(z)}dxdz Let q_{\varphi}(z|x) = q_{\varphi}(x|z)q(z)/q_{\varphi}(x), which represents the exact posterior distribution. Then for any approximate posterior distribution p(z|x), we have: \begin{aligned} I_{\varphi}(X,Z)=&\iint q_{\varphi}(x|z)q(z)\log \frac{q_{\varphi}(z|x)}{q(z)}dxdz\\ =&\iint q_{\varphi}(x|z)q(z)\log \frac{p(z|x)}{q(z)}dxdz + \iint q_{\varphi}(x|z)q(z)\log \frac{q_{\varphi}(z|x)}{p(z|x)}dxdz\\ =&\iint q_{\varphi}(x|z)q(z)\log \frac{p(z|x)}{q(z)}dxdz + \int q_{\varphi}(x)KL\Big(q_{\varphi}(z|x) \Big\Vert p(z|x)\Big)dz\\ \geq &\iint q_{\varphi}(x|z)q(z)\log \frac{p(z|x)}{q(z)}dxdz\\ =& \iint q_{\varphi}(x|z)q(z)\log p(z|x) - \underbrace{\iint q_{\varphi}(x|z)q(z)\log q(z) dxdz}_{=\int q(z)\log q(z)dz \text{ is a constant}} \end{aligned} In other words, mutual information is greater than or equal to \iint q_{\varphi}(x|z)q(z)\log p(z|x) plus a constant. To maximize mutual information, we can consider maximizing this lower bound. Since p(z|x) is arbitrary, we can simply assume p(z|x)=\mathcal{N}\left(z;E(x),\sigma^2\right), where E(x) is an encoder with parameters. Substituting this and omitting redundant constants, we find it is equivalent to adding a loss term to the generator: \mathbb{E}_{z\sim q(z)} \big[\Vert z - E(G(z))\Vert^2\big]

Therefore, based on the InfoGAN information lower bound idea, equation [eq:gan-energy] becomes: \begin{aligned} \theta =& \mathop{\text{argmin}}_{\theta} \mathbb{E}_{x\sim p(x)}\big[U_{\theta}(x)\big] - \mathbb{E}_{z\sim q(z)}\big[U_{\theta}(G_{\varphi}(z))\big] + \lambda_1 \mathbb{E}_{x\sim p(x)}\big[\Vert \nabla_x U_{\theta}(x)\Vert^2\big]\\ \varphi,E =& \mathop{\text{argmin}}_{\varphi,E} \mathbb{E}_{z\sim q(z)}\big[U_{\theta}(G_{\varphi}(z)) + \lambda_2 \Vert z - E(G_{\varphi}(z))\Vert^2\big] \end{aligned} \label{eq:gan-energy-3}

By this point, we have completed the treatment of H_{\varphi}(X) from two perspectives, thereby completing the derivation of the entire GAN and energy model.

Introduction to MCMC

Actually, I only slightly understand the meaning of MCMC and do not fully grasp its methods and essence. This “introduction” is merely a basic overview of its concepts. MCMC stands for “Markov Chain Monte Carlo.” In my understanding, it is roughly this: it is difficult for us to directly sample from a given distribution q(x), but we can construct the following stochastic process: x_{n+1} = f(x_n, \alpha) \label{eq:suijidigui} where \alpha is an easily implemented random process, such as sampling from a binary or normal distribution. In this way, starting from some x_0, the resulting sequence \{x_1, x_2, \dots, x_n, \dots\} is random.

If it can be further proved that the stationary distribution of equation [eq:suijidigui] is exactly q(x), then it means the sequence \{x_1, x_2, \dots, x_n, \dots\} is a batch of samples sampled from q(x). Thus, sampling from q(x) is achieved, although the sampling results are arranged in a certain order.

Langevin Equation

A special case of equation [eq:suijidigui] is the Langevin equation: x_{t+1} = x_t - \frac{1}{2}\varepsilon \nabla_x U(x_t) + \sqrt{\varepsilon}\alpha,\quad \alpha \sim \mathcal{N}(\alpha;0,1) \label{eq:sde} It is also called a stochastic differential equation. When \varepsilon \to 0, its stationary distribution is exactly the energy distribution: p(x) = \frac{e^{-U(x)}}{Z} That is to say, given the energy function U(x), we can achieve sampling from the energy distribution through equation [eq:sde]. This is the original idea of MCMC sampling for energy distributions.

Of course, directly sampling x from the energy function and equation [eq:sde] might not be realistic because the dimension of x (in common scenarios, x represents an image) is too large, and controllability is difficult to guarantee. On the other hand, the last term of equation [eq:sde] is Gaussian noise, so as long as \varepsilon \neq 0, the result will inevitably be noisy, and image authenticity is hard to guarantee.

An interesting transformation is: we can consider MCMC sampling of z instead of x. Because in the previous model, we finally obtained both the energy function U_{\theta}(x) and the generative model G_{\varphi}(z), which means the energy function for z is: U_{\theta}(G_{\varphi}(z))

Note: This result is not strictly established; it can only be considered an empirical formula. Strictly speaking, it only holds when the Jacobian determinant of G is 1. I once discussed this with the author on GitHub, and he also pointed out that there is no rigorous theoretical derivation for this; it just comes from intuition. For details, refer to: https://github.com/ritheshkumar95/energy_based_generative_models/issues/4

With the energy function for z, we can implement MCMC sampling for z through equation [eq:sde]: z_{t+1} = z_t - \frac{1}{2}\varepsilon \nabla_z U_{\theta}(G_{\varphi}(z_t)) + \sqrt{\varepsilon}\alpha,\quad \alpha \sim \mathcal{N}(\alpha;0,1) \label{eq:sde-2} In this way, all the problems mentioned earlier disappear because the dimension of z is generally much smaller than x, and there is no need to worry about noise caused by \varepsilon \neq 0, because z itself is noise.

Better Truncation Techniques

At this point, readers whose minds are not yet confused might realize: isn’t the distribution of z just the standard normal distribution? Isn’t sampling easy? Why bother with a whole set of MCMC sampling?

In an ideal situation, the energy distribution corresponding to the energy function U_{\theta}(G_{\varphi}(z)) of z: q_{\theta,\varphi}(z)=\frac{e^{-U_{\theta}(G_{\varphi}(z))}}{Z} should indeed be the standard normal distribution q(z) we originally passed to it. But in fact, there is always a gap between ideal and reality. When we train a generative model using a standard normal distribution, the noise that produces real samples often tends to be narrower. This requires some truncation techniques, or rather, screening techniques.

For example, flow-based generative models often use the “annealing” technique after training, which means setting the variance of the noise smaller during generation to produce more stable samples; refer to “NICE in Detail: Basic Concepts and Implementation of Flow Models.” BigGAN, released last year, also discussed truncation techniques for noise in GANs.

If we believe in our model, and believe that both the energy function U_{\theta}(x) and the generative model G_{\varphi}(z) are valuable, then we have reason to believe that e^{-U_{\theta}(G_{\varphi}(z))}/Z will be a better distribution for z than the standard normal distribution (a distribution of z that can generate more realistic x, because it incorporates G_{\varphi}(z) into the definition of the distribution). Therefore, sampling from e^{-U_{\theta}(G_{\varphi}(z))}/Z will be superior to sampling from q(z). In other words, MCMC sampling [eq:sde-2] can improve the generation quality after sampling, as verified in the original paper. We can understand it as a better truncation technique.

More Efficient MALA

The sampling process [eq:sde-2] is actually still relatively inefficient. The original paper actually uses an improved version called MALA (Metropolis-adjusted Langevin algorithm), which introduces a screening process on top of [eq:sde-2]: \begin{aligned} \tilde{z}_{t+1} =& z_t - \frac{1}{2}\varepsilon \nabla_z U_{\theta}(G_{\varphi}(z_t)) + \sqrt{\varepsilon}\alpha,\quad \alpha \sim \mathcal{N}(\alpha;0,1)\\ \\ z_{t+1} =& \left\{\begin{aligned}&\tilde{z}_{t+1}, \quad \text{if } \beta < \gamma\\ &z_t, \quad \text{otherwise}\end{aligned}\right.,\quad \beta \sim U[0,1]\\ \\ \gamma =& \min\left\{1, \frac{q(\tilde{z}_{t+1})q(z_t|\tilde{z}_{t+1})}{q(z_t)q(\tilde{z}_{t+1}|z_t)}\right\} \end{aligned} Here: \begin{aligned} q(z)\propto&\, \exp\Big(-U_{\theta}(G_{\varphi}(z))\Big)\\ q(z'|z)\propto&\, \exp\left(-\frac{1}{2\varepsilon}\Vert z' - z + \varepsilon \nabla_z U_{\theta}(G_{\varphi}(z))\Vert^2\right) \end{aligned} That is, z_{t+1}=\tilde{z}_{t+1} is accepted with probability \gamma, and it remains unchanged with probability 1-\gamma. According to Wikipedia, such an improvement allows the sampling process to have a better chance of sampling high-probability samples, which means generating more realistic samples. (I don’t fully understand this set of theories, so I can only copy it!)