How does the nn.Embedding module relate intuitively to the idea of an embedding in general?

Viewed 762

So, I'm having a hard time understanding nn.Embedding. Specifically, I can't connect the dots between what I understand about embeddings as a concept and what this specific implementation is doing.

My understanding of an embedding is that it is a smaller dimension representation of some larger dimension data point. So it maps data in N-d to a M-d latent/embedding space such that M < N.

As I understand it, this mapping is achieved through the learning process, as in an auto-encoder. The encoder learns the optimal embedding so that the decoder can reconstruct the original input.

So my question is, how does this relate to nn.Embedding module:

A simple lookup table that stores embeddings of a fixed dictionary and size. This module is often used to store word embeddings and retrieve them using indices. The input to the module is a list of indices, and the output is the corresponding word embeddings.

does this layer "learn" a lower dimensional representation of a larger input space? Or is it something else entirely?

What I'm looking for is to take the very abstract language of the documentation to something real:

Let's say I have some input x. This input might be a vectorized image or maybe some sequence daily temperature data. In any case, this input x has 100 elements (100 days of temperature, or a 10x10 image).

How can you explain the use of nn.Embedding() in this case?

What does each argument mean in a real world context?

1 Answers

As you said, the aim when using an embedding is to reduce the dimension of your data. However, it does not learn a lower dimensional representation of a larger input space on its own. Starting from a random initialization you can improve this embedding through a learning process. This requires finding a suitable task to train the embedding on, I think, for another question. I believe it's called a "pretext task", where ultimately the objective is to have an accurate embedding matrix.

You can check the parameters of any nn.Module with .parameters(). It will return an generator.

<< [x for x in nn.Embedding(10, 2).parameters()][0].shape
>> torch.Size([10, 2])

Here, there are 10*2 parameters (i.e. dimension_input*dimension_output or by PyTorch's naming num_embeddings*embedding_dims). However it is, still, a lookup table: given an index it will return an embedding of size embedding_dims. But you these embeddings (the values of this matrix) can be changed.

Here's a little experiment:

E = nn.Embedding(10, 2)
optim = optim.SGD(E.parameters(), lr=0.01)
X = torch.randint(0, 10, size=(100,))

loss_before = E(X).mean()
loss_before.backward()
optim.step()

loss_after = E(X).mean()

As expected, loss_before and loss_after are different which shows nn.Embedding's parameters are learnable.


Edit: your question comes down to, "how do I encode my data?".

For those examples you gave precisely:

Let's say I have some input x. This input might be a vectorized image or maybe some sequence daily temperature data. In any case, this input x has 100 elements (100 days of temperature, or a 10x10 image).

You can't use a nn.Embedding to solve these cases. Embedding layers are different to a reduction matrix. The latter can be used to reduce every single vector of dimension d into dimension n where n<<d. The prerequisite to using an embedding layer is having a finite dictionnary of possible elements. For example, you might want to represent a word with a vector of size n then you would use a embedding of nb_possible_words x n. This way, for any given word in the dictionnary the layer will produce the corresponding n-size vector.

As I said in the comments below, num_embeddings is the number of unique elements you are working with and embedding_dim is the size of the embedding, i.e. the size of the output vector.

nn.Embedding is usually used at the head of a network to cast encoded data into a lower dimensionality space. It won't solve your problem by magically reducing your dimensions.

If you have a sequence of temperatures you want to analyse. You could encode each temperature into a one-hot-encoding. But this vector representation might be very large (depending on the number of different temperatures). Using an embedding layer would allow to reduce the size of these vectors. This is important when the aim is to analyse the data with a RNN any other MLP for that matter. Since the bigger your input size, the more paramaters you will have!

Related