Possible bug in LFQ distance calculation

[MattMcPartlon]

vector-quantize-pytorch/vector_quantize_pytorch/lookup_free_quantization.py

Line 222 in e244d47

distance = -2 * einsum('... i d, j d -> ... i j', original_input, self.codebook)

Hey Phil, Thanks for another great implementation 😄 . Regarding the distance calculation in LFQ (linked) I think this only holds if you're comparing self.codebook to the quantization of x (i.e. both have constant norm).

import torch
xs = torch.randn(10,3)
ys = torch.randn(10,3)
#xs,ys = map(lambda x: x/torch.norm(x,dim=-1,keepdim=True), (xs,ys))
print(torch.argmin(torch.cdist(xs,ys),dim=-1))
print(torch.argmin(-torch.einsum("i d, c d -> ... i c", xs, ys),dim=-1))

# output:
#tensor([8, 4, 7, 2, 3, 0, 4, 7, 1, 4])
#tensor([3, 9, 3, 3, 3, 2, 5, 9, 9, 5])

import torch
xs = torch.randn(10,3)
ys = torch.randn(10,3)
xs,ys = map(lambda x: x/torch.norm(x,dim=-1,keepdim=True), (xs,ys))
print(torch.argmin(torch.cdist(xs,ys),dim=-1))
print(torch.argmin(-torch.einsum("i d, c d -> ... i c", xs, ys),dim=-1))

# output:
#tensor([4, 7, 7, 7, 4, 7, 7, 8, 8, 8])
#tensor([4, 7, 7, 7, 4, 7, 7, 8, 8, 8])

[lucidrains]

@MattMcPartlon hey Matt 😄 ! yes i think you are right, and digging into it, @Lijun-Yu made the change here a while back https://github.com/lucidrains/vector-quantize-pytorch/pull/89/files#diff-3488501716ef0c4a1b84127592be7e1015e0faf0d7d3b630b69d88e55a134701R213

i can prepare a PR and merge it once he gives the final word

[MattMcPartlon]

He would certainly know :). Thanks, Phil!

[lucidrains]

@MattMcPartlon going to just merge it for now so you can go train your potentially life-saving model 😉

@Lijun-Yu let me know if you see any issues with this revert

[RobertLuo1]

I think whether normalization on the original_input and the codebook or not perform the same. In LFQ, the codebook is the combination of {-1, 1}. Therefore, it seems to be equaliviant to the commit loss here. The distance of the same sign between original input and the token in codebook will be the largest.

[lucidrains]

@RobertLuo1 ah yes, you are right 🤦 reverted the revert 😆

that's really neat!

import torch

xs = torch.randn(10,3)
ys = torch.randn(10,3)

ys = ys / ys.norm(dim = -1, keepdim = True)

print(torch.argmin(torch.cdist(xs, ys),dim=-1))
print(torch.argmin(-torch.einsum("i d, c d -> ... i c", xs, ys),dim = -1))

# output:
# tensor([1, 1, 1, 1, 6, 9, 8, 6, 5, 5])
# tensor([1, 1, 1, 1, 6, 9, 8, 6, 5, 5])

[MattMcPartlon]

yup, great catch @RobertLuo1. thank you! (and thanks for the quick (yet unneeded) patch Phil!)

[RobertLuo1]

@MattMcPartlon huh, You are welcome! I have been trying it for a while. Yet, I still could not exploit its power as I have tried many times I found that the following loss term avg_codebook_entropy can not be optimized to be larger, It seems to be contradicted to the per_sample_entropy term. What about you? @MattMcPartlon

[MattMcPartlon]

I have not had luck yet @RobertLuo1. The per_sample_entropy should be a n_tokens (rows) x n_codewords (cols) matrix. The first term term says that each column should have low entropy. The second term says that if you average over the rows then the entropy should be large (i.e. each codeword should have uniform probability of being sampled for some token ). I agree that these can be somewhat opposing objectives, especially when the number of tokens is much smaller than the number of codewords...

Anyways, LFQ has not worked at all for me compared to the standard VAE objective. Have you tried any other quantizers?

[RobertLuo1]

@MattMcPartlon Yeah! I see in magvit2 paper, they only utilize 256 token in each batch. The flattened batch operation makes the value of avg_codebook_entropy loss increase along with the batch. LFQ seems does not work for me either. I assume that maybe there is something missing when implementing it.

[MattMcPartlon]

yup, I agree. I can start to overfit with a large enough codebook size (64k), but can't find a setting that generalizes well. It's a shame because I really liked the magvit2 paper 😓.

I'm planning to try out the other quantizers tomorrow. I'll post back if I have any luck :).

[MattMcPartlon]

@RobertLuo1 This is not a very satisfying fix, but I think I've managed to get LFQ working well by cranking down the commitment + entropy loss weights (0.001) and adding another loss term:

aux_loss=torch.mean(
                ((original_input**2 - torch.ones_like(original_input)) ** 2)[mask]
            ),

this should already be enforced in the comitment loss, but it appears to be useful for convergence (still very early in training, but it's the first run I've seen that behaves normally).

The yellow and light blue runs have the same configuration as two of the others, except for the addition of that aux _loss term. the distogram_loss is my proxy for reconstruction loss.

[lucidrains]

another thing I wanted to try is to make the projection before the scalar quantization use cosine similarity with some learned gamma