Add new tokens to tokenizer before distillation with Model2Vec

[aoezdTchibo]

Hello!

Is it possible to incorporate new tokens into the tokenizer before I distill the corresponding model with Model2Vec? Also, in the context of Zipf's Law, at which position in the vocabulary?

What I could do is the following:

tokens = ["my_special_token", "my_special_token2"]
tokenizer = AutoTokenizer.from_pretrained("intfloat/multilingual-e5-large")
tokenizer.add_tokens(tokens)
static_embedding = StaticEmbedding(
    tokenizer,
    embedding_dim=1024,
)
model = SentenceTransformer(
    modules=[static_embedding],
    device="cpu",
)
model.tokenizer.encode("my_special_token my_special_token2").tokens
# --> ['<s>', 'my_special_token', 'my_special_token2', '</s>']

However, this would be without distillation and corresponding weighting. The context is that we use the distilled model for product search, and we don't want to break down specific brand names further, but rather take them as they are, especially while fine-tuning.

[aoezdTchibo]

My first try looks like this (see #744 for adding new tokens):

tokens = ["my_special_token", "my_special_token2"]
model_name = "intfloat/multilingual-e5-large"
base_model = SentenceTransformer(model_name)
base_model_path = f"./models/{model_name}"

# Add new tokens to the base model and save this custom model version
base_model_module = base_model._first_module()
base_model_module.tokenizer.add_tokens(tokens, special_tokens=True)
base_model_module.auto_model.resize_token_embeddings(len(base_model_module.tokenizer))
base_model.save(base_model_path)

# For distillation load the custom base model version
static_embedding = StaticEmbedding.from_distillation(
    base_model_path,
    device="cpu",  # I am testing it locally
    pca_dims=256,
    apply_zipf=True,
)
static_embedding.embedding.weight.data = (
    static_embedding.embedding.weight.data.contiguous()
)
model = SentenceTransformer(
    modules=[static_embedding],
    device="cpu",
)

At file level, I can see that the new tokens are present in the added_tokens.json file. However, during the distillation process, I get the following warning: "Reported vocab size 501153 is inconsistent with the vocab size 501155." The fine-tuning afterwards fails.

Currently, the only solution is to manually integrate the tokens into the tokenizer.json file at the file level, but this seems to be way to hacky. After doing this, both the distillation and subsequent fine-tuning work properly.

Any suggestions?

[tomaarsen]

Hello!

I feel like this might not work as you'd like it to. Notably, from_distillation uses a pretrained Sentence Transformer model to determine the static embeddings, but if you add new tokens to the pretrained Sentence Transformer, then that model will not know anything about those tokens yet. The corresponding embeddings will be awful. Instead, we don't want to add new tokens, but tell the distillation function to also distil with texts that aren't just individual tokens.

I believe you should instead be using the vocabulary argument from from_distillation, and set it to e.g. list(model[0].tokenizer.get_vocab()) + ["my_special_token", "my_special_token2"]. With "my_special_token" being e.g. "Star Wars", I believe Model2Vec (which is used in from_distillation) will embed the multiple tokens in "Star Wars", and then mean them together. The idea is that this produces a meaningful embedding for the new token, which will then be a part of the StaticEmbedding. However, I'm not 100% sure the model[0].tokenizer will be correct afterwards.
Please do experiment with this, and let me know if it works as you hope it does!

• Tom Aarsen

[aoezdTchibo]

Hi Tom!

We had already experimented internally with the vocabulary argument of the from_distillation function, but came to the conclusion that this couldn't be a suitable solution due to the fact that this feature only works with WordPiece tokenizers (yet), and all foundation models relevant to our use case are based on subword tokenizers:

ValueError: You passed a vocabulary, but the model you are using does not use a WordPiece tokenizer. This is not supported yet. Feel free to open an issue if this is a blocker: https://github.com/MinishLab/model2vec/issues

But your uncertainty is justified, because I also wonder how the inference should run if the tokenizer doesn't know about the new tokens. This way, no look-up to the corresponding static embedding is possible...

Your point regarding more valuable embeddings is completely understandable and was also confirmed after fine-tuning because the embedding model determined lower similarity values for the new tokens, even though direct matches were present in the comparison text. So generally, the embedding model was less confident regarding the new tokens, which, as you describe yourself, is also understandable. Essentially, it doesn't know the new tokens as well yet. Do you think with enough examples while fine-tuning, the model would eventually be able to generate valuable embeddings?

[tomaarsen]

Before I dive into considering finetuning, I am aware that the MinishLab team behind Model2Vec recently reworked their backend to allow these multi-word vocabulary options for more tokenizers, I believe.

Could you install the bleeding edge version of model2vec and try again?

pip install git+https://github.com/MinishLab/model2vec.git

cc @stephantul @Pringled

• Tom Aarsen

[stephantul]

Hello! Thanks for the ping @tomaarsen!

As tom said, this should work from main now. Please let us know if it doesn't, and open an issue in model2vec itself. We're super happy to solve this correctly.

I've confirmed that it works with e5 and modernbert-based models. But I don't know, we might run into other issues.

[aoezdTchibo]

Hi @stephantul and @tomaarsen,

Thank you both for your valuable assistance. We greatly appreciate your help!

I've created a new issue in the Model2Vec repository: MinishLab/model2vec#209

Lets us continue the discussion there. Consequently, I'm closing this issue.