codebook initialization

[ramyamounir]

Hi, Thank you for this great work. It's quite useful!

I have been having problems with index collapse and I'm not sure where it's coming from. But upon digging into the code, it seems that when we're not using k-means to initialize the codebook vectors, randn (normal distribution) is used to initialize them. The vqvae paper specifically uses uniform distribution for initialization, which allows the authors to ignore KL divergence when training.

This is from the vqvae paper: "Since we assume a uniform prior for z, the KL term that usually appears in the ELBO is constant w.r.t. the encoder parameters and can thus be ignored for training."

Is there any reason why you changed to Normal distribution here?

Thanks!

[lucidrains]

@ramyamounir Hi Ramy! I checked the paper as well as deepmind's original sonnet code (from which the implementation is based) and you are correct, I should have been using uniform init

I've corrected it with this commit Thank you for raising this issue!

[ramyamounir]

Awesome! Thanks for the quick reply and fix.

[gorold]

From the VQGAN paper: " The codebook variables are initialized from a normal distribution."

Should CosineSimCodebook continue to use normal distribution for initialisation?