Q-Loss calculation between normalized emb and unnormalized z

[emergenz]

jafar/train_tokenizer.py

Line 69 in 5ff9fc7

q_loss = jnp.square(jax.lax.stop_gradient(outputs["emb"]) - outputs["z"]).mean()

In the q-loss calculation, the normalized embedding and the unnormalized latent is used. Is this a typo?

In the following code, shouldn't we return the unnormalized x?

jafar/utils/nn.py

Lines 115 to 129 in 5ff9fc7

	def __call__(self, x: jax.Array, training: bool) -> Dict[str, jax.Array]:
	# --- Compute distances ---
	x = normalize(x)
	codebook = normalize(self.codebook)
	distance = -jnp.matmul(x, codebook.T)
	if training:
	distance = self.drop(distance)
	# --- Get indices and embeddings ---
	indices = jnp.argmin(distance, axis=-1)
	z = self.codebook[indices]
	# --- Straight through estimator ---
	z_q = x + jax.lax.stop_gradient(z - x)
	return z_q, z, x, indices

[emergenz]

(Or do the loss calculation with all parts normalized)