PaperlayerView exact revision on arXiv

Source-linked AI summary

Jina Embeddings 2: 8192-Token General-Purpose Text Embeddings for Long Documents

Michael Günther, Jackmin Ong, Isabelle Mohr, Alaeddine Abdessalem, Tanguy Abel, Mohammad Kalim Akram, Susana Guzman, Georgios Mastrapas, Saba Sturua, Bo Wang, Maximilian Werk, Nan Wang, Han Xiao

arXiv:2310.19923v4cs.CLcs.AIcs.LG

TL;DR

Many embedding models have limited input lengths, and chunking documents creates fragmented representations and more vectors. Jina Embeddings v2 uses bidirectional ALiBi in a modified BERT architecture to encode up to 8192 tokens; it performs competitively on MTEB, and extended sequence lengths offer an advantage over models without that capability.

  • Problem

    Many embedding models have limited input lengths, so chunking documents yields fragmented representations and more vectors, increasing memory use, search demands, and latency.

    Source: §1 · p. 1

  • Method

    Jina Embeddings v2 uses bidirectional ALiBi in a modified BERT architecture to create open-source embedding models that encode texts up to 8192 tokens.

    Source: §1 · p. 1, §7 · p. 9

  • Results

    Jina Embeddings v2 is competitive on MTEB, and extended sequence lengths offer an advantage over models without that capability.

    Source: §1 · p. 2, §7 · p. 10

  • Takeaways & Limitations

    Extended context can amplify the efficacy of numerous downstream tasks tied to embeddings.

    Source: §1 · p. 2

    MTEB omits very long texts, prompting the authors to add extended-document retrieval and clustering tasks.

    Source: §6.2 · p. 7

Abstract

from arXiv · show

Text embedding models have emerged as powerful tools for transforming sentences into fixed-sized feature vectors that encapsulate semantic information. While these models are essential for tasks like information retrieval, semantic clustering, and text re-ranking, most existing open-source models, especially those built on architectures like BERT, struggle to represent lengthy documents and often resort to truncation. One common approach to mitigate this challenge involves splitting documents into smaller paragraphs for embedding. However, this strategy results in a much larger set of vectors, consequently leading to increased memory consumption and computationally intensive vector searches with elevated latency. To address these challenges, we introduce Jina Embeddings 2, an open-source text embedding model capable of accommodating up to 8192 tokens. This model is designed to transcend the conventional 512-token limit and adeptly process long documents. Jina Embeddings 2 not only achieves state-of-the-art performance on a range of embedding-related tasks in the MTEB benchmark but also matches the performance of OpenAI's proprietary ada-002 model. Additionally, our experiments indicate that an extended context can enhance performance in tasks such as NarrativeQA.

1 Introduction

The introduction presents long-document embedding as a problem of semantic fragmentation and resource costs, then introduces an ALiBi-enhanced BERT encoder for texts up to 8192 tokens. Its fine-tuned embeddings outperform Jina Embeddings v1 and rival E5 on MTEB, while longer context can improve downstream embedding tasks.

1 Introduction

  • Chunking documents to work around sequence limits fragments meaning and increases vector counts, memory use, search computation, and latency.The introduction notes that memory-intensive database index structures can exacerbate these costs.Source: §1 · p. 1
  • 8192 tokens: bidirectional ALiBi in BERT enables open-source encoder embeddings for long texts beyond conventional sequence limits.The model is fine-tuned to encode texts into single embedding representations.Source: §1 · p. 2
  • The fine-tuned embeddings outperform the Jina Embeddings v1 suite [Günther et al., 2023] on MTEB and rival E5 [Wang et al., 2022].The model was fine-tuned using hundreds of millions of text samples.Source: §1 · p. 2
  • Longer context can amplify performance on numerous downstream embedding tasks, motivating evaluation datasets with long text values.The authors curated these datasets to address the predominance of short passages in existing embedding evaluations.Source: §1 · p. 2

2 Related Work

Embedding research has advanced from foundational statistical methods to pretrained and contrastively trained models, but long-document encoding remains constrained and inadequately evaluated by common benchmarks.

  • Embedding methods progressed from LSA and LDA to Sentence-BERT [Reimers and Gurevych, 2019] and unsupervised contrastive learning [Gao et al., 2022, Wang et al., 2022].Condenser and RetroMAE introduced architectures and pretraining methods tailored to dense encoding and retrieval.Source: §2 · p. 2
  • E5 [Wang et al., 2022], Jina Embeddings v1 [Günther et al., 2023], and GTE use multi-stage contrastive training and achieve state-of-the-art MTEB outcomes.Training first uses weak pairs from public data to improve domain generalization, then supervised fine-tuning on annotated text triples spanning diverse tasks.Source: §2 · p. 2
  • BERT-derived models retain a 512-token input limit, while ALiBi supports sequences beyond 2048 tokens but is typically used in generative models with unidirectional bias.Source: §2 · p. 2
  • BEIR evaluates retrieval tasks and settings, while MTEB spans tasks and languages, but both give limited attention to long-document encoding.Source: §2 · p. 2

3 Training Paradigm Overview

Jina Embeddings v2 is trained in three stages: long-context BERT pre-training, text-pair fine-tuning, and further fine-tuning with hard negatives for embedding tasks.

I Pre-training a Modified BERT

  • The modified BERT backbone is trained from scratch on a full-text corpus with masked language modeling to encode documents up to 8192 tokens.Source: §3 · p. 2

II Fine-tuning with Text Pairs

  • Text-pair fine-tuning trains the model to encode a passage as a single vector representation.Source: §3 · p. 2

III Fine-tuning with Hard Negatives

  • Hard-negative fine-tuning helps the model distinguish relevant passages from related but irrelevant ones, especially improving retrieval and classification.The paper identifies this stage as especially critical for those tasks.Source: §3 · pp. 2–3

4 Pre-training a Modified BERT

The section introduces a BERT-based backbone modified to encode extended sequences and strengthen language modeling, with pre-training largely following Liu et al. [2019a] plus performance optimizations.

4.1 Model Architecture

  • ALiBi replaces positional embeddings with symmetric per-head linear attention-score biases, enabling bidirectional encoder attention; each head uses a distinct bias scalar.Unlike the original causal variant, the encoder variant mirrors attention biases so tokens can attend in both directions.Source: §4.1 · p. 3, Figure 1 · p. 4
  • Feedforward layers use GEGLU in small/base and ReGLU in large because large-model GEGLU training was unstable; attention blocks use post-layer normalization after pre-layer tests failed to improve stability or performance.The large-model GEGLU variant showed promising initial MLM accuracy but unstable training.Source: §4.1 · p. 3

4.2 Training Data

  • Pre-training uses English C4 (about 365 million web documents and 170 billion tokens), making the models English-only; 1% is held out for validation loss and MLM accuracy.C4 is cleaned Common Crawl data retaining English content and removing webpages with inappropriate content.Source: §4.2 · p. 3

4.3 Training Algorithm

  • Pre-training uses whole-word masked language modeling, masking 30% of tokens with an 80/10/10 replacement split; it omits NSP and predicts masked tokens with a vocabulary-wide decoder and cross-entropy.NSP is excluded because it was considered to contribute little to downstream performance.Source: §4.3 · p. 4
  • Pre-training caps sequences at 512 tokens, takes only each document’s first 512 tokens without document packing or multiple samples per document, and uses global batches of 4096.Because sequence lengths vary, each batch contains a different number of masked tokens for loss calculation; ALiBi also makes position-embedding training unnecessary, allowing short-sequence pre-training and later adaptation to longer sequences.Source: §4.3 · p. 4
  • Optimization uses AdamW (β1 = 0.9, β2 = 0.98, ϵ = 1e−6, weight decay 0.01), dropout 0.1, and FP16 dynamic mixed precision; BF16 underperformed on MLM accuracy and GLUE.The learning rate warms up linearly for 10,000 steps, then decays to zero after 100,000 steps; peak rates are 1e−3, 6e−4, and 4e−4 for small, base, and large models, respectively.Source: §4.3 · p. 4

5 Fine-Tuning for Embeddings

The models are fine-tuned to produce single-vector text representations through an unsupervised text-pair stage followed by supervised training with negatives to improve ranking.

5 Fine-Tuning for Embeddings

  • Mean pooling averages token embeddings into one representation without adding trainable parameters, following Sentence-BERT [Reimers and Gurevych, 2019].The models are fine-tuned in an unsupervised phase followed by a supervised one.Source: §5 · pp. 4–5

5.1 Fine-tuning with Text Pairs

  • About 40 diverse sources supply text pairs; title–abstract pairs improve clustering, while consistency filtering and source-random batch sampling shape the corpus.Pairs are pre-shuffled, with sampling rates adjusted to source quality and quantity.Source: §5.1 · p. 5
  • InfoNCE pulls paired query-target embeddings together relative to other batch targets; empirical testing found τ = 0.05 effective.The loss aims to give paired texts analogous representations and unpaired texts distinct representations.Source: §5.1 · p. 5

5.2 Fine-tuning with Hard Negatives

  • The supervised stage targets better ranking with batches of one positive and 15 negatives; retrieval uses validated hard negatives, while non-retrieval uses random negatives.Hard negatives diminished quality on non-retrieval datasets, which were retained to preserve performance on those tasks.Source: §5.2 · p. 5
  • The hard-negative stage uses a bidirectional modified InfoNCE loss that incorporates additional negatives when pairing queries with passages.The batch is represented as r = (q, p, n1, . . . , n15).Source: §5.2 · pp. 5–6

5.3 Memory Optimizations

  • Mixed precision, DeepSpeed, and activation checkpointing after each BERT layer reduce memory overhead to enable larger training batches.Larger batches increase the number of texts each example is compared against in the InfoNCE losses.Source: §5.3 · p. 6

6 Evaluation

Jina BERT maintains masked-language-model accuracy when extended to long sequences, while Jina Embeddings v2 performs strongly on MTEB and generally benefits from longer inputs, with task-specific exceptions.

6.1 Evaluation of Jina BERT

  • GLUE evaluation fine-tunes pretrained Jina BERT models on task-specific training data, selects hyperparameters on development data, and reports test-set results through the benchmark server.For RTE, STS, and MRPC, fine-tuning starts from the MNLI single-task model rather than the baseline pretrained models.Source: §6.1 · p. 6
  • 8192 tokens: Jina BERT’s MLM accuracy does not drop when evaluated at this sequence length, despite training only at 512 tokens.Other BERT and RoBERTa models with absolute positional embeddings trained at 512 tokens cannot be evaluated beyond that length in this comparison.Source: Figure 2 · p. 6, §6.1 · p. 7, §6.2.2 · p. 8

6.2 Evaluation of Jina Embeddings v2

  • Because MTEB omits very long texts, the evaluation adds extended-document retrieval and clustering tasks, including PatentClustering, WikiCitiesClustering, NarrativeQA, and LoCo.Patent documents average 6,376 tokens, WikiCities articles average 2,031 tokens, and NarrativeQA documents average 74,843 tokens; LoCo comprises five retrieval tasks selected partly for document length and whole-document understanding.Source: §6.2 · pp. 7–8

6.2.1 Results on MTEB

  • On MTEB, jina-small-v2 and jina-base-v2 show standout performance in Classification and Retrieval relative to other models, with jina-base-v2 leading in Classification and both models excelling in Retrieval.The text attributes some of the retrieval advantage to second-stage fine-tuning, which all-MiniLM-L6-v2 and all-mpnet-base-v2 omit.Source: §6.2.1 · p. 8, Table 3 · p. 9

6.2.2 Impact of Maximum Sequence Length

  • Longer inputs generally improve retrieval and clustering results, especially on NarrativeQA, where models restricted to the beginning of each document often underperform.Only jina-base-v2 and OpenAI’s text-embedding-ada-002 support 8K inputs; other models’ results at 8191 tokens are truncated to their intrinsic maximum, typically 512.Source: §6.2.2 · p. 8, Figure 3 · p. 9
  • Longer inputs improve BigPatent clustering but slightly reduce performance on WikiCities in most cases, where article openings often contain the country information needed for grouping.Later article content may be less relevant and can alter attributes that influence city-embedding clustering.Source: §6.2.2 · pp. 8–9

7 Conclusion

Jina Embeddings v2 is an open-source embedding suite for short and long texts, supporting sequences up to 8192 tokens. It uses a modified BERT architecture and performs competitively on MTEB.

  • The modified BERT architecture replaces positional embeddings with bidirectional ALiBi slopes to capture positional information.Source: §7 · p. 9
  • The open-source model suite supports up to 8192 tokens, a 16x increase over leading open-source embedding models, and performs competitively on MTEB.Source: §7 · pp. 9–10
  • Trained on the C4 Web document corpus and then fine-tuned, the models encode short and long texts as meaningful vector representations.Source: §7 · p. 9
  • Using extended sequence lengths gives the models an advantage over models without that capability.Source: §7 · p. 10

A Appendix: MTEB and LoCo Becnharmk

The appendix provides detailed performance tables for six MTEB task categories and a separate performance table for the LoCo dataset.

Related papers

Loading 2310.19923v4…