mxbai-embed-2d-large-v1

Model Description

mxbai-embed-2d-large-v1 is the world's first 2D-Matryoshka embedding model. The 2D-Matryoshka model introduces a novel approach that enables you to reduce both the number of layers and the dimensions of embeddings within the model. This dual reduction strategy allows for a more compact model size while still delivering performance on par with that of leading models such as Nomic's embedding model. Specifically, reducing the model's layers by approximately 50% retains up to 85% of its original performance, even without additional training.

The model was pretrained using contrastive training on over 700 million pairs, covering a wide variety of topics across the internet. It was then fine-tuned with over 30 million high-quality triplets using novel loss functions. mxbai-embed-2d-large-v1 allows users to get multiple models out of one and use different embedding sizes, providing full control over the trade-offs between speed, storage consumption, and model performance.

On the Massive Text Embedding Benchmark (MTEB), mxbai-embed-2d-large-v1a performs at the level of current embedding models of different sizes. The model's performance remains competitive even when the embedding size is reduced by a factor of 16. Additionally, the model retains about 75% of its performance after cutting half of its layers, demonstrating the effectiveness of the 2D-Matryoshka approach.

Using a Prompt

Adding a domain-specific prompt to a text can help the model understand how the embedding will be used.

For retrieval tasks, the query can be preceded by the prompt: Represent this sentence for searching relevant passages:. For other tasks, the text can be used as-is without any additional prompt.

Suitable Scoring Methods

• Cosine Similarity: Ideal for measuring the similarity between text vectors, commonly used in tasks like semantic textual similarity and information retrieval.
• Euclidean Distance: Useful for measuring dissimilarity between embeddings, especially effective in clustering and outlier detection.
• Dot Product: Appropriate when embeddings are normalized; used in tasks where alignment of vector orientation is critical.

Limitations

• Language: mxbai-embed-2d-large-v1 is trained on English text and is specifically designed for the English language.
• Sequence Length: The suggested maximum sequence length is 512 tokens. Longer sequences may be truncated, leading to a loss of information.

Examples

Semantic Search

The following code illustrates how to compute similarities between sentences using the cosine similarity score function. The number of dimensions can be adjusted using the dimensions parameter of the SDK.

import torch
from mixedbread import Mixedbread
from sentence_transformers.util import semantic_search
 
mxbai = Mixedbread(api_key="YOUR_API_KEY")
model = "mixedbread-ai/mxbai-embed-2d-large-v1"
 
prompt = 'Represent this sentence for searching relevant passages:'
query = "A man is eating a piece of bread"
 
docs = [
    "A man is eating food.",
    "A man is eating pasta.",
    "The girl is carrying a baby.",
    "A man is riding a horse.",
]
 
query_res = mxbai.embed(
    model=model,
    prompt=prompt,
    input=[query],
    normalized=True,
    encoding_format='float',
    dimensions=512
)
 
docs_res = mxbai.embed(
    model=model,
    input=docs,
    normalized=True,
    encoding_format='float',
    dimensions=512
)
 
query_embedding_list = query_res.data[0].embedding
docs_embeddings_list = [item.embedding for item in docs_res.data]
 
query_tensor = torch.tensor([query_embedding_list])
docs_tensor = torch.tensor(docs_embeddings_list)
 
hits = semantic_search(query_tensor, docs_tensor, top_k=len(docs))
 
print(f"Query: {query}\n")
print("Results (sorted by relevance):")
for hit in hits[0]:
    doc_index = hit['corpus_id']
    score = hit['score']
    print(f"Score: {score:.4f}\tDocument: {docs[doc_index]}")