[All-about-embedding 2] Implementing Word2Vec — From Data to Meaning Axes

2026-08-26 → 2026-08-26

02. Implementing Word2Vec — From Data to Meaning Axes#

All About Embedding · Part 1: Word2Vec, SGNS, and Gravity · Part 2: Implementation, Relations, and Meaning Axes · Part 3: Culture and Semantic Change · Part 4: Research Companion

© 2026 Byunghwee Lee. Created and edited in collaboration with OpenAI Codex. Unauthorized copying, reproduction, distribution, and use are prohibited.

Part 1 established the mathematical story: a context window creates center–context events, SGNS compares those events with noise, and the learned dot product approximates a low-rank data-versus-noise log-odds matrix. This note turns that story into code and then explains the most useful geometry that emerges from it.

The central progression is:

corpus
→ positive and negative pairs
→ SGNS training
→ relation profiles
→ repeated difference directions
→ analogies and meaning axes

After reading this note#

  • You can build positive pairs and a negative-sampling distribution from a corpus.
  • You can implement the SGNS objective in NumPy or PyTorch and explain every tensor shape.
  • You can connect the learned vectors to relation profiles rather than treating coordinates as named features.
  • You can explain why vector subtraction can reveal a repeated relationship.
  • You can construct, score, and validate a meaning axis from matched pairs or two pole sets.
  • You can connect Word2Vec to matrix factorization, recommendation, node embedding, and energy models.

Table of Contents#

  1. Ⅰ. Implementing and Creating Data — 1–3
  2. Ⅱ. Deeper Interpretation of SGNS — 4–6
  3. Ⅲ. A Minimal Training Experiment — 7–9
  4. Ⅳ. Relationship Direction and Meaning Axis — 10–16
  5. Ⅴ. Connections to Other Models — 17–18
  6. Ⅵ. Core Formulas and Mental Model — 19–21

Ⅰ. Implementing and Creating Data#

↑ Table of Contents

1. From text to training events#

SGNS does not train directly on sentences. It trains on positive (center, context) pairs produced by a context window.

For the sentence

the cat likes warm food

with window=2, the center cat produces:

(cat, the)
(cat, likes)
(cat, warm)

The window therefore defines what “related” means. Small windows emphasize local syntactic roles; larger windows include broader topical context.

from collections import Counter


def build_vocabulary(sentences, min_count=1):
    counts = Counter(token for sentence in sentences for token in sentence)
    words = [word for word, count in counts.items() if count >= min_count]
    word_to_id = {word: i for i, word in enumerate(words)}
    return word_to_id, counts


def positive_pairs(sentences, word_to_id, window=2):
    for sentence in sentences:
        ids = [word_to_id[word] for word in sentence if word in word_to_id]
        for i, center in enumerate(ids):
            left = max(0, i - window)
            right = min(len(ids), i + window + 1)
            for j in range(left, right):
                if i != j:
                    yield center, ids[j]

A dynamic window samples a radius between 1 and the maximum window for each center. This gives nearby contexts more weight without changing the learning objective.

High-frequency words can dominate the event stream. Word2Vec commonly subsamples them before pair generation. Conceptually, this changes the effective corpus: the model learns from a deliberately reweighted set of encounters, not raw text alone.

2. Negative sampling and the noise baseline#

For every positive pair (w, c), draw several context IDs from a noise distribution:

$$ P_n(c) \propto \operatorname{count}(c)^{3/4}. $$

The $3/4$ exponent softens the raw unigram distribution. Frequent contexts still appear often as negatives, but not in direct proportion to their overwhelming frequency.

import torch


def make_noise_distribution(word_to_id, counts, exponent=0.75):
    weights = torch.tensor(
        [counts[word] ** exponent for word in word_to_id],
        dtype=torch.float,
    )
    return weights / weights.sum()


def draw_negatives(noise_probs, batch_size, n_negatives):
    sampled = torch.multinomial(
        noise_probs,
        batch_size * n_negatives,
        replacement=True,
    )
    return sampled.view(batch_size, n_negatives)

Negative sampling does two jobs at once:

  1. It avoids a full softmax over the vocabulary.
  2. It defines the baseline against which an observed pair must be surprising.

Thus (cat, the) may be common but not especially informative because the is also common under the noise distribution. (cat, pet) can receive a stronger relative score because it occurs more often than its baseline would predict.

3. The SGNS model in code#

Each token has two learned rows:

  • Q[w]: the word as a center, or query.
  • K[c]: the word as a context, or key.

For one positive context $c^+$ and negative contexts $c_i^-$, the loss is

$$ \mathcal L = -\log\sigma(q_w^\top k_{c^+}) -\sum_i\log\sigma(-q_w^\top k_{c_i^-}). $$

The following NumPy function exposes one update step explicitly.

import numpy as np


def sigmoid(x):
    return 1.0 / (1.0 + np.exp(-x))


def sgns_step(Q, K, center_id, positive_id, negative_ids, lr=0.01):
    q = Q[center_id].copy()
    k_pos = K[positive_id].copy()
    k_neg = K[negative_ids].copy()

    pos_score = q @ k_pos
    neg_scores = k_neg @ q

    loss = -np.log(sigmoid(pos_score) + 1e-12)
    loss -= np.log(sigmoid(-neg_scores) + 1e-12).sum()

    pos_error = sigmoid(pos_score) - 1.0
    neg_error = sigmoid(neg_scores)

    grad_q = pos_error * k_pos
    grad_q += (neg_error[:, None] * k_neg).sum(axis=0)
    grad_k_pos = pos_error * q
    grad_k_neg = neg_error[:, None] * q

    Q[center_id] -= lr * grad_q
    K[positive_id] -= lr * grad_k_pos
    for context_id, gradient in zip(negative_ids, grad_k_neg):
        K[context_id] -= lr * gradient

    return float(loss)

The copied rows matter because all gradients must be computed from the same forward state. Automatic differentiation handles this bookkeeping in PyTorch:

import torch.nn as nn
import torch.nn.functional as F


class SGNS(nn.Module):
    def __init__(self, vocab_size, dim):
        super().__init__()
        self.input_embedding = nn.Embedding(vocab_size, dim)
        self.output_embedding = nn.Embedding(vocab_size, dim)

    def forward(self, center, positive, negatives):
        # center, positive: [batch]
        # negatives:       [batch, n_negatives]
        q = self.input_embedding(center)       # [B, d]
        k_pos = self.output_embedding(positive) # [B, d]
        k_neg = self.output_embedding(negatives) # [B, n, d]

        pos_score = (q * k_pos).sum(dim=1)                    # [B]
        neg_score = torch.bmm(k_neg, q.unsqueeze(2)).squeeze(2) # [B, n]

        pos_loss = -F.logsigmoid(pos_score)
        neg_loss = -F.logsigmoid(-neg_score).sum(dim=1)
        return (pos_loss + neg_loss).mean()

Training then follows the standard pattern:

model = SGNS(vocab_size=len(word_to_id), dim=100)
optimizer = torch.optim.Adam(model.parameters(), lr=1e-3)

for center, positive in loader:
    negatives = draw_negatives(
        noise_probs,
        batch_size=center.shape[0],
        n_negatives=5,
    )

    optimizer.zero_grad()
    loss = model(center, positive, negatives)
    loss.backward()
    optimizer.step()

The entire pipeline is now visible:

tokens → window pairs → noise samples → embedding lookup
       → dot products → logistic loss → shared-vector updates

Ⅱ. Deeper Interpretation of SGNS#

↑ Table of Contents

4. The pairwise optimum#

Let $x_{wc}$ be the number of positive occurrences of pair $(w,c)$ and $x_w=\sum_c x_{wc}$. With $n_{\mathrm{neg}}$ negatives per positive event, the expected number of negative occurrences of the pair is

$$ n_{\mathrm{neg}}x_wP_n(c). $$

If each pair score could be optimized independently, its optimum would be

$$ \boxed{ s_{wc}^* = \log \frac{P_{\mathrm{data}}(c\mid w)} {n_{\mathrm{neg}}P_n(c)} }. $$

With unigram noise, $P_n(c)=P(c)$, this becomes

$$ q_w^\top k_c \approx PMI(w,c)-\log n_{\mathrm{neg}}. $$

The actual model cannot assign every pair an independent score. It must use the shared, low-dimensional factorization $s_{wc}=q_w^\top k_c$. The embedding is therefore a compressed approximation to the full log-odds relation matrix.

Increasing $n_{\mathrm{neg}}$ lowers the idealized score baseline by $\log n_{\mathrm{neg}}$. This is why a genuinely observed pair need not have a positive dot product: it only needs to score appropriately relative to the alternatives.

5. A word vector is a compressed relation profile#

For a center word $w$, consider its scores against every context key:

$$ r_w = q_w^\top K^\top = [q_w^\top k_1,\ldots,q_w^\top k_V]. $$

$r_w$ is the word’s full relation profile. If cat and dog occur in similar contexts, then

$$ r_{cat}\approx r_{dog}. $$

The $V$-dimensional profiles are too large to store and compare directly. Word2Vec represents each one with a $d$-dimensional vector. Similar vectors arise because similar words repeatedly solve similar prediction problems.

The training score depends on both direction and magnitude:

$$ q^\top k=\|q\|\,\|k\|\cos\theta. $$

Cosine similarity is useful after training because it removes vector magnitude and focuses on direction. It is an evaluation geometry, not the SGNS training loss.

6. From log-odds to gravity#

Rearranging the pairwise optimum gives

$$ P_{\mathrm{data}}(w,c) \propto P_{\mathrm{data}}(w)P_n(c) \exp(q_w^\top k_c). $$

This has the structure of a gravity model:

observed interaction
= source activity
× destination baseline
× latent pair attraction

The endpoint frequencies explain how active the two sides are. The exponential dot product explains pair-specific attraction beyond that baseline.

The same idea extends to a graph flow $F_{ij}$:

$$ F_{ij}\propto m_i m_j\exp(q_i^\top k_j). $$

Traditional gravity models begin with a known space, such as geographic distance. Embedding models instead learn a latent interaction space from observed flows.


Ⅲ. A Minimal Training Experiment#

↑ Table of Contents

7. What to inspect while training#

Loss alone does not explain the learned geometry. A small diagnostic should track three levels:

  1. Optimization: Does the average SGNS loss decrease?
  2. Pair scores: Do observed pairs score above mismatched pairs for the same center?
  3. Word geometry: Do words with similar context profiles have higher cosine similarity?

For a tiny corpus containing animal and vehicle sentences, inspect values such as:

def cosine(a, b, eps=1e-12):
    return float(a @ b / max(np.linalg.norm(a) * np.linalg.norm(b), eps))


animal_score = Q[id_cat] @ K[id_pet]
mismatch_score = Q[id_cat] @ K[id_road]
animal_similarity = cosine(Q[id_cat], Q[id_dog])
cross_similarity = cosine(Q[id_cat], Q[id_car])

A successful mechanism-level result should usually satisfy

$$ s(cat,pet)>s(cat,road) $$

and

$$ \cos(q_{cat},q_{dog})>\cos(q_{cat},q_{car}). $$

The precise numbers are not meaningful on a tiny corpus. The relative pattern is.

8. Four implementation checks#

Only a few checks are essential:

  • Train with dot products and the logistic objective; use cosine for downstream similarity.
  • Keep tensor shapes explicit: q [B,d], k_pos [B,d], k_neg [B,n,d].
  • In manual NumPy updates, compute all gradients from copied pre-update vectors.
  • Sample negatives from the intended distribution with replacement; repeated IDs should accumulate gradients.

Everything else is an engineering choice. Larger corpora may need faster samplers, streaming pair generation, or sparse optimizers, but these do not change the conceptual model.

9. What an experiment can and cannot show#

A toy corpus is a mechanism microscope. It can show that shared contexts produce shared directions and that the model distinguishes observed pairs from noise. It cannot demonstrate high-quality general semantics, because the vocabulary and contexts are too small.

When changing the window, noise exponent, number of negatives, or dimensionality, ask one question at a time:

Which definition of relation or baseline did this change?

This keeps hyperparameters connected to their statistical meaning.


Ⅳ. Relationship Direction and Meaning Axis#

↑ Table of Contents

10. Meaning is relational, not attached to one coordinate#

It is tempting to label coordinates directly:

dimension 17 = animalness
dimension 42 = formality

But the SGNS factorization does not identify individual coordinates. For an orthogonal matrix $R$,

$$ Q'=QR,\qquad K'=KR $$

preserves every query–key score because

$$ Q'K'^\top=QRR^\top K^\top=QK^\top. $$

The coordinate labels can change under rotation while the learned relationships remain identical. What is interpretable is the relative geometry within a fixed representation:

  • nearby directions: similar relation profiles;
  • a difference between two vectors: a change in relation profile;
  • a direction repeated across pairs: a recurring semantic relationship.

Before analyzing geometry, choose one word representation and use it consistently:

$$ v_w=Q[w] $$

or, after confirming compatible preprocessing,

$$ v_w=\frac{Q[w]+K[w]}{2}. $$

The first choice is simplest and keeps the analysis in the query space.

11. What vector subtraction measures#

For two query vectors, define the relationship direction

$$ d_{a\rightarrow b}=q_b-q_a. $$

Projecting this difference onto a context key $k_c$ gives

$$ \begin{aligned} (q_b-q_a)^\top k_c &=q_b^\top k_c-q_a^\top k_c\\ &\approx \log\frac{P(c\mid b)}{n_{\mathrm{neg}}P_n(c)} -\log\frac{P(c\mid a)}{n_{\mathrm{neg}}P_n(c)}\\ &= \log\frac{P(c\mid b)}{P(c\mid a)}. \end{aligned} $$

The noise baseline cancels. The difference vector summarizes how associations with many contexts change when moving from $a$ to $b$.

Suppose several pairs satisfy

$$ v_{b_1}-v_{a_1} \approx v_{b_2}-v_{a_2} \approx \cdots. $$

Then the corpus has produced a recurring transformation. That shared transformation is a relationship direction.

This is more precise than saying that one coordinate “contains” the relationship. The relationship is expressed by a direction across the whole space.

12. Analogy as parallel translation#

The familiar analogy

$$ v_{king}-v_{man} \approx v_{queen}-v_{woman} $$

says that two pairwise changes are approximately parallel. Rearranging gives

$$ v_{king}-v_{man}+v_{woman} \approx v_{queen}. $$

The operation is a translation:

1. Estimate the movement from man to king.
2. Apply the same movement from woman.
3. Search for the word nearest the predicted destination.

It does not literally remove a pure “male component” from king. The safer interpretation is that the two pairs exhibit similar changes in their compressed context profiles.

The standard 3CosAdd search is:

$$ t=v_b-v_a+v_c, \qquad \hat{x}=\arg\max_x\cos(v_x,t). $$

E_unit = E / np.linalg.norm(E, axis=1, keepdims=True).clip(min=1e-12)

target = E[id_b] - E[id_a] + E[id_c]
target /= max(np.linalg.norm(target), 1e-12)

scores = E_unit @ target
scores[[id_a, id_b, id_c]] = -np.inf
answer_ids = np.argsort(scores)[-10:][::-1]

Analogy quality depends on how consistently the relationship repeats in the corpus. It is an empirical property of the learned geometry, not a constraint explicitly guaranteed by SGNS.

13. From repeated pair differences to a relationship direction#

One pair difference can mix the intended relationship with word-specific details. Multiple matched pairs give a more stable estimate.

For aligned pairs $(a_i,b_i)$, compute

$$ d_i=v_{b_i}-v_{a_i}. $$

The orientation must be consistent. For a female-directed relationship, use

$$ v_{woman}-v_{man},\quad v_{queen}-v_{king},\quad v_{actress}-v_{actor}, $$

not a mixture of forward and reversed differences. Otherwise the directions cancel.

A simple shared direction is the normalized average:

$$ a_{rel} = \frac{\frac{1}{n}\sum_i d_i} {\left\|\frac{1}{n}\sum_i d_i\right\|}. $$

Before averaging, check whether the pair directions are actually coherent:

offsets = E[right_ids] - E[left_ids]
offsets_unit = offsets / np.linalg.norm(
    offsets, axis=1, keepdims=True
).clip(min=1e-12)

agreement = offsets_unit @ offsets_unit.T
mean_offset = offsets.mean(axis=0)
relation_axis = mean_offset / max(np.linalg.norm(mean_offset), 1e-12)

If the pairwise cosine matrix agreement contains many negative or near-zero values, there may be no single shared relationship to summarize.

14. A meaning axis turns a direction into a question#

A relationship direction describes a recurring transformation. A meaning axis uses such a direction as a ruler for all words.

Examples include:

informal <--------------------------> formal
concrete <--------------------------> abstract
negative <--------------------------> positive

Let $a$ be a unit axis and $\mu_0$ a reference point. A word’s signed projection is

$$ z_w=(v_w-\mu_0)^\top a. $$

The vector decomposes into its axis component and residual meaning:

$$ v_w-\mu_0 = z_wa+r_w, \qquad r_w\perp a. $$

$z_w$ answers one specified question. It is not the complete meaning of the word. Two words can have the same axis score and still differ greatly in the residual space.

This distinction matters:

analogy:
    apply a relationship direction to one point

meaning axis:
    project many points onto a relationship direction

Both operations use the same geometry.

15. Building an axis from two poles#

A single contrast such as formal - informal may be sensitive to the idiosyncrasies of those two words. Use several seed words for each pole:

$$ S^+=\{formal,professional,official\}, \qquad S^-=\{informal,casual,colloquial\}. $$

Compute the pole centroids:

$$ \mu^+ = \frac{1}{|S^+|}\sum_{w\in S^+}v_w, \qquad \mu^- = \frac{1}{|S^-|}\sum_{w\in S^-}v_w. $$

The axis points from the negative pole to the positive pole:

$$ \boxed{ a = \frac{\mu^+-\mu^-}{\|\mu^+-\mu^-\|} }. $$

A natural reference point is the midpoint

$$ \mu_0=\frac{\mu^++\mu^-}{2}. $$

E_work = E - E.mean(axis=0, keepdims=True)  # optional global centering

mu_pos = E_work[positive_ids].mean(axis=0)
mu_neg = E_work[negative_ids].mean(axis=0)

axis = mu_pos - mu_neg
axis /= max(np.linalg.norm(axis), 1e-12)
midpoint = 0.5 * (mu_pos + mu_neg)

projection_scores = (E_work - midpoint) @ axis
ranked_ids = np.argsort(projection_scores)[::-1]

There are two common scoring choices:

  • Signed projection $(v_w-\mu_0)^\top a$: retains magnitude along the axis.
  • Cosine alignment $\cos(v_w,a)$: focuses on direction and removes word-vector norm.

Neither is universally correct. Choose the score that matches the research question, document preprocessing, and use the same procedure for seeds and evaluated words.

Matched-pair averaging and pole centroids are related but not identical. Matched pairs preserve correspondence between contrasts; pole centroids compare two groups as wholes. Use matched pairs when pair identity matters, and pole centroids when the two conceptual poles are primary.

16. When is an axis convincing?#

An axis is useful when it generalizes beyond the words used to construct it. A compact validation set is enough:

  1. Held-out order: Do unused contrast pairs appear in the expected order?
  2. Seed stability: Does removing one seed or resampling seeds preserve the ranking?
  3. Coherent extremes: Do the highest- and lowest-scoring words fit the stated contrast?
  4. Run stability: Does the conclusion survive retraining with another random seed?

When comparing separately trained spaces, align them before comparing coordinates or directions. Within one fixed space, words and their axis rotate together, so orthogonal rotation does not change projections. Across independent spaces, arbitrary rotations make raw coordinate comparisons meaningless.

Finally, the axis measures how a corpus organizes language between chosen poles. It does not establish an essential property or a causal effect. The accurate claim is:

In this corpus, under this embedding and pole definition, the word is more aligned with one side of the specified semantic contrast.


Ⅴ. Connections to Other Models#

↑ Table of Contents

17. One family of interaction models#

The same low-rank scoring pattern appears in several domains:

Model Source Target Score
Word2Vec center word context word $q_w^\top k_c$
Recommendation user item $u_i^\top v_j$
node2vec source node context node $q_i^\top k_j$

SGNS can therefore be read as a low-rank logistic factorization of a sparse interaction matrix. DeepWalk and node2vec create their interaction events from random walks rather than sentences, but the training structure is closely related.

Separate query and key embeddings are particularly natural for directional data: a word or node can behave differently as a source and as a target.

18. Energy and softmax#

Define energy as

$$ E(w,c)=-q_w^\top k_c. $$

Then full softmax has the Boltzmann form

$$ P(c\mid w) = \frac{e^{-E(w,c)}}{\sum_j e^{-E(w,j)}}. $$

Observed pairs should have low energy, or high dot products. Negative sampling learns this compatibility structure contrastively without evaluating the full normalization term for every update.

Transformer attention also uses query–key dot products, but its queries and keys are dynamically computed from contextual token states. Word2Vec attaches static query and key parameters to word IDs. The scoring motif is shared; the representation mechanism is different.


Ⅵ. Core Formulas and Mental Model#

↑ Table of Contents

19. Five equations to remember#

Score

$$ s(w,c)=q_w^\top k_c $$

SGNS objective for one positive pair

$$ \log\sigma(q_w^\top k_c) + \sum_{i=1}^{n_{\mathrm{neg}}} \log\sigma(-q_w^\top k_{c_i^-}) $$

Pairwise optimal score

$$ q_w^\top k_c \approx \log\frac{P(c\mid w)}{n_{\mathrm{neg}}P_n(c)} $$

Shifted PMI under unigram noise

$$ q_w^\top k_c \approx PMI(w,c)-\log n_{\mathrm{neg}} $$

Gravity form

$$ P(w,c) \propto P(w)P_n(c)\exp(q_w^\top k_c) $$

20. The compact mental model#

window defines positive relations
        ↓
noise defines the comparison baseline
        ↓
SGNS learns low-rank log-odds
        ↓
similar relation profiles become similar vectors
        ↓
similar changes in profiles become parallel offsets
        ↓
parallel offsets support analogies and meaning axes

The crucial distinction is:

word similarity     = similarity between relation profiles
relationship        = change between two relation profiles
meaning-axis score  = projection onto a repeated relationship direction

21. From geometry to research#

Implementation gives us vectors; interpretation turns their geometry into measured questions. The next note uses meaning axes, WEAT, temporal alignment, and external validation to study cultural association and semantic change.

Continue to Part 3: Embedding Applications — Reading Culture Through Semantic Space.

Text is licensed under CC BY-NC 4.0. Cited images remain the property of their respective rights holders.