Introduction

Word2Vec is a group of related models that are used to produce word embeddings, or vector representations of words. Developed by Mikolov et al. at Google in 2013, Word2Vec has become a foundational piece in the field of natural language processing (NLP). Apache Spark's machine learning library is equipped with an implementation of Word2Vec, leveraging distributed computing capabilities for scalability and efficiency.

In this article, we'll explore the mathematics behind Spark's Word2Vec, technical explanations, and examples to illustrate its functionalities.

Fundamentals of Word2Vec

Word2Vec operates in two primary modes:

1. Continuous Bag of Words (CBOW): This approach predicts the target word from context words within a given window surrounding the target.
2. Skip-gram: In this method, the model tries to predict the context words based on the target word.

Both methods learn to map words into a continuous vector space where semantically similar words are mapped to nearby points.

Vector Mathematics in Word2Vec

Notations and Definitions

• $V$: Vocabulary Size • $d$: Dimension of Word Vectors • $X_{VX}$: Context matrix with size $V \times d$ • $Y_{dV}$: Target matrix with size $d \times V$

Objective Function

For CBOW, given a context $C$:

$$\max , P(\text{target} \mid \text{context}) = \prod\_{t=1}^{T} P(w\_t \mid w\_{t-k}, \ldots, w\_{t+k})$$

For Skip-gram:

$$\max , P(\text{context} \mid \text{target}) = \prod\_{t=1}^{T} \prod\_{-k \leq j \leq k, j \neq 0} P(w\_{t+j} \mid w\_t)$$

Softmax Function

To solve these probabilities, we use the softmax function:

$$P(o \mid c) = \frac{e^{\textbf{v}\_o^{\top} \cdot \textbf{v}*c}}{\sum*{w=1}^{V} e^{\textbf{v}\_w^{\top} \cdot \textbf{v}\_c}}$$

Where $\textbf{v}_o$ and $\textbf{v}_c$ are the vector representations.

Negative Sampling

Negative sampling refines the calculation by approximating the softmax function, focusing on observing observed word pairs and a small sample of random word pairs:

$$\log \sigma(\textbf{v}*o^{\top} \textbf{v}c) + \sum{i=1}^{k} \mathbb{E}*{w\_i \sim P\_n(w)}[\log \sigma(-\textbf{v}\_{w\_i}^{\top} \textbf{v}\_c)]$$

Optimizing the Objective

The optimization of the objective functions is performed through stochastic gradient descent (SGD). Gradients are computed and used to update vector representations:

• Modify vectors $\textbf{v}_o$ and $\textbf{v}_c$: • $\textbf{v}_o \leftarrow \textbf{v}_o - \eta \cdot \nabla_{\textbf{v}_o} J(\theta)$ • $\textbf{v}_c \leftarrow \textbf{v}_c - \eta \cdot \nabla_{\textbf{v}_c} J(\theta)$

Where $\eta$ represents the learning rate.

Technical Implementation in Spark

Spark's implementation of Word2Vec allows for scaling across large datasets by using its distributed processing features. Here is an outline:

1. Input Processing: Text data is tokenized into words.
2. Word Pair Generation: For skip-gram, create pairs of target and context words.
3. Model Training: Using SGD and negative sampling to iterate and optimize word vectors.

Example Code

Below is an example of Word2Vec initialization and training in Spark: