What My Word2Vec Model Learned

Abstract

Previously, we looked into how a neural probabilistic language model differed from counting models. This time, I wanted to see what the embeddings themselves learned. Following Mikolov et al.'s word2vec paper, I built skip-gram in NumPy and trained it on a small slice of text8. I compared training time with full softmax and hierarchical softmax, then evaluated the saved vectors using nearest neighbours and a few analogy questions.

Contents

• Initializing the experiment
• Can we remove the hidden layer?
• Softmax being the culprit
• Taking the path to a word
• Running all the pairs
• How does C learn the representation?
• What did the vectors learn?
• The relationship we wanted but couldn't have Pt.2
• Conclusion and what's next
• References

Initializing the experiment

For this experiment, I used the first 2 million tokens of text8. Here is the setup for the completed run:

Setting	This run
Corpus	First 2 million tokens of text8
Vocabulary filter	Keep words seen at least 5 times
After filtering	21,681 unique words; 1,910,050 tokens
Training objective	Skip-gram
Context window	Radius sampled from 1 to 5 for each center position
Training pairs	11,457,894
Embedding dimension	100
Full training run	Hierarchical softmax, one pass
Learning rate	Linear decay from 0.025, with a floor of 0.0000025

I used skip-gram, where the center word predicts each neighbour separately. CBOW goes the other way, using the neighbouring words to predict the center word. The paper shows both architectures in Figure 1 and Section 3.

Let's take an example:

My name is [Pragalva] and I live in Nepal.

With a window radius of 2, the pairs are:

Center: our input	Context: the target
Pragalva	name
Pragalva	is
Pragalva	and
Pragalva	I

Every token position gets its turn as the center. A radius of 2 gives up to four pairs, fewer at the edges. In the actual run, sampling the radius from 1 to 5 gives about six pairs per position on average. That is how 1.91 million tokens became 11.46 million pairs.

The code drops rare words before making windows, so words on either side of a removed word become closer. It also treats this slice as one stream, without sentence boundaries. These pairs are what the model learns from; the next question is how it turns a center word into a prediction.

Can we remove the hidden layer?

The paper was trying to make learning useful word representations much cheaper, so the models could learn from far more text. It looked at both the quality of the vectors and the computation needed to train them (Sections 1.1 and 2).

The expensive parts were passing the embeddings through a nonlinear hidden layer and calculating an output across the entire vocabulary.

In the previous model, the context embeddings passed through a nonlinear hidden layer before reaching the output. For these skip-gram pairs, I wanted to understand whether a simpler model could still learn useful relationships.

The paper writes the feedforward language model's cost per example as:

$$Q = ND + NDH + HV$$

Here, $N$ is the number of context words, $D$ the embedding dimension, $H$ the number of hidden units, and $V$ the vocabulary size.

Cost	What the model is doing
$ND$	Looking up the context embeddings
$NDH$	Connecting those embeddings to the hidden layer
$HV$	Connecting the hidden layer to every vocabulary output

The $HV$ output term dominates with full softmax; after making the output cheaper, $NDH$ becomes the bottleneck. The new architectures remove that nonlinear hidden layer (Sections 2.1 and 3). These are two separate savings: removing the hidden layer simplifies the model, and hierarchical softmax makes each prediction cheaper.

In my skip-gram baseline, the path is simply:

center word → look up C[center] → score the vocabulary → softmax

The embedding matrix C is still learned, but its vectors now go directly into the output calculation. That removes the nonlinear hidden layer's cost. Scoring the vocabulary is still expensive, as the first run showed.

Softmax being the culprit

I started with only 20,000 pairs.

Full-softmax baseline	Recorded result
Reported average loss, early → late	9.984 → 9.470
Training time	178 seconds
Estimated time for all 11.46M pairs	About 28 hours

At the measured rate, all 11.46 million pairs would take about 28 hours. I didn't complete that run; I looked at what each pair was costing first.

The baseline gives each word two vectors: C[word] for the input role and W_out[word] for the output role. For every pair, it computes:

h = C[center]
scores = W_out @ h

That scores all 21,681 vocabulary entries, even though the pair has only one target. Full softmax then normalizes those scores, and the update touches every output row. The expensive part is the whole vocabulary-wide calculation, not just the softmax function in isolation.

Taking the path to a word

Hierarchical softmax changes that output calculation. It represents a word by a path through a binary tree, so training a pair only needs the forks on the target word's path.

The implementation builds a Huffman tree using a min-heap: repeatedly take the two smallest counts, join them, and put their combined count back. Words are leaves; each internal fork gets a trainable vector. Frequent words receive shorter paths.

These fork vectors replace W_out's per-word output vectors. The input embedding is still C[center].

Here is a five-word example with made-up counts, just to see the structure:

Figure 1. One valid Huffman tree for these illustrative counts. Left and right are code choices, not categories of meaning.

At each fork, the model computes:

$$\begin{aligned} &P(\text{right}\mid\text{center, fork}) \\ &= \sigma(C[\text{center}]\cdot W_{\text{fork}}) \end{aligned}$$

Going left has probability $1-p$ . Multiply the probabilities along a word's path to get its probability. During training, we already know the target, so we follow its path. We aren't greedily choosing whichever branch looks more likely.

At initialization, every fork is 50/50. In the toy tree, the leaf probabilities are:

Word	Path length	Initial probability
the	1	1/2
of	2	1/4
fox	3	1/8
cat	4	1/16
zebra	4	1/16
Total		1

Each fork splits the probability reaching it between its children. That is why the leaves sum to 1, even after the fork probabilities change.

My actual tree had 21,680 forks for 21,681 words:

Target	Forks evaluated per pair
the	4
of	5
king	11
zebra	18

A balanced tree would have depths around 14 or 15. Huffman coding gives common words cheaper paths at the expense of some rare words. The work now depends on the target's path length, rather than scoring the entire vocabulary.

I then trained on the same 20,000 pairs using this output structure:

Method	Pairs	Recorded training time
Full softmax	20,000	178 s
Hierarchical softmax	20,000	0.57 s

That was about 300 times faster in this small NumPy comparison, making a full pass practical. This compares training time for the same number of pairs; embedding quality still needs to be evaluated.

Running all the pairs

I used hierarchical softmax for all 11.46 million pairs, decreasing the learning rate linearly from 0.025 toward zero as it worked through them. The full training loop took about 350 seconds.

Stage	Loss recorded in my notes
Early hierarchical-softmax run	Around 7.4
First 500,000 pairs of the full run	Around 6.6
Later reporting chunks	Fluctuated around 6.4–6.8

The script printed the average loss every 500,000 pairs. I didn't save that log, so these are the rounded values I noted while it ran.

The starting loss was already lower than the full-softmax baseline's 9.98, before the vectors had learned anything. To understand that difference, look at the two initializations. With full softmax, W_out starts at zero, so all words have the same score and probability $1/21681$ :

$$-\ln(1/21681) = \ln(21681) \approx 9.98$$

With zero-initialized forks, a word at depth $d$ instead has probability $2^{-d}$ . So the, at depth 4, starts at $1/16$ . The tree gives common words an advantage through its shape, which explains the lower starting loss.

Later, the loss fluctuated around 6.4–6.8. A center word can have many different neighbours, so I wouldn't expect zero loss. But this plateau doesn't tell me the best achievable loss either: each reported chunk contains different text, while the parameters and learning rate keep changing.

What I wanted to know was whether words used in similar ways had learned similar vectors. For that, I compared the embeddings directly.

How does C learn the representation?

We start with random numbers in C. How does predicting a neighbour turn those numbers into a useful representation?

When three is the center and years is the neighbour, the model uses C[three] to predict the turns on the path to years. It compares those predictions with the correct turns, then updates both the fork vectors and C[three] to reduce the error.

When four is the center and years is the neighbour, C[four] gets updated using the same fork vectors and the same target turns.

In the training loop, the update to the input embedding comes from:

error = p - turns
grad_h = error @ W_fork[rows]
C[center] -= lr * grad_h

p contains the predicted probabilities of going right, and turns contains the correct choices, 0 for left and 1 for right. A right turn pushes the input vector in the fork vector's direction; a left turn pushes it in the opposite direction. The size of each correction depends on the prediction error. The actual loop calculates grad_h before updating the fork vectors.

So it isn't simply pulling the input toward every vector on the path. It is adjusting it to make that path more likely. C[years] is not the target vector here; the output vectors belong to the forks.

Across many shared neighbours, three and four receive similar kinds of corrections. Words used in similar contexts can develop similar representations. Nobody tells the model that these two words are numbers, and no part of the loss directly asks their embeddings to become close. That similarity can emerge because they need to make similar predictions. It doesn't mean every individual update brings them closer.

What did the vectors learn?

I kept C and compared its rows using cosine similarity. Training uses an input vector and output or fork vectors; here, I am comparing two input embeddings. These are the first five neighbours from the saved vectors, in rank order:

Word	Nearest neighbours	Top cosine
three	four, five, six, seven, eight	0.944
france	spain, italy, provence, portugal, toulouse	0.764
king	castile, alfonso, prince, trebizond, afonso	0.828
computer	iigs, design, computers, software, cpu	0.768

three was the cleanest result. Its five nearest neighbours were all numbers, even though training only asked it to predict surrounding words.

Were they already close at initialization? I compared some of these pairs before and after training:

Figure 2. Cosine similarity between input embeddings, before and after one training pass.

three and four went from 0.007 to 0.944. It isn't only numbers: france and spain, king and prince, and computer and software all started near zero and ended above 0.7. For comparison, three and computer went from 0.126 to 0.108. The model didn't simply make all of these vectors point in the same direction.

I reconstructed the starting C using the code's fixed seed and both random draws: the first initializes the unused baseline, and the second initializes hierarchical softmax. The final values come from the saved C.npy. The comparison script reproduces these numbers, and with --plot this chart, from that run's C.npy and vocab.txt.

These are neighbours selected from the trained vectors, plus one comparison pair. This shows how those particular pairs changed; it doesn't establish how well the whole vocabulary learned, or isolate the effect of years from all the other neighbours.

For france, countries and French places are mixed together. So “similar” doesn't strictly mean “interchangeable.” It can include a shared topic as well as a shared role in a sentence.

king sits near names and places associated with particular histories, and computer sits near iigs. These neighbours suggest the influence of the text used for training, though I haven't traced them back to individual articles.

The relationship we wanted but couldn't have Pt.2

Finding countries near each other doesn't yet tell me whether the vectors capture the relationship between a country and its capital. That needs a more specific test: does the offset from france to paris also take us from italy to rome?

For an analogy of the form “a is to b as c is to what?”, I calculated:

$$x = v_b - v_a + v_c$$

The paper evaluates this kind of offset by looking for the closest word and excluding the three inputs (Section 4). My eval.py normalizes each input vector before combining them, then ranks candidates by cosine similarity.

Question	Expected	Actual top three
man : king :: woman : ?	queen	epirus, sancho, tsar
france : paris :: italy : ?	rome	strasbourg, ethelred, taganrog
big : bigger :: small : ?	smaller	safe, humid, glaciers
two : four :: three : ?	six, if interpreted as doubling	six, seven, five

The first three failed. The number example returned six, but its lead over the other candidates was small:

six     0.876
seven   0.874
five    0.868

The number vectors are already packed together, and seven is barely behind six. I can't conclude from this example that the model learned doubling. Across these questions, the useful neighbourhoods didn't translate into reliable relationship offsets.

Limited data is likely part of the explanation. In Table 2 of the paper, analogy accuracy only keeps climbing when the amount of training data and the vector dimensions grow together, and even its smallest setting uses 24 million words; the strongest results come from hundreds of millions to billions. I trained on 1.91 million tokens. Training duration and the other settings could matter too, and I didn't vary any of them, so I can't isolate the cause. Four hand-picked questions are a sanity check, not an analogy benchmark.

Conclusion and what's next

This run learned recognisable neighbourhoods without the earlier model's nonlinear hidden layer. That answers part of the question I started with, while the failed analogies show where this particular run fell short.

There are limits built into this setup too. Each word has one vector regardless of its meaning in a sentence. The pair objective doesn't record whether a neighbour was on the left or the right.

What stayed with me was how the computational problem changed the experiment. With full softmax, I was watching 20,000 pairs. With the tree, I could train on all 11.46 million and actually inspect the result.

word2vec uses prediction to learn the representations we keep. It builds on the idea we explored with Bengio: learned continuous representations can capture semantic and syntactic relationships that one-hot vectors alone don't express. In my run, some of that structure was visible in the neighbours. Testing what it would take to get consistent analogy offsets is the next question.

References

• Implementation and evaluation: word2vec.py and eval.py. The full-softmax baseline remains commented out; the executable training loop runs hierarchical softmax.
• Mikolov et al., Efficient Estimation of Word Representations in Vector Space: the paper followed here.

Originally published on pragalva.me on September 27, 2026.