Word2Vec for User Behavior: From Sequences to Training Pairs
Before you start: an embedding is a learned table lookup. Each item ID selects a vector. This chapter explains where the learning signal comes from.
Skip-gram learns an item’s vector by asking it to predict nearby items. In text, the items are words. In a behavior log, they might be product views in a shopping session. We will follow one pair, jacket → socks, from the session that creates it to the update that changes its vectors.
A session becomes several training examples
Consider boots → socks → jacket → boots → hat. Choose jacket as the center and a context radius of one position. Its two neighbors create the positive pairs (jacket, socks) and (jacket, boots). “Positive” means observed together under this window rule; it does not mean purchased, liked or interchangeable.
With jacket at position 3 and radius 1, the context positions are 2 (socks) and 4 (boots). These create two separate positive training pairs.
The window is clipped at the session boundary. Applying the same procedure to every center creates the training data. This symmetric window uses items on both sides; the resulting pair does not tell the model whether socks occurred before or after jacket. A next-action task needs a construction that respects that different objective.
Two tables give each item two roles
Maintain an input table V and a separate context table U. Each has one row per vocabulary item and the same number of coordinates per row. For the pair (jacket, socks), look up jacket’s input vector $v_{\text{jacket}}$ and socks’ context vector $u_{\text{socks}}$. Their dot product is the pair’s score:
A two-coordinate dot product multiplies the first coordinates, multiplies the second coordinates, and adds those products. The score is an unrestricted real number. It is not yet a probability, a cosine similarity, or a guarantee that jacket and socks are close in the input table.
Every item can appear in either role, but its two rows are separate parameters. For this pair, jacket supplies the input row and socks supplies the context row. When socks later becomes a center, training uses socks’ input row instead.
How this relates to full-softmax skip-gram
One way to predict a context ID is to score every item in a vocabulary $\mathcal V$, exponentiate the scores and normalize:
Here $w$ is a center ID and $c$ a context ID. The denominator makes probabilities sum to one across all candidate contexts, but it also involves the entire vocabulary. The negative-sampling objective below trains a binary discriminator instead. Its sigmoid scores do not form this normalized distribution.
Turn one observed pair into a small classification task
Keep (jacket, socks) as an observed pair with label 1. Draw two context IDs from a noise distribution, say kettle and mug, and give (jacket, kettle) and (jacket, mug) label 0. These are sampled training comparisons, not evidence that a user dislikes those products.
Convert each pair score $s$ to a binary probability $D$ with the sigmoid function:
$D$ estimates “observed pair rather than sampled noise” under this constructed task. A score of 0 gives $D=0.5$; increasing the score increases $D$. It is not a purchase probability, and different pairs’ values do not have to sum to one.
Let $s_+$ be the observed pair’s score and $s_{n_j}$ the score for the $j$th sampled context, with $K$ negatives. Use negative log probability for the desired label in each comparison:
The positive term rewards a large $D$. A negative term rewards a large $1-D$, which equals $\sigma(-s)$. We sum the $K+1$ terms here; averaging them would rescale this batch’s gradient. At three zero scores, all three desired-label probabilities are 0.5, so the loss is $3\log 2\approx2.0794$.
Watch an actual update change the selected rows
Choose input vector jacket = (1, 0), and context vectors socks = (0, 1), kettle = (0, −1), mug = (0, −0.5). These illustrative starting vectors make all three scores zero while allowing nonzero gradients.
For one labeled comparison, the loss derivative with respect to its score is D − label. Initially, this is −0.5 for socks and +0.5 for each noise item. With a learning rate of 0.2, the first simultaneous update moves the jacket input vector to (1, 0.25). The context vectors also change, giving scores 0.35, −0.35, −0.225.
Each bar shows D: probability assigned to the observed-pair label.
After one update the loss is about 1.6537. Only jacket’s input row and the three participating context rows change. In corpus training, many pairs share rows, so their updates interact; not every individual pair score must improve on every step.
Follow the first gradient through the vectors
Let $y_r$ be comparison $r$’s label, $u_r$ its context vector, and $v$ the shared center vector. Since the score is $u_r^\top v$, the chain rule gives:
At initialization, the three contributions to the center gradient are (0, −0.5), (0, −0.5), and (0, −0.25). Their sum is (0, −1.25). Subtract 0.2 times this gradient from (1, 0), giving (1, 0.25).
For the context rows, subtract 0.2 times the respective gradients from the old vectors. Socks becomes (0.1, 1), kettle (−0.1, −1), and mug (−0.1, −0.5). Use the old parameters for every derivative in this simultaneous update. If a context ID appears more than once in the batch, add its contributions before applying the update.
The noise distribution changes the comparisons
A common negative-sampling choice draws item $i$ with probability proportional to its training frequency $f_i$ raised to the power 3/4. Divide by the sum of those weights to obtain probabilities:
The following miniature catalog has counts [64, 64, 8, 8, 1, 1]. Lowering the exponent from 1 to 0.75 increases each rare item’s draw probability from about 0.68% to 1.76%. Frequent items remain more likely, but their dominance is reduced. Exponent 0 makes this six-item distribution uniform.
Fixed counts: boots 64, socks 64, jacket 8, hat 8, kettle 1, mug 1.
Drawing negatives and discarding some frequent positive training tokens are different operations; Word2Vec can use both. Also record whether repeated negative IDs or collisions with the positive ID are kept, skipped or resampled. Filtering changes the effective sampling process.
The training pairs define the meaning of “related”
For activity logs, session boundaries, event types and time cutoffs are part of the model. Sort events within the intended user/session key. Investigate accidental duplicate logging separately from meaningful repeated views. Joining the end of one user’s session to another user’s first event invents a relationship.
A purchase, an impression and a comparison click convey different information. Boots and socks may be useful complements while two boots are substitutes. A co-occurrence objective does not choose between those meanings for you. State the intended retrieval task before judging the neighbors, and build training data using only information available before the evaluation cutoff.
A small implementation you can inspect
The Python below keeps pair generation inside each session and implements the update used in the figure. It deliberately repeats a fixed pair and fixed negatives; a full training loop would sample many pairs and negatives, initialize the full tables, and evaluate on held-out data.
Python: session pairs and the computed 20-update example
import numpy as np
def session_pairs(sessions, radius=1):
if not isinstance(radius, int) or isinstance(radius, bool) or radius < 1:
raise ValueError("radius must be a positive integer")
for tokens in sessions: # Never concatenate separate sessions.
for i, center in enumerate(tokens):
for j in range(max(0, i-radius), min(len(tokens), i+radius+1)):
if i != j:
yield center, tokens[j]
def sgns_step(V, U, center, context, negatives, rate=0.2):
# V and U are separate floating-point arrays: [vocabulary, dimension].
# IDs must be valid rows. This function uses a SUM over this fixed batch.
ids = np.array([context, *negatives], dtype=int)
labels = np.zeros(len(ids)); labels[0] = 1
v_old = V[center].copy()
scores = U[ids] @ v_old
probabilities = np.exp(-np.logaddexp(0, -scores)) # Stable sigmoid.
errors = probabilities - labels
loss = np.logaddexp(0, -scores[0]) + np.logaddexp(0, scores[1:]).sum()
grad_v = errors @ U[ids] # All derivatives use the old parameters.
V[center] -= rate * grad_v
# Repeated context/noise IDs must accumulate into the same output row.
np.add.at(U, ids, -rate * errors[:, None] * v_old)
return float(loss)
# IDs: jacket=0, socks=1, kettle=2, mug=3.
V = np.array([[1., 0.], [.2, .1], [.1, -.2], [-.1, .2]])
U = np.array([[.3, .2], [0., 1.], [0., -1.], [0., -.5]])
for step in range(21):
scores = U[[1, 2, 3]] @ V[0]
loss = np.logaddexp(0, -scores[0]) + np.logaddexp(0, scores[1:]).sum()
if step in [0, 1, 20]:
print(step, np.round(scores, 6), round(float(loss), 6))
if step < 20:
sgns_step(V, U, center=0, context=1, negatives=[2, 3])
# 0 [0. 0. 0.] 2.079442
# 1 [0.35 -0.35 -0.225] 1.653726
# 20 [3.566834 -3.566834 -3.225431] 0.094677
np.add.at accumulates updates when a sampled context ID repeats. The stable sigmoid and log-loss expressions also avoid directly exponentiating a very large positive score.
Check a library’s actual training objective before equating its results with this example. Spark’s documented Word2Vec implementation uses hierarchical softmax. Its DataFrame transformer aggregates a sequence into a document vector; that output is different from exporting an individual item’s learned row.
Choose the vectors that the retrieval task will consume
The trained input table is one possible representation of items. Input-to-input cosine similarity, context-to-context similarity, and the trained input-to-context dot product are different scores. An observed-pair update directly trains the last of these; it does not make the other choices equivalent.
Save the item-to-row vocabulary with the selected vectors, preprocessing decisions and model version. Decide how unseen items are handled: a table lookup cannot infer a new item’s row without another source of features or a fallback. Continue to embedding evaluation to compare useful neighbors, held-out rankings and dimension choices.
References and further reading
- Mikolov et al., Distributed Representations of Words and Phrases — skip-gram, negative sampling, the frequency-power noise distribution and frequent-token subsampling; Sections2–2.3.
- Goldberg and Levy, word2vec Explained — the distinction between predicting a vocabulary distribution and discriminating observed pairs from noise.
- Spark MLlib: Word2Vec — the documented hierarchical-softmax training implementation.
- Spark ML: Word2Vec transformer — sequence inputs and averaged document-vector outputs.