Unconditional to conditional generation via distributional level matching

Motivation

Recently, priming models have shown a very intuitive way of building an image generation with the two stage set up. In

Unconditional models:
JiT

Distributional level matching：
FD Loss' advantages:
It can precompute the entire image dataset's feature statistics, $\mu$ nd $\Sigma$, making comparison with the online EMA network very easy and fast.

What's the difference

Unconditional Base Model

Creating the unconditional base model steps:

First, we make the Imagenet 100 dataset. Take every 10th id, our list of images to produce is [0, 10...990]

The image generation model is the JiT-B/16. It is a transformer model trained on direct image pixels.

The generator accepts 3 inputs, and outputs an predicted clean image $x_\text{pred}$
$z_t \quad$: Image-shaped tensor. The shape is $3 \times 256 \times 256$
$t \quad$: Timestep indicating noise level
$c \quad$: Class ID as an integer

How to make it unconditional?

Embedding table: Created inside JiT as a trainable model parameter. Randomly initialized, containing a total of 1001 rows, each row being a vector of 768 numbers. Normally, in a conditional generator, each row is a learned class embedding vector. Each row when passed into the generator will make it output a different class based on the learned embedding value updated during training.

So the 1st row's vector, embedding_table[0] #shape [1, 768], when passed into the model, makes it output a tench fish.

The 1000th row's vector, embedding_table[999] #shape [1, 768], makes model output toilet paper.

Null Label: This is the 1001st row of the embedding table. When training a conditional generator, we pass this into the model 10% of the time for classifier-free guidance training.

To make a fully unconditional model that doesn't learn any fundamental class relations, we use the null label for every training step. While the null vector's value is updated each training step, it doesn't make the model conditional because the model doesn't receive a different, identifiable embedding saying this should be "dog" or "car".

self.embedding_table = nn.Embedding(num_classes + 1, hidden_size)
num_classes = 1000
hidden_size = 768
embedding table shape = [1001, 768]

EMA weights: "exponential moving average" of model's weights. So with EMA1:

EMA1 weights = 0.9999 × previous EMA1 weights + 0.0001 × current model weights

This helps the model maintain its weights during training, as each batch's gradient affects all weights each time. We prefer a smoothing to improve sampling quality in the end.

Unconditional base model

This is the unconditional baseline model. To generate each image, it takes 50 de-noising steps, flowing from pure Gaussian noise to the clean image.

FD Loss

Using only the FD objective, the model remains unconditional. Starting from the pretrained checkpoint above, FD Loss can improve two things:

1. Turn the model from 50 step generation into one step generation.

2. Improves image quality by distributional level matching of generator output distribution with

Distribution level matching

The following are methodologies of doing distribution level matching. Basically, all the goal is to match the generators feature density distribution, within a given class c, to the real dataset's class feature distribution.
$${q_\theta(z\mid c)\approx p(z\mid c)}$$Where
$p(z\mid c)=\text{real class feature distribution}$
$q_\theta(z\mid c)=\text{generated class feature distribution}$

Sinkhorn Divergence

Step 1: Generate 3 sets of samples.
$$\begin{aligned}G&=\{g_1,\ldots,g_B\} &&\text{generated features, with gradients}\\ G'&=\{g'_1,\ldots,g'_B\} &&\text{generated reference features, detached}\\ R&=\{r_1,\ldots,r_N\} &&\text{real features, detached}\end{aligned}$$In the priming set up, they did used 64 samples of G and G', and 128 samples (real images) of R.
Both generated sets comes from the same generator with independent Gaussian noises.

Step 2: Compute pairwise distance matrices
We want to find a generator-real, generator-generator pairing.
For generator-real:
For every generated feature $g_i$, real feature $r_j$, compute the squared Euclidean distance between the two feature vectors:$$g_i=[g_{i1},g_{i2},\dots,g_{id}]\\r_j=[r_{j1},r_{j2},\dots,r_{jd}]$$Then:
$$\boxed{\|g_i-r_j\|_2^2=\sum_{k=1}^{d}(g_{ik}-r_{jk})^2}$$Each $g_i$ is compared with all 128 $r_j$, and we compute the squared distance between, saving it to a total distance matrix, with 64 x 128 = 8192 pairwise distances.$$ C^{GR}\in\mathbb{R}^{64\times128} $$Similarly, calculate generated-to-generated reference distances:
$$\boxed{C_{ik}^{GG'}=\frac12\|g_i-g'_k\|_2^2}$$We have 64 x 64 = 4096 pairwise distances between the generated features. $$ C^{GG'}\in\mathbb R^{64\times64}$$
Step 3: solve the optimal transport