Skip to content

add benchmarks #20

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
stephantul merged 6 commits into from
Nov 4, 2025
Merged

add benchmarks #20

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
6 commits
Select commit Hold shift + click to select a range
e40f0be
add benchmarks
stephantul Nov 4, 2025
5fe7f57
fix rounding
stephantul Nov 4, 2025
054286f
fix number
stephantul Nov 4, 2025
a5bda41
ToC
stephantul Nov 4, 2025
53337eb
fix docs
stephantul Nov 4, 2025
b651ef0
fix e-mail
stephantul Nov 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
203 changes: 52 additions & 151 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,14 +16,30 @@
</a>
</div>

NIFE compresses large embedding models into static, drop-in replacements with up to 200x faster query embedding [see benchmarks]().

<div align="center">
<h2>
<a href="https://huggingface.co/collections/stephantulkens/nife-models"><strong>Models</strong></a> |
<a href="https://huggingface.co/collections/stephantulkens/nife-data"><strong>Datasets</strong></a> |
<a href="./benchmarks/"><strong>Benchmarks</strong></a> |
<a href="./docs/README.md"><strong>Create your own model</strong></a>
</div>

NIFE compresses large embedding models into static, drop-in replacements with up to 200x faster query embedding [see benchmarks](#benchmarks).

## Features

- 200x faster CPU query embedding
- 400-900x faster CPU query embedding
- Fully aligned with their teacher models
- Re-use your existing vector index

## Table of contents

1. [Quickstart](#quickstart)
2. [Installation](#installation)
3. [Usage](#usage)
4. [Rationale](#rationale)

## Introduction

Nearly Inference Free Embedding (NIFE) models are [static embedding](https://huggingface.co/blog/static-embeddings) models that are fully aligned with a much larger model. Because static models are so small and fast, NIFE allows you to:
Expand Down Expand Up @@ -113,6 +129,13 @@ print(model.similarity(query, docs))

```

## Pretrained models

I have two pretrained models:

* [`stephantulkens/NIFE-mxbai-embed-large-v1`](https://huggingface.co/stephantulkens/NIFE-mxbai-embed-large-v1): aligned with [`mxbai-embed-large-v1`](https://huggingface.co/mixedbread-ai/mxbai-embed-large-v1).
* [`stephantulkens/NIFE-gte-modernbert-base`](https://huggingface.co/stephantulkens/NIFE-gte-modernbert-base): aligned with [`gte-modernbert-base`](https://huggingface.co/Alibaba-NLP/gte-modernbert-base).

## Rationale

For retrieval using dense models, the normal mode of operation is to embed your documents, and put them in some index. Then, using that same model, also embed your queries. In general, larger embedding models are better than smaller models, so you're often better off by making your embedder as large as possible. This however, makes inference more difficult; you need to host a larger model, and embedding queries might take longer.
Expand All @@ -130,6 +153,31 @@ See this table:

As in doc-SPLADE, you lose performance. No real way about it, but as with other fast models, the gap is smaller than you might think.

## Benchmarks

I benchmark our models on [NanoBEIR](https://huggingface.co/collections/zeta-alpha-ai/nanobeir). I use two trained models:

* [`stephantulkens/NIFE-mxbai-embed-large-v1`](https://huggingface.co/stephantulkens/NIFE-mxbai-embed-large-v1)
* [`stephantulkens/NIFE-gte-modernbert-base`](https://huggingface.co/stephantulkens/NIFE-gte-modernbert-base)

For all models, I report NDC@10 and queries per second. I do this for the student model and teacher model, to show how much performance you lose when switching between them. Detailed benchmark performance can be found in [the benchmarks folder.](./benchmarks/). The query timings were performed on the first 1000 queries of the msmarco dataset, and averaged over 7 runs. The benchmarks were run on an Apple M3 pro.

### `gte-modernbert-base`

| | Queries per second (CPU) | NDCG@10 |
|---------|--------------------------|---------|
| NIFE | 71400 (14ms/1k queries) | 59.2 |
| Teacher | 237 (4210ms/1k queries) | 66.34 |

### `mxbai-embed-large-v1`

| | Queries per second (CPU) | NDCG@10 |
|---------|--------------------------|---------|
| NIFE | 65789 (15ms/1k queries) | 59.2 |
| Teacher | 108 (9190ms/1k queries) | 65.6 |

It is interesting that both NIFE models get the same performance, even with different teacher models. This could point towards a ceiling effect, where a certain percentage of queries can be answered correctly by static models, while others require contextualization.

## How does it work?

We use knowledge distillation from an initialized static model to the teacher we want to emulate. Some special things:
Expand All @@ -147,156 +195,9 @@ NIFE can't do the following things:
1) Ignore words based on context: the query "What is the capital of France?" the word "France" will cause documents containing the term "France" to be retrieved. There is no way for the model to attenuate this vector and morph it into the answer ("Paris").
2) Deal with negation: for the same reason as above; there is no interaction between tokens, so the similarity between "Cars that aren't red" and "Cars that are red" will be really high.

# Creating a NIFE model

To create a NIFE model, you can run the scripts in `scripts`, or directly use the code from the repository. First, you should create a corpus of embeddings for your embedder. You can also use pre-computed collections of embeddings I created:

* [mixedbread-ai/mxbai-embed-large-v1](https://huggingface.co/collections/stephantulkens/mxbai-large-v1-embedpress)
* [Alibaba-NLP/gte-modernbert-base](https://huggingface.co/collections/stephantulkens/gte-modernbert-embedpress)

Broadly construed, training a NIFE model has 5 separate steps.

### 1. Create a set of embeddings using the teacher

Let's assume we want to create embeddings on [trivia QA](https://huggingface.co/mandarjoshi/trivia_qa), using `mxbai-embed-large-v1` as a teacher.

```python
from datasets import load_dataset
from pynife.distillation.infer import generate_and_save_embeddings
from sentence_transformers import SentenceTransformer

model_name = "mixedbread-ai/mxbai-embed-large-v1"
model = SentenceTransformer(model_name)

dataset_name = "mandarjoshi/trivia_qa"
dataset = load_dataset(dataset_name, "rc", split="train")
dataset_iterator = (x['question'] for x in dataset)

output_directory = "my-trivia-qa"

generate_and_save_embeddings(
model=model,
records=dataset_iterator,
output_folder=output_directory,
limit_batches=None,
batch_size=8,
save_every=512,
max_length=512,
model_name=model_name,
dataset_name=dataset_name,
lowercase=False,
make_greedy=False,
)

```

This piece of code loads the model, the dataset and then starts inference. Inference takes a while, and will stream snippets to disk as .txt files and torch tensor files. After the whole dataset has been inferenced, the .txt and tensor files are converted into parquet files, and the .txt and torch tensor files are deleted.

Your dataset will be ready and saved as parquet files in `output_directory`. If you want to upload these, please use the `HfAPI`, not `dataset.push_to_hub`, because we rely on some metadata embedded in the README to infer the base model later on. Note that the dataset iterator can be anything, and does not need to be a Hugging Face dataset. For example, it could also work with a stream from your database.

For a simple inference script with a lot of pre-made datasets, see [the infer_datasets script](./scripts/infer_datasets.py).

### 2. (optional) Expanding a tokenizer

NIFE models work really well if you create a custom tokenizer for your domain. Empirically, it also works really well if you just expand the tokenizer of your teacher model with additional words. We call this _tokenizer expansion_. We have a pre-defined corpus to work on:

```python
from transformers import AutoTokenizer

from datasets import load_dataset
from pynife.tokenizer.expand_tokenizer import expand_tokenizer


dataset = load_dataset("stephantulkens/msmarco-vocab", split="train")
print(dataset.tolist()[:5])
# [{'token': '.', 'frequency': 36174594, 'document_frequency': 8701009},
# {'token': 'the', 'frequency': 28806701, 'document_frequency': 7712172},
# {'token': ',', 'frequency': 25825435, 'document_frequency': 7411743},
# {'token': 'of', 'frequency': 15196930, 'document_frequency': 6562023},
# {'token': 'a', 'frequency': 13702107, 'document_frequency': 6064770},

tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
# Function expects an iterator over dictionaries with "token" and "frequency" as keys.
new_tokenizer = expand_tokenizer(tokenizer, data, new_vocabulary_size=30000)
new_tokenizer.save_pretrained("my_tokenizer")

```

This will do a couple of things:
1) It will remove all tokens from the original tokenizer that aren't present in your data.
2) It will then add the most frequent tokens until the size of the tokenizer == `new_vocabulary_size`.

This works a lot better than training a tokenizer from scratch on equivalent data. For a runnable version, see [the expand_tokenizer script](./scripts/expand_tokenizer.py).

To get frequency counts, you can use `count_tokens_in_dataset`, as follows:

```python
from datasets import load_dataset, Dataset

from pynife.tokenizer.count_vocabulary import count_tokens_in_dataset

dataset = load_dataset("sentence-transformers/msmarco", "corpus", split="train", streaming=True)
dataset_iterator = (item["passage"] for item in dataset)
counts = count_tokens_in_dataset(dataset_iterator)

# Save the counts as a dataset if you want.
dataset = Dataset.from_list(counts, split="train")
dataset.push_to_hub("my_hub")

```

This dataset can be used directly to expand your tokenizer, above. For a runnable version, see [the create_vocabulary script](./scripts/create_vocabulary.py)

### 3. Train

Given a dataset and optionally a tokenizer, there's two steps to complete for a successful training.

#### 3a Initialize a static model using your teacher

Using *your teacher model*, initialize a static model. For example, when using [`mixedbread-ai/mxbai-embed-large-v1`](https://huggingface.co/mixedbread-ai/mxbai-embed-large-v1):

```python
from sentence_transformers import SentenceTransformer
from transformers import AutoTokenizer

from pynife.initialization import initialize_from_model

teacher = SentenceTransformer("mixedbread-ai/mxbai-embed-large-v1")
# The tokenizer you trained in step 2. or an off-the-shelf tokenizer.
tokenizer = AutoTokenizer.from_pretrained("my_tokenizer")
model = initialize_from_model(teacher, tokenizer)

```

#### 3b Actually train

Now you can train, just like a regular sentence transformer. In my experiments, I found that using the cosine distance as a loss function was superior to using MSE, so I recommend using that, find it in `pynife.losses`. In addition, I also recommend using [Matryoshka Representation Learning](https://arxiv.org/abs/2205.13147). There's a bunch of helper functions in `pynife` to make training easier. In general, I recommend using hyperparameters like the following:

* `batch_size`: 128
* `learning rate`: 0.01
* `scheduler`: "cosine_warmup_with_min_lr"
* `warmup_ratio`: 0.1
* `weight_decay`: 0.01
* `epochs`: 5

It can be tempting to move to very high batch sizes, but this has a very large detrimental effect on performance, even with higher learning rates. As a consequence, GPU usage during training is actually pretty low, because there's very little actual computation happening. For a complete runnable training loop, including model initialization, see [the training script](./scripts/experiment_distillation.py).

```python
from pynife.losses import CosineLoss
from pynife.data import get_datasets

# Fill with datasets you trained yourself.
datasets_you_made = [""]
train_dataset = get_datasets(datasets_you_made)

# Model is initialized in step 3a.
loss = CosineLoss(model=model)

# Train as usual.

```
## Inquiries

This will train a model and report the result to wandb. The `experiment_distillation` script is otherwise completely the same as a regular sentence transformers training loop, so there's very little actual code involved.
If you think NIFE could be interesting for your business let me know, I am open to consulting jobs regarding training models and fast inference. Just reach out to me [via e-mail](mailto:stephantul@gmail.com).

## License

Expand Down
66 changes: 66 additions & 0 deletions benchmarks/README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,66 @@
# Benchmarks

This document contains detailed benchmarks for all nanoBEIR datasets, and some additional benchmarks with batch sizes.

## Full nanoBEIR results

Here are the full nanoBEIR results for each teacher/NIFE pair.

### `NIFE-mxbai-embed-large-v1`

| dataset | ndcg@10 teacher | ndcg@10 NIFE |
|:-------------------|------------------:|---------------:|
| NanoArguAna | 0.67 | 0.64 |
| NanoClimateFEVER | 0.40 | 0.32 |
| NanoDBPedia | 0.64 | 0.61 |
| NanoFEVER | 0.92 | 0.87 |
| NanoFiQA2018 | 0.56 | 0.44 |
| NanoHotpotQA | 0.87 | 0.75 |
| NanoMSMARCO | 0.66 | 0.62 |
| NanoNFCorpus | 0.38 | 0.34 |
| NanoNQ | 0.71 | 0.60 |
| NanoQuoraRetrieval | 0.96 | 0.90 |
| NanoSCIDOCS | 0.45 | 0.38 |
| NanoSciFact | 0.79 | 0.74 |
| NanoTouche2020 | 0.53 | 0.47 |

The NIFE model performs worse on most datasets, but performs closely to its teacher on DBPedia and Arguana. Note that the NIFE model was trained on MSMARCO, so this represents a case of in-domain results: if you have many documents and can train on them, your NIFE model can approach the performance of your teacher model.

### `NIFE-gte-modernbert-base`

| dataset | ndcg@10 teacher | ndcg@10 NIFE |
|:-------------------|------------------:|---------------:|
| NanoArguAna | 0.77 | 0.61 |
| NanoClimateFEVER | 0.46 | 0.38 |
| NanoDBPedia | 0.60 | 0.61 |
| NanoFEVER | 0.94 | 0.81 |
| NanoFiQA2018 | 0.62 | 0.54 |
| NanoHotpotQA | 0.77 | 0.65 |
| NanoMSMARCO | 0.65 | 0.64 |
| NanoNFCorpus | 0.35 | 0.35 |
| NanoNQ | 0.72 | 0.63 |
| NanoQuoraRetrieval | 0.97 | 0.90 |
| NanoSCIDOCS | 0.47 | 0.37 |
| NanoSciFact | 0.82 | 0.75 |
| NanoTouche2020 | 0.48 | 0.46 |

As you can see the NIFE model performs worse on most datasets, but outperforms the base model on DBPedia and is pretty close on NanoMSMARCO. Note that the NIFE model was trained on MSMARCO, so this represents a case of in-domain results: if you have many documents and can train on them, your NIFE model can approach the performance of your teacher model.

## Effect of batch size on speed

One interesting issue is that most query embedders are not run on very large batch sizes. In real-life workloads, embedders are often run in microservices that ingest a single query at a time. Therefore, reporting QPS at high batch sizes is actually not a realistic estimate of performance.

For this experiment, I use `NIFE-mxbai-embed-large-v1`, all timings are done on CPU (Macbook Pro M3)

| Batch size | NIFE QPS | Teacher QPS | x Speedup |
|----:|-------:|----------:|------------:|
| 1 | 12925 | 13 | 1007 |
| 2 | 14090 | 26 | 552 |
| 4 | 19670 | 45 | 435 |
| 8 | 31326 | 64 | 486 |
| 16 | 46346 | 82 | 567 |
| 32 | 62077 | 99 | 628 |
| 64 | 75661 | 107 | 706 |
| 128 | 93948 | 101 | 931 |

As you can see both NIFE and the teacher benefit a lot from batching, although NIFE is already very fast to begin with. This shows that, on CPU, there is a very large gain on QPS by switching from the teacher to NIFE.
Loading