Token Embedding Pruning for Sentence Transformers

Introduction

Embedding models are usually deployed in many applications for semantic search, clustering, deduplication, recommendation, and other retrieval-heavy systems. For English, we have many options and people have trained and released models ranging from a few million parameters to billions. For most use cases, the smaller models typically in the range of 20-40 million have very good performance and can be deployed without too much worry about compute and memory. Once we move from English to other languages, there are not as many model choices. This is even worse for low resource languages like Nepali where lack of research and lack of data has made the choices very small. Fortunately, big companies and universities have trained multilingual models on massive corpora that work quite well for many languages. However, the side effect is that the model must know about every language hence the vocabulary can become very large and a surprisingly large share of the model parameters can live in the token embedding layer.

So the question is: do we really need the model to store the data about languages we do not care about? Another question is, does the model’s vocabulary have redundancies? If some token embeddings are already very similar to each other, can we merge them without hurting downstream retrieval quality too much?

In this post, I explore two ways to shrink Sentence Transformer style embedding models by pruning the token embedding layer. The first approach clusters similar token embeddings and replaces each cluster with a shared embedding. The second approach trims the vocabulary to the tokens actually needed for a target language or domain. Both methods only modify the tokenizer and embedding layer, so the rest of the transformer remains unchanged.

With those approaches, the models can be made significantly smaller with an acceptable drop in performance. In my experiments, the best trade-off came from trimming multilingual-e5-small for Nepali, reducing it from 117.7 million parameters to 29.7 million parameters while retaining most of its retrieval performance on NanoBEIR-ne.

The table below shows how much we can reduce for different models and approaches. When adapting multilingual-e5-small and EmbeddingGemma-300m to Nepali only, I was able to shrink them by 74% and 58% respectively. The potential memory usage reduction in fp32 is 335 MB and 685 MB respectively - a significant reduction!

	experiment	Base Vocab	Pruned Vocab	Vocab reduction %	Base params (mil)	Pruned params (mil)	Parameter reduction %	Memory Reduction (mb)
0	MiniLM-L6-v2 — Clustering	30522	20808	31.83	22.71	18.98	16.42	14.23
1	E5-small-v2 — Clustering	30522	22225	27.18	33.36	30.17	9.55	12.15
2	Multilingual E5-small — Clustering	250037	172569	30.98	117.65	87.91	25.28	113.48
3	EmbeddingGemma 300M — Clustering	262144	188191	28.21	307.58	250.79	18.47	216.66
4	Multilingual E5-small — English trim	250037	71547	71.39	117.65	49.11	58.26	261.46
5	Multilingual E5-small — Nepali trim	250037	20928	91.63	117.65	29.68	74.78	335.61
6	EmbeddingGemma 300M — Nepali trim	262144	28195	89.24	307.58	127.91	58.41	685.4

I’ve uploaded several pruned models on HuggingFace that you can try as well.

• Pruned multilingual-e5-small
• Pruned embeddinggemma-300m
• Pruned all-MiniLM-L6-v2
• English only trim of multilingual-e5-small
• Nepali only trim of multilingual-e5-small
• Nepali only trim of embeddinggemma-300m

The source code for the approaches I discuss here is also available in Github: jangedoo/model-pruner.

Motivation

I maintain a Nepali news aggregator which uses pre-trained embedding models for clustering, de-duplication and even in a personalized recommender system as one of the candidate generators. When I started the project, I evaluated several models which had to work with both English and Nepali. Naturally I wanted the model with the best performance in my evaluation dataset for ndcg, recall and other metrics but I was also constrained with how much memory the model would take when deployed. However, multi lingual models were quite big so I decided to look for separate models for each language.

For English, the choice was obvious to me: sentence/transformers/all-miniLM-L6-v2 - a small model with decent performance but such a model (at similar size) for Nepali didn’t exist. So I fine-tuned that model in two iterations jangedoo/all-MiniLM-L6-v2-nepali and jangedoo/all-MiniLM-L6-v3-nepali. However, I realized that it takes way too many tokens to represent Nepali text than English since all-miniLM-L6-v2 was trained on an English dataset and the tokenizer doesn’t really know about Nepali words.

For the first iteration of my news aggregator project, I went with two deployed models - one for English and one for Nepali. Each of those models has about 23 million parameters. This introduced complexity in my code and deployment setup as well. Again after a second round of model evaluation, I ended up with intfloat/multilingual-e5-small which is multi-lingual and has around 100 million parameters which is a bit larger than the previous two models combined. But I was happy with the offline evaluation metrics and decided to deploy this even though it meant very little memory was left in my server.

Lack of available memory really bugged me and I started looking for ways to reduce the memory usage. Then I had a couple of ideas:

1. I use the embeddings of news articles to cluster and de-duplicate them so why not take the same approach for tokens as well. We just run a clustering algorithm on the weights of the embedding layer and for each cluster with more than X tokens, merge the embeddings of each cluster together. Surely there must be some tokens which are very similar to each other. Based on my research, I did not find any previous literature discussing this approach, so it might be a novel approach for pruning the models.

2. Multi-lingual models typically have a large vocabulary size to account for different languages they are trained on. A large vocab size means a large embedding layer which accounts for most of the model parameters. e.g. google/embeddinggemma-300m has around 307 million total parameters out of which 201 million are just in the Embedding layer. So an obvious approach is why not remove the embeddings of tokens which are from languages I do not care about. There are a few papers discussing this idea already. TextPruner: A Model Pruning Toolkit for Pre-Trained Language Models proposes vocabular pruning (same idea) and transformer pruning which removes attention heads with least importance. An Efficient Multilingual Language Model Compression through Vocabulary Trimming also proposes vocabular pruning and has a Python package available in Github as well.

Pruning

There are many approaches to prune the model weights, but here I’m specifically focused on reducing the parameters in the embedding layer. For example, if we consider a typical vocabulary size of 32,000 and an Embedding layer with 512 dimensions, then total parameters of the Embedding layer alone is about 16.3 million. Google’s embeddinggemma-300m has a vocabulary size of approx. 262K with 768 embedding dimensions which leads to about 201 million parameters (66% of total weights) just for the Embedding layer. The figure below shows the share of Embedding layer parameters in four different models. As expected, multi-lingual models have high number of parameters in Embedding layer compared to mono-lingual ones.

In my framework, I structure this in a four step pipeline:

1. Grouping: This stage assigns a cluster_id to each of the tokens. Tokens with the same cluster_id are considered to be redundant. This abstraction works for both clustering the token embeddings and domain adaptation. Following scikit learn’s convention, a cluster_id=-1 means all items are distinct even though they have the same cluster_id.

2. Reducing: This stage extracts embeddings for each cluster_id. A single cluster_id can have multiple tokens, so we need to decide what the embedding of this new cluster_id will be. An obvious choice is to simply take the average of all embeddings in this cluster.

3. Updating the weights: We simply replace the original embedding weights with new embeddings. The number of rows in the embedding matrix will change from Vocab Size x Embed Dim to Cluster Size x Embed Dim.

4. Updating the tokenizer: Since the embedding layer has changed, which now has embeddings per cluster rather than per token, we need to update the tokenizer as well. An additional mapping from token_id to cluster_id is stored in the tokenizer. During the encoding process, we extract the token_ids and then use the mapping to lookup cluster_ids for each of those tokens. There could be a better way to do this, but I ended up with a new subclass PrunedTokenizer which maintains the mapping and uses this mapping during encoding. During the pruning process, the original tokenizer is replaced with this tokenizer.

Evaluation setup

To evaluate the model’s performance after pruning I will be using the following datasets:

• stsb: Semantic Textual Similarity Benchmark
• NanoBEIR: A condensed version of standard BEIR (Benchmarking Information Retrieval) and covers diverse tasks such as web search, duplicate detection, question answering etc.
• NanoBEIR-ne: Nepali translation of NanoBEIR dataset which I created. The translation was done using DeepSeek V4 Flash.

There are many metrics that can be computed such as NDCG, precision, recall etc. However, for my use cases where these models are mainly used for candidate generation, recall is the metric I care about most. Also the absolute performance of the pruned model is less important to me than its relative performance compared to the base model. If a model was pruned by 30% but still retains 95% of base model’s recall performance, then it is an acceptable trade-off for me.

Clustering token embeddings

The idea is to simply cluster the tokens based on their embedding weights. Instead of per-token embedding, we will have per-cluster embedding. Depending on the parameters we choose for the clustering algorithm, the reduction can be quite significant. For clustering, I chose DBSCAN, simply because I don’t need to choose the number of clusters to create. There are two main hyperparameters eps and distance metric used. You can read the scikit learn’s documentation about the details of the algorithm. For all my experiments, I set the min_samples parameter to 2 - because I want the clusters to have at least 2 items. Although this will probably result in many clusters with few items, I want to capture as much redundant tokens as possible.

As for the metric, I used cosine since it is always bounded between 0 and 2. DBSCAN uses distance rather than similarity so cosine distance has bound 0 to 2 since cosine distance = 1 - cosine similarity and cosine similarity is between -1 and 1. So the final knob to tweak is the eps parameter which indicates the max. distance between two tokens to be considered as neighbors. Any pair with more than eps distance won’t be in the same cluster.

Domain adaptation

Multi-lingual language models have a large vocabulary to account for different languages. Vocabulary size directly contributes to the overall number of parameters of a model. Assuming that the total number of tokens related to Nepali language is around 40k, embeddinggemma’s Embedding layer would just need 30 million parameters. This is a significant reduction compared to the original 201 million parameters.

As mentioned before, these two papers TextPruner: A Model Pruning Toolkit for Pre-Trained Language Models and An Efficient Multilingual Language Model Compression through Vocabulary Trimming explore this idea.

Implementation is actually quite straightforward. We take a multilingual model M with original vocabulary size V and a target language corpus C. Then we tokenize the corpus C with the model M’s tokenizer and collect unique token_ids. The size of this unique token_ids will be our new vocabulary size. From the model M’s embedding layer, we only select the embeddings of the collected token_ids and remove the weights of unrelated tokens from the embedding layer.

Below is a pseudocode of the implementation.

tokenizer = m.tokenizer
# extract unique tokens of this domain
domain_token_ids = set(tokenizer.tokenize(corpus.text).token_ids)
# selct embeddings of these tokens only 
pruned_embeddings = m.embedding_layer.weights[domain_token_ids, :]
# replace the weights of embedding layer with pruned embeddings
m.embedding_layer.weights = pruned_embeddings

Evaluation

First, let’s take a look at the reduction in total model parameters using two different methods: Clustering and Domain adaptation. The figure below shows the reduction of different models. In this setup, Clustering approach just merged very similar tokens and Nepali trim/English trim indicates Domain adaptation by only keeping Nepali or English tokens.

As shown in the figure above, clustering approach reduced the model size quite a bit for all models - multilingual-e5-small seeing the best reduction of about 25%. Even MiniLM-L6-V2, which is already pretty small and compact, dropped from 22.7 million to 19 million parameters.

The biggest drop was seen from domain adaptation (indicated by English or Nepali trim in the figure). I was able to make multilingual-e5-small 75% smaller by adapting it to Nepali text only. Similarly, English trim of multilingual-e5-small saw a reduction of about 60%.

Embedding Clustering

But how much do we lose by pruning the models like this? Actually, not that much. The figure below shows the mean recall@10 across all NanoBEIR and NanoBEIR-ne tasks.

The best outcome was Nepali trim of multilingual-e5-small model. The Nepali trim has 29.7 million parameters and still has very respectable 0.458 mean recall@10 compared to 0.485 of the base model with 117.7 million parameters.

The scores above are averaged across all NanoBEIR tasks, but not every task might be of interest to you. I use these models for semantic search and deduplication. Semantic search is asymmetric - a short query like a phrase or question should match a longer passage that answers it. For clustering/deduplication however we’d like the model to recognize that one question is similar to another question or one text is a paraphrase of another. Because these use cases have different goals, I focus mainly on DBPedia and MSMARCO for semantic search and QuoraRetrieval for duplicate detection.

The figure below shows the difference in recall@10 for base vs pruned via clustering method. For English (NanoBEIR), the performance drop (recall@10) on MSCARCO, DBPedia and QuoraRetrieval is negligible. For Nepali (NanoBEIR-ne) MiniLM-L6-v2 is just used as a baseline and can be ignored since it is an English-only model. I didn’t evaluate E5-small-v2. For EmbeddingGemma and multilingual-e5-small the performance drop can be seen in most of the tasks, and surprisingly for some tasks, the pruned models had better performance than the base model.

	Model	Metric	Base	Pruned	Absolute change	Change (pp)	Retained (%)
0	MiniLM-L6-v2	Pearson cosine	0.8274	0.8183	-0.0091	-0.91	98.9
1	MiniLM-L6-v2	Spearman cosine	0.8203	0.8095	-0.0108	-1.08	98.68
2	E5-small-v2	Pearson cosine	0.8317	0.8234	-0.0084	-0.84	99
3	E5-small-v2	Spearman cosine	0.8492	0.8386	-0.0106	-1.06	98.76
4	Multilingual E5-small	Pearson cosine	0.8092	0.7925	-0.0166	-1.66	97.94
5	Multilingual E5-small	Spearman cosine	0.8359	0.8014	-0.0344	-3.44	95.88
6	EmbeddingGemma 300M	Pearson cosine	0.8446	0.8113	-0.0333	-3.33	96.06
7	EmbeddingGemma 300M	Spearman cosine	0.8485	0.8074	-0.0412	-4.12	95.15

Domain adaptation

The figure below shows the performance difference for base vs pruned via Nepali domain adaptation. EmbeddingGemma suffered a bit more loss than multilingual-e5-small. The performance of pruned multilingual-e5-small on DBPedia, MSMARCO and QuoraRetrieval is excellent given that we shrunk the model by more than 75%.

In STSB-ne dataset, which is used to measure pairwise semantic similarity, the performance of base vs pruned (Nepali trim) is almost exactly the same. After removing 91.6% of vocabulary and 75% of all parameters, multilingual-e5-small retains 99.71% of its STS-ne Pearson correlation. Similarly, EmbeddingGemma removes 89% of its vocabulary and 58% of all parameters while retaining 99.31% of Pearson correlation. For tasks like clustering and de-duplication, domain adaptation pruning has almost no impact.

	Model	Metric	Base	Pruned	Absolute change	Retained (%)
0	Multilingual E5-small	Pearson cosine	0.7227	0.7205	-0.0021	99.71
1	Multilingual E5-small	Spearman cosine	0.7182	0.7156	-0.0026	99.64
2	EmbeddingGemma 300M	Pearson cosine	0.6941	0.6893	-0.0048	99.31
3	EmbeddingGemma 300M	Spearman cosine	0.6831	0.6786	-0.0046	99.33

For English domain trim, I evaluated only multilingual-e5-small and the performance on STSB is exactly the same.

	Model	Metric	Base	Pruned	Absolute change	Retained (%)
0	Multilingual E5-small	Pearson cosine	0.8092	0.8091	-0	100
1	Multilingual E5-small	Spearman cosine	0.8359	0.8359	-0	100

Clustering + Domain adaptation

We can go one step further and combine both clustering and domain adaptation. Even after domain adaptation, remaining tokens could be redundant. We take a two step process by first performing domain adaptation and then clustering on remaining tokens. I’ve already implemented this here but didn’t get the chance to evaluate it yet. I think we can shrink the model even more by combining.

Conclusion

For my use case of having a good embedding model but with a lower memory footprint, this approach has worked quite well. I showed the evaluation results from public datasets. In my private datasets on typical information retrieval tasks and even end to end recommender systems, I did not see any decrease in the metrics and sometimes they even improved for a few tasks.

Also, since the proposed solution only modifies the tokenizer and the embedding layer, this is naturally applicable to generative language models as well. Although I suspect generative models are more sensitive to the identity of each token compared to the embedding models. This is something I wish to explore in the future.

I hope you enjoyed the post. If you found any mistakes, please let me know.