ML coding · Embeddings and retrieval
Skip-gram with negative sampling: loss and gradients
Given the centre-word vector v, the true context vector u and k negative
vectors U_neg (k × d), with s = u·v and sj = uj·v, the loss is
L = −log σ(s) − Σj log σ(−sj).
Return [L, dL/dv, dL/du, dL/dU_neg], the last a k × d list. Dot products may reach ±1000.
Examples
Input v = [1, 0], u = [1, 0], U_neg = [[0, 1]]
Output [1.0064, [-0.2689, 0.5], [-0.2689, 0], [[0.5, 0]]]
s = 1, one negative orthogonal to v (s₁ = 0).
Input v = [10, 0], u = [-100, 0], U_neg = [[100, 0]]
Output [2000, [200, 0], [-10, 0], [[10, 0]]]
Dot products of −1000 and 1000: the naive log σ gives −inf.
Constraints
- 1 ≤ d ≤ 300, 1 ≤ k ≤ 20
- Checked to a relative tolerance of 10⁻⁹
Hints
Hint 1
d/dx [−log σ(x)] = σ(x) − 1 and d/dx [−log σ(−x)] = σ(x).
Hint 2
So dL/dv = (σ(s) − 1)u + Σ σ(s_j)u_j, dL/du = (σ(s) − 1)v and dL/du_j = σ(s_j)v.
Hint 3
log σ(x) = −softplus(−x), and softplus(y) = max(y, 0) + log1p(exp(−|y|)).
Where candidates go wrong
- np.log(sigmoid(x)) for large negative x: σ rounds to 0 and the log is −inf.
- Sign errors on the negative terms: their gradient pushes σ(s_j) towards 0, so it is +σ(s_j), not σ(s_j) − 1.
- Forgetting that v receives gradient from every negative.
The idea
Compute the k + 1 dot products, then the loss with a stable softplus: −log σ(s) = softplus(−s) and −log σ(−sj) = softplus(sj). The gradients follow from σ′ = σ(1 − σ): the coefficient on the positive pair is σ(s) − 1, on each negative σ(sj), and each gradient is that coefficient times the other vector of the pair.
Time O(kd) · Space O(kd)
Reference solution
import numpy as np
def softplus(y):
return np.maximum(y, 0) + np.log1p(np.exp(-np.abs(y)))
def sigmoid(z):
return np.exp(-softplus(-z)) # σ(z) = exp(log σ(z)), never overflows
def sgns_loss_grad(v, u, U_neg):
v = np.asarray(v, dtype=float)
u = np.asarray(u, dtype=float)
N = np.asarray(U_neg, dtype=float)
s = u @ v
sn = N @ v
loss = softplus(-s) + softplus(sn).sum() # -log σ(s) - Σ log σ(-s_j)
gp = sigmoid(s) - 1.0 # coefficient on the positive pair
gn = sigmoid(sn) # coefficients on the negatives
grad_v = gp * u + gn @ N
grad_u = gp * v
grad_N = gn[:, None] * v[None, :]
return [float(loss), grad_v.tolist(), grad_u.tolist(), grad_N.tolist()]Follow-ups an interviewer asks
- Check the gradients with central finite differences on a random example.
- Sample negatives from the unigram distribution raised to the 3/4 power: why that exponent?
- Subsampling of frequent words, and why it speeds training and improves rare-word vectors.
Python runs in your browser with NumPy. The first run downloads the runtime, which takes a few seconds.
Why interviewers ask it
Deriving and coding these gradients is a classic research-engineer question. It tests the derivative of the log-sigmoid, the direction of each term, and numerical care.
The theory behind it is in the lesson Cosine search, contrastive loss and negative sampling.
Keep going
Top-k search by cosine similarity
Return the k rows most similar to a query by cosine similarity, with zero vectors and ties handled by stated rules, vectorised in NumPy..
Open the problem → Medium15 minThe InfoNCE contrastive loss
The contrastive loss behind CLIP and SimCLR in NumPy: normalised embeddings, a temperature, and a log-sum-exp that survives tiny temperatures..
Open the problem → Hard20 minSkip-gram with negative sampling: loss and gradients
The word2vec skip-gram loss with k negative samples and its exact gradients, written with a log-sigmoid that never overflows..
Open the problem →