I’m quite new to using Pytorch and deep learning. What size of unique categories of a categorical variable is appropriate for applying the nn.Embedding ideally (best practices)? for example, if a feature has only two unique values, does that mean that it’s better to use one-hot encoding instead of teaching nn.Embedding (1)? or I should just initialize nn.Embedding with one-hot matrix [[0, 1], [1,0]] (2)?

The question is which of three approaches ((3) - learn nn.Embedding without my initialization) are preffered for small category space


1 Answer 1


nn.embedding layer is a sytactic sugar equivalent to one hot vector+ linear layer.

Suppose you have 2 distinct variables. and you want your model to learn their vector representation of size 3 based on your task at hand.

You can either describe them as one hot vector of 1,2 and a linear layer of 2,3 When you multiple one hot vector with weight matrix of linear layer, it is equivalent to picking up a column vector from that matrix.

import numpy as np
first_one_hot = np.array([1,0])
second_one_hot = np.array([0,1])
word_vectors = np.array([[3,4],[5,6]])
print(f"first_vector: {np.dot(first_one_hot, word_vectors)}")
print(f"second_vector: {np.dot(second_one_hot, word_vectors)}")

Equivalent pytorch nn.Embeddings code:

    import torch
    import torch.nn as nn

    embedding = nn.Embedding(2, 2)
    print(f"embedding.weight: {embedding.weight}")
    # Compute the word vectors using matrix-vector multiplication
    first_vector = embedding(torch.tensor([0]))
    second_vector = embedding(torch.tensor([1]))

    print(f"first_vector: {first_vector}")
    print(f"second_vector: {second_vector}")

nn.Embedding saves you from manually picking corresponding embedding by using matrix multiplication with one-hot. Instead you can just pass the index and get the corresponding embedding vector.

If at all you need to initialize the Embedding layer with existing vectors it"s usually not one-hot vectors but rather pretrained vectors known to work well for that task. For example, we can initialize embedding layer with word2vec vectors and continue training.

Using following:

embedding = nn.Embedding.from_pretrained(word_vectors, freeze=True)

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service and acknowledge that you have read and understand our privacy policy and code of conduct.

Not the answer you're looking for? Browse other questions tagged or ask your own question.