# Contrastive Learning and Mixture of Experts Enables Precise Vector
Embeddings

Rohan Kapur  Logan Hallee  Arjun Patel  The College of the University of
Chicago Bohdan Khomtchouk

###### Abstract

The advancement of transformer neural networks has significantly elevated the
capabilities of sentence similarity models, particularly in creating effective
vector representations of natural language inputs. However, these models face
notable challenges in domain-specific contexts, especially in highly
specialized scientific sub-fields. Traditional methods often struggle in this
regime, either overgeneralizing similarities within a niche or being overly
sensitive to minor differences, resulting in inaccurate text classification
and subpar vector representation. In an era where retrieval augmentation and
search are increasingly crucial, precise and concise numerical representations
are essential. In this paper, we target this issue by assembling niche
datasets using co-citations as a similarity metric, focusing on biomedical
domains. We employ two key strategies for fine-tuning state-of-the-art models:
1. Domain-specific Fine-Tuning, which tailors pretrained models to a single
domain, and 2. Universal Applicability with Mixture of Experts (MoE), adapting
pretrained models with enforced routing for multiple domains simultaneously.
Our training approach emphasizes the use of abstracts for faster training,
incorporating Multiple Negative Rankings loss for efficient contrastive
learning. Notably, our MoE variants, equipped with $N$ experts, achieve the
efficacy of $N$ individual models, heralding a new era of versatile, One-Size-
Fits-All transformer networks for various tasks. This methodology marks
significant advancements in scientific text classification metrics and holds
promise for enhancing vector database search and compilation.

##  1 Introduction

The remarkable success of transformer-based large language models (LLMs) since
2017 [1] has significantly increased our confidence in their abilities and
outputs. Nowadays, LLMs are treated as a de facto knowledge base for many and
have been adopted on a mass scale since the release of services like ChatGPT
and open-source counterparts like Llama and Mistral [2, 3]. However, despite
their widespread use, challenges persist, particularly in terms of the
accuracy and reliability of these models. Common issues like LLM
hallucinations [4, 5] highlight the ongoing need for improvement, especially
in areas like text classification, where sentence similarity models often
struggle to accurately discern nuances in specialized scientific domains. The
ability to generate reliable vector embeddings and perform precise
classification is crucial, especially for technologies that rely on
information retrieval and web search.

In this paper, we present our efforts to enhance the precision and reliability
of LLMs in identifying similar or niche intradisciplinary texts, a key
component in improving text classification at the document level. Our approach
encompasses two primary strategies:

  1. 1.

Domain-Specific Fine-Tuning: We apply standard fine-tuning methods to
pretrained BERT (Bidirectional Encoder Representation Transformers) models,
tailoring them to learn and understand specific scientific domains.

  2. 2.

Universal Applicability through Mixture of Experts (MoE): We introduce a
scalable method of seeding MoE models for fine-tuning pretrained BERT models
across multiple domains, aiming for a versatile, “One-Size-Fits-All” model.

Central to our framework is the use of co-citation data, which allows for the
rapid compilation of related papers within specific scientific fields. This is
complemented by contrastive learning techniques, which further refine the
pretrained models to better understand and categorize domain-specific content.
Our research seeks to address the gaps in current LLM applications,
particularly in enhancing the accuracy and utility of sentence similarity
models in scientific text classification.

##  2 Related Work

Sentence similarity models [6, 7, 8] have been developed by researchers for
years due to their broad range of applications, which include search retrieval
and vector databases, textual classification, ranking systems, among others.
Despite their popularity, current state-of-the-art (SOTA) LLMs and sentence
similarity models, including those built on scientific datasets [9, 10], are
not capable of consistently identifying sub-field specific similar and
dissimilar papers due to their generality or narrow definition of similar.

The bulk of citation network research has been done in the metascience
community, where citation networks are used to track and identify highly
influential papers, as well as identify future citations [11, 12, 13, 14, 15,
16, 17, 18, 19].

In this paper, we take advantage of prior work in both fields. Our models are
built by fine tuning available SOTA general-purpose models. To fine tune these
models, we used co-citations as our similarity metric, whereby co-citations
were used to help our model identify similar and dissimilar work.

##  3 Methods

###  3.1 Data Compilation

To generate sufficiently large and accurate training datasets, we took
advantage of co-citation networks. Co-citations represent instances where two
papers are cited together in another paper and enable the production of large
training datasets from small amounts of data due to their nonlinear nature.
For a dataset of 10,000 individual papers, for example, over 125,000 co-
citation pairs can be produced. Co-citations also imply a high degree of
similarity. Representing occurrences where two papers are cited together, any
two co-cited papers are not only individually relevant due to their individual
citations, but also are relatively similar given their appearance in the same
paper. While this degree of similarity may not be perfect, co-citations have
generally been shown to imply a high degree of similarity between papers [18].
Thus, by labeling any given co-citation pair as similar, we produced
sufficiently large and accurate datasets for model training.

We chose to use the cardiovascular disease (CVD) and chronic obstructive
pulmonary disease (COPD) sub-fields of the biomedical sciences as example case
studies for our approach, chosen based on our familiarity with biomedicine.
The CVD and COPD sub-fields represent two domains that contrast significantly
in cocitation network size and thus, allowed us to compare the scalability of
our approach as it relates to data availability. Since 2010, on PubMed
Central, around 290,000 articles relating to CVD have been published, while
just 14,000 articles relating to COPD have been published.

We obtained all papers, including their respective abstracts and citation
information, used in this study from PubMed Central, a database of open-access
papers in the biomedical sciences. We queried papers by utilizing Medical
Subject Heading (MeSH) terms for both cardiovascular diseases (CVD) and
chronic obstructive pulmonary disease (COPD), querying specifically for open-
access papers having at least one citation and an abstract. For CVD, we
queried specifically for papers published between 2017 - 2022, while querying
for papers published between 2010 - 2022 for COPD. Using these queries, we
obtained all relevant citation and abstract information for our study. Our
decision to use a larger time range for COPD was based on the size of the sub-
field and the inability to create a large enough dataset from querying from
just 2017 onwards, as in the case of CVD. In total, we queried 99,490 papers
for our CVD dataset and 10,051 papers for our COPD dataset.

As described earlier, we generated our training dataset for both CVD and COPD
datasets by compiling all pairs of papers that had been co-cited together in
our dataset. Our test dataset for both sub-fields was constructed by taking
all similar abstract pairs in the training dataset where at least one paper
had been published in 2022, the most recent year in our dataset. Our test
dataset was constructed such that at least one paper was from 2022, as we are
interested in building models capable of distinguishing between new relevant
and new irrelevant research. In this case, research from 2022 represented our
source of new research.

Our validation dataset was constructed from the remaining training dataset.
The remaining training dataset was split randomly in a 99:1 ratio without
duplicates, with the larger of the new datasets being used as our final
training dataset and the smaller of the two being used as the initial
validation dataset. On top of this initial validation dataset made up of
similar abstract pairs, an equal amount of dissimilar abstract pairs were
added. We generated pairs of dissimilar papers by compiling pairs of papers,
which had been cited at least 15 times each, but had never been co-cited
together. While it is impossible to guarantee any two non-co-cited papers
won’t be cited together in the future, this approach, by requiring papers to
be cited individually 15+ times, the maximum citation requirement that allowed
us to have enough dissimilar pairs to match the number of similar pairs in
each validation dataset, minimized this possibility. Note that while we
produced pairs of dissimilar papers for model evaluation, the production of
dissimilar paper pairs is not necessary for model generation.

After constructing the validation and test datasets, we accounted for co-
citation frequency in our training dataset by duplicating co-citation pairs
that had been co-cited multiple times in our training dataset. If two papers
had been co-cited together 5 times, for example, this duplication would result
in this pair of papers occurring 5 times in our dataset. This duplication
allowed us to weigh pairs that had been co-cited more frequently more heavily
than pairs that had been co-cited less frequently.

###  3.2 Transformer Neural Networks

The transformer architecture is adept at sequential processing and is SOTA for
NLP tasks [20, 21, 22, 23, 24, 25, 26]. A key component of its performance is
the token embedding process, which turns natural language inputs into integers
through tokenization. These integers serve as the indices of a lookup table
which maps to representative vectors. These vector representations or token
embeddings capture the syntactic and semantic nuances of the language through
extensive semi-supervised denoising [27, 28]. Other additions like position or
token type embeddings are also common to improve the proficiency of these
networks [29].

Beyond the base embeddings, a transformer model is made up of $L$ transformer
blocks, where the $i^{th}$ block maps a sequence of $T$-length vectors
$(h_{1}^{(i-1)},...,h_{T}^{(i-1)})$ to another $T$-length vector sequence
$(h_{1}^{(i)},...,h_{T}^{(i)})=H^{(i)}$, where all vectors are dimension $d$
[30]. This mapping is done using two main steps: the self-attention mechanism
which mixes information across the vectors, and a multi-layer perceptron (MLP)
that mixes information across $d$ [31].

GPT (Generative Pretrained Transformers) models or causal models, like
OpenAI’s GPT series (GPT-2, GPT-3, etc.), focus on generative tasks and use a
variant called transformer decoders [32, 33, 34]. They use unidirectional
attention when processing text. This means they can predict the next word in a
sentence but cannot modify their understanding based on words that come later
in the text. BERT models or transformer encoders, utilize bidirectional
attention, capture more context and word relationships, and are better suited
for tasks like text classification [29].

####  3.2.1 Mixture of Experts

Mixture of Experts (MoE) models add a linear layer or router network to each
transformer block, which outputs logits from $H^{(i)}$. These logits route
$H^{(i)}$ to multiple unique copies of the MLP section in the transformer
block called experts [35]. For causal GPT-like models, this is typically done
on a per-token basis, allowing for experts to specify in language classes like
punctuation, nouns, numbers, etc [36]. For BERT-like models this allows for
sentence-wise routing of the entire $H^{(i)}$ [37]. While there are many ways
to route MoE networks, two main approaches involve either calling one expert
per block or the top $k$ experts per block. Our variants call one expert per
block, otherwise known as the switch transformer [38].

Controlling the routing of $H^{(i)}$, allows for a one-size-fits-all approach
to text classification where one expert per transformer layer is an expert in
a specific domain. To control this routing we add special tokens for each
domain, like [CVD] and [COPD], and replace the standard [CLS] token with these
upon tokenization. An additional cross-entropy loss is added which compares
the router logits to the correct domain identity. For faster finetuning, we
utilized pretrained models for this novel MoE approach, seeding the experts
with the learned MLP layers and seeding the new special domain tokens with the
learned [CLS] embedding; the only randomly initialized weights are the
routers.

####  3.2.2 Models of choice

We minimized model training time by taking advantage of pretrained BERT
models. While these models were initially poor at consistently identifying
similar papers within each sub-field, they provided us with adequate
infrastructure to build our model, preventing us from having to train from
scratch. We chose to fine-tune two models. 1.) all-MiniLM-L6-v2 model (Mini)
[39, 40, 41, 42, 43, 44], a popular sentence transformer model trained on 1+
billion sentence pairs and is highly effective for sentiment analysis and text
classification. Mini is a standard BERT-like model with a hidden size of 384,
an intermediate size of 1536, 12 attention heads, and 6 hidden layers for a
total of 23 million parameters. 2.) SciBERT is a general-purpose BERT model
finetuned on entire papers from the semantic scholar: over 1.14 million
diverse papers [45]. SciBERT is SOTA for many domain-specific relation
extraction and sentence classification tasks with a hidden size of 786, an
intermediate size of 3072, 12 attention heads, and 12 hidden layers for a
total of 110 million parameters. Compared to many newer BERT architectures
with up to (or more than) 1 billion parameters, these modest model sizes allow
for effective fine-tuning given minimal training data and reduced
computational cost during training iterations.

MoE versions of our fine-tuned models have a larger parameter count, 30
million and 167 million respectively, due to two experts (two domains) per
layer. However, the effective parameter account is the same as the original
due to only one expert being called at a time, resulting in fast training and
inference.

Our fine-tuned models were compared to a basic term frequency-inverse document
frequency (TF-IDF) model alongside the following popular base models without
any fine-tuning: Mini, BERT, Mpnet [46], Declustr [47], SciBERT, BiomedBERT
[48], and ClincalBERT [49]. This large model variety and pretraining strategy
allows for a more general comparison of how effective our fine-tuning
framework is.

The TF-IDF model, representing the most basic and straightforward sentence
similarity model, acts as a baseline for expected performance. Finally, we
also prompted GPT3.5 with sets of abstracts to have them assess qualitatively
if the papers were similar or not.

###  3.3 Training Strategy

To minimize training time, we also chose to use abstracts rather than entire
papers as the text input to the model. Abstracts represent a human-generated
summarized version of a paper and as a a result, include almost all relevant
textual information contained in a paper. They are significantly smaller than
actual papers and are devoid of the pre-processing complications that come
with using full papers such as figure and data pre-processing. Thus, by using
abstracts, we were able to significantly condense our dataset, allowing for
faster training, without significant information loss. We fine-tune Mini and
SciBERT on the individual datasets for CVD and COPD as well as combined in the
MoE version. These are referred to as single expert (SE) and MoE versions.

To train our models, we feed similar abstracts from the CVD or COPD domain to
the model independently, extracting the vector embedding from the pooler
output for each abstract. The pooler output is generated by a small neural
network from the [CLS] token (our special domain tokens in our case) embedding
from the last hidden state $H^{(L)}$, a standard practice when using the
Huggingface transformers package. These vector embeddings are compared with a
variant of the Multiple Negative Rankings (MNR) loss used to train cdsBERT
[28, 50]. MNR Loss is a loss function that has seen significant success with
sentence embedding problems [51] and was highly successful in our local
experiments. Our variant utilizes dot products as a similarity heuristic and
scales the similarity by a learned temperature parameter. Furthermore, MNR
loss only requires positive/similar text pairs, generating negative/dissimilar
text pairs from the other positive pairs in a given mini-batch. As a result,
MNR loss removed the need to generate dissimilar text pairs for our training
dataset under the assumption that the random chance of finding a similar paper
randomly with a batch size of 20 is sufficiently small, significantly reducing
the potential resources needed for the training dataset generation. During
training, we randomly switch the order of the two input abstract pairs to
prevent any bias in the way they are fed to the loss function.

We performed hyperparameter optimization on the CVD dataset, using random
search to approximate the best batch size alongside warmup steps and total
training length that optimized the model’s F1 score on the validation dataset.
Due to resource constraints, we decided to use a smaller dataset than the
actual training dataset, with the hyperparameter optimization dataset being a
random 10% sample of the training dataset. Our tested range for each
hyperparameter was also limited by resource constraints. After trying batch
sizes ranging from 5-20 we found that 20 supported the best model performance,
offering a large enough batch size to require nuanced understanding but not
too large to hinder model training; as the contrastive loss chosen is
significantly harder to minimize as the batch size increases. For our final
training runs we utilized a learning rate of $1e^{-5}$, a one cycle learning
rate scheduler [52] with 500 warmup steps, and periodic validation every 5000
steps. Training was conducted for 10 epochs and halted early when a patience
of 3 was exceeded regarding the validation metrics; the best model was loaded
at the end. SciBERT training was conducted on a single A100 GPU, Mini training
was done on a single A10 GPU, and all training runs took less than 24 hours.

###  3.4 Evaluation Strategy

All models were evaluated on the validation and test sets separately for each
domain. We utilize cosine similarity between two vectors extracted from an
abstract pair to classify the abstracts as similar or nonsimlar given a cosine
similarity threshold. Cosine similarity is a common measure of vector
similarity that ranges from -1 to 1, where -1 is exactly opposite and 1 occurs
for a pair of the same vector. This binary threshold is determined by an F1
variant called $F1_{max}$ [53, 54]. $F1_{max}$ is simply the maximum F1 score
calculated for all possible thresholds for a reported metric. While typically
used for imbalanced multilabel classification, randomly choosing a binary
threshold for the reported F1 would not be a fair comparison of different
models. For example, one model may perform much better with a cosine
similarity threshold at 0.5 for abstract text similarity compared to 0.4, or
vice versa. We also report average distance and accuracy. The average distance
was calculated by taking the absolute value of the difference between the
similarity score between two abstract pairs and their label, 0 or 1. For
example, if a model gave two co-cited (similar) abstracts a cosine similarity
of 0.75, the distance would be $(|1-0.75|)$ = 0.25. Full details of the
conversion of any two abstract pairs into binary similar and dissimilar
metrics are summarized in Figure 1. Because the test dataset contains no
negative examples, we limit the $F1_{max}$ threshold search between 0.5 and 1
to prevent the trivial -1 threshold that always leads to $F1_{max}=1$, and
report the accuracy using the found threshold.

Figure 1: Determination of abstract pair similarity for model evaluation

##  4 Results

The performance of our models on the validation and test datasets is
summarized and compared to other leading sentence similarity models, as well
as a standard term frequency inverse document frequency model (TF-IDF) in
Tables 1 & 2\. We also evaluated the performance of GPT-3.5 on our generated
validation and test datasets. Upon evaluation, it was found that that GPT-3.5
represented the prototypical example of a model with a definition of similar
that was too narrow, as GPT-3.5 regarded close to all inputted examples as
dissimilar when asked whether the inputted examples were similar due to
(correctly) identified minor differences between each inputted text. Given
this, GPT-3.5 was excluded from the below tables, but an example of its output
when given two similar papers can be seen in Figure 3. We also excluded the
MoE version of Mini from the results as they did not score better than random
on either validation set with a low 0.67 $F1_{max}$.

Model | Cutoff | F1Max${}_{Valid}\uparrow$ | Dist.${}_{Valid}\downarrow$ | Acc.${}_{Valid}\uparrow$ | Dist.${}_{Test}\downarrow$ | Acc.${}_{Test}\uparrow$
---|---|---|---|---|---|---
Our models
Mini-SE | 0.51 | 0.94 | 0.23 | 0.93 | 0.19 | 0.97
SciBERT-SE | 0.98 | 0.97 | 0.48 | 0.97 | 0.01 | 1.00
SciBERT-MoE | 0.98 | 0.97 | 0.48 | 0.97 | 0.01 | 1.00
Base models
TF-IDF | 0.00 | 0.67 | 0.50 | 0.50 | 0.02 | 1.00
Mini | 0.46 | 0.90 | 0.33 | 0.90 | 0.32 | 0.95
BERT | 0.89 | 0.71 | 0.48 | 0.65 | 0.07 | 0.89
Mpnet | 0.49 | 0.94 | 0.28 | 0.94 | 0.26 | 0.96
Declutr | 0.65 | 0.84 | 0.42 | 0.83 | 0.25 | 0.89
SciBERT | 0.44 | 0.89 | 0.34 | 0.90 | 0.39 | 0.87
BiomedBERT | 0.99 | 0.72 | 0.50 | 0.73 | 0.01 | 0.76
ClinicalBERT | 0.92 | 0.71 | 0.49 | 0.68 | 0.06 | 0.81
Table 1: Evaluation metrics of our fine tuned models compared to base models for the CVD datasets. Single expert (SE) and Mixture of Expert (MoE) models are compared, showcasing near-identical performance despite MoE models exhibiting mastery of multiple datasets. $F1_{max}$ is not shown for the test set due to the trivial -1 threshold resulting in all models performing perfectly due to no negative data.n Thus, a high $F1_{max}$ is key for the validation data and a low distance is key for the test data. Model | Cutoff | F1Max${}_{Valid}\uparrow$ | Dist.${}_{Valid}\downarrow$ | Acc.${}_{Valid}\uparrow$ | Dist.${}_{Test}\downarrow$ | Acc.${}_{Test}\uparrow$
---|---|---|---|---|---|---
Our models
Mini-SE | 0.36 | 0.83 | 0.33 | 0.82 | 0.47 | 0.56
SciBERT-SE | 0.99 | 0.81 | 0.49 | 0.80 | 0.01 | 1.00
SciBERT-MoE | 0.98 | 0.80 | 0.49 | 0.79 | 0.01 | 1.00
Base models
TF-IDF | 0.00 | 0.67 | 0.50 | 0.50 | 0.02 | 1.00
Mini | 0.48 | 0.69 | 0.45 | 0.61 | 0.42 | 0.76
BERT | 0.87 | 0.69 | 0.49 | 0.58 | 0.09 | 0.85
Mpnet | 0.53 | 0.68 | 0.45 | 0.62 | 0.38 | 0.71
Declutr | 0.62 | 0.68 | 0.47 | 0.56 | 0.27 | 0.91
SciBERT | 0.46 | 0.67 | 0.47 | 0.55 | 0.42 | 0.82
BiomedBERT | 0.98 | 0.67 | 0.50 | 0.51 | 0.01 | 0.97
ClinicalBERT | 0.91 | 0.69 | 0.49 | 0.61 | 0.07 | 0.80
Table 2: Evaluation metrics of our fine tuned models compared to base models
for the COPD datasets. Single expert (SE) and Mixture of Expert (MoE) models
are compared, showcasing near-identical performance despite MoE models
exhibiting mastery of multiple datasets. $F1_{max}$ is not shown for the test
set due to the trivial -1 threshold resulting in all models performing
perfectly due to no negative data.n Thus, a high $F1_{max}$ is key for the
validation data and a low distance is key for the test data. Figure 2: Typical
ChatGPT Response to a Set of Similar Papers, qualitatively classifying all
similar papers as dissimilar.

The results presented in Tables 1 & 2 elucidate the effectiveness of our fine-
tuning strategy, as our models demonstrated a pronounced proficiency in
identifying similar or dissimilar papers within highly specific domains. On
the CVD dataset, our models achieved superior $F1_{max}$ and accuracy scores
than all base models. Particularly, the SciBert variants exhibited a near
perfect 0.97 $F1_{max}$. While the $F1_{max}$ and accuracy scores were
marginally lower on the COPD dataset, the relative performance of our models
on this dataset is even more clear: On the COPD validation dataset, our models
surpassed every base model by at least 10% in both $F1_{max}$ and accuracy
scores. This performance highlights the capability of our approach to yield
high-quality results, even with limited training data.

On the COPD test dataset, our SciBert models continued to outshine the base
models, albeit by a narrower margin. Viewed in isolation, this might suggest
only a modest enhancement by our models. However, in conjunction with the
validation dataset results, it becomes evident that our models are uniquely
proficient at distinguishing between similar and dissimilar papers—a
capability not mirrored by the base models, which exhibited a tendency to
erroneously classify nearly all papers as similar. This distinction is
especially notable considering that all base models with over 0.90+ accuracy
on the test dataset scored $\approx$ 0.67 $F1_{max}$ and $\approx$ 0.50
accuracy on the validation dataset, highlighting a pronounced inability to
differentiate between similar and dissimilar papers beyond random chance.
Consequently, the prowess of our models with high performance on the test and
validation sets is not merely a product of a bias toward classifying papers as
similar but is indicative of a genuine capacity to discern nuanced
distinctions.

In a broader context, the combined results from the COPD and CVD datasets
underscore our models’ adeptness at comprehending similarities between papers
specific to particular domains. The Sci-Bert models showcased remarkable
generalizable performance, surpassing that of Mini on average. However, it is
worth noting that Mini demonstrated superior performance on the smaller COPD
dataset, aligning with established fine-tuning trends that suggest smaller
models require less nonredundant data for optimal performance than their
larger counterparts [55]. Interestingly, the MoE models, adopting a one-size-
fits-all approach, did not exhibit inferior performance compared to the SE
models, which points to the robustness and versatility of the MoE approach
through enforced routing.

##  5 Discussion

Our research advances the use of Large Language Models (LLMs) by focusing on
improving their domain-specific understanding and document-wide comprehension.
As vector databases become increasingly prevalent, the summarization abilities
of LLMs gain importance. Our innovative framework, which incorporates co-
citation networks, contrastive learning, custom MNR and router losses, novel
special tokens, MoE seeding, and enforced routing, significantly enhances
sentence similarity metrics compared to base models. This allows for more
accurate differentiation between similar and dissimilar scientific papers
based on abstracts, overcoming the limitations of previous SOTA models that
were either too sensitive or too general.

More specifically, our SciBERT SE and MoE variants performed equally on the
CVD validation set with an $F1_{max}$ of 0.97, an average distance of 0.48,
and an accuracy of 0.97. The Mini SE variant performed similarly with a
$F1_{max}$ of 0.94, an average distance of 0.23, and an accuracy of 0.93. On
the CVD validation set, Mini SE performed best with a $F1_{max}$ of 0.83, an
average distance of 0.33, and an accuracy of 0.82. Following that is SciBERT
SE with $F1_{max}$ 0.81 and SciBERT MoE with $F1_{max}$ 0.80. Both SciBERT
variants reuslted in average distance of 0.49 and accuracies of 0.80 and 0.79
respectively.

The test sets require a slightly more nuanced look. We withhold $F1_{max}$
from the tables since there is a trivial cutoff of -1 to get 1.00 $F1_{max}$,
as there is no negative data. By limiting the threshold search and reporting
the accuracy we get a nontrivial representation of performance. We found that
SciBERT variants had more precise vector representations with 1.00 accuracy
all around. Mini SE variants had lower accuracies of 0.97 for CVD and 0.56 for
COPD. Interestingly, the Mini SE CVD variant performed better than the Mini SE
COPD variant on the COPD test data; with 0.85 accuracy. This suggests that the
cosine similarity threshold limit of 0.5 may artificially hinder the metrics
here. The average distance metric gives more context to test dataset
performance. SciBERT variants excelled at placing similar abstract vector
representations close in space with an average distance of 0.01 for CVD and
COPD test sets. On the other hand, the smaller Mini SE had a high distance
even compared to base models. Notably, the BiomedBERT base model also had an
average test set distance of 0.01 on both test sets, which is unsurprising
given the possibility of training data overlap.

Importantly, our MoE approach performs equally to our SE models on multiple
domains. On top of that, the MoE seeding is completely scalable, enabling $N$
experts and $N$ new special tokens given $N$ different datasets. Because our
MoE model only calls one expert per transformer block, the development of a
one-size-fits-all MoE model is approachable and computationally affordable;
offering a promising solution to building reliable vector embeddings and
improved text classification. While adding new special tokens for each domain
may enhance domain-specific understanding when used for the pooler output,
using them for enforced routing is certainly an optional approach. We chose
this so that the model could correctly route examples on their own.
Alternatively, one could just enforce the expert routing by passing a list of
routing indices desired for a specific input.

Despite the MoE engineering, our attempts with a Mini MoE variant were less
successful, suggesting a minimum size requirement in the base BERT model for
effective performance. This may be due to the need for sufficiently capable
shared attention layers that can generalize to support diverse experts and
domains.

In summary, our use of co-citation networks enables rapid and efficient
dataset compilation for training transformers in niche domains. SE finetuning
of base BERT models through contrastive learning with MNR-inspired loss
significantly improves sentence similarity. The MoE approach further expands
these capabilities, suggesting the feasibility of a universal model for text
classification and vector embeddings across various domains through MoE
seeding and enforced routing. We anticipate future research to extend these
methodologies, developing effective MoE BERT models with specialized knowledge
across multiple fields, vocabularies, or tasks.

##  6 Data availability

Our data is available for download and use with the Huggingface datasets
package at huggingface.co/datasets/lhallee/abstract_domain_cvd and
huggingface.co/datasets/lhallee/abstract_domain_cvd. The code and framework
used to train the models as well as data compilation can be found at
github.com/lhallee/Mixture-of-Experts-Sentence-Similarity.

##  7 Funding and acknowledgments

This work was partly supported by the University of Delaware Graduate College
through the Unidel Distinguished Graduate Scholar Award (LH), Indiana
University (BBK), and NIH R01DK132090 (BBK).

##  8 Conflict of interest

BBK is the co-founder of Dock Therapeutics, Inc. The remaining authors have
nothing to disclose.

##  9 Notes

## References

  * [1] A. Vaswani, N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. u. Kaiser, and I. Polosukhin, “Attention is all you need,” in Advances in Neural Information Processing Systems (I. Guyon, U. V. Luxburg, S. Bengio, H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett, eds.), vol. 30, Curran Associates, Inc., 2017.
  * [2] H. Touvron, L. Martin, K. Stone, P. Albert, A. Almahairi, Y. Babaei, N. Bashlykov, S. Batra, P. Bhargava, S. Bhosale, D. Bikel, L. Blecher, C. C. Ferrer, M. Chen, G. Cucurull, D. Esiobu, J. Fernandes, J. Fu, W. Fu, B. Fuller, C. Gao, V. Goswami, N. Goyal, A. Hartshorn, S. Hosseini, R. Hou, H. Inan, M. Kardas, V. Kerkez, M. Khabsa, I. Kloumann, A. Korenev, P. S. Koura, M.-A. Lachaux, T. Lavril, J. Lee, D. Liskovich, Y. Lu, Y. Mao, X. Martinet, T. Mihaylov, P. Mishra, I. Molybog, Y. Nie, A. Poulton, J. Reizenstein, R. Rungta, K. Saladi, A. Schelten, R. Silva, E. M. Smith, R. Subramanian, X. E. Tan, B. Tang, R. Taylor, A. Williams, J. X. Kuan, P. Xu, Z. Yan, I. Zarov, Y. Zhang, A. Fan, M. Kambadur, S. Narang, A. Rodriguez, R. Stojnic, S. Edunov, and T. Scialom, “Llama 2: Open foundation and fine-tuned chat models.”
  * [3] A. Q. Jiang, A. Sablayrolles, A. Mensch, C. Bamford, D. S. Chaplot, D. d. l. Casas, F. Bressand, G. Lengyel, G. Lample, L. Saulnier, L. R. Lavaud, M.-A. Lachaux, P. Stock, T. L. Scao, T. Lavril, T. Wang, T. Lacroix, and W. E. Sayed, “Mistral 7b.”
  * [4] R. Branco, A. Branco, J. António Rodrigues, and J. R. Silva, “Shortcutted commonsense: Data spuriousness in deep learning of commonsense reasoning,” in Proceedings of the 2021 Conference on Empirical Methods in Natural Language Processing (M.-F. Moens, X. Huang, L. Specia, and S. W.-t. Yih, eds.), (Online and Punta Cana, Dominican Republic), pp. 1504–1521, Association for Computational Linguistics, Nov. 2021.
  * [5] N. M. Guerreiro, E. Voita, and A. Martins, “Looking for a needle in a haystack: A comprehensive study of hallucinations in neural machine translation,” in Proceedings of the 17th Conference of the European Chapter of the Association for Computational Linguistics (A. Vlachos and I. Augenstein, eds.), (Dubrovnik, Croatia), pp. 1059–1075, Association for Computational Linguistics, May 2023.
  * [6] N. Reimers and I. Gurevych, “Sentence-BERT: Sentence embeddings using Siamese BERT-networks,” in Proceedings of the 2019 Conference on Empirical Methods in Natural Language Processing and the 9th International Joint Conference on Natural Language Processing (EMNLP-IJCNLP) (K. Inui, J. Jiang, V. Ng, and X. Wan, eds.), (Hong Kong, China), pp. 3982–3992, Association for Computational Linguistics, Nov. 2019.
  * [7] Z. Quan, Z.-J. Wang, Y. Le, B. Yao, K. Li, and J. Yin, “An efficient framework for sentence similarity modeling,” IEEE/ACM Transactions on Audio, Speech, and Language Processing, vol. 27, no. 4, pp. 853–865, 2019.
  * [8] H. Yao, H. Liu, and P. Zhang, “A novel sentence similarity model with word embedding based on convolutional neural network,” Concurrency and Computation: Practice and Experience, vol. 30, no. 23, p. e4415, 2018.  e4415 cpe.4415.
  * [9] J. Giorgi, O. Nitski, B. Wang, and G. Bader, “DeCLUTR: Deep contrastive learning for unsupervised textual representations,” in Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 11th International Joint Conference on Natural Language Processing (Volume 1: Long Papers) (C. Zong, F. Xia, W. Li, and R. Navigli, eds.), (Online), pp. 879–895, Association for Computational Linguistics, Aug. 2021.
  * [10] P. Deka, A. Jurek-Loughrey, and Deepak, “Unsupervised keyword combination query generation from online health related content for evidence-based fact checking,” The 23rd International Conference on Information Integration and Web Intelligence, 2021.
  * [11] A. Abrishami and S. Aliakbary, “Predicting citation counts based on deep neural network learning techniques,” Journal of Informetrics, vol. 13, no. 2, pp. 485–499, 2019.
  * [12] T. Zhang, H. Chi, and Z. Ouyang, “Detecting research focus and research fronts in the medical big data field using co-word and co-citation analysis,” in 2018 IEEE 20th International Conference on High Performance Computing and Communications; IEEE 16th International Conference on Smart City; IEEE 4th International Conference on Data Science and Systems (HPCC/SmartCity/DSS), pp. 313–320, 2018.
  * [13] D. G. Brizan, K. Gallagher, A. Jahangir, and T. Brown, “Predicting citation patterns: Defining and determining influence,” Scientometrics, vol. 108, no. 1, p. 183–200, 2016.
  * [14] R. Malkawi, M. Daradkeh, A. El-Hassan, and P. Petrov, “A semantic similarity-based identification method for implicit citation functions and sentiments information,” Information, vol. 13, no. 11, 2022.
  * [15] O. Rodriguez-Prieto, L. Araujo, and J. Martinez-Romo, “Discovering related scientific literature beyond semantic similarity: A new co-citation approach,” Scientometrics, vol. 120, no. 1, p. 105–127, 2019.
  * [16] M. Eto, “Extended co-citation search: Graph-based document retrieval on a co-citation network containing citation context information,” Information Processing & Management, vol. 56, no. 6, p. 102046, 2019.
  * [17] L. Galke, F. Mai, I. Vagliano, and A. Scherp, “Multi-modal adversarial autoencoders for recommendations of citations and subject labels,” in Proceedings of the 26th Conference on User Modeling, Adaptation and Personalization, UMAP ’18, (New York, NY, USA), p. 197–205, Association for Computing Machinery, 2018.
  * [18] G. Colavizza, K. W. Boyack, N. J. van Eck, and L. Waltman, “The closer the better: Similarity of publication pairs at different cocitation levels,” Journal of the Association for Information Science and Technology, vol. 69, no. 4, pp. 600–609, 2018.
  * [19] B. Gipp and J. Beel, “Citation proximity analysis (cpa) : A new approach for identifying related work based on co-citation analysis,” 2009.
  * [20] E. Beeching, C. Fourrier, N. Habib, S. Han, N. Lambert, N. Rajani, O. Sanseviero, L. Tunstall, and T. Wolf, “Open llm leaderboard.” https://huggingface.co/spaces/HuggingFaceH4/open_llm_leaderboard, 2023.
  * [21] P. Clark, I. Cowhey, O. Etzioni, T. Khot, A. Sabharwal, C. Schoenick, and O. Tafjord, “Think you have solved question answering? try arc, the ai2 reasoning challenge,” 2018.
  * [22] R. Zellers, A. Holtzman, Y. Bisk, A. Farhadi, and Y. Choi, “Hellaswag: Can a machine really finish your sentence?,” 2019.
  * [23] D. Hendrycks, C. Burns, S. Basart, A. Zou, M. Mazeika, D. Song, and J. Steinhardt, “Measuring massive multitask language understanding,” 2021.
  * [24] S. Lin, J. Hilton, and O. Evans, “Truthfulqa: Measuring how models mimic human falsehoods,” 2022.
  * [25] K. Sakaguchi, R. L. Bras, C. Bhagavatula, and Y. Choi, “WINOGRANDE: an adversarial winograd schema challenge at scale,” 2019.
  * [26] K. Cobbe, V. Kosaraju, M. Bavarian, M. Chen, H. Jun, L. Kaiser, M. Plappert, J. Tworek, J. Hilton, R. Nakano, C. Hesse, and J. Schulman, “Training verifiers to solve math word problems,” 2021.
  * [27] L. Hallee and J. P. Gleghorn, “Protein-protein interaction prediction is achievable with large language models.”
  * [28] L. Hallee, N. Rafailidis, and J. P. Gleghorn, “cdsBERT - extending protein language models with codon awareness.”
  * [29] J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova, “BERT: Pre-training of deep bidirectional transformers for language understanding.”
  * [30] P. Sharma, J. T. Ash, and D. Misra, “The truth is in there: Improving reasoning in language models with layer-selective rank reduction.”
  * [31] D. Y. Fu, S. Arora, J. Grogan, I. Johnson, S. Eyuboglu, A. W. Thomas, B. Spector, M. Poli, A. Rudra, and C. Ré, “Monarch mixer: A simple sub-quadratic GEMM-based architecture.”
  * [32] A. Radford, J. Wu, R. Child, D. Luan, D. Amodei, and I. Sutskever, “Language models are unsupervised multitask learners,” 2019.
  * [33] T. B. Brown, B. Mann, N. Ryder, M. Subbiah, J. Kaplan, P. Dhariwal, A. Neelakantan, P. Shyam, G. Sastry, A. Askell, S. Agarwal, A. Herbert-Voss, G. Krueger, T. Henighan, R. Child, A. Ramesh, D. M. Ziegler, J. Wu, C. Winter, C. Hesse, M. Chen, E. Sigler, M. Litwin, S. Gray, B. Chess, J. Clark, C. Berner, S. McCandlish, A. Radford, I. Sutskever, and D. Amodei, “Language models are few-shot learners.”
  * [34] OpenAI, J. Achiam, S. Adler, S. Agarwal, L. Ahmad, I. Akkaya, F. L. Aleman, D. Almeida, J. Altenschmidt, S. Altman, S. Anadkat, R. Avila, I. Babuschkin, S. Balaji, V. Balcom, P. Baltescu, H. Bao, M. Bavarian, J. Belgum, I. Bello, J. Berdine, G. Bernadett-Shapiro, C. Berner, L. Bogdonoff, O. Boiko, M. Boyd, A.-L. Brakman, G. Brockman, T. Brooks, M. Brundage, K. Button, T. Cai, R. Campbell, A. Cann, B. Carey, C. Carlson, R. Carmichael, B. Chan, C. Chang, F. Chantzis, D. Chen, S. Chen, R. Chen, J. Chen, M. Chen, B. Chess, C. Cho, C. Chu, H. W. Chung, D. Cummings, J. Currier, Y. Dai, C. Decareaux, T. Degry, N. Deutsch, D. Deville, A. Dhar, D. Dohan, S. Dowling, S. Dunning, A. Ecoffet, A. Eleti, T. Eloundou, D. Farhi, L. Fedus, N. Felix, S. P. Fishman, J. Forte, I. Fulford, L. Gao, E. Georges, C. Gibson, V. Goel, T. Gogineni, G. Goh, R. Gontijo-Lopes, J. Gordon, M. Grafstein, S. Gray, R. Greene, J. Gross, S. S. Gu, Y. Guo, C. Hallacy, J. Han, J. Harris, Y. He, M. Heaton, J. Heidecke, C. Hesse, A. Hickey, W. Hickey, P. Hoeschele, B. Houghton, K. Hsu, S. Hu, X. Hu, J. Huizinga, S. Jain, S. Jain, J. Jang, A. Jiang, R. Jiang, H. Jin, D. Jin, S. Jomoto, B. Jonn, H. Jun, T. Kaftan, . Kaiser, A. Kamali, I. Kanitscheider, N. S. Keskar, T. Khan, L. Kilpatrick, J. W. Kim, C. Kim, Y. Kim, H. Kirchner, J. Kiros, M. Knight, D. Kokotajlo, . Kondraciuk, A. Kondrich, A. Konstantinidis, K. Kosic, G. Krueger, V. Kuo, M. Lampe, I. Lan, T. Lee, J. Leike, J. Leung, D. Levy, C. M. Li, R. Lim, M. Lin, S. Lin, M. Litwin, T. Lopez, R. Lowe, P. Lue, A. Makanju, K. Malfacini, S. Manning, T. Markov, Y. Markovski, B. Martin, K. Mayer, A. Mayne, B. McGrew, S. M. McKinney, C. McLeavey, P. McMillan, J. McNeil, D. Medina, A. Mehta, J. Menick, L. Metz, A. Mishchenko, P. Mishkin, V. Monaco, E. Morikawa, D. Mossing, T. Mu, M. Murati, O. Murk, D. Mély, A. Nair, R. Nakano, R. Nayak, A. Neelakantan, R. Ngo, H. Noh, L. Ouyang, C. O’Keefe, J. Pachocki, A. Paino, J. Palermo, A. Pantuliano, G. Parascandolo, J. Parish, E. Parparita, A. Passos, M. Pavlov, A. Peng, A. Perelman, F. d. A. B. Peres, M. Petrov, H. P. d. O. Pinto, Michael, Pokorny, M. Pokrass, V. Pong, T. Powell, A. Power, B. Power, E. Proehl, R. Puri, A. Radford, J. Rae, A. Ramesh, C. Raymond, F. Real, K. Rimbach, C. Ross, B. Rotsted, H. Roussez, N. Ryder, M. Saltarelli, T. Sanders, S. Santurkar, G. Sastry, H. Schmidt, D. Schnurr, J. Schulman, D. Selsam, K. Sheppard, T. Sherbakov, J. Shieh, S. Shoker, P. Shyam, S. Sidor, E. Sigler, M. Simens, J. Sitkin, K. Slama, I. Sohl, B. Sokolowsky, Y. Song, N. Staudacher, F. P. Such, N. Summers, I. Sutskever, J. Tang, N. Tezak, M. Thompson, P. Tillet, A. Tootoonchian, E. Tseng, P. Tuggle, N. Turley, J. Tworek, J. F. C. Uribe, A. Vallone, A. Vijayvergiya, C. Voss, C. Wainwright, J. J. Wang, A. Wang, B. Wang, J. Ward, J. Wei, C. J. Weinmann, A. Welihinda, P. Welinder, J. Weng, L. Weng, M. Wiethoff, D. Willner, C. Winter, S. Wolrich, H. Wong, L. Workman, S. Wu, J. Wu, M. Wu, K. Xiao, T. Xu, S. Yoo, K. Yu, Q. Yuan, W. Zaremba, R. Zellers, C. Zhang, M. Zhang, S. Zhao, T. Zheng, J. Zhuang, W. Zhuk, and B. Zoph, “GPT-4 technical report.”
  * [35] N. Shazeer, A. Mirhoseini, K. Maziarz, A. Davis, Q. Le, G. Hinton, and J. Dean, “Outrageously large neural networks: The sparsely-gated mixture-of-experts layer.”
  * [36] M. AI, “Mixtral of experts.”  Section: news.
  * [37] S. Zuo, Q. Zhang, C. Liang, P. He, T. Zhao, and W. Chen, “MoEBERT: from BERT to mixture-of-experts via importance-guided adaptation.”
  * [38] W. Fedus, B. Zoph, and N. Shazeer, “Switch transformers: Scaling to trillion parameter models with simple and efficient sparsity.”
  * [39] W. Wang, F. Wei, L. Dong, H. Bao, N. Yang, and M. Zhou, “MiniLM: Deep self-attention distillation for task-agnostic compression of pre-trained transformers.”
  * [40] P. Lewis, Y. Wu, L. Liu, P. Minervini, H. Küttler, A. Piktus, P. Stenetorp, and S. Riedel, “PAQ: 65 million probably-asked questions and what you can do with them.”
  * [41] D. Khashabi, A. Ng, T. Khot, A. Sabharwal, H. Hajishirzi, and C. Callison-Burch, “GooAQ: Open question answering with diverse answer types.”
  * [42] M. Dunn, L. Sagun, M. Higgins, V. U. Guney, V. Cirik, and K. Cho, “SearchQA: A new q&a dataset augmented with context from a search engine.”
  * [43] M. Koupaee and W. Y. Wang, “WikiHow: A large scale text summarization dataset.”
  * [44] M. Henderson, P. Budzianowski, I. Casanueva, S. Coope, D. Gerz, G. Kumar, N. Mrkšić, G. Spithourakis, P.-H. Su, I. Vulić, and T.-H. Wen, “A repository of conversational datasets.”
  * [45] I. Beltagy, K. Lo, and A. Cohan, “SciBERT: A pretrained language model for scientific text.”
  * [46] K. Song, X. Tan, T. Qin, J. Lu, and T.-Y. Liu, “MPNet: Masked and permuted pre-training for language understanding.”
  * [47] J. Giorgi, O. Nitski, B. Wang, and G. Bader, “DeCLUTR: Deep contrastive learning for unsupervised textual representations.”
  * [48] Y. Gu, R. Tinn, H. Cheng, M. Lucas, N. Usuyama, X. Liu, T. Naumann, J. Gao, and H. Poon, “Domain-specific language model pretraining for biomedical natural language processing,” vol. 3, no. 1, pp. 1–23.
  * [49] K. Huang, J. Altosaar, and R. Ranganath, “ClinicalBERT: Modeling clinical notes and predicting hospital readmission.”
  * [50] M. M. Shariatnia, “Simple CLIP,” 4 2021.
  * [51] M. Henderson, R. Al-Rfou, B. Strope, Y.-H. Sung, L. Lukács, R. Guo, S. Kumar, B. Miklos, and R. Kurzweil, “Efficient natural language response suggestion for smart reply,” ArXiv, vol. abs/1705.00652, 2017.
  * [52] L. N. Smith and N. Topin, “Super-convergence: Very fast training of neural networks using large learning rates.”
  * [53] “DeepGraphLearning/torchdrug.”  original-date: 2021-08-10T03:51:24Z.
  * [54] J. Su, C. Han, Y. Zhou, J. Shan, X. Zhou, and F. Yuan, “SaProt: Protein language modeling with structure-aware vocabulary.”  Pages: 2023.10.01.560349 Section: New Results.
  * [55] J. Hoffmann, S. Borgeaud, A. Mensch, E. Buchatskaya, T. Cai, E. Rutherford, D. d. L. Casas, L. A. Hendricks, J. Welbl, A. Clark, T. Hennigan, E. Noland, K. Millican, G. van den Driessche, B. Damoc, A. Guy, S. Osindero, K. Simonyan, E. Elsen, J. W. Rae, O. Vinyals, and L. Sifre, “Training Compute-Optimal large language models,” 2022.