PyTorch - Word Embedding

PyTorch - Modern Word Embeddings with nn.Embedding and Skip-Gram

In Natural Language Processing (NLP), word embeddings are dense vector representations of words that capture semantic meaning. This chapter focuses on understanding and implementing a foundational word embedding model, Word2Vec, specifically its Skip-Gram variant, using modern PyTorch practices. While Word2Vec itself was originally implemented in C, PyTorch allows for flexible and efficient custom model creation.

The Word2Vec Skip-Gram model learns word embeddings by predicting context words (surrounding words) given a target word. We will implement this using PyTorch’s nn.Embedding layer, which is the standard way to handle embeddings.

Core Concepts of Skip-Gram

Target Word: The input word for which we want to learn an embedding.
Context Words (Positive Samples): Words that appear nearby the target word in a corpus.
Negative Samples: Words that do not appear in the context of the target word, randomly sampled from the vocabulary. These are crucial for efficient training.
Two Embedding Layers: Typically, Word2Vec uses two sets of embeddings for each word: one when it’s a target word (u_embeddings) and another when it’s a context word (v_embeddings). Often, only the u_embeddings are used as the final word vectors.

Step 1: Essential Imports

First, let’s import the necessary PyTorch modules. Note that torch.autograd.Variable is no longer needed as tensors can track gradients directly.

import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim # For optimizers

# Example: Check if CUDA is available and set device
# device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

Step 2: Implementing the Skip-Gram Model

We’ll define a class SkipGramModel that inherits from nn.Module. This model will include two nn.Embedding layers: one for target words and one for context words.

class SkipGramModel(nn.Module):
    def __init__(self, vocab_size, embedding_dim):
        super(SkipGramModel, self).__init__()
        self.vocab_size = vocab_size
        self.embedding_dim = embedding_dim

        # Embedding layer for target words (center words)
        self.u_embeddings = nn.Embedding(vocab_size, embedding_dim, sparse=True)
        # Embedding layer for context words
        self.v_embeddings = nn.Embedding(vocab_size, embedding_dim, sparse=True)

        self.init_embeddings()

    def init_embeddings(self):
        init_range = 0.5 / self.embedding_dim
        self.u_embeddings.weight.data.uniform_(-init_range, init_range)
        # Initialize v_embeddings similar to u_embeddings, not to zero
        self.v_embeddings.weight.data.uniform_(-init_range, init_range)

    def forward(self, target_word_indices, positive_context_indices, negative_context_indices):
        """
        target_word_indices: Tensor of shape (batch_size)
        positive_context_indices: Tensor of shape (batch_size)
        negative_context_indices: Tensor of shape (batch_size, num_negative_samples)
        """
        # Get embeddings for the target word
        # Shape: (batch_size, embedding_dim)
        emb_target = self.u_embeddings(target_word_indices)

        # Get embeddings for the positive context word
        # Shape: (batch_size, embedding_dim)
        emb_positive_context = self.v_embeddings(positive_context_indices)

        # Calculate score for positive pairs (target * positive_context)
        # Higher score means more similar
        # torch.mul performs element-wise multiplication
        # torch.sum along dim=1 sums up the dot product components
        positive_score = torch.sum(torch.mul(emb_target, emb_positive_context), dim=1)
        positive_loss = F.logsigmoid(positive_score) # Maximize this score

        # Get embeddings for negative context words
        # Shape: (batch_size, num_negative_samples, embedding_dim)
        emb_negative_context = self.v_embeddings(negative_context_indices)

        # Calculate scores for negative pairs (target * negative_contexts)
        # torch.bmm (batch matrix multiplication) is efficient here
        # emb_target needs to be unsqueezed to (batch_size, embedding_dim, 1) for bmm
        # Or, more commonly, (batch_size, 1, embedding_dim) and multiply with (batch_size, num_negative_samples, embedding_dim).transpose(1,2)
        # Simpler: iterate or use broadcasting if num_negative_samples is small, or use element-wise product and sum carefully.
        # For this example, let's use batch matrix multiplication after reshaping target embedding.
        # emb_target_expanded shape: (batch_size, 1, embedding_dim)
        emb_target_expanded = emb_target.unsqueeze(1)

        # emb_negative_context shape: (batch_size, num_negative_samples, embedding_dim)
        # We want dot product for each negative sample with the target embedding.
        # (batch_size, 1, embedding_dim) @ (batch_size, embedding_dim, num_negative_samples)
        # -> (batch_size, 1, num_negative_samples)
        negative_scores = torch.bmm(emb_target_expanded, emb_negative_context.transpose(1, 2)).squeeze(1)
        negative_loss = F.logsigmoid(-negative_scores) # Minimize these scores (maximize -scores)

        # Total loss is the sum of negative log likelihoods
        # We want to maximize positive_loss and sum of negative_loss elements
        # So we minimize -(sum of positive_loss + sum of all negative_loss elements)
        total_loss = -(torch.sum(positive_loss) + torch.sum(negative_loss))
        return total_loss

    def get_word_embeddings(self):
        # Usually, the u_embeddings are used as the final word vectors
        return self.u_embeddings.weight.data.cpu().numpy()

    def save_embedding_vectors(self, id_to_word_map, output_file_path, use_cuda_device=False):
        if use_cuda_device:
            embeddings = self.u_embeddings.weight.cpu().data.numpy()
        else:
            embeddings = self.u_embeddings.weight.data.numpy()

        with open(output_file_path, 'w', encoding='utf-8') as fout:
            fout.write(f'{len(id_to_word_map)} {self.embedding_dim}\n')
            for word_id, word_string in id_to_word_map.items():
                embedding_vector = embeddings[word_id]
                embedding_str = ' '.join(map(str, embedding_vector))
                fout.write(f'{word_string} {embedding_str}\n')
        print(f"Embeddings saved to {output_file_path}")

# Example usage (conceptual - full training loop not shown here)
def example_training_step(model, optimizer, target, pos_context, neg_context):
    model.train() # Set model to training mode
    optimizer.zero_grad() # Clear previous gradients
    loss = model(target, pos_context, neg_context)
    loss.backward() # Compute gradients
    optimizer.step() # Update model parameters
    return loss.item()

Step 3: Using the Model (Conceptual)

To train this model, you would need a large text corpus, a vocabulary (mapping words to integer IDs), and a way to generate (target, positive_context, negative_context) batches. An optimizer like torch.optim.SGD or torch.optim.Adam would be used to update the model’s weights.

if __name__ == '__main__':
    # --- Hyperparameters (Example) ---
    VOCAB_SIZE = 10000  # Size of your vocabulary
    EMBEDDING_DIM = 100 # Desired dimensionality of word vectors
    # device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

    # --- Instantiate model ---
    model = SkipGramModel(VOCAB_SIZE, EMBEDDING_DIM) # .to(device)
    # print(model)

    # --- Dummy data for illustration (replace with actual data loading) ---
    # BATCH_SIZE = 64
    # NUM_NEGATIVE_SAMPLES = 5
    # dummy_target = torch.randint(0, VOCAB_SIZE, (BATCH_SIZE,)) # .to(device)
    # dummy_pos_context = torch.randint(0, VOCAB_SIZE, (BATCH_SIZE,)) # .to(device)
    # dummy_neg_context = torch.randint(0, VOCAB_SIZE, (BATCH_SIZE, NUM_NEGATIVE_SAMPLES)) # .to(device)

    # --- Optimizer ---
    # optimizer = optim.Adam(model.parameters(), lr=0.001)

    # --- Training loop (conceptual) ---
    # for epoch in range(num_epochs):
    #     for target, pos_context, neg_context in data_loader:
    #         loss = example_training_step(model, optimizer, target, pos_context, neg_context)
    #         # print(f'Epoch {epoch}, Loss: {loss}')

    # --- Saving embeddings ---
    # id_to_word = {i: str(i) for i in range(VOCAB_SIZE)} # Dummy id_to_word map
    # model.save_embedding_vectors(id_to_word, 'word_embeddings.txt')
    print("SkipGramModel defined. Fill in data loading and training loop to use.")

Applications and Further Learning

Word embeddings learned via Word2Vec are foundational for many NLP tasks, including:

Sentiment Analysis
Machine Translation
Text Classification
Question Answering

While implementing Word2Vec from scratch is a great learning exercise, for many applications, using pre-trained word embeddings (like GloVe, FastText, or those from transformer models like BERT) or fine-tuning nn.Embedding layers as part of a larger task-specific neural network is common practice. You can find more details on nn.Embedding and related PyTorch functionalities in the official PyTorch documentation.