VQ-VAE uniform prior over the latent variables VQ-VAE propose to

*

*make the posterior categorical distribution $q(z|x)$ one-hot = defined as $1$ for the closest embedding $e$ in the codebook to the output of the encoder $z_e(x)$;

*define a simple uniform prior over $z$ (the latent variables).

The authors claim that this obtains a KL divergence which is discrete and equal to $\log K$, where $K$ is the number of discrete latent embeddings.
If I understand correctly, that means that the range of the uniform prior over $z$ should be so that every embedding has a probability equal to $1/K$ to be the closest to the encoder output.
Am I correct?
If that's so, how should we initialize the embedding vectors? Should their weights be initialized to a uniform probability - $U[\min(z_e(x)), \max(z_e(x))] ; x\in X$?

Several confusions here:
The VQ-VAE ELBO objective can be written as the KL divergence between the true and approximate posterior, as in Rocca which can be shown to be written as the expectation of the conditional likelihood - the KL divergence between the approximate posterior and the prior:
$$ argmax ELBO = argmin KL(q_x(z), p(z|x)) = $$
$$argmax(E_{z: q_x}(logp(x|z)) - KL(q_x(z),p(z)))$$
Now if the prior is uniform over the K possible values of z, and $q_x(z)$ is a one-hot distribution, it's clear that the KL divergence is a constant $logK$, so we're left with only maximizing $logp(x|z_q(x))$, which is exactly what's presented in VQ-VAE.
Also, if we assume that p(x|z) is a Gaussian whose mean is a deterministic function of z (the decoder), and the covariance is $cI$, we can arrive at the familiar reconstruction MSE loss.
All this is unrelated to the initialization of the $e_k$ embedding values, which is also uniform but might as well be from another distribution.
