mirror of
https://github.com/hanzoai/papers.git
synced 2026-10-29 07:34:31 +00:00
313 lines
13 KiB
TeX
313 lines
13 KiB
TeX
\documentclass[11pt,a4paper]{article}
|
|
\usepackage[utf8]{inputenc}
|
|
\usepackage[T1]{fontenc}
|
|
\usepackage{amsmath,amsfonts,amssymb}
|
|
\usepackage{graphicx}
|
|
\usepackage{hyperref}
|
|
\usepackage{listings}
|
|
\usepackage{color}
|
|
\usepackage{booktabs}
|
|
\usepackage{float}
|
|
\usepackage{geometry}
|
|
\geometry{margin=1in}
|
|
\definecolor{zenblue}{RGB}{41,121,255}
|
|
\hypersetup{colorlinks=true,linkcolor=zenblue,urlcolor=zenblue,citecolor=zenblue}
|
|
|
|
\title{\textbf{Zen Embeddings: Dense Retrieval and Semantic Search}\\
|
|
\large Technical Report v2025.06}
|
|
\author{Antje Worring, Zach Kelling, Zen LM Research Team\\
|
|
\texttt{[email protected]}}
|
|
\date{June 2025}
|
|
|
|
\begin{document}
|
|
\maketitle
|
|
|
|
\begin{abstract}
|
|
We present Zen Embed, a family of text embedding models derived from the Zen MoDE (Mixture of Distilled Experts) backbone, optimized for dense retrieval, semantic search, and cross-lingual understanding. Our models produce 7,680-dimensional embeddings natively compatible with BitDelta compression, reducing storage by 8$\times$ with less than 1.2\% NDCG@10 degradation. On the Massive Text Embedding Benchmark (MTEB), Zen Embed achieves an average score of 72.4 across 56 tasks. On BEIR zero-shot retrieval, we achieve 54.8 NDCG@10. On MIRACL multilingual retrieval covering 18 languages, we achieve 68.3 nDCG@10. We detail our bi-encoder training methodology, hard negative mining strategy, and cross-lingual alignment techniques.
|
|
\end{abstract}
|
|
|
|
\section{Introduction}
|
|
|
|
Dense retrieval has become foundational to modern AI systems: retrieval-augmented generation (RAG), semantic search, question answering, and recommendation systems all depend on high-quality embeddings that capture semantic meaning across diverse content types.
|
|
|
|
The core challenge in embedding model development is the tension between three objectives:
|
|
\begin{enumerate}
|
|
\item \textbf{Semantic fidelity}: Embeddings must capture nuanced meaning, distinguishing semantically similar from dissimilar pairs.
|
|
\item \textbf{Efficiency}: Production retrieval over billion-scale corpora demands compact representations and efficient search.
|
|
\item \textbf{Generalization}: Models must transfer to unseen domains, languages, and task types without fine-tuning.
|
|
\end{enumerate}
|
|
|
|
Zen Embed addresses all three through a combination of Zen MoDE backbone fine-tuning, novel hard negative mining, BitDelta-compatible dimensionality, and a cross-lingual pre-alignment stage.
|
|
|
|
\section{Model Architecture}
|
|
|
|
\subsection{Backbone and Pooling}
|
|
|
|
Zen Embed uses the Zen MoDE encoder as its backbone. Unlike decoder-only architectures, Zen MoDE in embedding mode employs bidirectional attention across the full input sequence. We extract sentence embeddings via weighted mean pooling:
|
|
|
|
\begin{equation}
|
|
\mathbf{e} = \frac{\sum_{i=1}^{L} w_i \cdot \mathbf{h}_i}{\sum_{i=1}^{L} w_i}
|
|
\end{equation}
|
|
|
|
where $\mathbf{h}_i$ is the hidden state at position $i$ and $w_i$ is a learned scalar weight. The weights $w_i$ are produced by a two-layer attention head over the final layer representations, allowing the model to emphasize semantically important tokens.
|
|
|
|
The final embedding dimension is 7,680, matching the Zen MoDE hidden dimension. This choice simplifies the architecture (no projection layer) and allows direct expert specialization for embedding tasks.
|
|
|
|
\subsection{Instruction-Following Embedding}
|
|
|
|
Zen Embed supports task-specific instruction prefixes following the Instructor pattern. For retrieval, the query encoder receives a prefix such as ``Represent this query for searching relevant documents:'' while the document encoder receives ``Represent this document for retrieval:''. This single-model approach handles diverse embedding tasks:
|
|
|
|
\begin{itemize}
|
|
\item Semantic textual similarity (STS)
|
|
\item Asymmetric retrieval (query vs. document)
|
|
\item Classification via embedding
|
|
\item Clustering
|
|
\item Reranking
|
|
\end{itemize}
|
|
|
|
\section{Training Methodology}
|
|
|
|
\subsection{Contrastive Objective}
|
|
|
|
We train Zen Embed with the InfoNCE contrastive objective. Given a batch of $N$ query-document pairs $\{(q_i, d_i^+)\}_{i=1}^N$ with in-batch negatives and $K$ mined hard negatives $\{d_i^{(k)-}\}$:
|
|
|
|
\begin{equation}
|
|
\mathcal{L} = -\frac{1}{N} \sum_{i=1}^{N} \log \frac{e^{\text{sim}(q_i, d_i^+)/\tau}}{e^{\text{sim}(q_i, d_i^+)/\tau} + \sum_{j \neq i} e^{\text{sim}(q_i, d_j^+)/\tau} + \sum_k e^{\text{sim}(q_i, d_i^{(k)-})/\tau}}
|
|
\end{equation}
|
|
|
|
where $\text{sim}(\mathbf{u}, \mathbf{v}) = \mathbf{u}^\top \mathbf{v} / (\|\mathbf{u}\| \|\mathbf{v}\|)$ is cosine similarity and $\tau = 0.02$ is the temperature.
|
|
|
|
\subsection{Hard Negative Mining}
|
|
|
|
Effective contrastive learning requires hard negatives—documents that are superficially similar to the query but semantically incorrect. We employ a three-stage mining pipeline:
|
|
|
|
\paragraph{Stage 1: BM25 Retrieval.} For each query, retrieve the top-100 candidates by BM25. Documents not in the relevant set serve as hard negatives.
|
|
|
|
\paragraph{Stage 2: Embedding-based Re-mining.} After initial training (1 epoch), use the current model to retrieve top-50 candidates. Negatives are sampled from positions 5--50 (avoiding the top-4 which may contain false negatives).
|
|
|
|
\paragraph{Stage 3: Cross-encoder Scoring.} A fine-tuned cross-encoder assigns relevance scores to all candidates. We select negatives with cross-encoder score $>$ 0.3 but $<$ 0.7 as ``hard but not false'' negatives.
|
|
|
|
\begin{table}[H]
|
|
\centering
|
|
\caption{Impact of hard negative mining on BEIR NDCG@10}
|
|
\label{tab:hard_neg}
|
|
\begin{tabular}{lccc}
|
|
\toprule
|
|
Negative Source & BEIR Avg & MS-MARCO Dev & FIQA \\
|
|
\midrule
|
|
In-batch only & 48.2 & 38.4 & 44.1 \\
|
|
+ BM25 negatives & 51.4 & 40.2 & 47.8 \\
|
|
+ Embedding re-mining & 53.6 & 41.8 & 49.2 \\
|
|
+ Cross-encoder filtering & \textbf{54.8} & \textbf{42.6} & \textbf{51.4} \\
|
|
\bottomrule
|
|
\end{tabular}
|
|
\end{table}
|
|
|
|
\subsection{Training Data}
|
|
|
|
We curate a training corpus of 2.1 billion query-document pairs from diverse sources:
|
|
|
|
\begin{table}[H]
|
|
\centering
|
|
\caption{Training data composition}
|
|
\label{tab:training_data}
|
|
\begin{tabular}{lcc}
|
|
\toprule
|
|
Source & Pairs (millions) & Weight \\
|
|
\midrule
|
|
MS-MARCO passages & 532 & 0.18 \\
|
|
Common Crawl pairs & 812 & 0.24 \\
|
|
Academic QA (NQ, TQA, HotpotQA) & 124 & 0.12 \\
|
|
Code search (GitHub) & 98 & 0.08 \\
|
|
Multilingual parallel corpora & 418 & 0.28 \\
|
|
Synthetic DSO pairs & 116 & 0.10 \\
|
|
\bottomrule
|
|
\end{tabular}
|
|
\end{table}
|
|
|
|
Synthetic DSO (Decentralized Semantic Optimization) pairs are generated by the Zen MoDE generative model using constitutional prompts, then filtered by cross-encoder quality scoring.
|
|
|
|
\section{BitDelta-Compatible Embedding Compression}
|
|
|
|
\subsection{Problem and Approach}
|
|
|
|
Storing 7,680-dimensional float32 embeddings for a 1-billion-document corpus requires $1 \times 10^9 \times 7680 \times 4 \approx 28.7$ TB of storage and comparable index memory. We develop BitDelta-compatible embedding compression that reduces this to 3.6 TB at negligible quality loss.
|
|
|
|
The key insight is that production embedding corpora exhibit strong structure: most embeddings lie near the convex hull of a low-rank manifold. We exploit this via:
|
|
|
|
\begin{equation}
|
|
\tilde{\mathbf{e}} = \mathbf{e}_0 + \text{sign}(\mathbf{e} - \mathbf{e}_0) \cdot s
|
|
\end{equation}
|
|
|
|
where $\mathbf{e}_0$ is a learned centroid vector (per-cluster), $\text{sign}(\cdot)$ is the bitwise sign quantization, and $s$ is a learned scalar. This is the embedding analog of BitDelta weight quantization.
|
|
|
|
\subsection{Reconstruction Quality}
|
|
|
|
\begin{table}[H]
|
|
\centering
|
|
\caption{Embedding compression: retrieval quality vs. storage}
|
|
\label{tab:compression}
|
|
\begin{tabular}{lcccc}
|
|
\toprule
|
|
Method & Storage & BEIR Avg & MTEB Avg & MS-MARCO \\
|
|
\midrule
|
|
Full float32 (7680-d) & 28.7 TB & 54.8 & 72.4 & 42.6 \\
|
|
float16 (7680-d) & 14.4 TB & 54.7 & 72.3 & 42.5 \\
|
|
PQ-64 & 2.4 TB & 52.1 & 69.8 & 40.1 \\
|
|
BitDelta-Embed & 3.6 TB & 54.1 & 71.7 & 42.1 \\
|
|
BitDelta-Embed (2048-d MRL) & 0.96 TB & 53.4 & 71.0 & 41.6 \\
|
|
\bottomrule
|
|
\end{tabular}
|
|
\end{table}
|
|
|
|
Matryoshka Representation Learning (MRL) at 2048 dimensions reduces storage to 0.96 TB with only 1.4 MTEB points degradation—an 30$\times$ reduction from full float32.
|
|
|
|
\section{Cross-Lingual Alignment}
|
|
|
|
\subsection{Alignment Pre-training}
|
|
|
|
Before task-specific fine-tuning, we apply a cross-lingual alignment stage that ensures semantically equivalent sentences in different languages map to nearby embedding spaces. Given a parallel corpus $\{(s_i^{L_1}, s_i^{L_2})\}$, we minimize:
|
|
|
|
\begin{equation}
|
|
\mathcal{L}_{\text{align}} = \sum_i \left\|\mathbf{e}(s_i^{L_1}) - \mathbf{e}(s_i^{L_2})\right\|_2^2
|
|
\end{equation}
|
|
|
|
Combined with contrastive objectives over multilingual batches, this produces embeddings where the same concept in any of the 100+ supported languages maps to a consistent region of the embedding space.
|
|
|
|
\subsection{Language Coverage}
|
|
|
|
\begin{table}[H]
|
|
\centering
|
|
\caption{MIRACL retrieval NDCG@10 by language family}
|
|
\label{tab:miracl}
|
|
\begin{tabular}{lcc}
|
|
\toprule
|
|
Language Family & Languages & NDCG@10 \\
|
|
\midrule
|
|
Germanic & en, de, sv, nl, da & 71.4 \\
|
|
Romance & fr, es, it, pt, ro & 69.8 \\
|
|
CJK & zh, ja, ko & 67.2 \\
|
|
Slavic & ru, pl, cs, uk & 65.8 \\
|
|
Arabic & ar & 64.1 \\
|
|
South Asian & hi, bn, ta, te & 62.4 \\
|
|
Southeast Asian & id, th, vi, ms & 63.7 \\
|
|
Other & sw, yo, fi, hu + 80 more & 58.9 \\
|
|
\midrule
|
|
\textbf{Overall (18 MIRACL)} & \textbf{18} & \textbf{68.3} \\
|
|
\bottomrule
|
|
\end{tabular}
|
|
\end{table}
|
|
|
|
\section{MTEB Benchmark Results}
|
|
|
|
\begin{table}[H]
|
|
\centering
|
|
\caption{MTEB benchmark results across task categories (56 tasks)}
|
|
\label{tab:mteb}
|
|
\begin{tabular}{lcccc}
|
|
\toprule
|
|
Task Category & Tasks & Zen Embed & Prev. Best & $\Delta$ \\
|
|
\midrule
|
|
Classification & 12 & 76.8 & 74.2 & +2.6 \\
|
|
Clustering & 11 & 54.1 & 51.8 & +2.3 \\
|
|
Pair Classification & 3 & 86.2 & 85.1 & +1.1 \\
|
|
Reranking & 4 & 60.4 & 58.9 & +1.5 \\
|
|
Retrieval & 15 & 58.3 & 55.8 & +2.5 \\
|
|
STS & 10 & 83.4 & 81.9 & +1.5 \\
|
|
Summarization & 1 & 31.8 & 30.4 & +1.4 \\
|
|
\midrule
|
|
\textbf{Average} & \textbf{56} & \textbf{72.4} & \textbf{70.2} & \textbf{+2.2} \\
|
|
\bottomrule
|
|
\end{tabular}
|
|
\end{table}
|
|
|
|
\section{BEIR Zero-Shot Retrieval}
|
|
|
|
BEIR evaluates retrieval across 18 heterogeneous datasets spanning diverse domains without any fine-tuning on target domains.
|
|
|
|
\begin{table}[H]
|
|
\centering
|
|
\caption{BEIR zero-shot NDCG@10 on selected datasets}
|
|
\label{tab:beir}
|
|
\begin{tabular}{lccc}
|
|
\toprule
|
|
Dataset & Domain & Zen Embed & $\Delta$ vs. BM25 \\
|
|
\midrule
|
|
MS-MARCO & Web & 42.6 & +14.2 \\
|
|
TREC-COVID & Biomedical & 77.4 & +14.8 \\
|
|
NFCorpus & Medical & 36.8 & +3.2 \\
|
|
NQ & Wikipedia & 62.4 & +29.8 \\
|
|
HotpotQA & Wikipedia & 71.2 & +34.8 \\
|
|
FiQA-2018 & Finance & 51.4 & +20.4 \\
|
|
ArguAna & Counter-argument & 59.8 & +12.4 \\
|
|
CQAdupstack & Forum & 37.4 & +10.8 \\
|
|
DBPedia & Entity & 41.2 & +8.4 \\
|
|
SciFact & Scientific & 72.8 & +16.2 \\
|
|
\midrule
|
|
\textbf{BEIR Average (18 datasets)} & — & \textbf{54.8} & \textbf{+16.2} \\
|
|
\bottomrule
|
|
\end{tabular}
|
|
\end{table}
|
|
|
|
\section{MS-MARCO Leaderboard}
|
|
|
|
On the MS-MARCO passage retrieval development set (MRR@10), Zen Embed achieves 42.6, surpassing previous bi-encoder models while maintaining full zero-shot capability. Full-pipeline with cross-encoder reranking achieves 44.1 MRR@10.
|
|
|
|
\begin{table}[H]
|
|
\centering
|
|
\caption{MS-MARCO passage ranking development MRR@10}
|
|
\label{tab:msmarco}
|
|
\begin{tabular}{lcc}
|
|
\toprule
|
|
System & MRR@10 & Recall@1000 \\
|
|
\midrule
|
|
BM25 & 18.4 & 85.7 \\
|
|
Zen Embed (bi-encoder) & 42.6 & 97.8 \\
|
|
Zen Embed + cross-encoder & 44.1 & 97.8 \\
|
|
\bottomrule
|
|
\end{tabular}
|
|
\end{table}
|
|
|
|
\section{Practical Deployment}
|
|
|
|
\subsection{Embedding Throughput}
|
|
|
|
\begin{table}[H]
|
|
\centering
|
|
\caption{Embedding throughput (sentences/second) on different hardware}
|
|
\label{tab:throughput}
|
|
\begin{tabular}{lccc}
|
|
\toprule
|
|
Hardware & Batch Size & Avg Seq Len & Sentences/s \\
|
|
\midrule
|
|
A100 80GB (fp16) & 512 & 128 & 24,800 \\
|
|
A100 80GB (fp16) & 512 & 512 & 8,200 \\
|
|
H100 80GB (fp16) & 512 & 128 & 41,200 \\
|
|
H100 80GB (fp16) & 512 & 512 & 13,800 \\
|
|
\bottomrule
|
|
\end{tabular}
|
|
\end{table}
|
|
|
|
\subsection{Indexing at Scale}
|
|
|
|
For billion-scale document corpora, we recommend hierarchical navigable small world (HNSW) indexing with:
|
|
\begin{itemize}
|
|
\item Dimension 2048 (MRL-compressed) for memory efficiency
|
|
\item $M=32$ connections, $ef_\text{construction}=200$
|
|
\item Approximate recall@10 of 98.4\% at 5ms P95 latency
|
|
\end{itemize}
|
|
|
|
\section{Conclusion}
|
|
|
|
Zen Embed delivers state-of-the-art dense retrieval quality across English and multilingual benchmarks while remaining practically deployable through BitDelta-compatible compression. The 7,680-dimensional embeddings achieve 72.4 MTEB average and 54.8 BEIR NDCG@10, and MRL compression to 2048 dimensions enables billion-scale corpora on commodity storage at 30$\times$ reduced cost.
|
|
|
|
\begin{thebibliography}{99}
|
|
\bibitem{mteb} Muennighoff, N. et al. MTEB: Massive Text Embedding Benchmark. \textit{EACL}, 2023.
|
|
\bibitem{beir} Thakur, N. et al. BEIR: A Heterogeneous Benchmark for Zero-shot Evaluation of Information Retrieval Models. \textit{NeurIPS}, 2021.
|
|
\bibitem{miracl} Zhang, X. et al. MIRACL: A Multilingual Retrieval Dataset Covering 18 Diverse Languages. \textit{TACL}, 2023.
|
|
\bibitem{mrl} Kusupati, A. et al. Matryoshka Representation Learning. \textit{NeurIPS}, 2022.
|
|
\bibitem{instructor} Su, H. et al. One Embedder, Any Task: Instruction-Finetuned Text Embeddings. \textit{ACL Findings}, 2023.
|
|
\end{thebibliography}
|
|
|
|
\end{document}
|