Word2Vec in Pure NumPy

This repository contains an implementation of the Word2Vec algorithm, specifically the Skip-Gram with Negative Sampling (SGNS) architecture, written entirely from scratch using only NumPy. No deep learning frameworks were used.

Dataset

The model is trained on three Wikipedia corpora covering major financial crises (~387,000 words total from 60 articles):

Corpus	Words	Topics
financial_crisis_corpus.txt	168K	2007–2008 crisis: subprime mortgages, Lehman Brothers, TARP, CDOs, quantitative easing
dotcom_crisis_corpus.txt	116K	2001 dot-com bubble: Enron, WorldCom, Sarbanes-Oxley, NASDAQ, venture capital, Silicon Valley
black_monday_corpus.txt	102K	1987 crash: Black Monday, program trading, circuit breakers, NYSE, futures, volatility

Code Structure

File	Description
model.py	Word2VecSGNS class — embedding lookup, forward pass, loss, analytical gradients, and SGD parameter updates.
dataset.py	Tokenizer, Vocab, and Word2VecDataset — text preprocessing, vocabulary building, sub-sampling of frequent words, and unigram-based negative sampling.
train.py	Training loop — loads the corpus, trains the model, evaluates with word similarity and analogies, and saves the trained model.
test_model.py	Numerical gradient checker (finite differences vs. analytical) and dataset sanity checks.
evaluate.py	Interactive script to load a saved model and query word similarities/analogies.

How to Run

1. Install requirements:

   pip install -r requirements.txt

2. Run training (saves model to trained_model/):

   python train.py

3. Run tests:

   python test_model.py

4. Evaluate saved model:

   python evaluate.py

Training Results

Trained with: embedding_dim=50, window=5, neg_samples=5, min_count=5, lr=1.0, batch_size=256, epochs=20.

Vocabulary size : 6,325
Tokens (post-subsample): 273,232
Training time : 214 seconds
Final loss : 2.22

Sample Word Similarities

Query	Most Similar Words
lehman	brothers (0.95), shearson (0.78), fuld (0.78), stearns (0.78)
silicon	valley (0.97), stanford (0.85), san (0.82)
dow	jones (0.95), djia (0.85), average (0.84)
enron	andersen (0.70), worldcom (0.70), dynegy (0.67)
venture	capital (0.83), capitalists (0.80), vc (0.77)
fraud	wire (0.85), charges (0.85), criminal (0.84), sec (0.81)
derivatives	isda (0.86), swaps (0.85), otc (0.83)
futures	contract (0.82), options (0.81), forwards (0.78)
panic	1907 (0.88), 1873 (0.86), 1819 (0.85)
margin	maintenance (0.85), requirement (0.83), call (0.82)

Gradient Derivation

For center word vector $v_c$, positive context vector $v_p$, and negative sample vectors $v_{n_k}$ ($k=1 \dots K$), the objective minimizes:

$$L = -\log \sigma(v_c \cdot v_p) - \sum_{k=1}^K \log \sigma(-v_c \cdot v_{n_k})$$

Using $\sigma'(x) = \sigma(x)(1 - \sigma(x))$, the analytical gradients are:

1. Positive context vector $v_p$: $$\frac{\partial L}{\partial v_p} = (\sigma(v_c \cdot v_p) - 1) , v_c$$

2. Negative sample vector $v_{n_k}$: $$\frac{\partial L}{\partial v_{n_k}} = \sigma(v_c \cdot v_{n_k}) , v_c$$

3. Center vector $v_c$: $$\frac{\partial L}{\partial v_c} = (\sigma(v_c \cdot v_p) - 1) , v_p + \sum_{k=1}^K \sigma(v_c \cdot v_{n_k}) , v_{n_k}$$

These are computed in batched form in model.py using vectorized NumPy operations and normalized by batch size.

Optimizations Implemented

• Sub-sampling of Frequent Words: Probability-based dropping of ultra-frequent tokens (e.g. "the", "a") using the formula from the original paper, improving training speed and rare-word representation quality.
• Negative Sampling: Approximates the full softmax by drawing $K$ negative words per context pair from a unigram distribution raised to the power of 0.75.
• Batched Matrix Operations: All forward/backward computations are fully vectorized using numpy.add.at for correct sparse gradient accumulation across repeated indices.
• Dynamic Window Size: Context window size is randomly sampled from $[1, W]$ per center word, effectively weighting closer words more heavily (as in the original paper).

Possible Alternative Approaches

These were not implemented but are worth discussing:

• CBOW (Continuous Bag of Words): Instead of predicting context from center word, CBOW predicts the center word from the average of context vectors. It is typically faster to train but slightly worse on rare words.
• Hierarchical Softmax: An alternative to negative sampling that uses a binary tree (Huffman coding) to decompose the softmax into $O(\log V)$ binary classifications.
• Adam / Adagrad Optimizer: SGD works but per-parameter adaptive learning rates (Adagrad) can improve convergence, especially for infrequent words whose embeddings are updated rarely.
• Learning Rate Scheduling: Linearly decaying the learning rate from an initial value to near-zero over training (as in the original Word2Vec C code) generally improves final embedding quality.
• Larger Embeddings: Increasing embedding_dim from 50 to 100-300 can capture finer-grained semantic relationships, at the cost of more compute.