What's the role of the commitment loss in VQ-VAE?

Question

I'm reading about VQ-VAE, and trying to understand the commitment loss $\beta||z_e(x) - sg(e)||^2$, described in the following sentence:

Finally, since the volume of the embedding space is dimensionless, it can grow arbitrarily if the embeddings $e_i$ do not train as fast as the encoder parameters. To make sure the encoder commits to an embedding and its output does not grow, we add a commitment loss, the third term in equation 3

what do they mean by the embedding space volume being dimensionless? and what would happen if we omit this term? aren't we covered by the reconstruction loss term?

Answer 1

ok I tried to remove the commitment loss and train without it and my loss blew up and the training diverged. this is what I think happens:

suppose we have a codebook of just two entries, $z=0$ for cats and $z=1$ for dogs. also let's assume that the embedding has just one dimension (a scalar), and initialize $e_0 = [-1]$ and $e_1 = [1]$.

now let's optimize the reconstruction error (the first term of our loss) - the decoder gets as input one of our two embeddings, and of course it will want to push cats and dogs away from each other. when we back-prop, we will straight-through copy this gradient to the encoder output, so the encoder will always push its output of cats towards $-\infty$, and the output of dogs towards $+\infty$, and there would be absolutely no punishment for that in the reconstruction error.

the problem would now be when we optimize for the vq-vae loss (the second term of our loss). it would get larger and larger and the embeddings would always try to chase the encoder outputs towards $-\infty$ and $+\infty$, thus creating an instable cycle, effectively ruining the training process.