Codebook diversity loss incorrect

[pfeatherstone]

At this line:

vector-quantize-pytorch/vector_quantize_pytorch/vector_quantize_pytorch.py

Line 1266 in 2b367e5

prob = (-distances * self.codebook_diversity_temperature).softmax(dim = -1)

distances is already negative from

vector-quantize-pytorch/vector_quantize_pytorch/vector_quantize_pytorch.py

Line 570 in 2b367e5

dist = -cdist(flatten, embed)

If distances were positive then the diversity loss would be correct i think. You can check this my comparing:

torch.testing.assert_close(embed_idx, prob.argmax(-1))

At line

vector-quantize-pytorch/vector_quantize_pytorch/vector_quantize_pytorch.py

Line 1266 in 2b367e5

prob = (-distances * self.codebook_diversity_temperature).softmax(dim = -1)

This will fail.

[lucidrains]

@pfeatherstone yes, you are right, thank you for catching this

[pfeatherstone]

No problem. I have found though that a diversity loss can very easily kill convergence if not careful with weighting and temperature. Though when using ResidualVQ or ResidualSimVQ without diversity i get pretty good codebook usage when training a soundstream style model. So not sure if an explicit diversity loss is really required. I am however having trouble with FSQ. I don't know if you've seen #226, but i get extremely bad codebook usage. I quite would like a diversity loss with FSQ but not quite sure how to formulate it. Any tips would be greatly appreciated.

[lucidrains]

@pfeatherstone yeah, i have a few ideas, will get back into it after i finish another project