Unintentional decay of embeddings towards 0?

By updating all embeddings regardless of if they are being used (in the current batch) you are decaying them towards 0. Is this intended?

https://github.com/deepmind/sonnet/blob/master/sonnet/python/modules/nets/vqvae.py

_I have mostly read re-implementations of your code in pytorch and it could be a bug on their side but it looks like you are doing the same._

I have tried removing the hidden decay and only update the embeddings that are used but this seems to lower perplexity when training.