PaperlayerView exact revision on arXiv

Source-linked AI summary

SGPT: GPT Sentence Embeddings for Semantic Search

Niklas Muennighoff

arXiv:2202.08904v5cs.CLcs.AIcs.IR

TL;DR

Large decoder models had not been usable for semantic search and sentence embeddings, leaving these tasks to smaller encoders. SGPT adapts decoders through prompting or fine-tuning, and its scaled models achieve state-of-the-art search and embedding results.

  • Problem

    Decoder-only models had not been effectively used for semantic search and sentence embeddings, whose methods remained dominated by comparatively small encoders.

  • Method

    SGPT applies decoder-only transformers to semantic search using Bi-Encoder embeddings with pooling and bias-only fine-tuning, or Cross-Encoder scoring from GPT log probabilities.

  • Results

    At scale, SGPT-BE sets state-of-the-art results on BEIR and USEB, while SGPT-CE achieves unsupervised state-of-the-art performance on BEIR.

  • Takeaways & Limitations

    SGPT-BE supports semantic search and other embedding tasks, while SGPT-CE provides unsupervised search with storage and latency adjustable through prompting and reranking.

  • Takeaways & Limitations

    SGPT-BE’s sequence-length limit may constrain performance on BioASQ, where longer documents could require additional compute.

Abstract

from arXiv · show

Decoder transformers have continued increasing in scale reaching hundreds of billions of parameters. Due to their scale the same decoder sets state-of-the-art results on various language tasks via prompting or fine-tuning. Yet, these large foundation models remain unusable for the related fields of semantic search and sentence embeddings. This prevents possibly new state-of-the-art results and forces organizations to train and maintain separate models. To this end, we propose SGPT to use decoders for sentence embeddings and semantic search via prompting or fine-tuning. At 5.8 billion parameters SGPT improves on the previously best sentence embeddings by a margin of 7% and outperforms a concurrent method with 175 billion parameters as measured on the BEIR search benchmark. Code, models and result files are freely available at https://github.com/Muennighoff/sgpt.

1 Introduction

SGPT applies decoder-only transformers to semantic search and sentence embeddings, addressing their prior exclusion from these fields. It introduces Cross-Encoder and Bi-Encoder approaches, with large models achieving strong benchmark results.

  • Semantic search finds top-k corpus answers while interpreting queries and documents beyond keywords, but has been limited mainly to BERT-like encoder-only transformers.
  • Decoder-only transformers have scaled to hundreds of billions of parameters, yet meaningful embeddings for semantic search remain unclear and dominated by smaller encoders.
  • SGPT applies decoder-only transformers to semantic search and extracts sentence embeddings through Cross-Encoder and Bi-Encoder settings.
  • Bi-Encoder: SGPT-BE combines position-weighted mean pooling with bias-only fine-tuning, changing less than 0.1% of pretrained parameters while remaining competitive with full fine-tuning.
  • Bi-Encoder: At 5.8B parameters, SGPT-BE produces the best natural language embeddings available by a margin of 7% for semantic search.
  • Cross-Encoder: At 6.1B parameters, SGPT-CE uses GPT log probabilities without fine-tuning and achieves the best unsupervised performance on BEIR by a margin of 8%.

2 Related Work

Related work distinguishes Cross-Encoders from faster, cacheable Bi-Encoders and separates symmetric from asymmetric search according to whether queries and documents are interchangeable.

  • Cross-Encoders vs Bi-Encoders: Cross-Encoders jointly encode each query-document pair, requiring k forward passes for a corpus of k documents and preventing document-vector caching.
  • Cross-Encoders vs Bi-Encoders: Bi-Encoders encode queries and documents separately, allowing document vectors to be cached and queried through one query pass and a similarity function.
  • Cross-Encoders vs Bi-Encoders: Cross-Encoders tend to outperform Bi-Encoders but are slower, motivating multi-stage systems that retrieve broadly before Cross-Encoder reranking.
  • Symmetric vs Asymmetric Search: Asymmetric search treats queries and documents as non-interchangeable, commonly involving short questions and longer documents; BEIR contains 19 such datasets.
  • Symmetric vs Asymmetric Search: Symmetric search treats queries and documents as interchangeable, as in duplicate-question retrieval; evaluation uses USEB, Quora from BEIR, and STS-B.

3 SGPT Cross-Encoder

SGPT-CE applies pre-trained decoder-only GPT models to semantic-search re-ranking through log-probability extraction and prompt selection. Its performance depends on model scale and BM25’s candidate set, while prompt adaptation improves symmetric-search results.

  • Method: SGPT-CE extracts document-ranking scores from GPT log probabilities and re-ranks the top k documents returned by BM25.Documents are truncated from the left to fit the model’s maximum sequence length while maintaining a constant query length.
  • Method: The experiments use publicly available decoder transformers ranging from 125M to 6.1B parameters and select the best prompt from 12 MSMARCO prompts.The selected prompt is denoted PG.
  • Results: SGPT-CE-6.1B has almost 15x more parameters than BM25+CE, increasing latency despite winning on more datasets while achieving a worse average score.BM25+CE is a supervised fine-tuned BERT-based Cross-Encoder, whereas SGPT-CE uses only the pre-trained GPT model.
  • Results: In Re-rank Top 10, SGPT-CE-6.1B reaches around 80% of the maximum possible performance because BM25 bottlenecks the candidate set.Moving to Re-rank Top 100 produces large gains on datasets including HotpotQA and TREC-COVID.
  • Symmetric Search: Prompt customization for symmetric Quora search improves SGPT-CE-6.1B by 6% and outperforms all Quora results in Table 1.The experiment re-ranks the top 100 BM25 documents using the sum of query-token log probabilities.

4 SGPT Bi-Encoder

SGPT-BE extracts decoder sentence embeddings with position-weighted pooling and contrastive fine-tuning, including bias-only BitFit updates. At scale, it achieves state-of-the-art results across USEB and BEIR, while trading greater embedding storage and compute for accuracy.

  • Method: Position-weighted mean pooling gives later decoder tokens higher weight because causal attention prevents earlier tokens from seeing future tokens.The resulting vector represents the sequence using hidden states weighted by position.
  • Method: BitFit fine-tunes only bias parameters, changing less than 0.1% of pretrained parameters while remaining competitive with full fine-tuning.Freezing the remaining parameters reduces GPU memory and time required per training step.
  • Results: After identical fine-tuning, 125M-parameter SGPT closely trails 110M-parameter SBERT at layer 12, while weighted mean pooling outperforms mean and last-token pooling.In the unsupervised setting, decoder transformers strongly underperform encoders; increasing SGPT size ten-fold raises final-layer performance beyond SBERT.
  • Results: BitFit models remain within +2 to -2% of fully fine-tuned models, although BitFit degrades decoder performance more than encoder performance.The authors attribute this difference partly to the absence of a query bias vector in SGPT models.
  • Asymmetric search: For asymmetric search, SGPT uses MS-MARCO training, 300-token sequences, and special brackets distinguishing queries from documents.The model embeds queries in square brackets and documents in curly brackets.
  • Results: SGPT-BE-5.8B achieves the best average nDCG@10 on both the selected BEIR subset and the full BEIR benchmark.It improves on GTR-XXL by 7% but produces 4096-dimensional embeddings requiring about 5x more storage than GTR-XXL.

5 Conclusion and Future Work

SGPT adapts GPT models as Cross- or Bi-Encoders for semantic search. SGPT-BE targets sentence embeddings, while SGPT-CE uses log probabilities for unsupervised search.

  • SGPT modifies GPT models for both Cross-Encoder and Bi-Encoder semantic-search settings.The approach builds on SBERT.
  • SGPT-BE combines position-weighted mean pooling with fine-tuning only bias tensors.At scale, it produces new state-of-the-art sentence embeddings and can support semantic search or other embedding tasks.
  • SGPT-CE extracts log probabilities from pre-trained GPT models to produce unsupervised state-of-the-art search results.Its prompt and maximum re-rank parameter can be adjusted for performance and latency requirements.
  • Future work includes fine-tuning GPT Cross-Encoders on MSMARCO, combining SGPT with generative search, and studying missing biases in large GPT models.The authors suspect MSMARCO fine-tuning could outperform the presented non-fine-tuned SGPT-CE and SGPT-BE with sufficient re-ranking.

A Additional results

The appendix reports additional SGPT Cross-Encoder and Bi-Encoder results on BEIR and other embedding benchmarks. It also documents OpenAI embedding results and experimental settings for smaller BEIR datasets.

  • Additional SGPT Cross-Encoder and Bi-Encoder results on BEIR are reported as nDCG@10, excluding MS MARCO from average scores.
  • Table 7 reports Max Re-rank=10 OpenAI search-endpoint scores of 0.545 and 0.539 on TREC-COVID, and 0.670 and 0.658 on SciFact.The values correspond respectively to Ada and Davinci for each dataset.
  • Additional USEB, Quora, and STS-B results use average precision, nDCG@10, and Spearman correlation, respectively.The table distinguishes out-of-domain from in-domain results and notes that some fragments may be in-domain because of large pre-training data.
  • OpenAI Embeddings Endpoint results from December 2021 use nDCG@10 and closely match results published in the cited prior work.The table marks in-domain performance with ‡.
  • The smallest-BEIR-dataset experiments report that SGPT-125M-weightedmean-msmarco-specb performs 3% worse than SBERT-base-msmarco on average.The table also defines settings including special brackets, BitFit plus word-token-embedding training, batch size, and separate query-document transformers.

B.1 Prompts

The prompt experiments order prompts by increasing complexity and evaluate them with nDCG@10 on MSMARCO using SGPT-CE-2.7B.

  • Prompts are searched and listed in order of increasing complexity for SGPT-CE-2.7B.
  • Prompt performance is measured by nDCG@10 on MSMARCO.
  • The sum of log probabilities from {query} provides the re-rank score, while overflowing tokens are truncated from the left of {doc}.

B.2 Licenses

The appendix documents licensing information for the BEIR, USEB, and OpenAI resources referenced in the work.

  • BEIR datasets are licensed under various licenses listed in Appendix E of the BEIR paper.
  • USEB datasets are licensed under Apache 2.0.
  • OpenAI models are provided to customers through a non-exclusive, non-sublicensable, non-transferable, non-assignable, revocable OpenAI API license.

B.3 Computational Cost

The SGPT experiments used substantial computational resources, with costs varying by evaluation suite, model size, and training configuration.

  • Around 5,000 USD in tokens were used to evaluate the OpenAI Search and Embeddings endpoints.
  • SGPT experiments ran on one node equipped with 8 NVIDIA A100 Tensor Core GPUs, each with 40GB memory.
  • The complete BEIR evaluation for the 5.8B SGPT-CE model took around two weeks.
  • For the 5.8B SGPT-BE model, symmetric search training took 21 hours, while asymmetric training took 60 hours.
  • Oracle provided the cluster used for the experiments.
Loading 2202.08904v5…