Discrete Latent Representation For Image Generation

Fabianaugie Jametoni

Contents

• Discrete Vs Continuous
• VQ-VAE And VQ-GAN
• BERT
• MaskedGIT
• Conclusions

​

2

Discrete Vs Continuous

Discrete Vs Continuous

4

Which one is bigger? the amount of number between 1 and ∞ or between 0 and 1?

Discrete Vs Continuous

5

Which one is bigger? the number between 1-Infinity or between 0-1?

• According to Cantor there is more number between 0 and 1 than there is between 1 and infinity, he proved this in hi famous Cantor diagonalization argument
• Then According to this we can assume that model that use continuous latent representation like diffusion model, theoretically can encode more data than any model that use discrete latent representation which is very important in generative model

​

Discrete Vs Continuous

6

So, case closed then. Continuous latent space are better than discrete latent space for generative model

Discrete Vs Continuous

7

So, case closed then. Continuous latent space are better than discrete latent space for generative model��Well….. Not Yet

Discrete Vs Continuous

8

Not Yet

• While it’s true that continuous latent representation can encode more data, but, what continuous latent space has a hard time was in disentanglement
• In continuous latent space, we need to create extra step in order to disentangle the latent space into various “explained” representative, discrete latent representation does not have this problem, infact, given enough data, discrete latent model like vq-vae can separate this representation on its own, in theory anyway. in fact, VQ-tokenizer has been used everywhere from text2speech to Dall-e!

Discrete Vs Continuous

9

So how does it work?

VQ-VAE And VQ-GAN

VQ-VAE And VQ-GAN

11

Introduction

• Discrete model use “Codebook” to generate data, you can think of this as a set of lookup table that the model would look up in order to generate data

VQ-VAE And VQ-GAN

12

Introduction

• Discrete model use “Codebook” to generate data, you can think of this as a set of lookup table that the model would look up in order to generate data

But what does exactly the lookup table, look up? What is the meaning of all this number?

VQ-VAE And VQ-GAN

13

Introduction

The codebook store the centroid of a data distribution, so in a way, discrete model still use continuous latent space for data representation, the different is that discrete model has a lot of small distribution representing some part of the image, unlike continuous model that only has one giant continuous representation to represent the data

C1

C2

C5

C4

C3

C

VQ-VAE And VQ-GAN

14

Introduction

happiness

contentment

pleasure

joy

satisfaction

C

VQ-VAE And VQ-GAN

15

VQ-VAE

• VQ-VAE itself are actually quite simple

VQ-VAE And VQ-GAN

16

VQ-VAE

• VQ-VAE itself are actually quite simple
• Problem!

because of the discrete codebook, the gradient will be lost, so how are we gonna train the model?

VQ-VAE And VQ-GAN

17

VQ-VAE

• VQ-VAE itself are actually quite simple
• Problem!

Instead, what VQ-VAE do is they copy the gradient from decoder and used it in the enoceder! skipping the codebook entirely

VQ-VAE And VQ-GAN

18

VQ-VAE

• Problem: If we skip the codebook gradient, then how are we gonna update it?

VQ-VAE And VQ-GAN

19

VQ-VAE Loss

• Problem: If we skip the codebook gradient, then how are we gonna update it?

Reconstruction Loss

This is just the normal autoencoder loss which is the MSE between the input and output

VQ-VAE And VQ-GAN

20

VQ-VAE Loss

• Problem: If we skip the codebook gradient, then how are we gonna update it?

Reconstruction Loss

This loss are used to move the centroid to the embedding

Codebook Loss

CA

CB

VQ-VAE And VQ-GAN

21

VQ-VAE Loss

• Problem: If we skip the codebook gradient, then how are we gonna update it?

Reconstruction Loss

This loss are used to move the centroid to the embedding

Codebook Loss

CA

CB

VQ-VAE And VQ-GAN

22

VQ-VAE Loss

• Problem: If we skip the codebook gradient, then how are we gonna update it?

Reconstruction Loss

This loss are used to move the centroid to the embedding

Codebook Loss

CA

CB

VQ-VAE And VQ-GAN

23

VQ-VAE Loss

• Problem: If we skip the codebook gradient, then how are we gonna update it?

Reconstruction Loss

This loss are used to move the centroid to the embedding

Codebook Loss

CA

CB

VQ-VAE And VQ-GAN

24

VQ-VAE Loss

• Problem: If we skip the codebook gradient, then how are we gonna update it?

Reconstruction Loss

Codebook Loss

CA

CB

Commitment Loss

You can think of this loss as the reverse of Codebook loss, this loss basically move the embedding to the centroid, this ensure that the embedding commit to 1 centroid and don't fluctuate between 2 or more centroid

VQ-VAE And VQ-GAN

25

VQ-VAE Loss

• Problem: If we skip the codebook gradient, then how are we gonna update it?

Reconstruction Loss

Codebook Loss

CA

CB

Commitment Loss

You can think of this loss as the reverse of Codebook loss, this loss basically move the embedding to the centroid, this ensure that the embedding commit to 1 centroid and don't fluctuate between 2 or more centroid

VQ-VAE And VQ-GAN

26

VQ-VAE Loss

• Problem: If we skip the codebook gradient, then how are we gonna update it?

Reconstruction Loss

Codebook Loss

CA

CB

Commitment Loss

You can think of this loss as the reverse of Codebook loss, this loss basically move the embedding to the centroid, this ensure that the embedding commit to 1 centroid and don't fluctuate between 2 or more centroid

VQ-VAE And VQ-GAN

27

VQ-VAE Loss

• Problem: If we skip the codebook gradient, then how are we gonna update it?

Reconstruction Loss

You can think of this loss as the reverse of Codebook loss, this loss basically move the embedding to the centroid, this ensure that the embedding commit to 1 centroid and don't fluctuate between 2 or more centroid

Codebook Loss

CA

CB

Commitment Loss

VQ-VAE And VQ-GAN

28

VQ-GAN

• VQ-GAN itself is not much different than VQ-VAE, the first is they use perceptual loss instead of MSE

VQ-VAE And VQ-GAN

29

Perceptual Loss

Input Image

Reconstructed Image

Pre-Trained Conv Network (CLiP, ResNet, MobileNet, DiNO, Etc.)

Label Image Conv Embedding

Reconstruction Image Conv Embedding

MSE

VQ-VAE And VQ-GAN

30

Patch GAN

BERT

BERT

32

BERT

• Created by google in 2018, BERT was a model made as a backbone for many NLP task like Sentiment analysis and emotion classification, it use the encoder part of the transformer and train it on a certain nlp task in hope for the model to be able to extract feature from text

​

BERT

33

Masked Token Prediction

• One of BERT pre-training task is Masked Token Prediction, by masking a few word in a sentence, the model are then tasked to predict those masked token. this task are made so that the model can understand the context behind each word in a sentence

​

MaskedGIT

MaskedGIT

35

MaskedGIT

• MaskedGIT took the idea of BERT, but instead of predicting text token, maskedGIT predict image token

​

MaskedGIT

36

MaskedGIT

In theory, maskedGIT can generate an image in a single pass. in practice however, the model do required multiple iteration in order to generate an image, first we create scheduler consist of a vector of float consisted of a value in range of 0-1, the size of ow which are the amount of our iteration, then, the generation algorithm goes as follow:

• First, the input are started as a blank canvas with all token masked
• Then, we use this initial input and use it to predict our initial prediction
• we then masked the output with ratio determine from our scheduler
• the masked output are then imputed back to the transformer
• we do it until n-iteration at which point the image are generated

​

MaskedGIT

37

MaskedGIT

MaskedGIT

38

Training Detail

In the original paper, during training the ratio of masked token are sampled randomly to simulate the various decoding scenario. in practice however, i found that, instead of sample the ratio randomly, the result are better if the sample ration are uniformly sampled across batch on every epoch, the training took 200 epoch using Adam optimizer and Cross-Entropy for our loss function

MaskedGIT

39

Speed Comparison

Because the nature of transformer, we can parallelize the prediction process of every mask which speed up the prediction process compared to other method

MaskedGIT

40

Zero-Shot Editing Capabilities

Another capabilities of this model in in it capabilities of doing zero-shot image editing by masking part of the image and predicting the masked part specifically

Q&A

Thank You