A GAN That Needs No Lipschitz Constraint and Yet Doesn't Suffer From Vanishing Gradients—Interested?
At some point, without quite realizing it, I discovered that I too had fallen into the giant pit that is GANs. Sigh. I'll try to climb out of it sooner rather than later...
This post introduces a new GAN framework I recently submitted to arXiv. It centers on a new way of understanding probability divergences, and from this understanding derives a new GAN. The whole piece leans fairly theoretical—I've tried to give a complete argument for all the relevant properties of this GAN, and I believe the result is theoretically self-contained.
Paper link: https://papers.cool/arxiv/1811.07296
Let me state the conclusions up front:
1. The paper offers a direct way of analyzing and constructing probability divergences, which simplifies the process of building new GAN frameworks.
2. It derives a GAN framework called GAN-QP $\eqref{eq:gan-gp-gd}$. This GAN needs no Lipschitz constraint the way WGAN does, and yet it does not suffer from the vanishing-gradient problem that SGAN has. Experiments show it performs at least as well as, and sometimes better than, WGAN.
GAN-QP results
The largest experiment in the paper is 512×512 face generation (CelebA HQ), which sufficiently demonstrates the model's effectiveness (the results aren't perfect, but the model is remarkably simple). If you're interested, read on. more
Facing the Dual Space Head-On
Building a GAN framework typically involves three steps:
1. Find a good probability divergence;
2. Work out its dual form;
3. Turn it into a min-max game.
The issue is: what actually matters for training is steps two and three—step one isn't really necessary.
In fact, it's quite hard to define a new divergence from scratch in the primal space, and even once you've defined one, it's not necessarily easy to convert it into a dual form. However, we can analyze things directly in the dual space, and by doing so we can uncover a whole family of new, well-behaved divergences. In other words, we can directly argue, within the dual space, whether a given expression satisfies the definition of a divergence, and thereby obtain an optimizable objective directly—without needing to care whether it happens to be the JS divergence or the Wasserstein distance.
Let's walk through some examples to illustrate this approach.
Divergence
First, let's state the definition of a divergence:
If $\mathcal{D}[p, q]$ is a scalar function of $p,q$, and it satisfies:
1. $\mathcal{D}[p, q]\geq 0$ always holds;
2. $\mathcal{D}[p, q]=0\Leftrightarrow p=q$.
then we call $\mathcal{D}[p, q]$ a divergence of $p,q$. The main difference between a divergence and a "distance" is that a divergence need not satisfy the triangle inequality, nor does it need to be symmetric. But a divergence retains the most basic property of measuring a discrepancy, so we can use it to measure how different $p,q$ are from each other.
SGAN
Basic Definition
Let's start with the discriminator loss in SGAN, defined as
\begin{equation}\mathcal{D}[p(x),q(x)] = \max_T\, \frac{1}{2}\mathbb{E}_{x\sim p(x)}[\log \sigma(T(x))] + \frac{1}{2}\mathbb{E}_{x\sim q(x)}[\log (1 - \sigma(T(x)))] + \log 2\label{eq:js-dual-2}\end{equation}
This is, in fact, the dual form of the JS divergence. But we can prove directly, from this definition alone, that it is a divergence, and then discuss the properties of this divergence in its own right—without ever needing to know that it's the JS divergence.
How do we prove it? We just need to show that this expression satisfies the two requirements for a divergence stated above. Note that, following our logic, we don't know it's the JS divergence—but we can prove mathematically that it is a divergence.
If the reader really understands what equation $\eqref{eq:js-dual-2}$ is saying, the proof isn't hard. Equation $\eqref{eq:js-dual-2}$ first defines an expectation, and then takes the maximum over $T$ (more precisely, the "supremum")—only the result after taking the maximum is the divergence. Let me stress this again: "the result after taking the maximum is the divergence"; the expression $\frac{1}{2}\mathbb{E}_{x\sim p(x)}[\log \sigma(T(x))] + \frac{1}{2}\mathbb{E}_{x\sim q(x)}[\log (1 - \sigma(T(x)))] + \log 2$ by itself is not the divergence.
The full proof is a bit lengthy, so I won't reproduce it in full here—please refer to the appendix of the original paper. Or look at the WGAN section below, since that proof is comparatively simpler.
Adversarial Network
Once we have a divergence, we can train a generative model by shrinking the divergence between the two probability distributions. That is, what we need to do next is
\begin{equation}\min_{G} \mathcal{D}[p(x),q(x)]\end{equation}
Note that $\mathcal{D}[p(x),q(x)]$ is realized through the $\max_{T}$ operation, so combined together this becomes a min-max process. For the example above, this is equivalent to:
\begin{equation}G,T = \mathop{\text{argmin}}_G\mathop{\text{argmax}}_T\, \mathbb{E}_{x\sim p(x)}[\log \sigma(T(x))]+ \mathbb{E}_{x=G(z),z\sim q(z)}[\log (1 - \sigma(T(x)))]\label{eq:js-min-max}\end{equation}
This is SGAN.
So we see that the GAN process really consists of just two steps: 1. define a divergence via $\max$; 2. shrink the divergence between the two distributions via $\min$. The new perspective here is treating $\max$ directly as part of the definition of the divergence.
Performance Analysis
We know SGAN can suffer from vanishing gradients—why is that? Let's consider an extreme case:
\begin{equation}p(x)=\delta(x-\alpha),q(x)=\delta(x-\beta)\end{equation}
where $\alpha\neq\beta$. In this case, the two distributions are each just a point mass, with absolutely no overlap. Substituting this into $\eqref{eq:js-dual-2}$, we get
\begin{equation}\mathcal{D}[p(x),q(x)] = \max_T\, \frac{1}{2}[\log \sigma(T(\alpha))] + \frac{1}{2}[\log (1 - \sigma(T(\beta)))] + \log 2\end{equation}
Note that we have no constraint at all on $T$, so to maximize this we can let $T(\alpha)\to +\infty,T(\beta)\to -\infty$, giving a supremum equal to the constant $\log 2$. That is, in this case $\mathcal{D}[p(x),q(x)]=\log 2$.
This says that for two distributions with essentially no overlap, the divergence defined by $\eqref{eq:js-dual-2}$ gives a constant value $\log 2$—and a constant means zero gradient, so there's nothing to optimize. And the two WGAN papers show that "no overlap" is theoretically quite common in GANs, so this is an inherent flaw of SGAN.
The General f-Divergence
The sections above have laid out the full procedure of this way of thinking:
1. We define a mathematical expression via $\max$, and then prove directly, from mathematics alone, that it is a divergence—without worrying about what it's called;
2. We minimize this divergence via $\min$, and combined together this forms a min-max process, giving us a GAN;
3. To check how this divergence behaves in extreme cases, we can test it with $p(x)=\delta(x-\alpha),q(x)=\delta(x-\beta)$.
The argument above for SGAN generalizes in parallel to all f-GANs (see f-GAN: A Production Line for GAN Models). The various f-divergences turn out not to differ in any essential way—they all share the same inherent flaw (either vanishing gradients or exploding gradients).
WGAN
Basic Definition
Now let's turn to a new class of divergence: the Wasserstein distance. Note that the Wasserstein distance is a rigorous distance satisfying an axiomatic definition, though here we only care about its divergence property. Define
\begin{equation}\mathcal{W}[p(x),q(x)] = \max_{T,\,\Vert T\Vert_L \leq 1}\, \mathbb{E}_{x\sim p(x)}[T(x)] - \mathbb{E}_{x\sim q(x)}[T(x)]\label{eq:wd-dual}\end{equation}
where
\begin{equation}\Vert T\Vert_L = \max_{x\neq y} \frac{|T(x)-T(y)|}{d(x,y)}\end{equation}
and $d(x, y)$ is any off-the-shelf distance metric.
We can prove directly that this is a divergence. This proof is fairly classic, so let me lay it out here:
1. No matter what $p(x),q(x)$ is, if we simply let $T(x)\equiv 0$, we get $\mathbb{E}_{x\sim p(x)}[T(x)] - \mathbb{E}_{x\sim q(x)}[T(x)]=0$. Since the definition of a divergence requires taking the maximum over all possible $T$, the divergence is at least no smaller than 0—this proves non-negativity.
2. To prove that when $p(x)=q(x)$, $\mathcal{W}[p(x),q(x)]=0$—that is, $\mathcal{W}[p(x),p(x)]=0$—this is almost trivially true.
3. To prove that when $p(x)\neq q(x)$ (strictly speaking, when the measure of the set where they differ is positive), $\mathcal{W}[p(x),q(x)] > 0$: this one is slightly harder, but still simple. Just let $T_0(x) = \text{sign}(p(x) - q(x))$, and then clearly
\begin{equation}\begin{aligned}&\mathbb{E}_{x\sim p(x)}[T_0(x)] - \mathbb{E}_{x\sim q(x)}[T_0(x)] \\ > =&\int (p(x)-q(x))\cdot \text{sign}(p(x) - q(x)) dx > 0\end{aligned}\end{equation}
This gives us a direct proof that $\mathcal{W}[p(x),q(x)]$ satisfies the definition of a divergence.
Adversarial Network
Likewise, once we have this new divergence, we can define a new GAN:
\begin{equation}G,T = \mathop{\text{argmin}}_G\mathop{\text{argmax}}_{T,\,\Vert T\Vert_L \leq 1}\, \mathbb{E}_{x\sim p(x)}[T(x)] - \mathbb{E}_{x=G(z),z\sim q(z)}[T(x)]\label{eq:wd-min-max}\end{equation}
This is WGAN. Relevant references on this blog include The Art of Mutual Adversary: From Scratch All the Way to WGAN-GP and WGAN-div: An Unsung Hero That Patched Up WGAN.
Performance Analysis
Likewise, testing the $\mathcal{W}[p(x),q(x)]$ divergence with $p(x)=\delta(x-\alpha),q(x)=\delta(x-\beta)$, we get
\begin{equation}\mathcal{W}[p(x),q(x)] = \max_{T,\,\Vert T\Vert_L \leq 1} T(\alpha) - T(\beta)\end{equation}
Note that we have the Lipschitz constraint $\Vert T\Vert_L \leq 1$, which means $|T(\alpha) - T(\beta)| \leq d(\alpha, \beta)$, and equality can be achieved, so
\begin{equation}\mathcal{W}[p(x),q(x)] = d(\alpha,\beta)\end{equation}
The result is not a constant, so even in this extreme case we can still shrink the distance between the two distributions. So from this perspective, WGAN is better than SGAN.
The Lipschitz Constraint
The problem WGAN leaves behind is how to impose the Lipschitz constraint on the discriminator. Currently there are three approaches: weight clipping, gradient penalty, and spectral normalization—see Lipschitz Constraints in Deep Learning: Generalization and Generative Models and WGAN-div: An Unsung Hero That Patched Up WGAN.
Weight clipping has basically fallen out of use. The gradient penalty is, in principle, just a heuristic method, with its own irrationalities, and computing gradients is usually slow. Spectral normalization looks the most elegant and currently performs quite well too, though it might be overly restrictive. For further discussion see WGAN-div: An Unsung Hero That Patched Up WGAN.
A New Divergence, A New GAN
So the situation is: SGAN can suffer from vanishing gradients, and while WGAN is good, it needs an extra Lipschitz constraint. Naturally we'd ask: is there a GAN that needs no Lipschitz constraint and yet does not suffer from vanishing gradients? Can we have our cake and eat it too?
It turns out we can. Let me help you find one. Actually, not just one—let's find a whole family.
The Quadratic Potential Divergence
Basic Definition
The divergence I'm about to present takes the following form:
\begin{equation}\begin{aligned}&\mathcal{L}[p(x),q(x)] \\ =& \max_{T}\, \mathbb{E}_{(x_r,x_f)\sim p(x_r)q(x_f)}\left[T(x_r,x_f)-T(x_f,x_r) - \frac{(T(x_r,x_f)-T(x_f,x_r))^2}{2\lambda d(x_r,x_f)}\right]\end{aligned}\label{eq:qp-dual}\end{equation}
where $\lambda > 0$ is a hyperparameter, and $d$ can be any distance metric.
This looks like it's just WGAN with a quadratic potential term added on top, hence I call it the quadratic potential divergence (QP-div).
The paper's appendix has already proven that equation $\eqref{eq:qp-dual}$ is indeed a divergence.
Performance Analysis
Testing this divergence with $p(x)=\delta(x-\alpha),q(x)=\delta(x-\beta)$, we get
\begin{equation}\mathcal{L}[p(x),q(x)] = \max_{T}\, T(\alpha,\beta)-T(\beta,\alpha) - \frac{(T(\alpha,\beta)-T(\beta,\alpha))^2}{2\lambda d(\alpha,\beta)}\end{equation}
Let $z = T(\alpha,\beta)-T(\beta,\alpha)$, and we get $z - \frac{z^2}{2\lambda d(\alpha,\beta)}$—looks familiar, doesn't it? This is just the maximum of a quadratic function, and the maximum value is $\frac{1}{2}\lambda d(\alpha,\beta)$, so we have
\begin{equation}\mathcal{L}[p(x),q(x)] = \frac{1}{2}\lambda d(\alpha,\beta)\end{equation}
This is basically the same as WGAN—even for extreme distributions, there's no risk of vanishing gradients. We really can have our cake and eat it too.
GAN-QP
Adversarial Network
Once we have the divergence, we can build the adversarial network. The final form we arrive at is
\begin{equation}\begin{aligned}&T= \mathop{\text{argmax}}_T\, \mathbb{E}_{(x_r,x_f)\sim p(x_r)q(x_f)}\left[T(x_r,x_f)-T(x_f,x_r) - \frac{(T(x_r,x_f)-T(x_f,x_r))^2}{2\lambda d(x_r,x_f)}\right] \\ &G = \mathop{\text{argmin}}_G\,\mathbb{E}_{(x_r,x_f)\sim p(x_r)q(x_f)}\left[T(x_r,x_f)-T(x_f,x_r)\right] \end{aligned}\label{eq:gan-gp-gd}\end{equation}
I call this GAN-QP in the paper.
Note that we should not add the quadratic term $-\frac{(T(x_r,x_f)-T(x_f,x_r))^2}{2\lambda d(x_r,x_f)}$ into the generator's loss (in principle this isn't an issue, but in practice, when optimizing with gradient descent, it causes trouble), because the denominator of this term is $d(x_r,x_f)$. Once you minimize the quadratic term, it's equivalent to minimizing $d(x_r,x_f)$, i.e., using $d(x_r,x_f)$ to measure the discrepancy between images—which is not sensible.
Analysis of the Solution
Using the calculus of variations, we can prove (again, in the appendix) that the discriminator's optimal solution is:
\begin{equation}\frac{p(x_r)q(x_f) - p(x_f)q(x_r)}{p(x_r)q(x_f) + p(x_f)q(x_r)} = \frac{T(x_r,x_f)-T(x_f,x_r)}{\lambda d(x_r, x_f)}\label{eq:opt-t}\end{equation}
From this optimal solution, we can draw two conclusions. First, it's not hard to prove that the optimal solution automatically satisfies
\begin{equation}-1 \leq \frac{T(x_r,x_f)-T(x_f,x_r)}{\lambda d(x_r, x_f)}\leq 1\end{equation}
That is, the optimal solution automatically satisfies the Lipschitz constraint. So we can think of GAN-QP as a scheme with a self-adaptive Lipschitz constraint.
Second, substituting the optimal solution into the generator's loss, we obtain the overall objective:
\begin{equation}\lambda\iint p(x_r)q(x_f)\frac{p(x_r)q(x_f) - p(x_f)q(x_r)}{p(x_r)q(x_f) + p(x_f)q(x_r)} d(x_r, x_f) dx_r dx_f\end{equation}
This is also a probability divergence, and we've also proven theoretically that it does not suffer from vanishing or exploding gradients (this is related to the Cauchy–Schwarz inequality). Furthermore, we can see that $\lambda$ is just a scaling factor and isn't actually important—so GAN-QP is robust to $\lambda$, and $\lambda$ won't have a noticeable effect on the model's performance.
Experimental Results
The paper compares GAN-QP against a number of GANs on the CelebA HQ dataset, showing that GAN-QP can match, or even surpass, current state-of-the-art models.
Note that in model $\eqref{eq:gan-gp-gd}$, $T$ is a bivariate function of $(x_r,x_f)$, but experiments show that the simplest univariate special case $T(x_r,x_f)\equiv T(x_r)$ suffices—i.e., using $T(x_r) - T(x_f)$ for $T(x_r,x_f) - T(x_f,x_r)$ is enough, and switching to a bivariate function doesn't bring any noticeable improvement (though it's possible I just didn't tune it well). In that case, the form ends up looking very similar to WGAN-GP, but with a more complete theoretical foundation.
Code is open-sourced at: https://github.com/bojone/gan-qp
128×128
At 128×128 resolution, we ran fairly comprehensive comparisons, using FID as the quantitative metric. The results are shown below:
FID curves for different GANs
And the table below:
$$\begin{array}{c|ccccc} \hline \hline & \text{GAN-QP-L1/L2} & \text{WGAN-GP} & \text{WGAN-SN} & \text{SGAN-SN} & \text{LSGAN-SN} \\ \hline \text{Best FID} & 45.0 / 44.7 & 55.5 & 47.8 & 44.5 & 45.8\\ \hline \text{Speed} & 1\text{x} / 1\text{x} & 1.5\text{x} & 1\text{x} & 1\text{x} & 1\text{x}\\ \hline \hline \end{array}$$
256 and 512
At 128 resolution, GAN-QP and SGAN-SN gave the best performance, but at 256×256 resolution, a gap opens up between them:
$$\begin{array}{c|ccccc} \hline \hline & \text{GAN-QP} & \text{SGAN-SN} \\ \hline \text{Best FID} & 22.7 & 27.9\\ \hline \hline \end{array}$$
The largest experiment I ran for GAN-QP was 512×512 face generation, and the results are still pretty good, with a final FID of 26.44:
512x512 face generation results
Summary of the Paper
This work grew out of my own thinking about probability divergences, in an attempt to find a more direct way of understanding them—and it was also inspired in part by WGAN-div.
Fortunately, I managed to push this line of thinking through to completion, and along the way arrived at some new results, which I've now submitted to arXiv for everyone's reference. I hope to get feedback from more experienced researchers. In fact, based on a similar line of thinking, we can construct many similar divergences—for instance, replacing the square with a fourth power, sixth power, and so on—though the theoretical analysis becomes progressively harder.
Given limited compute, and the fact that I'm not a dedicated GAN researcher, the experiments here may not be as thorough as they could be—I've kept them just enough to support the conclusions, so please bear with me. Of course, I'd also welcome any guidance from readers.
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.