Tensorflow
Word2vec
CBOW
Machine Learning
Natural Language Processing

Tensorflow Word2vec CBOW model

Master System Design with Codemia

Enhance your system design skills with over 120 practice problems, detailed solutions, and hands-on exercises.

Introduction

Continuous Bag of Words, usually shortened to CBOW, is one of the two classic Word2Vec training objectives. Instead of predicting surrounding words from one center word, CBOW predicts the center word from its surrounding context, which makes it compact, fast, and still useful for learning simple word embeddings.

How CBOW Works

Imagine the sentence:

we love natural language processing

With a context window of 2, the target word natural can be predicted from the context words:

  • 'we'
  • 'love'
  • 'language'
  • 'processing'

The model pipeline is conceptually simple:

  1. map each context token ID to an embedding vector
  2. combine the context vectors, usually by averaging or summing
  3. predict the target word across the vocabulary

CBOW is called "bag of words" because the context is treated as a set of nearby words rather than an ordered sequence. That makes the model lightweight, but it also means it does not encode word order explicitly.

Building Training Pairs

Before defining the model, you need context-target examples. For CBOW, each training example is:

  • input: the context words around the center
  • label: the missing center word

Here is a tiny runnable data-preparation example:

python
1import numpy as np
2
3sentences = [
4    "we love machine learning",
5    "we love natural language processing",
6    "tensorflow makes word embeddings easy",
7    "we build word embeddings with tensorflow",
8]
9
10tokens = sorted({word for sentence in sentences for word in sentence.split()})
11word_to_id = {word: idx + 1 for idx, word in enumerate(tokens)}
12id_to_word = {idx: word for word, idx in word_to_id.items()}
13
14window = 2
15contexts = []
16targets = []
17
18for sentence in sentences:
19    ids = [word_to_id[word] for word in sentence.split()]
20    if len(ids) < window * 2 + 1:
21        continue
22    for i in range(window, len(ids) - window):
23        context = ids[i - window:i] + ids[i + 1:i + window + 1]
24        contexts.append(context)
25        targets.append(ids[i])
26
27x = np.array(contexts, dtype="int32")
28y = np.array(targets, dtype="int32")
29
30print(x.shape)
31print(y.shape)

The main thing to verify is direction: in CBOW, the center word is the label.

A Minimal TensorFlow CBOW Model

Once the data is ready, the Keras model is compact. An Embedding layer produces vectors for the context tokens, and GlobalAveragePooling1D averages them into one context representation.

python
1import tensorflow as tf
2from tensorflow import keras
3from tensorflow.keras import layers
4
5vocab_size = len(word_to_id) + 1
6embed_dim = 16
7
8inputs = keras.Input(shape=(window * 2,), dtype="int32")
9embeddings = layers.Embedding(input_dim=vocab_size, output_dim=embed_dim)(inputs)
10pooled = layers.GlobalAveragePooling1D()(embeddings)
11outputs = layers.Dense(vocab_size, activation="softmax")(pooled)
12
13model = keras.Model(inputs, outputs)
14model.compile(
15    optimizer="adam",
16    loss="sparse_categorical_crossentropy",
17    metrics=["accuracy"],
18)
19
20model.fit(x, y, epochs=200, verbose=0)

This version uses a full softmax over the entire vocabulary. That is perfectly fine for a small demo.

Inspecting the Learned Embeddings

After training, the embedding matrix contains the learned word vectors.

python
embedding_matrix = model.layers[1].get_weights()[0]
print(embedding_matrix.shape)
print(embedding_matrix[word_to_id["tensorflow"]])

You can also test the model by giving it a context and asking for the predicted center word:

python
1test_context = np.array([
2    [
3        word_to_id["we"],
4        word_to_id["love"],
5        word_to_id["language"],
6        word_to_id["processing"],
7    ]
8], dtype="int32")
9
10pred = model.predict(test_context, verbose=0)[0]
11predicted_id = int(np.argmax(pred))
12print(id_to_word[predicted_id])

On a tiny toy corpus the quality will be limited, but the mechanics are correct and runnable.

Scaling Beyond the Demo

The simple dense softmax output becomes expensive when the vocabulary is large. Real Word2Vec-style training often uses alternatives such as sampled softmax or negative sampling to avoid computing a full vocabulary distribution for every example.

Even if you move to a more scalable loss later, the training-pair logic and the basic embedding idea stay the same. The important shift is computational efficiency, not the underlying CBOW objective.

This is why a small full-softmax demo is still a good learning step before moving to larger corpora.

Common Pitfalls

The most common mistake is reversing the training objective and accidentally preparing skip-gram pairs instead of CBOW pairs. In CBOW, the context predicts the center word.

Another issue is forgetting that CBOW ignores word order. Averaging embeddings is intentional, but it means two contexts with the same words in different orders look identical to the model.

People also often run into scalability limits with large vocabularies because a full softmax output layer becomes expensive. That is expected and is one reason larger Word2Vec systems use sampled objectives.

Finally, watch token indexing carefully. Off-by-one mistakes around reserved IDs, padding, or vocabulary size can silently produce incorrect targets.

Summary

  • CBOW predicts a center word from surrounding context words.
  • In TensorFlow, a minimal CBOW model uses Embedding, pooling, and a softmax output layer.
  • The hardest part is often generating the correct context-target training pairs.
  • Full softmax is fine for small demos, but large vocabularies usually need more efficient training methods.
  • The learned embedding matrix can be reused in downstream NLP tasks after training.

Course illustration
Course illustration

All Rights Reserved.