Mastering Multi-Vector Embedding Models with Sentence Transformers

Multi-Vector Models: A Powerful Tool for Retrieval Tasks

Multi-vector models are shaking things up in information retrieval. Unlike traditional dense embeddings that squeeze text into a single vector, these models keep a vector for every token. This lets you match tokens more precisely using the MaxSim operator. Instead of mushing everything together, it compares each query token with document tokens, keeping all those subtle signals that dense embeddings tend to lose. It's a big step forward for retrieval accuracy.

Why Finetune Multi-Vector Models?

If you're working in specialized domains, like medical or legal fields, finetuning these models becomes vital. General-purpose models often miss the unique vocabulary and query style those fields demand. By finetuning, you align the model with domain specifics, drastically boosting retrieval performance. Especially in medicine, where the average document might be 941 tokens long, conventional models that chop off text at 256 or 512 tokens miss out on critical info.

Finetuning helps the model cope with longer documents, which is a plus for fields needing detailed information. During training, configuring document lengths ensures your model processes the full context, bumping up retrieval accuracy.

Setting Up for Finetuning

Choosing the Right Starting Point

Picking a solid starting point for finetuning is crucial. The source article points to various checkpoints, like lightonai/mLateOn-unsupervised, which adapts well to new domains. From what I've seen, unsupervised checkpoints—those hanging out after contrastive pretraining but before supervised finetuning—tend to adapt better than fully supervised ones.

Loading and Configuring the Model

To get rolling with finetuning, load up a multi-vector model through the Sentence Transformers library. Here's how you can configure one:

from sentence_transformers import MultiVectorEncoder

# Load the model with configuration for longer documents
model = MultiVectorEncoder(
    "lightonai/mLateOn-unsupervised",
    model_kwargs={"torch_dtype": "float32"},
    processor_kwargs={"model_max_length": 8192}
)

# Unset document length caps if present
model[0].query_length = None
model[0].document_length = None

# Add a punctuation skiplist
import string
model[2].skiplist_words = list(string.punctuation)
model[2].resolve_with_tokenizer(model.tokenizer)

This setup means your model can tackle documents up to 8192 tokens long, perfect for domains with lengthy texts.

Dataset Preparation

To finetune, you need the right datasets. The article suggests looking on the Hugging Face Datasets Hub or using local datasets in CSV or JSON format. Typically, you'll have query-passage pairs where each passage answers a query. Here's how you could load a dataset from the Hub:

from datasets import load_dataset

train_dataset = load_dataset("tomaarsen/miriad-4.4M-split", split="train")
print(train_dataset)

Make sure your dataset format fits your chosen loss function. MultiVectorMultipleNegativesRankingLoss works well for question-answer pairs, optimizing the model with in-batch negatives.

Training Process and Considerations

Training these models involves a few key pieces: the model, dataset, loss function, training settings, evaluator, and trainer. They all come together to fine-tune the model for your specific domain needs.

Training Components

• Model: Decide whether to fine-tune an existing one or start from scratch with a base transformer.
• Dataset: Your dataset should be compatible, with clear query-passage pairs.
• Loss Function: Pick a function that suits your data, like MultiVectorMultipleNegativesRankingLoss for in-batch negatives.
• Training Arguments: Set parameters that impact both performance and debugging.
• Evaluator: Optionally, keep tabs on model performance during training.
• Trainer: This component pulls everything together for the training process.

What I'd Do on Monday Morning

First thing Monday, I'd check out our current retrieval system to spot any weak points. If we're handling domain-specific data, especially in sectors like healthcare or law, finetuning a multi-vector model is on the table. Here's how I'd tackle it:

1. Identify what's unique about our domain and find datasets that match those needs.
2. Pick a good starting model, likely an unsupervised checkpoint for better domain fit.
3. Load both the model and dataset, ensuring they're configured for the domain's document length.
4. Kick off the finetuning, watching for performance gains and tweaking settings where necessary.
5. Compare the tuned model against our current setup to see if we've made real improvements.

Limitations and Trade-offs

Sure, multi-vector models can boost retrieval accuracy, but they have their downsides. More vectors mean a bigger index, which can slow down storage and retrieval speeds. Also, finetuning demands domain-specific data, and that's not always on hand. Organizations need to weigh these issues against the retrieval benefits they might see.

And don't forget the computational resources. Even though you might pull it off with a single consumer GPU, the time and cost still matter.

Worked Example: Improving Medical Retrieval

Let's get practical. Say we're out to improve a medical retrieval system. We start with the lightonai/mLateOn-unsupervised model, known for adapting well to new domains. First, our dataset of medical queries and passage pairs should be ready and work seamlessly with MultiVectorMultipleNegativesRankingLoss. Once we’ve loaded the model with the right configurations for long texts and added the punctuation skiplist, we dive into finetuning.

During this stage, keeping an eye on the model’s performance via a validation set is crucial. You might need to tweak training parameters like the learning rate or batch size for optimal results. Post-finetuning, you assess the model’s new performance, stacking it up against the current system with metrics like NDCG@10.

Finetuning with Sentence Transformers isn't just a task—it's a strategy for enhancing domain-specific retrieval work. Dig into the components and steps, and you're well on your way to better pull what you need from your data.