Paper Overview

While surveying work on generative models, I came across an interesting idea called GAN-supervised learning. Earlier reviews covered several papers at once, but this article focuses on a single paper. The name might sound as though the work simply trains a generative network. Instead, GAN-supervised learning means using images produced by a GAN as supervision while optimizing another model. To understand the idea, let us first look briefly at classical congealing.

What Is Congealing?

Consider the familiar MNIST handwritten-digit dataset, simplified to contain only zeros and ones. Even when the same person writes the same digit, the result varies in pixel position, tilt, width, height, and many other ways. In other words, the distribution associated with a single label is spread across many forms.

In the figure, the $j^{th}$ image $I^j$ is transformed by matrix $U^j$ into $I^{j\prime}$. Indexing each pixel location by $i$ lets us denote the same location in every image as $x_i^j$. We stack all transformed digit images and estimate the empirical entropy of the pixels at each shared location.

$$\hat{H}(x_i) = -\left( \frac{N_0}{N} \log _2 \frac{N_0}{N} + \frac{N_1}{N} \log _2 \frac{N_1}{N} \right)$$

Here $N_0$ and $N_1$ are the frequencies of black ($0$) and white ($1$) at the same pixel location across the images. Dividing each by the total number of images maps it to a probability. The objective minimizes this entropy so that the transformations make all images resemble the same canonical form as closely as possible.

Returning to the Abstract

Extending that framework leads to the paper's central question: if we can generate a consistent target mode for an unaligned dataset, could that mode provide supervision for training a spatial transformer network (STN)?

The method improved precise correspondence—the agreement between aligned regions—by roughly $3\times$ over the previous state of the art. The resulting preprocessing can support augmented reality, image editing, and a variety of downstream GAN tasks.

Why GANgealing?

This section summarizes the paper's introduction. The title asks why GANgealing is necessary because the classical method used for MNIST does not readily extend to complex RGB images.

Visual alignment means finding correspondence or registration across still images and video—for example, locating the same object or matching positions between different images. It underlies computer-vision problems such as optical flow, 3D matching and depth estimation, medical imaging, tracking, and augmented reality.

Recent research has made pairwise alignment from image A to image B easier, but aligning an entire dataset has received less attention. Problems that need a generalized reference frame—automatically annotating a whole dataset or finding keypoints in every image, for instance—remain difficult.

Other work has shown that generative models trained on already aligned images, such as FFHQ, AFHQ, and CelebA-HQ after face detection and cropping, can produce higher-quality samples. This creates a need for more global keypoint reasoning and alignment.

GANgealing expands the congealing idea beyond simple modalities such as binary digits. At a high level, it begins with a pretrained GAN generator learned from an unaligned dataset. Random generated samples and a target sample supervise an STN, while the STN parameters and the target sample are optimized together. The STN is trained independently of the GAN itself.

Related Work

Training with GANs

GAN-based generative models are used throughout computer vision, from classification and segmentation to representation learning and 3D graphics. This paper aims to exploit the pretrained joint distribution of a deep generative network. Unlike the familiar pipeline that first generates images with a GAN and then trains a discriminative network on them, GANgealing learns the discriminative model and optimizes its GAN-generated supervision together.

To avoid ambiguity: the method does not train the generative network. It optimizes the “GAN-generated image” itself. Generated images serve as supervision without pixel-level augmentation or postprocessing based on domain knowledge.

Joint Image-Set Alignment

The idea of an image average has been used to visualize joint alignment within a set sharing the same semantic content. Classical congealing, on which this paper builds, performs unsupervised joint alignment by minimizing entropy across images. It works well for structured data such as binary digits but not for more diverse data.

Later approaches simplified the problem by projecting high-dimensional datasets into low-rank subspaces or factorizing images into color, appearance, and shape features. These methods assume that every image can be represented by one mode—another subspace onto which the dataset is projected.

Other work combined joint visual alignment and clustering with user-provided data. Rather than relying on unsupervised alignment, it used bounding-box supervision and clustered multiple modes within an object category. Automated variants exist, but tend to be restricted to particular domains. Some approaches trained a network to predict a warp for each image, enabling clustering and alignment of large-scale collections, yet remained limited to simple color transformations. For complex datasets, those constraints undermine the underlying assumptions.

Spatial Transformer Networks

The main framework trained with GAN supervision is the spatial transformer network, a well-known example of deep-learning-based data processing. An STN learns parameters for a geometric transformation. It predicts a small set of warp parameters, builds a sampling grid over an input image, and defines sampling and warping as differentiable operations so the entire process can be trained.

STN modules have performed well in discriminative tasks and have been used for robust filter learning, view synthesis, and 3D representation learning. Those methods generally use an STN as a constraint that makes some other task easier. This paper differs by using a GAN to train the STN itself.

GAN-Based Supervised Learning

We can view the supervision pair $(x,~y)$ as the usual source–target relationship in supervised learning. Because the model being trained is an STN, the source image is unaligned—the casually positioned cat in the figure—while the target image is aligned, like the cat on the right. A GAN generator provides the images used to construct this aligned dataset.

First, $x$ is a random sample from a pretrained GAN generator. The target $y$ is generated by further manipulating the latent code used for $x$. Together they supervise an STN $f_\theta:x\rightarrow y$:

$$\mathcal{L}(f_\theta,~y) = l(f_\theta(x),~y)$$

Here $l$ is a reconstruction loss: the goal is to turn unaligned data $x$ into aligned data $y$. In conventional supervised learning, both $x$ and $y$ are fixed. In GAN-supervised learning, however, the target $y$ is also optimized end to end alongside the STN $f_\theta$. After training, the STN is evaluated on real inputs rather than GAN-generated images.

Dense Visual Alignment

How can congealing be formulated so that GAN supervision becomes possible? This brings us to the algorithm called GANgealing.

GANgealing begins by training a latent-variable generative model $G$ on an unaligned input dataset. Its latent vector is written $w\in\mathbb{R}^{512}$. The use of $w$ rather than $z$ reflects the StyleGAN-based generator used in the work. Given the trained $G$, a source latent is sampled from $w\sim\mathcal{W}$ and the source image is $x=G(w)$.

The target image is $G(c)$, where $c$ has the same dimensionality as $w$ and serves as the target for the STN. Although $w$ can be drawn randomly, we do not initially know which $c$ will be an appropriate target for $G(w)$. The solution is to optimize $c$ jointly with the STN. Because generator $G$ is differentiable with respect to $c$, gradient descent can do so.

$$\mathcal{L}_\text{align}(T,~c) = l(T(G(w)), G(c))$$

Here $T(\cdot)=f_\theta(\cdot)$, and $l$ is a distance between two images—apparently a perceptual loss in the figure. Minimizing the objective with respect to $c$ asks the latent $c$ to find a sample that every generated image $x$ can approach through spatial transformer $T$. Thus, $G(c)$ supplies the canonical alignment for the image set. If the initial $c$ generates an unreachable target, the loss automatically adjusts it; after enough iterations, $c$ learns to generate a useful average image.

This simple approach is plausible but works only when the dataset has limited diversity. Every input image $x$ is trained toward the same image $G(c)$ generated from one constant latent, which imposes a hard optimization limit.

Anyone who has experimented with GAN inversion will know that optimizing from a sample close to the input generally produces better images than optimizing a completely random latent. GANgealing applies the same idea. Rather than use the identical target $G(c)$ for every random $G(w)$, it preserves pose and orientation across samples while taking detailed appearance from $G(w)$.

StyleGAN constructs images from low resolution upward, applying latents to a learned constant and progressing from coarse style to fine detail. Combining target vector $c$ with random input latent $w$ at the appropriate layers makes $\text{mix}(c,w)\in\mathbb{R}^{512}$ retain the broad structure determined by $c$ while preserving details from the image generated by $w$.

$$\mathcal{L}_\text{align}(T,~c) = l(T(G(w)), G(\text{mix}(c,~w)))$$

The experiments use StyleGAN2 and exploit its style–pose disentanglement. Style mixing creates a separate target $G(\text{mix}(c,w))$ for each sample: $c$ controls the early, coarse generator features and determines pose, while $w$ enters later layers to determine texture. The paper also ablates the choice of mixing point.

STN Parameterization

Recall the STN architecture. The spatial-transformer function $T$ receives an image and produces a sampling grid $g\in\mathbb{R}^{H\times W\times2}$, whose parameterization depends on the intended transformation. GANgealing tests transformations for rotation, scaling, and horizontal and vertical translation, as well as unconstrained deformation. The final $T$ combines an unconstrained STN with a similarity STN for rotation, scale, and translation. It can even learn transformations such as a horizontal flip.

For the unconstrained $T$, a total-variation (TV) regularizer is effective. Total variation is commonly defined through the supremum of the sum or average of input gradients; in a simplified discrete form,

$$TV(x) = \sum_n \vert y_{n+1} - y_n \vert$$

Definitions of total-variation distance vary. In this paper it regularizes sampling grid $g$ so that the grid does not deform excessively:

$$\mathcal{L}_\text{TV} (T) = \mathcal{L}_\text{Huber} (\Delta_x g) + \mathcal{L}_\text{Huber} (\Delta_y g)$$

As the equation shows, it calculates directional differences along $x$ and $y$. An additional norm discourages $g$ from moving too far from the identity transformation:

$$\mathcal{L}_I (T) = \parallel g \parallel^2_2$$

Parameterization of $c$

The target latent vector $c$ is not optimized as a completely unconstrained input through differentiable generator $G$. Instead, it is parameterized as a linear combination of the top $N$ principal directions in $\mathcal{W}$ space:

$$c = \bar{w} + \sum_{i=1}^N \alpha_i d_i$$

These principal directions are the familiar axes extracted by PCA. Here $\bar{w}$ is the empirical mean latent, $d_i$ is the $i$th principal direction, and each $\alpha_i$ is a learned scalar coefficient.

This parameterization matters because StyleGAN's $\mathcal{W}$ space is highly expressive. Optimizing $c$ without constraints tends to degrade target-image quality and move the result away from the manifold of natural images. Reducing the number $N$ of principal axes restricts the feasible directions of $c$, effectively constraining it around the mean of $\mathcal{W}$ and mitigating this problem.

The final GANgealing objective is

$$\mathcal{L}(T,~c) = \mathbb{E}_{w \sim \mathcal{W}} (\mathcal{L}_\text{align} (T, c) + \lambda_\text{TV} \mathcal{L}_\text{TV} (T) + \lambda_I \mathcal{L}_I (T))$$

The experiments use $\lambda_\text{TV}=1000$ or $2500$, and $\lambda_I=1$.

Joint Alignment and Clustering

The description so far works well for several multimodal datasets, including LSUN Bicycles and Cats. Some datasets, such as LSUN Horses, contain poses too diverse to represent with a single mode. To handle them, GANgealing can be extended from one target latent $c$ to a clustering algorithm.

Suppose we learn $K$ vectors $c$. Each latent constant captures a particular data mode, so learning $(c_k)_{k=1}^K$ models several modes in a dataset. As before, each $c$ is a linear combination of the top $N$ principal directions in $\mathcal{W}$, with a separate set of learned $\alpha$ coefficients for every $c_k$. The paper similarly uses $K$ distinct spatial transformers $T_k$, one per mode. The resulting model learns dense correspondence among all images within each cluster.

The remaining question is how to guide each pair $(c_k,T_k)$ toward a particular mode. The method uses the hard-assignment step

$$\mathcal{L}_\text{align}^K (T, c) = \min_k \mathcal{L}_\text{align} (T_k, c_k)$$

so that each unaligned image is assigned to and aligned under one mode. The familiar unimodal case corresponds to $K=1$.

At test time, a generated image $G(w)$ can be assigned to

$$k^\ast = \arg \min_k \mathcal{L}_\text{align}(T_k,~c_k)$$

and warped with the corresponding transformer $T_{k^\ast}$. Real images are harder: calculating $\mathcal{L}_\text{align}$ requires not only the input image, but the latent $w$ that generated it.

Clustering over $K$ modes optimizes multiple modalities, and each $c$ is mixed and optimized with the input latent $w$, so applying the same process to a real image would require recovering a latent variable. GAN inversion is the intuitive solution, as discussed in my earlier post on GAN-based image manipulation: find a latent $w$ whose generated sample resembles image $x$. Accurate GAN inversion is nevertheless slow and particularly challenging for non-face datasets, since StyleGAN is commonly trained on FFHQ faces.

The authors therefore train a network that predicts the cluster assignment of an input image as a classification result. For generated images, the target cluster $k^\ast$ is already available, so pairs $(G(w),k^\ast)$ provide straightforward supervision. This permits assignment of an arbitrary image without GAN inversion. The classifier reuses the spatial transformer's weights but replaces its warping head with a classification head. Both the spatial transformer and classifier, although optimized on fake samples, generalize well to real images.

Closing Thoughts

That concludes the explanation of GANgealing; I will stop without covering the experiments and results in detail. In reading recent papers, I have become more interested in the insights and mathematical background used to develop an idea—the related work and preliminaries—than in cataloging every experiment and number.

The paper's main contribution is not a method for training a GAN. It uses a GAN to train an STN and thereby align unaligned real images effectively. I found the subject exceptionally interesting, and it seems broadly useful.