Word embedding is not getting better

Question

I created a simple neural network to train the word embeddings.

I have 6 tokens only: ["apple", "banana", "lime", "red", "yellow", "green"].

My training strategy is to predict the color of a fruit.

For example, "apple" should predict "red".

Since the expected outcome is straightforward, I thought the neural network would be able to predict the corresponding color of a fruit at 100% with enough training.

However, I noticed that the neural network's prediction is not improving. Increasing the token dimension, number of nodes in the hidden layer, and epoch didn't help.

Would you guys help me understand why the neural network cannot predict the color?

And what kind of approach should I take to solve this kind of problem?

Here is my code:

import torch
import torch.nn as nn
import torch.optim as optim

seed = 3
torch.manual_seed(seed)

class OurModel(nn.Module):

    def __init__(self, token_to_index):
        super().__init__()
        self.token_to_index = token_to_index
        self.vocab_size = len(token_to_index)
        token_vec_dim = 2
        self.embeds = nn.Embedding(self.vocab_size, token_vec_dim)
        self.fc = nn.Sequential(
            nn.Linear(token_vec_dim, token_vec_dim),
            nn.ReLU(),
            nn.Linear(token_vec_dim, self.vocab_size),
            nn.ReLU()
        )

    def forward(self, token_indexes):
        """
        For predicting the next token.
        :param token_indexes: List of token indexes.
        :return: A token with the size: (batch_size, vocab_size).
                 Each element in the batch represents the logits for the next token.
        """
        input_token_vectors = self.embeds(torch.LongTensor(token_indexes))
        return self.fc(input_token_vectors)

def main():
    token_to_index = {
        "apple": 0,
        "banana": 1,
        "lime": 2,
        "red": 3,
        "yellow": 4,
        "green": 5
    }
    model = OurModel(token_to_index)

    loss_function = nn.CrossEntropyLoss()
    optimizer = optim.Adam(model.parameters())
    max_epoch = 10000
    for epoch in range(max_epoch):
        logits = model([0, 1, 2])
        correct_indexes = [3, 4, 5]
        loss = loss_function(logits, torch.LongTensor(correct_indexes))
        if epoch % 1000 == 0:
            print(f"epoch: {epoch}, loss: {loss}")
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

main()

Answer 1

You need to remove the final relu layer. The relu is setting negative logit values to zero which inflates their predicted probability. This is essentially forcing the model to assign higher probabilities to incorrect values.

Also a more minor issue - it also doesn't make sense to output a vector of vocab_size if you only intend to predict color values. The output size should reflect what you actually intend to predict.