Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Contrastive Learning and Mixture of Experts Enables Precise Vector Embeddings

Logan Hallee, Rohan Kapur, Arjun Patel, Jason P. Gleghorn, Bohdan Khomtchouk

TL;DR

This work tackles the challenge of producing precise vector embeddings for highly discriminative scientific text, with a focus on biomedical domains. It combines co-citation-based data curation and a Mixture of Experts–augmented BERT framework trained via a contrastive objective to learn domain-specific representations, using domain tokens and targeted routing. Notably, MoE models with $N$ domain-specific experts match the performance of $N$ independently trained models, and extending just a single transformer block to MoE yields a large portion of the full multi-domain benefit, indicating the feasibility of universal, scalable vector embeddings. The approach shows promise for improved retrieval and vector databases, enabling more accurate cross-domain scientific literature search and curation.

Abstract

The advancement of transformer neural networks has significantly elevated the capabilities of sentence similarity models, but they still struggle with highly discriminative tasks and may produce sub-optimal representations of important documents like scientific literature. With the increased reliance on retrieval augmentation and search, representing diverse documents as concise and descriptive vectors is crucial. This paper improves upon the vectors embeddings of scientific text by assembling niche datasets using co-citations as a similarity metric, focusing on biomedical domains. We apply a novel Mixture of Experts (MoE) extension pipeline to pretrained BERT models, where every multi-layer perceptron section is enlarged and copied into multiple distinct experts. Our MoE variants perform well over $N$ scientific domains with $N$ dedicated experts, whereas standard BERT models excel in only one domain at a time. Notably, extending just a single transformer block to MoE captures 85% of the benefit seen from full MoE extension at every layer. This holds promise for versatile and efficient One-Size-Fits-All transformer networks for numerically representing diverse inputs. Our methodology marks advancements in representation learning and holds promise for enhancing vector database search and compilation.

Contrastive Learning and Mixture of Experts Enables Precise Vector Embeddings

TL;DR

This work tackles the challenge of producing precise vector embeddings for highly discriminative scientific text, with a focus on biomedical domains. It combines co-citation-based data curation and a Mixture of Experts–augmented BERT framework trained via a contrastive objective to learn domain-specific representations, using domain tokens and targeted routing. Notably, MoE models with domain-specific experts match the performance of independently trained models, and extending just a single transformer block to MoE yields a large portion of the full multi-domain benefit, indicating the feasibility of universal, scalable vector embeddings. The approach shows promise for improved retrieval and vector databases, enabling more accurate cross-domain scientific literature search and curation.

Abstract

The advancement of transformer neural networks has significantly elevated the capabilities of sentence similarity models, but they still struggle with highly discriminative tasks and may produce sub-optimal representations of important documents like scientific literature. With the increased reliance on retrieval augmentation and search, representing diverse documents as concise and descriptive vectors is crucial. This paper improves upon the vectors embeddings of scientific text by assembling niche datasets using co-citations as a similarity metric, focusing on biomedical domains. We apply a novel Mixture of Experts (MoE) extension pipeline to pretrained BERT models, where every multi-layer perceptron section is enlarged and copied into multiple distinct experts. Our MoE variants perform well over scientific domains with dedicated experts, whereas standard BERT models excel in only one domain at a time. Notably, extending just a single transformer block to MoE captures 85% of the benefit seen from full MoE extension at every layer. This holds promise for versatile and efficient One-Size-Fits-All transformer networks for numerically representing diverse inputs. Our methodology marks advancements in representation learning and holds promise for enhancing vector database search and compilation.
Paper Structure (10 sections, 2 equations, 2 figures, 3 tables)

This paper contains 10 sections, 2 equations, 2 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Methods
  3. Data Compilation
  4. Transformer Neural Networks
  5. Mixture of Experts
  6. Models of choice
  7. Training Strategy
  8. Evaluation Strategy
  9. Results
  10. Discussion

Figures (2)

  • Figure 1: Method for determination of abstract pair similarity for model evaluation.
  • Figure 2: A typical ChatGPT response to a set of similar papers, qualitatively classifying all similar papers as dissimilar.