Luke Metz, Ben Poole, David Pfau, Jascha Sohl-Dickstein
8 min
Generative Adversarial Networks (GANs) are notorious for being difficult to train. Common failure modes include mode collapse—where the generator produces only a small subset of the target distribution—and training instability, where the generator and discriminator oscillate without converging. The researchers seek to stabilize the training process and improve the diversity of the generated samples by changing how the generator perceives the discriminator's response.
Standard GAN training involves alternating gradient steps for the generator and discriminator. The authors introduce a "surrogate objective" that effectively allows the generator to "look ahead" at how the discriminator will respond to its updates. By unrolling the discriminator's optimization for K steps (differentiating through the optimization process itself), the generator's objective is redefined as the result of this unrolled optimization. This creates a smoother, more informative training signal that discourages the generator from collapsing into suboptimal, narrow points of the data distribution.
The researchers demonstrate that unrolling is highly effective at stabilizing GAN training across various complex tasks, including a synthetic mixture of Gaussians and recurrent neural networks (RNNs) trained on MNIST digits. Experiments consistently show that increasing the number of unrolling steps leads to better coverage of the data distribution, effectively mitigating mode collapse. Furthermore, quantitative metrics such as reconstruction MSE and pairwise distance distributions on the CIFAR-10 dataset confirm that unrolled GANs produce more diverse samples that better reflect the statistical characteristics of the training data compared to standard GANs.
This approach bridges a gap between theoretical expectations of GAN convergence and practical training difficulties. By acknowledging the iterative, game-theoretic nature of GANs and incorporating that dynamic into the gradient calculation, it offers a robust tool for training more stable generative models. While computationally more expensive than standard training, the technique provides a path toward higher-quality, more diverse generative outputs without needing to rely solely on heuristic tricks or rigid architectural constraints.
We introduce a method to stabilize Generative Adversarial Networks (GANs) by defining the generator objective with respect to an unrolled optimization of the discriminator. This allows training to be adjusted between using the optimal discriminator in the generator's objective, which is ideal but infeasible in practice, and using the current value of the discriminator, which is often unstable and leads to poor solutions. We show how this technique solves the common problem of mode collapse, stabilizes training of GANs with complex recurrent generators, and increases diversity and coverage of the data distribution by the generator.
Alex: That's exactly how the paper describes the extra term. If the generator starts collapsing onto one point, the unrolled discriminator reacts by marking that point as fake. That raises the generator's loss, which discourages the collapse. The real discriminator is still updated the ordinary way. Only the generator gets the look-ahead.
Sam: What happens at the extremes of that dial?
Alex: With zero unrolling steps, you recover standard GAN training. As the number of steps grows without limit, you approach the ideal objective, training against the optimal discriminator. Everything in between interpolates, and the cost per step rises with each extra unrolling step.
Sam: Let's get to evidence. What's the clearest demonstration?
Alex: A toy one. The data is eight Gaussian blobs arranged in a ring. A standard GAN rotates around the ring, putting its mass on one blob at a time, and never settles. With ten unrolling steps, the generator spreads out quickly and converges to all eight.
Sam: Toy problems are where every method looks good. What about a model people actually struggle with?
Alex: They tried a deliberately mismatched pair. The generator is a recurrent network, an LSTM, drawing handwritten digits one column at a time. The discriminator is convolutional. Without unrolling, it collapses and cycles through blob-like proto-digits. With twenty unrolling steps, the samples disperse and appear to cover the digit distribution. That judgment is visual, though.
Sam: Visual judgment is exactly what misses mode dropping. Do they count modes anywhere?
Alex: Yes, with a dataset called stacked MNIST. They stack three random handwritten digits into the three color channels of one image. That gives 1,000 distinct combinations, each one a mode. A pretrained digit classifier labels the generated samples, so you can count how many combinations ever appear.
Sam: And how many combinations do the models find?
Alex: It depends heavily on discriminator size. With a discriminator a quarter the generator's size, the standard GAN averaged about 31 modes. With ten unrolling steps, about 327. With a discriminator half the generator's size, the gap narrows. The standard GAN found about 628 modes, against about 817 with unrolling.
Sam: So unrolling helps most when the discriminator is weak. Is that their reading too?
Alex: It is. They believe unrolling effectively increases the discriminator's capacity, because it can react to whatever specific trick the generator is using. There's a catch, though. They report that reasonably sized models covered all 1,000 modes even without unrolling. So they had to use smaller models to see an effect at all.
Sam: That's a real scoping caveat. Does it hold up on natural images, where you can't count modes?
Alex: On CIFAR-10, a standard natural-image dataset, the samples look about the same whatever the number of unrolling steps. So they test diversity another way. They take real training images and search the generator's input space for its closest possible reconstruction. If a model has dropped a region of the data, it can't reproduce images from that region.
Sam: And which models reconstruct best?
Alex: Each configuration was run with five random seeds. Average reconstruction error falls from about 0.023 with no unrolling to about 0.018 with ten steps. The more telling comparison is head to head. The standard GAN gave the best reconstruction on under one percent of images. The ten-step model won on about 61 percent.
Sam: So the standard GAN rarely wins outright. Is there anything independent of that search procedure?
Alex: Yes. They compare distances between random pairs of samples. Collapse squeezes samples together, so distances shrink. The standard GAN skews toward short distances, and with more unrolling the distribution looks more like the real data's.
Sam: Is it monotonic, though? More unrolling always better?
Alex: No. On a colored-digit task, where they measure how well the generator recovers the distribution of digit colors, one unrolling step actually hurt. They suspect it introduces oscillation. Five to ten steps did best there. And the authors are candid overall. They say each measure is suggestive, none is conclusive alone, but together they find them compelling.
Sam: What's the main price someone pays to use this?
Alex: Compute. Each training step costs more in proportion to the number of unrolling steps. The authors say one step can be enough for some architectures, while unstable ones, like the recurrent generator, need more. They mention early hints of a cheaper approximation, but say that needs further investigation.
Sam: So who should sit down with the full paper, and where should they start?
Alex: Anyone training GANs, or any two-player learning setup. Start with the method section on the missing gradient term, where they split the generator's gradient into the standard piece and the opponent-response piece. Then read the stacked MNIST results across discriminator sizes. That shows where the method helps and where it barely matters.
Sam: And for everyone else, the one line to carry around?
Alex: A generator that only sees its opponent's current move collapses. One that anticipates the reply spreads out.
Sam: Thinking a few moves ahead, it turns out, is worth the compute.