WordLlama is a utility for NLP and word embedding that repurposes components from large language models (LLMs) to generate efficient and compact word representations, similar to GloVe, Word2Vec, or FastText. It starts by extracting the token embedding codebook from a state-of-the-art LLM (e.g., LLaMA3 70B) and trains a small, context-free model within a general-purpose embedding framework.
WordLlama outperforms other word models like GloVe 300d on all MTEB benchmarks, while being significantly smaller (16MB for the default 256-dimension model compared to over 2GB).
WordLLama Key Features
Key features of WordLlama include:
- Matryoshka Representations: Flexible truncation of the embedding dimensions as needed.
- Low Resource Requirements: Simple token lookup with average pooling allows it to run efficiently on CPUs.
- Binarization: Models trained with the straight-through estimator can be packed into small integer arrays, enabling faster Hamming distance calculations (coming soon).
- Numpy-only Inference: Lightweight and simple implementation.
WordLlama uses Matryoshka representation learning to offer flexibility. The largest model (1024 dimensions) can be truncated to 64, 128, 256, or 512 dimensions. For binary embedding models, straight-through estimators are employed during training. Dense embeddings perform well with 256 dimensions, while binary embeddings see validation accuracy nearing saturation at 512 dimensions (64 bytes packed).
The final weights are saved after weighting, projection, and truncation of the entire tokenizer vocabulary, reducing the model to a single embedding matrix (nn.Embedding) that is far smaller than the gigabyte-scale LLM codebooks. The original tokenizer preprocesses text into tokens, and the reduced-size embeddings are average pooled, requiring minimal computation. Resulting model sizes range from 16MB to 250MB for the 128k LLaMA3 vocabulary.
WordLlama is ideal for lightweight NLP tasks such as training sklearn classifiers, performing semantic matching, fuzzy deduplication, ranking, and clustering. It’s well-suited for creating LLM output evaluators or multi-hop workflows, and with its compact and portable design, you can train models on consumer GPUs in just a few hours. Its speed and versatility make it a great “Swiss Army knife” for exploratory analysis and utility applications.
MTEB Results (l2_supercat)
| Metric | WL64 | WL128 | WL256 (X) | WL512 | WL1024 | GloVe 300d | Komninos | all-MiniLM-L6-v2 |
|---|---|---|---|---|---|---|---|---|
| Clustering | 30.27 | 32.20 | 33.25 | 33.40 | 33.62 | 27.73 | 26.57 | 42.35 |
| Reranking | 50.38 | 51.52 | 52.03 | 52.32 | 52.39 | 43.29 | 44.75 | 58.04 |
| Classification | 53.14 | 56.25 | 58.21 | 59.13 | 59.50 | 57.29 | 57.65 | 63.05 |
| Pair Classification | 75.80 | 77.59 | 78.22 | 78.50 | 78.60 | 70.92 | 72.94 | 82.37 |
| STS | 66.24 | 67.53 | 67.91 | 68.22 | 68.27 | 61.85 | 62.46 | 78.90 |
| CQA DupStack | 18.76 | 22.54 | 24.12 | 24.59 | 24.83 | 15.47 | 16.79 | 41.32 |
| SummEval | 30.79 | 29.99 | 30.99 | 29.56 | 29.39 | 28.87 | 30.49 | 30.81 |
HuggingFace
License
This project is licensed under the MIT License.