DATASCI 447 Lecture 20: Conditioning, Editing, and the Bridge to Sequences

Administrative Stuff

REVIEW: WHAT THE VQ-VAE DOES

How does the VQ-VAE differ from the traditional VAE?

• Latent space?
• Regularization?
• Reconstruction/Adversarial Loss?

If VAEs are fancy nonlinear PCA, what would be the VQ-VAE’s equivalent?

REVIEW: WHAT THE VQ-VAE DOES

The encoder compresses an image into a spatial feature map \((B, D, H', W')\).

Each spatial position is snapped to the nearest entry in a learned codebook.

The decoder reconstructs the image from the quantized features.

REVIEW: WHAT THE VQ-VAE DOES

REVIEW: THE SNAP UP CLOSE

The encoder outputs a \(K \times K \times D\) tensor of the image. Each of the \(K \times K\) positions are spatial codes.

For a single spatial position:

• The encoder outputs a continuous vector \(\mathbf{z}_e \in \mathbb{R}^D\).

• We compute \(\|\mathbf{z}_e - \mathbf{e}_j\|_2\) for every codebook entry \(j = 1, \ldots, K\).

• Replace with the winner: \(\mathbf{z}_q = \mathbf{e}_{k^*}\) where \(k^* = \arg\min_j \|\mathbf{z}_e - \mathbf{e}_j\|_2\)

REVIEW: THE SNAP UP CLOSE

REVIEW: REGULARIZATION

How did standard VAEs regularize the latent space to prevent mass collapse and yield more generalizable representations that could support generation?

How does the VQ-VAE differ?

REVIEW: The DIFFERENTIABILITY PROBLEM

The argmin is not differentiable. Gradients from the decoder can’t flow through the discrete selection!

z_q = z_e + (z_q - z_e).detach()

REVIEW: The DIFFERENTIABILITY PROBLEM

REVIEW: WHY TWO LOSSES, NOT ONE?

Naive idea: just minimize \(\|\mathbf{z}_e - \mathbf{e}_{k^*}\|_2^2\) and let both the encoder and codebook update.

Why this fails: the encoder and codebook would chase each other. The encoder moves \(\mathbf{z}_e\) left, so the codebook moves \(\mathbf{e}_k\) left. Now \(\mathbf{z}_e\) moves further left. They drift together in an arbitrary direction, never converging.

REVIEW: WHY TWO LOSSES, NOT ONE?

WHAT DID THE CODEBOOK LEARN?

One of the big advantages of the VQ approach is interpretability of the latent space

The codebook has \(K\) entries, each a \(D\)-dimensional vector.

• The SNAP operation ensures that all images are using some combination (usually 16, 64, or 256 entries depending on richness of the training set) of the same visual vocabulary

• Say that \(K = 512\). That’s not that many words! (I can say 512 unique words right now!)

WHAT DID THE CODEBOOK LEARN?

The issue is that each codebook entry is a vector of length \(D\)

• In weird latent visual space…

• These are abstract — we can’t directly visualize a 64-dimensional vector.

Basically, the same problem as intepreting features from a CNN!

SHARED CODEBOOK USAGE

The codebook entries are shared visual primitives. A golden retriever and a German shepherd might use the same “grass background” entry in their top-left positions, but different “fur texture” entries in their center positions.

The Issue

We’re designing latent vectors that are intended to capture specific attributes of the input images.

• Pose, background, age of animal, etc.

But, we’re also asking the latent space to learn about:

• Species

• Breed

• Color

Why waste capacity learning when we already know these values for each image!

ADDING CONTROL

We can reconstruct sharply with good spatial structure. But every generation is random — we can’t say “generate a golden retriever” or “make this dog black.”

CATEGORICAL EMBEDDINGS

Each categorical attribute gets a learned embedding — a lookup table that maps a discrete category to a dense vector:

self.breed_embed  = nn.Embedding(num_breeds, 16)    # 120 breeds → R¹⁶
self.pcolor_embed = nn.Embedding(num_pcolors, 8)     # 9 colors → R⁸
self.scolor_embed = nn.Embedding(num_scolors, 8)     # 9 colors → R⁸

• Dimensionality scales with the complexity of what’s encoded. Breed (120 classes with subtle distinctions) needs more dimensions than color (9 broad categories).

Concatenated into a single conditioning vector: \(\mathbf{c} = [\mathbf{c}_{\text{breed}};\; \mathbf{c}_{\text{pcolor}};\; \mathbf{c}_{\text{scolor}}] \in \mathbb{R}^{32}\)

WHAT IS AN EMBEDDING?

An embedding is just a matrix \(\mathbf{W} \in \mathbb{R}^{K \times D}\) where row \(k\) is the learned vector for class \(k\).

# nn.Embedding(120, 16) creates a 120×16 matrix
# Input: integer index 42 (golden retriever)
# Output: row 42 — a learned 16-dim vector

• The vectors are initialized randomly and trained end-to-end. The network learns whatever representation of “golden retriever” or “black” makes reconstruction easiest.

• Why not a one-hot vector?

CONDITIONING THE ENCODER

The encoder needs to know the class so it can organize the codebook indices efficiently.

Spatial broadcast: reshape \(\mathbf{c}\) to a spatial tensor and concatenate as extra input channels.

c_spatial = c.view(B, 32, 1, 1).expand(B, 32, H, W)   # (B, 32, 128, 128)
x_cond = torch.cat([x, c_spatial], dim=1)               # (B, 35, 128, 128)

CONDITIONING THE ENCODER

CONDITIONING THE DECODER

The decoder needs to know the class so it can interpret the codebook indices correctly.

Broadcast \(\mathbf{c}\) to the bottleneck spatial size and concatenate to the quantized features:

c_small = c.view(B, 32, 1, 1).expand(B, 32, H_prime, W_prime)
z_q_cond = torch.cat([z_q, c_small], dim=1)
reconstruction = self.decoder(z_q_cond)

• The same codebook entry can mean different things for different conditions. An entry that captures “pointy ear shape” gets decoded differently when \(\mathbf{c}\) says “husky” vs “chihuahua.”

THE CODEBOOK DOESN’T KNOW ABOUT CONDITIONING

The codebook \(\mathbf{E} \in \mathbb{R}^{K \times D}\) is shared across all breeds and colors.

• Quantization happens on the encoder’s output after it has processed the conditioned input. The codebook learns universal visual primitives useful across all classes.

\(\mathbf{c}\) flows around the codebook, not through it:

\[\mathbf{x}, \mathbf{c} \;\to\; \text{Encoder} \;\to\; \mathbf{z}_e \;\to\; \text{Quantize} \;\to\; \mathbf{z}_q \;\to\; [\mathbf{z}_q \,;\, \mathbf{c}] \;\to\; \text{Decoder} \;\to\; \hat{\mathbf{x}}\]

THE CODEBOOK DOESN’T KNOW ABOUT CONDITIONING

THE LOSS DOESN’T CHANGE

\[\mathcal{L} = \text{L1} + \lambda \cdot \text{LPIPS} + w_{\text{gan}} \cdot \mathcal{L}_{\text{GAN}} + \|\text{sg}[\mathbf{z}_e] - \mathbf{e}\|^2 + \beta_c\|\mathbf{z}_e - \text{sg}[\mathbf{e}]\|^2\]

Exactly the same loss as the unconditional VQ-VAE. No new terms!

Conditioning flows through the architecture. The loss handles the rest.

THE SEPARATION PRINCIPLE

Because \(\mathbf{c}\) and the codebook indices encode different things, we can swap parts of \(\mathbf{c}\) while keeping the indices fixed.

The LAZY SKIP

As we’ve seen before, optimizers are lazy

• Like my QTM 110ers, given the option to take the path of least resistance, they will!

Any guesses about what can happen when training a conditioned VQ model?

The LAZY SKIP

If the bottleneck is given too much capacity (e.g. \(K\) is too high), it will learn that the most efficient strategy is to completely ignore the conditions and learn everything through the bottleneck!

• The Codebook loss prioritizes reconstruction quality over conditioning fidelity.

Like with the GAN loss, we’ll need to put gutter-guards on the optimizer to prevent this.

The Strict Proctor

The problem: the encoder hides the labels (e.g. “Cat” or “Golden Retriever”) inside the 16 x 16 latents for an image so the decoder doesn’t have to read our condition vector

The solution: Hire a strict proctor to prevent cheating - a tiny latent classifier network

• Look at the encoder latents before quantization. Use those to try to predict the condition vector (e.g. “Cat” or “Dog”, “Black”, “Brown” or “White”)

• If the classifier can use the encoder outputs to predict the conditions, CHEATING

• If not, we’re all good

We want to train a classifier that will fail the cheating test.

Gradient Reversal

The classifier is independent of the VQ-VAE. How do we give that info back to the encoder?

• Have a separate optimizer for the proctor network. In each batch step the classifier.

• In the forward pass, let the latents flow as normal and let the classifier make its guesses

• In the backward pass, reverse the gradients (multiply by -1) for the classifier. We want to penalize the classifier when it gets the conditions right. Since the encoder gradient flows through the jointly trained classifier, the encoder combines its good gradients with the bad gradients!

• In equilibrium, the classifier should be unable to predict the conditions from the latents and send a gradient signal of 0.

Gradient Reversal

4D chess, baby!

• Reward the snitch network for reporting cheaters and punish the cheater.

• The encoder is punished for leaking information and learns to not do that!

• In theory, this leads to better conditioning. In theory…

The Strict Proctor

EDITING IMAGES

This now admits a very powerful property - editing

1. Train the Conditional VQ-VAE

2. Take an image and its conditions and encode it to the latent space (e.g. the 16 x 16 matrix of codebook values)

3. Concatenate new conditioning values to the existing codebook entries

4. Decode

EDITING IMAGES

Take a photo of a golden retriever. Make it black.

Step 1: Encode with original labels \(\to\) get codebook indices

z_q, indices = model.encode(image, breed=GOLDEN_RETRIEVER,
                             pcolor=GOLDEN, scolor=WHITE)

EDITING IMAGES

Take a photo of a golden retriever. Make it black.

Step 2: Decode with edited condition — only color changes

c_edited = get_condition(breed=GOLDEN_RETRIEVER,
                          pcolor=BLACK, scolor=BLACK)
edited = model.decode(z_q, c_edited)

Same codebook indices. Same pose. Same background. Different fur color.

EDITING: BREED CHANGE

Same codebook indices, different breed:

c_original = get_condition(GOLDEN_RETRIEVER, GOLDEN, WHITE)
c_edited   = get_condition(SIBERIAN_HUSKY, GRAY, WHITE)

original_image = model.decode(z_q, c_original)
edited_image   = model.decode(z_q, c_edited)

Same spatial layout. The animal’s pose, the background, the lighting — all preserved. But the breed shifts.

WHAT WE’VE (ALMOST) BUILT: DALL-E 1

Take stock of what you have:

• A VQ-VAE that tokenizes images into discrete codebook indices
• Perceptual loss + GAN loss for quality and sharpness
• Categorical conditioning through learned embeddings
• Editing by swapping conditions while keeping tokens fixed

EDITING RESULTS

Not that great yet.

There’s something that we’re still missing:

Conditioning on features can’t completely remove the identity information from the latent space.

• DALL-E 1.0 worked well in this way because they had millions of images and hundreds upon hundreds of GPUs

We’re still missing the ability to learn important structure in the latent space - the order in which the codebook entries appear!

RANDOM CODEBOOK INDICES → GARBAGE

Sample a random index at each of the 16 spatial positions. Decode.

• It will be garbage!

We’ve designed the VQ-VAE to preserve spatial information.

• But we then ignore the spatial information when we’re trying to do new things with images.

• With a large enough sample size (millions of images), this all comes out in the wash.

We don’t have the computational resources to process that!

ORDER MATTERS A LOT

“Dog ear” + “grass background” + “sky” in the right spatial positions = a dog photo.

The same entries in random positions = garbage.

• Neighboring positions are correlated. If the top-left says “indoor background,” the top-right should probably also be “indoor background” — not “blue sky.”

ORDER MATTERS

The codebook indices form a sequence (raster scan: top-left to bottom-right):

\[s_1, s_2, s_3, \ldots, s_{16}\]

Each index depends on the previous ones:

\[P(s_1, s_2, \ldots, s_{16}) \neq \prod_i P(s_i)\]

Sampling independently ignores all dependencies.

ORDER MATTERS

THE AUTOREGRESSIVE SOLUTION

Factor the joint distribution into a chain of conditionals:

\[P(s_1, s_2, \ldots, s_N) = \prod_i P(s_i \,|\, s_1, \ldots, s_{i-1})\]

Predict each token given all previous tokens.

CONDITIONAL AUTOREGRESSIVE SAMPLING

Add \(\mathbf{c}\) as a prefix or bias to the sequence model:

\[P(s_1, \ldots, s_N \,|\, \mathbf{c}) = \prod_i P(s_i \,|\, s_1, \ldots, s_{i-1}, \mathbf{c})\]

“Given that this is a golden retriever with golden primary color, what codebook indices should I produce?”

WE NEED A NEW TOOL

Predicting the next token given all previous tokens. Does this sound like a familiar task? Maybe a tool that most of you probably have open on your laptop right now?

WHY DISCRETE EDITS ARE CLEAN

In a continuous VAE, the latent vector \(\mathbf{z}\) is continuous. Even when you keep \(\mathbf{z}\) “fixed” and swap \(\mathbf{c}\), the decoder can pick up on subtle continuous signals in \(\mathbf{z}\) that leak identity information.

BABY’S FIRST DEEPFAKE

We’ve just shown that with a conditional VQ-VAE, anyone with this notebook can:

• Take a real photo of a real dog and change its color
• Take a real photo and change the breed entirely
• Generate photorealistic images of dogs that don’t exist

THE ETHICAL CONUNDRUM

The same architecture applied to faces produces deepfakes.

• Change someone’s expression
• Put words in their mouth
• Generate a photo of someone in a place they’ve never been
• Age someone, change their hair, swap their identity entirely

# Face version of what we just did
z_q, indices = model.encode(face_photo, identity=PERSON_A)
deepfake = model.decode(z_q, identity=PERSON_B)

DEEPFAKES: A DOUBLE-EDGED SWORD

The same technology that enables creative expression can also be misused for deception.

• Change someone’s expression
• Put words in their mouth
• Generate a photo of someone in a place they’ve never been
• Age someone, change their hair, swap their identity entirely

THE COMPLETE PICTURE

THE OTHER PATH

VQ-VAE path: discrete tokens → autoregressive transformer → DALL-E

Alternative: keep continuous spatial latents. Instead of autoregressive prediction, use iterative denoising — start from noise, gradually refine.

• This is latent diffusion \(\to\) Stable Diffusion.

SUMMARY

• VQ-VAE discretizes the latent space through codebook clustering — explicit regularization via \(K\)
• The codebook learns shared visual primitives — interpretable through usage analysis and PCA
• Conditioning enters through learned embeddings at encoder and decoder — no change to the loss
• Editing swaps conditioning while keeping codebook indices fixed
• Generation and high quality editing requires modeling sequential dependencies among codebook indices — a bridge to transformers
• Ethical warning: the same architecture applied to faces enables deepfakes
• You have built every component of DALL-E 1 (minus the autoregressive transformer adjuster)