Equivalent of embeddings_regularizer in pyTorch

Question

tf.keras.layers.Embedding has parameter embeddings_regularizer. What would be equivalent of this in pyTorch or nn.Embedding?

Answer 1

There is no direct equivalent for PyTorch as PyTorch only supports L2 regularization on parameters via torch.optim optimizers.

For example torch.optim.SGD has weight_decay parameter. If you set it up and you optimize your nn.Embedding it will be regularized by L2 with specified strength (you can pass only nn.Embedding for weight_decay, see per-parameter-options of optimizers).

If you wish to use L1 regularization you would have to:

• code it on your own
• use available third party solutions

Coding on your own

Usually we add L1 regularization to loss and backpropagate, but this is inefficient approach.

It is better to populate gradients of our parameters (there are some edge cases though) with the derivative of regularization (for L1 is it sign value). Something along those lines:

import torch

# Do this in your optimization loop AT THE TOP
embedding = torch.nn.Embedding(150, 100)
embedding.weight.grad = torch.sign(embedding.weight)

# Do the rest of optimization AND clear gradients!
...

Though it is a little harder to make it work in general (stuff like batch accumulation etc.) and pretty unclear IMO. You could apply L2 on top of that also.

torchlayers third party library

Disclaimer: I'm the author of this project

You can install torchlayers-nightly and get per-layer L1 and L2 regularization.

Install via pip:

pip install -U torchlayers-nightly

In your code you could do:

import torchlayers as tl
import torch

embedding = torch.nn.Embedding(150, 100)
regularized_embedding = tl.L1(embedding)

# Do the rest as per usual

This feature is experimental for now, but should work and I've used it with success previously.

Also, you should be able to use tl.L2 the same way, see docstrings about usage of this particular layer.

For more info check github repository and read documentation here.