word2vec

Word2Vec paper implementation in pytorch to replicate the paper results and generate a proper embedding structure.

Paper Refs. :

• Mikolov et al 2013 - Distributed Representations of Words and Phrases (SGNS)
• Rong, Xin 2014 - word2vec Parameter Learning Explained
• Levy & Goldberg 2014 - Neural Word Embedding as Implicit Matrix Factorization

Configuration used

min_word_count = 5
neg_sampling_power = 0.75
epochs = 10

params = InitializationParams(
    D=200,  # Embedding dimension
    window_size=5,  # Context window size
    k=5,  # Number of negative samples
    min_word_count=min_word_count,  # Minimum word freq required to include in the vocabulary
    t=1e-4,  # Subsampling threshold
    learning_rate=1,  # Initial learning rate (looks large but for batch size with meaned loss it is fine)
    epochs=epochs,  # Number of training epochs
    batch_size=1024,  # Batch size
    seed = 101,
    vocab_size=71290
)

Model

• Model Training Pipeline
• SPPMI + SVD baseline
• Trained Model - 200 vec. embedding size

Hardware & wall-clock

Device	CPU (Apple Silicon; MPS benchmarked ~3-4x slower on this gather/scatter workload)
Vocabulary	71,290 (min_count = 5)
Throughput	~128k tokens/sec
Wall-clock	~66 s/epoch (~11 min for 10 epochs)
Peak memory	~1.96 GB

Deviations from the paper / reference C implementation

Precise enough to reproduce the numbers. Most stem from batching (the paper/word2vec.c do pure online SGD).

• Minibatching + mean loss → learning rate retuned. Training uses batch_size=1024 with the loss meaned over the batch. Mean-reduction divides each example's gradient by the batch size, so the paper's online lr=0.025 becomes ineffective (~2.4e-5 per example). lr is retuned to 1.0 to recover a comparable per-example step. This is the reason the learning rate looks large.
• Optimizer: SGD with a LinearLR schedule decaying 1.0 → 1.0*1e-4 over training (the C decays 0.025 → 0.025*1e-4).
• Subsampling: the paper formula P(discard) = 1 - sqrt(t/f) (chosen over the C variant), applied to the token stream each epoch, t = 1e-4.
• Negative sampling: unigram P(w) ∝ count(w)^0.75, using counts taken after min_count filtering and before subsampling; k = 5; collisions with the true context are not filtered (matches the C).
• Dynamic window: symmetric reduced reach window - b, b ~ U{0..window-1} (matches word2vec.c); context count varies per center, down-weighting distant words.
• Boundaries: text8 is a single token stream with no sentence markers, so windows cross "sentence" boundaries.
• Init: input vectors U(-0.5/D, +0.5/D), output vectors zeros (from word2vec.c, not specified in the paper). Only syn0 is exported as the final vectors.

Seeds: seed = 101 (single seed; multi-seed variance not yet reported).

Results

Skip Gram implementation

Nearest Neighbours

king -> [('viii', 0.953), ('vii', 0.949), ('elizabeth', 0.947), ('constantine', 0.947), ('crowned', 0.941), ('queen', 0.94), ('prince', 0.939), ('xiv', 0.939)]

france -> [('austria', 0.96), ('spain', 0.96), ('germany', 0.96), ('italy', 0.957), ('hungary', 0.951), ('portugal', 0.948), ('netherlands', 0.948), ('poland', 0.948)]

water -> [('salt', 0.93), ('grain', 0.916), ('fresh', 0.914), ('warm', 0.908), ('soil', 0.907), ('vegetation', 0.902), ('mild', 0.901), ('coal', 0.896)]

music -> [('dance', 0.956), ('musical', 0.928), ('hop', 0.92), ('hip', 0.919), ('pop', 0.915), ('folk', 0.913), ('lindy', 0.903), ('jazz', 0.897)]

war -> [('wars', 0.879), ('invasion', 0.875), ('vietnam', 0.873), ('civil', 0.873), ('battles', 0.872), ('allied', 0.867), ('battle', 0.867), ('fought', 0.866)]

computer -> [('digital', 0.934), ('graphics', 0.923), ('desktop', 0.921), ('computing', 0.921), ('interactive', 0.919), ('interface', 0.919), ('unix', 0.918), ('bsd', 0.917)]

physics -> [('chemistry', 0.928), ('mathematical', 0.917), ('mathematics', 0.908), ('quantum', 0.901), ('theory', 0.9), ('mechanics', 0.896), ('theoretical', 0.887), ('analysis', 0.868)]

god -> [('spirit', 0.956), ('divine', 0.937), ('eternal', 0.927), ('allah', 0.926), ('heaven', 0.917), ('jesus', 0.914), ('yahweh', 0.908), ('baptism', 0.906)]

island -> [('coast', 0.923), ('harbour', 0.921), ('shore', 0.919), ('cape', 0.918), ('northwest', 0.917), ('strait', 0.915), ('southwest', 0.913), ('hills', 0.909)]

president -> [('minister', 0.945), ('deputy', 0.944), ('chairman', 0.942), ('vice', 0.94), ('secretary', 0.929), ('presidential', 0.928), ('attorney', 0.926), ('senator', 0.926)]

3CosAdd

man:king :: woman:? -> [('emperor', 0.887), ('empress', 0.881), ('augustus', 0.871), ('alfonso', 0.869), ('sigismund', 0.867)]
france:paris :: germany:? -> [('berlin', 0.959), ('munich', 0.911), ('moscow', 0.905), ('vienna', 0.904), ('bonn', 0.887)]

3CosMul

man:king :: woman:? -> [('emperor', 0.973), ('empress', 0.967), ('augustus', 0.96), ('alfonso', 0.959), ('julius', 0.958)]
france:paris :: germany:? -> [('berlin', 0.983), ('munich', 0.958), ('moscow', 0.955), ('vienna', 0.954), ('bonn', 0.946)]

Notes:

• france:paris :: germany:? gets berlin right at rank 1.
• man:king :: woman:? misses queen (it's outside the top 5). Instead it returns empress at rank 2, because king's nearest words are mostly other rulers and regnal numbers (viii, vii), not queen.
• 3CosAdd and 3CosMul give the same ranking here (3CosMul just rescales the scores). 3CosMul mainly helps on weaker, less-trained vectors.

SGNS vs Shifted-PPMI + SVD

Levy & Goldberg (2014) showed that SGNS is secretly doing matrix factorization: it factorizes the matrix PMI(w,c) - log k. So we can build that matrix directly (same corpus, same window, same k=5) and factorize it with SVD instead, then compare the neighbours. See svd-baseline.ipynb.

Where they agree : for clear topics, both methods give the same neighbours. This is the theory holding up: they're approximating the same matrix.

music  SGNS: musical, dance, folk, pop        SVD: musical, dance, folk, jazz, pop
god    SGNS: spirit, divine, eternal, heaven  SVD: divine, eternal, heaven, spirit
physics SGNS: chemistry, quantum, mechanics   SVD: quantum, mechanics, electrodynamics

Where they diverge : SGNS returns words of the same kind (countries next to countries). SVD instead surfaces rare words that happen to appear right next to the target. This is because plain SVD treats every cell of the matrix equally, so rare word pairs with a very high PMI score dominate. SGNS weights pairs by how often they actually occur, so those rare spikes get damped:

germany SGNS: france, russia, italy, finland, hungary          (other countries)
germany SVD : hauptbahnhof, neubrandenburg, hbf, magdeburg     (rare German-specific terms)

island  SGNS: coast, harbour, shore, cape                      (common geography)
island  SVD : archipelago, uninhabited, atoll, tutuila, lihou  (rare specific islands)

hauptbahnhof (German for "central station") almost only appears next to German place names, so its PMI score is huge and it dominates the SVD result. SGNS weights by frequency and subsamples common words, so it filters this out. In short: SVD and SGNS factorize the same matrix, but SGNS weights it by word frequency and SVD does not.

Future Work

• CBOW model implementation
• Binary tree implementation for embeddings (Heirarchial Softmax)