Question about applying Cautious Optimizer trick to nn.Embedding parameters

[jootanehorror]

Hi, thanks for releasing the code and the paper! I have a question regarding the use of the cautious update trick with nn.Embedding parameters.
As you know, during training with large vocabularies, only a small subset of tokens usually appears in each batch, so many rows in the embedding matrix do not receive gradients on every step. In AdamW (or similar optimizers), this means that the first and second moment buffers for those tokens gradually decay to near zero until they are updated again.
In the cautious update, the mask is defined as:

mask = (grad * update) > 0

I was wondering: for embeddings, where many tokens have grad = 0 most of the time, this mask might prevent certain updates or behave differently compared to dense parameters. I could not find any explicit mention of how embeddings were handled in the paper.
So my question is: should the cautious optimizer trick be applied to nn.Embedding parameters as well, or is it recommended to exclude them?
Thanks a lot for clarifying!

[kyleliang919]

so far we simplify apply the same mask = (grad * update) > 0 to the embedding and didn't do anything special.

But now that you mentioned it, it might be a good call to set mask = (grad * update) >= 0 instead to give more updates to the embedding layers. It should be a simple modification.

Overall though, since the batch size was large enough in our experiments, the gradient of the embedding might be less sparse than what we are imagining here.