Learning codebooks also updates the encoder ?

[gdoras]

In the case of a learnable codebook, it seems to me that the encoder outputs are not completely detached before they are used to compute the codebook loss because they are still connected to the loss indirectly via distance matrix and quantized vectors . The encoder still accumulates gradients (even though they are not updated by the in-place optimizer). Should not the encoder outputs be detached also before they are used to compute the distances ?

[lucidrains]

@gdoras hey, thanks for pointing this out

think this is only an issue if you have straight_through turned on? could you write a short script to confirm this?

[gdoras]

The assertion you added in 1.27.14 actually raised another question. If we use STE (no matter if we update the codebook via EMA or optimizer), it is set both in the VQ class and in the codebook class. The encoder gets gradients from two paths: 1) from z_q to z here but also 2) from the logits to the one-hots here. It can be checked easily by making VQ STE optional here:

    elif self.straight_through:
        # standard STE to get gradients through VQ layer.
        quantize = straight_through(x, quantize)

and toggling the STE in the VQ and/or the codebook.
Is this behavior really intended ? Why should the encoder accumulate somewhat redundant gradients from both paths ?

[lucidrains]

@gdoras yeah you are right, probably never makes sense to use it