Hi,
Why do you think it is too slow?
Running the backward should be between 1 and 2x the forward pass.
Then the gradient step depends on the size of your weights.
If your Embedding layers are very large compared to the rest of the net, you can do sparse updates by using sparse=True for it. See doc 254. And by using an optimizer that supports sparse updates like SGD or SparseAdam.