Model2Vec: Distill a Small Fast Model from any Sentence Transformer

Model2Vec is a method to distill a small, fast model from any Sentence Transformer model.

Model2vec allows you to create really fast small models that still perform well.

Table of Contents

• Quickstart
• What is Model2Vec?
• Main Features
• Who is this for?
• Usage
  • Distilling a Model2Vec model
  • Inferencing a Model2Vec model
  • Evaluating a Model2Vec model
• Model List
• Results
• Citing

Quickstart

Install the package with:

pip install model2vec

The easiest way to get started with Model2Vec is to download one of our flagship models from the HuggingFace hub. These models are pre-trained and ready to use. The following code snippet shows how to load a model and make embeddings:

from model2vec import StaticModel

# Load a model from the HuggingFace hub (in this case the M2V_base_output model)
model_name = "minishlab/M2V_base_output"
model = StaticModel.from_pretrained(model_name)

# Make embeddings
embeddings = model.encode(["It's dangerous to go alone!", "It's a secret to everyone."])

Alternatively, you can distill your own Model2Vec model from a Sentence Transformer model. The following code snippet shows how to distill a model:

from model2vec.distill import distill

# Choose a Sentence Transformer model
model_name = "BAAI/bge-base-en-v1.5"

# Distill the model
m2v_model = distill(model_name=model_name, pca_dims=256)

# Save the model
m2v_model.save_pretrained("m2v_model")

What is Model2Vec?

Model2Vec is a simple and effective method to distill any sentence transformer into static embeddings. It works by inferencing a vocabulary with the specified Sentence Transformer model, reducing the dimensionality of the embeddings using PCA, weighting the embeddings using zipf weighting, and storing the embeddings in a static format. When a vocabulary is passed, a word-level tokenizer is created on the fly based on the vocabulary. When output embeddings are used, the subword tokenizer from the Sentence Transformer is used.

This technique creates a small, fast, and powerful model that outperforms other static embedding models by a large margin on a a number of relevent tasks, while being much faster to create than traditional static embedding models such as GloVe, without need for a dataset.

Main Features

• Small: Model2Vec can reduce the size of a Sentence Transformer model by a factor of 15 *.
• Fast distillation: Model2Vec can distill a Sentence Transformer model in ~5 minutes on CPU *.
• Fast inference: Model2Vec creates static embeddings that are up to 500 times * faster than the original model.
• State-of-the-art static embedding performance: Model2Vec outperforms traditional static embeddings by a large margin on a number of benchmarks.
• No data needed: Distillation happens directly on a token leven, so no dataset is needed.
• Simple to use: Model2Vec provides an easy to use interface for distilling and inferencing Model2Vec models.
• Bring your own model: Model2Vec can be applied to any Sentence Transformer model.
• Bring your own vocabulary: Model2Vec can be applied to any vocabulary, allowing you to use your own domain-specific vocabulary.
• Multi-lingual: Model2Vec can easily be applied to any language.
• Tightly integrated with HuggingFace hub: Model2Vec models can be easily shared and loaded from the HuggingFace hub. Our models can be found here.
• Easy Evaluation: Model2Vec comes with a set of evaluation tasks to measure the performance of the distilled model.

* Based on the bge-base-en-v1.5 model.

Who is this for?

Model2Vec allows anyone to create their own static embeddings from any Sentence Transformer model in minutes. It can easily be applied to other languages by using a language-specific Sentence Transformer model and vocab. Similarly, it can be applied to specific domains by using a domain specific model, vocab, or both. This makes it an ideal tool for fast prototyping, research, and production use cases where speed and size are more important than performance.

Usage

Distilling a Model2Vec model

Distilling a model from the output embeddings of a Sentence Transformer model:

from model2vec.distill import distill

# Choose a Sentence Transformer model
model_name = "BAAI/bge-base-en-v1.5"

# Distill the model
m2v_model = distill(model_name=model_name, pca_dims=256)

# Save the model
m2v_model.save_pretrained("m2v_model")

Distilling with a custom vocabulary:

from model2vec.distill import distill

# Load a vocabulary as a list of strings
vocabulary = ["word1", "word2", "word3"]
# Choose a Sentence Transformer model
model_name = "BAAI/bge-base-en-v1.5"

# Distill the model with the custom vocabulary
m2v_model = distill(model_name=model_name, vocabulary=vocabulary, pca_dims=None)

# Save the model
m2v_model.save_pretrained("m2v_model")

Alternatively, the command line interface can be used to distill a model:

python3 -m model2vec.distill --model-name BAAI/bge-base-en-v1.5 --vocabulary-path vocab.txt --device mps --save-path model2vec_model

Inferencing a Model2Vec model

Inferencing with one of our flagship Model2Vec models:

from model2vec import StaticModel

# Load a model from the HuggingFace hub
model_name = "minishlab/M2V_base_output"
model = StaticModel.from_pretrained(model_name)

# Make embeddings
embeddings = model.encode(["It's dangerous to go alone!", "It's a secret to everyone."])

Inferencing with a saved Model2Vec model:

from model2vec import StaticModel

# Load a saved model
model_name = "m2v_model"
model = StaticModel.from_pretrained(model_name)

# Make embeddings
embeddings = model.encode(["It's dangerous to go alone!", "It's a secret to everyone."])

Evaluating a Model2Vec model

Model2Vec models can be evaluated using our evaluation package. To run this, first install the optional evaluation package:

pip install evaluation@git+https://github.com/MinishLab/evaluation@main

Then, the following code snippet shows how to evaluate a Model2Vec model:

from model2vec import StaticModel

from evaluation import CustomMTEB, get_tasks, parse_mteb_results, make_leaderboard, summarize_results
from mteb import ModelMeta

# Get all available tasks
tasks = get_tasks()
# Define the CustomMTEB object with the specified tasks
evaluation = CustomMTEB(tasks=tasks)

# Load the model
model_name = "m2v_model"
model = StaticModel.from_pretrained(model_name)

# Optionally, add model metadata in MTEB format
model.mteb_model_meta = ModelMeta(
            name=model_name, revision="no_revision_available", release_date=None, languages=None
        )

# Run the evaluation
results = evaluation.run(model, eval_splits=["test"], output_folder=f"results")

# Parse the results and summarize them
parsed_results = parse_mteb_results(mteb_results=results, model_name=model_name)
task_scores = summarize_results(parsed_results)

# Print the results in a leaderboard format
print(make_leaderboard(task_scores))

Model List

Model	Language	Description	Vocab	Sentence Transformer	Params
M2V_base_glove	English	Flagship embedding model based on GloVe vocab.	GloVe	bge-base-en-v1.5	102M
M2V_base_output	English	Flagship embedding model based on bge-base-en-v1.5 vocab. Uses a subword tokenizer.	Output	bge-base-en-v1.5	7.5M

Results

Main Results

Model2Vec is evaluated on MTEB, as well as two additional tasks: PEARL (a phrase representation task) and WordSim (a word similarity task). The results are shown in the table below.

Model	Avg (All)	Avg (MTEB)	Class	Clust	PairClass	Rank	Ret	STS	Sum	PEARL	WordSim
all-MiniLM-L6-v2	56.08	56.09	62.62	41.94	82.37	58.04	41.95	78.90	30.81	60.83	49.91
M2V_base_glove	48.58	47.60	61.35	30.52	75.34	48.50	29.26	70.31	31.50	50.28	54.29
M2V_base_output	46.79	45.34	61.25	25.58	74.90	47.63	26.14	68.58	29.20	54.02	49.18
GloVe_300d	42.84	42.36	57.31	27.66	72.48	43.30	22.78	61.90	28.81	45.65	43.05
WL256*	48.88	49.36	58.98	33.34	74.00	52.03	33.12	73.34	29.05	48.81	45.16

Task Abbreviations

For readability, the MTEB task names are abbreviated as follows:

• Class: Classification
• Clust: Clustering
• PairClass: PairClassification
• Rank: Reranking
• Ret: Retrieval
• STS: Semantic Textual Similarity
• Sum: Summarization

* WL256, introduced in the WordLlama package is included for comparison due to its similarities to Model2Vec. However, we believe it is heavily overfit to the MTEB dataset since it is trained on datasets used in MTEB itself. This can be seen by the fact that the WL256 model performs much worse on the non-MTEB tasks (PEARL and WordSim) than our models. The results shown in the Classification and Speed Benchmarks further support this.

Classification and Speed Benchmarks

In addition to the MTEB evaluation, Model2Vec is evaluated on a number of classification datasets. These are used as additional analysis to avoid overfitting to the MTEB dataset and to benchmark the speed of the model. The results are shown in the table below.

model	Average	sst2	imdb	trec	ag_news
bge-base-en-v1.5	90.00	91.54	91.88	85.16	91.45
all-MiniLM-L6-v2	84.10	83.95	81.36	81.31	89.77
M2V_base_output	82.23	80.92	84.56	75.27	88.17
M2V_base_glove	80.76	83.07	85.24	66.12	88.61
WL256	78.48	76.88	80.12	69.23	87.68
GloVe_300d	77.77	81.68	84.00	55.67	89.71

As can be seen, the Model2Vec models outperforms the GloVe and WL256 models on all classification tasks, and is competitive with the all-MiniLM-L6-v2 model while being much faster.

The scatterplot below shows the relationship between the number of sentences per second and the average classification score. The bubble sizes correspond to the number of parameters in the models (larger = more parameters), and the colors correspond to the sentences per second (greener = more sentences per second). This plot shows that the Model2Vec models are much faster than the other models, while still being competitive in terms of classification performance with the all-MiniLM-L6-v2 model.

Figure: The average accuracy over all classification datasets plotted against sentence per second. The circle size indicates model size.

Citing

If you use Model2Vec in your research, please cite the following:

@software{minishlab2024word2vec,
  authors = {Stephan Tulkens, Thomas van Dongen},
  title = {Model2Vec: Turn any Sentence Transformer into a Small Fast Model},
  year = {2024},
  url = {https://github.com/MinishLab/model2vec},
}