GitHub - benjismith/word2vec-mojo · GitHub
/" data-turbo-transient="true" />
Skip to content
Type / to search
Sign in<br>Sign upAppearance settings
You signed in with another tab or window. Reload to refresh your session.<br>You signed out in another tab or window. Reload to refresh your session.<br>You switched accounts on another tab or window. Reload to refresh your session.
Dismiss alert
{{ message }}
benjismith
word2vec-mojo
Public
Notifications<br>You must be signed in to change notification settings
Fork
Star
main
BranchesTags
Go to file
CodeOpen more actions menu
Folders and files<br>NameNameLast commit message<br>Last commit date<br>Latest commit
History<br>5 Commits<br>5 Commits
docs
docs
scripts
scripts
src
src
tests
tests
.gitattributes
.gitattributes
.gitignore
.gitignore
LICENSE
LICENSE
README.md
README.md
pixi.lock
pixi.lock
pixi.toml
pixi.toml
View all files
Repository files navigation
word2vec-mojo
A faithful, fast port of the classic<br>word2vec tool (Mikolov et al.,<br>Google, 2013) to the Mojo programming<br>language — same command-line interface, same file formats, same training<br>algorithm, rewritten for clarity and modern hardware.
Vectors trained by this port are interchangeable with vectors trained by the<br>original C tool: each implementation's distance/readbin utilities read<br>the other's binary output byte-for-byte.
Why Mojo
The original implementation is ~700 lines of dense C built around global<br>variables, hand-rolled hash tables, and compiler-dependent auto-vectorization.<br>Mojo lets this codebase say what it means:
Explicit SIMD. The two kernels all training arithmetic reduces to —<br>dot product and scaled vector addition — are written directly against the<br>CPU's vector registers (SIMD[DType.float32, width]), sized automatically<br>for whatever machine compiles them (src/word2vec/simd_math.mojo).
Structured parallelism. The reference's pthread boilerplate becomes a<br>closure handed to parallelize; the lock-free "Hogwild" weight sharing the<br>algorithm depends on is expressed in one deliberate unsafe_origin_cast<br>with a comment, instead of being implicit everywhere<br>(src/word2vec/trainer.mojo).
Safety by default, escape hatches where earned. Bounds-checked<br>collections and the compiler's exclusivity checking cover all of the code<br>except the hot loops that opt out on purpose. Porting under those checks<br>caught a real latent indexing hazard at the sigmoid table's boundary that<br>the C code only avoids through an integer-division accident.
A vocabulary that reads like the algorithm. syn0 is<br>input_embeddings, neu1e is input_gradient, vocab[word].point is<br>ancestor_rows — every module carries the explanation of both the<br>algorithm and its data layout.
Two implementation upgrades over the reference (both documented in-source,<br>neither changing the learned model):
Negative sampling draws noise words with Vose's alias method — exact<br>O(1) sampling from a few megabytes of tables — replacing the C tool's<br>quantized 100-million-entry array (400 MB of RAM)<br>(src/word2vec/noise_sampler.mojo).
The corpus tokenizer streams the training file in 1 MiB slabs instead of<br>one fgetc() per byte (src/word2vec/corpus_stream.mojo).
Performance
Benchmarked against the reference C implementation on an Apple M2 Pro<br>(8 performance + 4 efficiency cores), training a 192 MB / 35.7M-word corpus<br>with production-typical settings (skip-gram, -negative 15 -sample 1e-4 -size 100), both tools loading an identical shared vocabulary:
~3.3x the single-thread throughput of the C tool (390K vs. 120K<br>words/sec) — the SIMD kernels, fused update pass, and streaming<br>tokenizer at work.
Saturates the machine's memory bandwidth at 8 threads (~1.6M<br>words/sec) and holds that plateau flat through 16 threads; the C tool<br>peaks at 0.90M words/sec at 16 threads, still short of that ceiling.<br>Net: about 1.8x end-to-end at high thread counts , approaching 3x at<br>low ones.
Work is claimed by threads dynamically, so on heterogeneous CPUs<br>(performance + efficiency cores) every thread count from 8 up performs<br>identically here — efficiency cores help rather than gate. -threads 8<br>is the energy-efficient choice on this machine.
Full tables, the optimization history (including one measured-and-rejected<br>idea), and measurement methodology: docs/benchmarks.md.<br>Rerun with scripts/benchmark.sh.
Related work
As of August 2026 this appears to be the only word2vec implementation<br>written in Mojo (a GitHub search for word2vec repositories in the language<br>returns exactly this repository). The nearest neighbor is<br>mojo-gensim, which pursues a<br>different goal: a gensim-workalike library (negative sampling only, no<br>hierarchical softmax, no compatibility with the C tools' CLI or file<br>formats) benchmarked against gensim rather than the reference.
Among faithful rewrites — implementations keeping the original's exact<br>SGD semantics — published comparisons cluster around parity to ~1.6x the<br>C tool's throughput:<br>gensim's Cython<br>core is roughly at parity (with known...