Skip to content
W3S101

TensorFlow - Word Embedding

TensorFlow - Word Embedding with Keras

Section titled “TensorFlow - Word Embedding with Keras”

Word embedding is a technique in Natural Language Processing (NLP) where words or phrases from a vocabulary are mapped to vectors of real numbers. These dense vector representations capture semantic relationships between words, such that words with similar meanings are closer in the vector space. This is crucial for providing meaningful input to machine learning models.

For example, after training, word vectors might look like this (illustrative, actual values depend on training data and model): blue: (0.013, 0.001, 0.246, ..., -0.252, 1.005, 0.063) blues: (0.014, 0.119, -0.490, ..., 0.033, -0.100, 0.116) orange: (-0.248, -0.124, 0.210, ..., 0.080, 0.239, -0.014) oranges: (-0.356, 0.219, 0.081, ..., -0.354, 0.385, -0.071) Notice how plural forms might be close to their singular counterparts, and colors might cluster together.

Word2Vec (Skip-gram with Negative Sampling)

Section titled “Word2Vec (Skip-gram with Negative Sampling)”

Word2Vec is a popular model for learning word embeddings in an unsupervised manner from a large corpus of text. One of its architectures is the Skip-gram model, which learns to predict the context words (surrounding words) given a target word. To make training more efficient, Negative Sampling is often used, where the model is trained to distinguish true context words from a few randomly selected ‘negative’ (incorrect) words.

TensorFlow, particularly with its Keras API, provides tools to implement Word2Vec. Below is an example demonstrating the Skip-gram model with negative sampling using tf.keras.layers.Embedding and a custom training loop for NCE loss. Note: tf.nn.nce_loss is a lower-level API. For simpler applications, one might use pre-trained embeddings or simpler Keras models.

import os
import math
import numpy as np
import tensorflow as tf
from tensorflow import keras
from tensorflow.keras import layers
import collections
import random
from tensorflow.keras.preprocessing.sequence import skipgrams


# Parameters
batch_size = 128
embedding_dimension = 128  # Increased for better representation
negative_samples = 5       # Number of negative samples for NCE loss
window_size = 2            # Context window size for skip-grams
num_epochs = 5
LOG_DIR = "logs/word2vec_keras"


if not os.path.exists(LOG_DIR):
    os.makedirs(LOG_DIR)


# Sample sentences (a larger, more diverse corpus is needed for meaningful embeddings)
digit_to_word_map = {
    1: "One", 2: "Two", 3: "Three", 4: "Four", 5: "Five",
    6: "Six", 7: "Seven", 8: "Eight", 9: "Nine", 0: "Zero"
}
sentences = []
for i in range(10000):
    rand_odd_ints = np.random.choice(list(range(1, 10, 2)) + [0], 3)
    sentences.append(" ".join([digit_to_word_map[r].lower() for r in rand_odd_ints]))
    rand_even_ints = np.random.choice(list(range(2, 10, 2)) + [0], 3)
    sentences.append(" ".join([digit_to_word_map[r].lower() for r in rand_even_ints]))


# Tokenize sentences and build vocabulary
words = [word for sent in sentences for word in sent.split()]
word_counts = collections.Counter(words)
vocabulary_size = len(word_counts)
word2index = {word: i for i, (word, _) in enumerate(word_counts.most_common())}
index2word = {i: word for word, i in word2index.items()}


# Generate skip-gram pairs
encoded_sentences = [[word2index[word] for word in sent.split()] for sent in sentences]


skip_gram_pairs = []
for sent_indices in encoded_sentences:
    pairs, _ = skipgrams(sequence=sent_indices,
                           vocabulary_size=vocabulary_size,
                           window_size=window_size,
                           negative_samples=0, # We'll handle negative sampling in the loss
                           shuffle=True)
    for pair in pairs:
        skip_gram_pairs.append(pair)


skip_gram_pairs = np.array(skip_gram_pairs)
target_words, context_words = skip_gram_pairs[:, 0], skip_gram_pairs[:, 1]


# Define the Word2Vec model using Keras Functional API
input_target = keras.Input((1,), name='target_word')
input_context = keras.Input((1,), name='context_word')


embedding_layer = layers.Embedding(vocabulary_size,
                                   embedding_dimension,
                                   name='word_embedding')


target_embedding = embedding_layer(input_target)
target_embedding = layers.Reshape((embedding_dimension,))(target_embedding)


context_embedding = embedding_layer(input_context)
context_embedding = layers.Reshape((embedding_dimension,))(context_embedding)


# Using dot product to calculate similarity for positive pairs
positive_similarity = layers.Dot(axes=1, normalize=False)([target_embedding, context_embedding])


# For negative sampling, we'd typically use nce_loss or sample_softmax_loss in a custom training loop,
# or approximate it with binary classification (true pair vs. false pair).
# The original example used tf.nn.nce_loss. A full Keras model with NCE is complex to integrate cleanly.
# Here's a conceptual model structure. Actual NCE training is better with custom train_step.


# Simplified approach: define a model that can be trained with binary cross-entropy
# by generating explicit negative samples and labels (1 for positive, 0 for negative).
# This is not exactly NCE but a common alternative.


# For this example, we stick to a structure that *could* use tf.nn.nce_loss in a custom loop.
# However, demonstrating the full custom loop with NCE is extensive for this format.
# Let's focus on the embedding layer and how to extract embeddings.


# We'll define a model that just has the embedding layer for visualization purposes,
# as full NCE training is verbose.
class Word2VecModel(keras.Model):
    def __init__(self, vocab_size, embed_dim):
        super(Word2VecModel, self).__init__()
        self.target_embedding = layers.Embedding(vocab_size, embed_dim,
                                                input_length=1,
                                                name="w2v_embedding")
        # NCE loss typically requires separate weights for context predictions
        self.nce_weights = layers.Dense(embed_dim, use_bias=False)
        # This structure is simplified. True NCE involves sampling.


    def call(self, pair):
        target, context = pair
        word_embed = self.target_embedding(target)
        # context_embed = self.target_embedding(context) # if using shared weights
        # In a full NCE setup, one would compute loss against sampled negative context words.
        return word_embed # Simplified for now


# Instantiate a simple embedding model for demonstration
simple_embedding_model = keras.Sequential([
    layers.InputLayer(input_shape=(1,)),
    layers.Embedding(vocabulary_size, embedding_dimension, name='embedding_layer')
])


# A full training example with tf.nn.nce_loss requires a custom training loop.
# Let's illustrate setting up the embedding matrix and NCE loss function.


# Embedding matrix (target embeddings)
embeddings_matrix = tf.Variable(
    tf.random.uniform([vocabulary_size, embedding_dimension], -1.0, 1.0),
    name='target_embeddings')


# NCE weights and biases (context/output embeddings)
nce_weights = tf.Variable(
    tf.random.truncated_normal([vocabulary_size, embedding_dimension],
                               stddev=1.0 / math.sqrt(embedding_dimension)))
nce_biases = tf.Variable(tf.zeros([vocabulary_size]))


optimizer = keras.optimizers.Adam(learning_rate=0.01)


@tf.function
def train_step(target_input, context_labels):
    with tf.GradientTape() as tape:
        # Look up embeddings for the target words
        embed = tf.nn.embedding_lookup(embeddings_matrix, target_input)


        # Compute NCE loss
        loss = tf.reduce_mean(
            tf.nn.nce_loss(weights=nce_weights,
                           biases=nce_biases,
                           labels=tf.expand_dims(context_labels, axis=1),
                           inputs=embed,
                           num_sampled=negative_samples,
                           num_classes=vocabulary_size))
    gradients = tape.gradient(loss, [embeddings_matrix, nce_weights, nce_biases])
    optimizer.apply_gradients(zip(gradients, [embeddings_matrix, nce_weights, nce_biases]))
    return loss


print("Starting training with custom NCE loop...")
num_steps = len(target_words) // batch_size
for epoch in range(num_epochs):
    epoch_loss = 0
    # Shuffle data each epoch
    indices = np.arange(len(target_words))
    np.random.shuffle(indices)
    shuffled_targets = target_words[indices]
    shuffled_contexts = context_words[indices]


    for step in range(num_steps):
        start = step * batch_size
        end = (step + 1) * batch_size
        batch_targets = shuffled_targets[start:end]
        batch_contexts = shuffled_contexts[start:end]


        current_loss = train_step(batch_targets, batch_contexts)
        epoch_loss += current_loss


        if step % 500 == 0:
            print(f"Epoch {epoch+1}, Step {step}, Loss: {current_loss.numpy():.4f}")
    print(f"Epoch {epoch+1} average loss: {epoch_loss.numpy()/num_steps:.4f}")


print("Training finished.")


# Normalize embeddings before using (optional, but common for cosine similarity)
final_embeddings = embeddings_matrix.numpy()
norm = np.sqrt(np.sum(np.square(final_embeddings), 1, keepdims=True))
normalized_embeddings = final_embeddings / norm


# Example: Find words similar to 'one'
ref_word_str = "one"
if ref_word_str in word2index:
    ref_word_idx = word2index[ref_word_str]
    ref_word_embedding = normalized_embeddings[ref_word_idx]


    # Calculate cosine similarity
    similarities = np.dot(normalized_embeddings, ref_word_embedding)


    # Get top N similar words
    sorted_indices = np.argsort(similarities)[::-1]
    print(f"\nWords similar to '{ref_word_str}':")
    for i in range(1, 6): # Skip the first one (itself)
        similar_word_idx = sorted_indices[i]
        similar_word_str = index2word[similar_word_idx]
        print(f"- {similar_word_str} (similarity: {similarities[similar_word_idx]:.3f})")
else:
    print(f"Word '{ref_word_str}' not in vocabulary.")


# For TensorBoard embedding visualization (this part requires more setup for TF2)
# You'd typically save the embeddings and metadata.tsv file.
# E.g., save metadata:
with open(os.path.join(LOG_DIR, 'metadata.tsv'), "w") as f:
    for i in range(vocabulary_size):
        f.write(index2word[i] + "\n")
# The embeddings themselves (normalized_embeddings) would also be saved,
# and TensorBoard configured to load them. This is more involved than can be shown briefly.
# See: https://www.tensorflow.org/tensorboard/tensorboard_projector_plugin


print(f"\nEmbeddings and metadata.tsv (for projector) can be found in {LOG_DIR}")
# To visualize, you can use the projector: http://projector.tensorflow.org/
# by uploading the saved embedding tensor and metadata file.
# Or configure TensorBoard with a projector_config.pbtxt.
# For Keras Embedding layers, TensorBoard callback can also log embeddings.

Output Interpretation

Section titled “Output Interpretation”

The script will:

  1. Generate skip-gram pairs from sample sentences.
  2. Define TensorFlow Variables for target embeddings, NCE weights, and biases.
  3. Train these embeddings using a custom loop with tf.nn.nce_loss and Adam optimizer.
  4. Print loss values during training.
  5. After training, it will normalize the learned embeddings.
  6. Finally, it will find and print words most similar to ‘one’ based on cosine similarity of their embeddings. Due to the small and synthetic dataset, the ‘similarity’ might not be deeply semantic but should reflect co-occurrence patterns in the generated sentences (e.g., other odd numbers or ‘zero’). With a large, real-world text corpus, Word2Vec can learn rich semantic relationships.

Modern approaches often involve using pre-trained embeddings (like Word2Vec, GloVe, FastText) or training more advanced models like Transformers (e.g., BERT, GPT) which learn contextual embeddings.