We know that for ordinary models, we simply build the architecture, define the loss, and pass it to the optimizer for training. However, GANs are different. Generally, they involve two different losses that need to be optimized alternately. The mainstream approach currently is to train the discriminator and the generator in a 1:1 alternating frequency (training each once, and if necessary, setting different learning rates for both, i.e., TTUR). Alternating optimization means we need to pass data twice (from RAM to VRAM) and perform forward and backward passes twice.
If we could merge these two steps into a single optimization step, it would certainly save time. This is known as the synchronous training of GANs.
(Note: This article does not introduce a new GAN, but rather a new way to write GANs. This is a programming problem, not an algorithmic one.)
If in TensorFlow
In TensorFlow, implementing synchronous training is not difficult
because we have already defined the training operators for the
discriminator and the generator (let’s assume they are
D_solver and G_solver). Then, we can simply
execute:
sess.run([D_solver, G_solver], feed_dict={x_in: x_train, z_in: z_train})This is possible because we can separately access the parameters of
the discriminator and generator and directly operate on
sess.run.
A More General Method
But what if we are using Keras? Keras encapsulates the workflow, and generally, we cannot operate with such fine granularity. Therefore, below we introduce a general trick: by defining only a single loss and passing it to the optimizer, we can achieve GAN training. At the same time, from this trick, we can learn how to manipulate losses more flexibly to control gradients.
Optimization of the Discriminator
Let’s take the Hinge Loss of GAN as an example. Its form is: \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} Note that \mathop{\text{argmin}}_D implies that G must be fixed, because G itself has optimizable parameters. If it weren’t fixed, it would be \mathop{\text{argmin}}_{D,G}.
To fix G, besides the method of
"removing G’s parameters from the
optimizer," we can also use stop_gradient to manually fix
it: 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} Here, G_{ng}(z)=\text{stop\_gradient}(G(z)) In
this way, in Equation [eq:dg-d], although we have released the
weights of both D and G, by continuously optimizing Equation [eq:dg-d], only D will change, while G will not. This is because we are using a
gradient-descent-based optimizer, and the gradient of G has been stopped. In other words, we can
understand it as the gradient of G
being forcibly set to 0, so its update amount is always 0.
Optimization of the Generator
Now that the optimization of D is
solved, what about G? While
stop_gradient can easily help us fix the gradients of
internal parts (such as G(z) in D(G(z))), the optimization of G requires us to fix the external D, for which there is no direct function. But
do not be discouraged; we can use a mathematical trick to transform
it.
First, we must be clear: we want the gradient of G inside D(G(z)), but we do not want the gradient of D. If we directly calculate the gradient of D(G(z)), we get the gradients of both D and G simultaneously. What if we calculate the gradient of D(G_{ng}(z))? We only get the gradient of D because G has been stopped. Now, here is the key point: by subtracting these two, don’t we get the pure gradient of G! 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} Now, by optimizing Equation [eq:dg-g], D will not change, and what changes is G.
Note: You don’t need to understand this formulation from the chain rule; rather, understand it through the inherent meaning of
stop_gradient. For L(D,G), regardless of the relationship between G and D, the full gradient is (\nabla_D L, \nabla_G L). After stopping the gradient of G, it is equivalent to forcibly setting the gradient of G to 0, meaning the gradient of L(D,G_{ng}) is actually (\nabla_D L, 0). Therefore, 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 is worth mentioning that if you directly output this expression, the result is identically 0 because the two parts are the same, and subtracting them naturally yields 0. However, its gradient is not 0. That is to say, this is a loss that is identically 0, but its gradient is not identically 0.
Synthesizing a Single Loss
Alright, now both Equation [eq:dg-d] and Equation [eq:dg-g] have released D and G simultaneously, and both are \text{argmin} operations. Therefore, we can combine the two steps into a single loss: \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] \end{aligned} \label{eq:dg-dg} By writing this loss, we can complete the optimization of both the discriminator and the generator simultaneously without alternating training. The effect is essentially equivalent to 1:1 alternating training. The introduction of \lambda acts to set the ratio of the learning rates of the discriminator and the generator to 1:\lambda.
Reference Code: https://github.com/bojone/gan/blob/master/gan_one_step_with_hinge_loss.py
Summary
This article introduced a small trick for implementing GANs, allowing
us to write a single model and use a single loss to achieve GAN
training. It is essentially a technique for manually controlling
gradients using stop_gradient, which may also be useful in
other tasks.
So, in the future, I will use this way of writing GANs—it saves effort and time. Of course, theoretically, this approach requires more video memory, which can be seen as sacrificing space for time.