Cleverly Cutting Off Gradients: Implementing a GAN with a Single Loss

巧断梯度：单个loss实现GAN模型 苏剑林 · 2019-02-22 · Machine Learning GAN generative model Keras gradient

We know that for ordinary models, the usual workflow is to build the architecture, define a loss function, and just hand it to the optimizer for training. But GANs are different: generally speaking, they involve two distinct losses that need to be optimized alternately. The mainstream approach nowadays is to train the discriminator and the generator alternately at a 1:1 ratio (each trained once per step, and if necessary you can assign them different learning rates, i.e. TTUR). This alternating optimization means we need to feed in data twice (transferring it from memory to GPU memory), and run forward and backward passes twice.

If we could merge these two steps into a single optimization step, we would surely save time — this is what's known as synchronous training for GANs.

(Note: this post is not about introducing a new GAN variant, but about introducing a new way of writing GANs. It's a programming exercise, not an algorithmic one.)

If We Were Using TensorFlow

If we were working in TensorFlow, implementing synchronous training wouldn't be hard, because we would have already defined separate training ops for the discriminator and the generator (say D_solver and G_solver), and we could simply run

sess.run([D_solver, G_solver], feed_dict={x_in: x_train, z_in: z_train})

and be done with it. This relies on being able to access the discriminator's and generator's parameters separately, and being able to operate directly on sess.run.

A More General Approach

But what if we're using Keras? Keras has already wrapped up the training loop, so generally speaking we can't manipulate things at such a fine-grained level. So below we'll introduce a general technique: by defining just a single loss and handing it to the optimizer, we can implement GAN training. Along the way, we'll also learn how to manipulate the loss more flexibly in order to control gradients.

Optimizing the Discriminator

Let's take the hinge loss of a GAN as our example. It takes the form:

\begin{equation}\begin{aligned}D =& \mathop{\text{argmin}}_D \mathbb{E}_{x\sim p(x)}\big[\max\big(0, 1 + D(x)\big)\big]+\mathbb{E}_{z\sim q(z)}\big[\max\big(0, 1 - D(G(z))\big)\big]\\ G =& \mathop{\text{argmin}}_G \mathbb{E}_{z\sim q(z)}\big[D(G(z))\big] \end{aligned}\end{equation}

Note that $\mathop{\text{argmin}}_D$ means we need to keep $G$ fixed, because $G$ itself also has trainable parameters — if we didn't fix it, it would instead be $\mathop{\text{argmin}}_{D,G}$.

To fix $G$, besides the approach of "removing $G$'s parameters from the optimizer," we can also use stop_gradient to manually fix it:

\begin{equation}D,G = \mathop{\text{argmin}}_{D,G} \mathbb{E}_{x\sim p(x)}\big[\max\big(0, 1 + D(x)\big)\big]+\mathbb{E}_{z\sim q(z)}\big[\max\big(0, 1 - D(G_{ng}(z))\big)\big]\label{eq:dg-d}\end{equation}

Here

\begin{equation}G_{ng}(z)=\text{stop_gradient}(G(z))\end{equation}

This way, in equation $\eqref{eq:dg-d}$, even though we've simultaneously unfrozen the weights of $D,G$, continually optimizing equation $\eqref{eq:dg-d}$ will only change $D$, while $G$ stays fixed. This is because we're using a gradient-descent-based optimizer, and the gradient of $G$ has already been stopped — in other words, we can think of the gradient with respect to $G$ as having been forcibly set to zero, so its update is always zero.

Optimizing the Generator

Now that we've handled the optimization of $D$, what about $G$? stop_gradient makes it very convenient to fix the gradient of some inner part (e.g. the $G(z)$ of $D(G(z))$), but optimizing $G$ requires us to fix the outer $D$ instead, and there's no function that directly implements that. But don't despair — we can use a mathematical trick to transform the problem.

First, we need to be clear about what we want: we want the gradient of $G$ inside $D(G(z))$, but not the gradient of $D$. If we directly take the gradient of $D(G(z))$, we'll get the gradient of $D,G$ as well. What if we directly take the gradient of $D(G_{ng}(z))$? Then we'd only get the gradient of $D$, since $G$ has already been stopped. Now here's the key idea: if we subtract these two from each other, don't we end up with purely the gradient of $G$!

\begin{equation}D,G = \mathop{\text{argmin}}_{D,G} \mathbb{E}_{z\sim q(z)}\big[D(G(z)) - D(G_{ng}(z))\big]\label{eq:dg-g}\end{equation}

Now, optimizing equation $\eqref{eq:dg-g}$ leaves $D$ unchanged, while $G$ changes.

Note: there's no need to understand this construction via the chain rule — instead, understand it through the meaning of stop_gradient itself. For $L(D,G)$, regardless of the relationship between $G,D$, the full gradient is $(\nabla_D L, \nabla_G L)$; but once we stop the gradient of $G$, it's equivalent to forcibly setting the gradient with respect to $G$ to zero — that is, the gradient of $L(D,G_{ng})$ is effectively $(\nabla_D L, 0)$, so the gradient of $L(D,G)-L(D,G_{ng})$ is $(\nabla_D L, \nabla_G L) - (\nabla_D L, 0) = (0, \nabla_G L)$.

It's worth pointing out that if you directly evaluate this expression, the result is identically zero, since the two parts are identical and subtracting them naturally gives zero — but its gradient is not zero. In other words, this is a loss that is identically zero in value, yet whose gradient is not identically zero.

Combining Into a Single Loss

Good — now both equation $\eqref{eq:dg-d}$ and equation $\eqref{eq:dg-g}$ have simultaneously unfrozen $D,G$, and both are $\text{argmin}$, so we can combine the two steps into a single loss:

\begin{equation}\begin{aligned}D,G = \mathop{\text{argmin}}_{D,G}&\,\mathbb{E}_{x\sim p(x)}\big[\max\big(0, 1 + D(x)\big)\big]+\mathbb{E}_{z\sim q(z)}\big[\max\big(0, 1 - D(G_{ng}(z))\big)\big]\\ &\, + \lambda\, \mathbb{E}_{z\sim q(z)}\big[D(G(z)) - D(G_{ng}(z))\big]\label{eq:dg-dg}\end{aligned}\end{equation}

Writing out this loss lets us optimize the discriminator and generator simultaneously, without needing alternating training, and the effect is essentially equivalent to 1:1 alternating training. The role of $\lambda$ here is to set the ratio between the discriminator's and generator's learning rates to $1:\lambda$.

Reference code: https://github.com/bojone/gan/blob/master/gan_one_step_with_hinge_loss.py

Summary

This post mainly introduced a small trick for implementing GANs, which lets us write a single model with a single loss to carry out GAN training. At its core, it's a technique for manually controlling gradients using stop_gradient, which may well be useful for other tasks too.

So, from now on, this is the way I'll be writing GANs — it saves both effort and time. Of course, in theory this approach does consume a bit more GPU memory, so you could call it trading space for time.