Appropriate input size for nn.Embedding

Question

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

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)