Add KL loss

[borisdayma]

KL loss can avoid using a codebook (with associated quantization loss).

[borisdayma]

Advantage of using codebook:

• can be used for AR models
• smaller pre-encoding (only small sequence of int)

Advantage of KL loss:

• no quantization loss from codebook
• no hyper-parameter for codebook size
• no issue with unused codes
• if we have a small codebook embedding, pre-encoding may also be feasible

What do you think @patil-suraj

[patil-suraj]

We can add KL loss for sure. But then it depends how do we want to use this model. If we are going to do diffusion, then we can train the KL model. But for auto-regressive model like dalle-mini we will need the VQ part.