When training word2vec, why is a new negative sampling process formulated instead of simply downsampling?

Question

(For background, see The Skip-Gram Model.¹ This question does not exactly use their notation, but you should be able to follow along.)

The original skip-gram log-likelihood of a single word, $w$, given a single word in its context, $c$, is—

$$ \begin{align*} \mathcal{L} &= \log \frac{ \exp(\mathbf{u}_c^T \mathbf{v}_w) }{ \sum_{i=1}^{V} \exp(\mathbf{u}_i^T \mathbf{v}_w) } \\ &= \mathbf{u}_c^T \mathbf{v}_w - \log \sum_{i=1}^{V} \exp(\mathbf{u}_i^T \mathbf{v}_w) \\ &= \mathbf{u}_c^T \mathbf{v}_w - \log \bigg( \exp(\mathbf{u}_c^T \mathbf{v}_w) + \sum_{i=1, i \neq w}^{V} \exp(\mathbf{u}_i^T \mathbf{v}_w) \bigg). \end{align*} $$

—where:

• $\mathbf{u}_c$ is a vector representing a word, $c$, which occurred in the context of the input word $w$
• $\mathbf{v}_w$ is a vector representing the input word $w$
• $i$ is a word in the vocabulary
• $V$ is the size of the vocabulary.

The sum is extremely expensive to compute, because a text corpus can contain 100s of thousands of unique words. The way skip-gram gets around this issue is by contriving a new probabilistic process based on negative sampling (NS). This process and its log-likelihood are derived here.¹ The log-likelihood turns out to be—

$$ \mathcal{L}_{\text{NS}} = \log \sigma (\mathbf{u}_c^T \mathbf{v}_w) + \sum_{k=1}^{K} \log (1 - \sigma(\mathbf{u}_{\bar{c}_k}^T \mathbf{v}_w)). $$

—where $\bar{c}_k$ denotes a randomly sampled word that isn't in the context of $w$. The likelihood is that of $1 + K$ independent classifications: $1$ for the positive sample, and $K$ for the negative samples. It's worth mentioning that maximizing this likelihood gives you an approximate Noise Contrastive Estimator.²

$\mathcal{L}_{\text{NS}}$ is well-motivated and clearly works. But I'm wondering why we don't simply use the negative samples to approximate the sum in $\mathcal{L}$. Define a new log-likelihood, $\mathcal{L}_{\text{DS}}$ (DS = downsampled), which does exactly that:

$$ \begin{align} \mathcal{L}_{\text{DS}} = \mathbf{u}_c^T \mathbf{v}_w - \log \bigg( \exp(\mathbf{u}_c^T \mathbf{v}_w) + \sum_{k=1}^{K} \exp(\mathbf{u}_{\bar{c}_k}^T \mathbf{v}_w) \bigg). \end{align} $$

$\mathcal{L}_{\text{DS}}$ seems like the most straightforward way to tackle the computational issue with $\mathcal{L}$. It does not require re-formulating the statistical model. It's also well-studied. If "word" is replaced with "text passage", "in-context" is replaced with "similar", and "out-of-context" is replaced with "not similar", then this likelihood was used to train the (currently) top 4 text similarity transformers on the MTEB leaderboard.

The word2vec paper³, and resources which explain its loss function, $-\mathcal{L}_{\text{NS}}$^2,4,5,6 (including this CV answer), don't bring up $\mathcal{L}_{\text{DS}}$. I'm wondering why that is. More practically, how does one decide between using $\mathcal{L}_{\text{NS}}$ and $\mathcal{L}_{\text{DS}}$ for their representation learning task? Below is my progress on answering this question.

Gradients

I'm not sure how useful it is to dig into the differences between the probabilistic processes underlying $\mathcal{L}_{\text{NS}}$ and $\mathcal{L}_{\text{DS}}$. Let's see if their gradients reveal something.

Staring at the two equations, let's extract common dot products so that the gradient computation is reduced to computing simpler derivatives:

$$ a = \mathbf{u}_c^T \mathbf{v}_w \\ b_k = \mathbf{u}_{\bar{c}_k}^T \mathbf{v}_w. $$

Re-express the likelihood in terms of these:

$$ \begin{align} \mathcal{L}_{\text{NS}} &= \log \sigma (a) + \sum_{k=1}^{K} \log( 1 - \sigma(b_k)) \\ \mathcal{L}_{\text{DS}} &= a - \log \bigg( \exp(a) + \sum_{k=1}^{K} \exp(b_k) \bigg). \end{align} $$

Partial derivatives wrt $a$—

$$ \begin{align} \frac{ \partial \mathcal{L}_{\text{NS}} }{ \partial a } &= 1 - \sigma(a) \\ \frac{ \partial \mathcal{L}_{\text{DS}} }{ \partial a } &= 1 - f(a). \end{align} $$

—where $f(x)$ is the lazily notated softmax function: the numerator of $f(x)$ is $\exp(x)$ and the denominator is the input of the $\log$ in $\mathcal{L}_{\text{DS}}$. ^{Later, you might catch why it's called $f$ in particular :-)}

And wrt $b$ :

$$ \begin{align} \frac{ \partial \mathcal{L}_{\text{NS}} }{ \partial b_k } &= -\sigma(b_k) \\ \frac{ \partial \mathcal{L}_{\text{DS}} }{ \partial b_k } &= -f(b_k). \end{align} $$

Now the partial derivatives wrt the vectors (all are parameters in skip-gram):

$$ \begin{align} \frac{ \partial a }{ \partial \mathbf{v}_w } &= \mathbf{u}_c \:\:\:\:\: \frac{ \partial b_k }{ \partial \mathbf{v}_w } = \mathbf{u}_{\bar{c}_k} \\ \frac{ \partial a }{ \partial \mathbf{u}_c } &= \mathbf{v}_w \:\:\:\:\: \frac{ \partial b_k }{ \partial \mathbf{u}_c } = \mathbf{0} \\ \frac{ \partial a }{ \partial \mathbf{u}_{\bar{c}_{k'}} } &= \mathbf{0} \:\:\:\:\:\:\: \frac{ \partial b_k }{ \partial \mathbf{u}_{\bar{c}_{k'}} } = I(k = k') \mathbf{v}_w. \end{align} $$

Applying the chain rule, e.g. for $\mathbf{v}_w$—

$$ \frac{ \partial \mathcal{L} }{ \partial \mathbf{v}_w } = \frac{ \partial \mathcal{L} }{ \partial a } \frac{ \partial a }{ \partial \mathbf{v}_w } + \sum_{k=1}^{K} \frac{ \partial \mathcal{L} }{ \partial b_k } \frac{ \partial b_k }{ \partial \mathbf{v}_w } $$

—the gradients are:

$$ \begin{align} \frac{ \partial \mathcal{L}_{\text{NS}} }{ \partial \mathbf{v}_w } &= (1 - \sigma(\mathbf{u}_c^T \mathbf{v}_w) \mathbf{u}_c) - \sum_{k=1}^{K} \sigma(\mathbf{u}_{\bar{c}_k}^T \mathbf{v}_w) \mathbf{u}_{\bar{c}_k} \\ \frac{ \partial \mathcal{L}_{\text{DS}} }{ \partial \mathbf{v}_w } &= (1 - f(\mathbf{u}_c^T \mathbf{v}_w) \mathbf{u}_c) - \sum_{k=1}^{K} f(\mathbf{u}_{\bar{c}_k}^T \mathbf{v}_w) \mathbf{u}_{\bar{c}_k} \\ \frac{ \partial \mathcal{L}_{\text{NS}} }{ \partial \mathbf{u}_c } &= (1 - \sigma(\mathbf{u}_c^T \mathbf{v}_w) \mathbf{v}_w \\ \frac{ \partial \mathcal{L}_{\text{DS}} }{ \partial \mathbf{u}_c } &= (1 - f(\mathbf{u}_c^T \mathbf{v}_w) \mathbf{v}_w \\ \frac{ \partial \mathcal{L}_{\text{NS}} }{ \partial \mathbf{u}_{\bar{c}_k} } &= -\sigma(\mathbf{u}_{\bar{c}_k}^T \mathbf{v}_w)\mathbf{v}_w \\ \frac{ \partial \mathcal{L}_{\text{DS}} }{ \partial \mathbf{u}_{\bar{c}_k} } &= -f(\mathbf{u}_{\bar{c}_k}^T \mathbf{v}_w)\mathbf{v}_w. \end{align} $$

All of these gradients pass sanity checks about how vectors should get pushed.⁶ The only difference between the gradients of $\mathcal{L}_{\text{DS}}$ vs $\mathcal{L}_{\text{NS}}$ is the use of softmax vs sigmoid. This implies that the only difference is whether dot products of other words are taken into account when computing the weight of the gradient. For $\mathcal{L}_{\text{DS}}$, they are. For $\mathcal{L}_{\text{NS}}$, they aren't. Does this reveal anything about when to use $\mathcal{L}_{\text{DS}}$ vs $\mathcal{L}_{\text{NS}}$?

References

1. Zhang, Aston, et al. "Dive into deep learning." arXiv preprint arXiv:2106.11342 (2021). [link]

2. Dyer, Chris. "Notes on noise contrastive estimation and negative sampling." arXiv preprint arXiv:1410.8251 (2014). [link]

3. Mikolov, Tomas, et al. "Distributed representations of words and phrases and their compositionality." Advances in neural information processing systems 26 (2013). [link]

4. Goldberg, Yoav, and Omer Levy. "word2vec Explained: deriving Mikolov et al.'s negative-sampling word-embedding method." arXiv preprint arXiv:1402.3722 (2014). [link]

5. Manning, Christopher, and Richard Socher. "Natural Language Processing with Deep Learning CS224N/Ling284." (2019). [link]

6. Rong, Xin. "word2vec parameter learning explained." arXiv preprint arXiv:1411.2738 (2014). [link]