Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Using text embedding models as text classifiers with medical data

Rishabh Goel

TL;DR

The paper addresses the need for accurate medical text classification without training large medical LLMs. It evaluates a pipeline that uses text embeddings and a vector database as a knowledge base to classify clinician notes, with ground-truth vectors generated by LLMs and queries produced by multiple LLMs. Across 8 ailments and 18 robustness permutations, higher embedding dimensionality and richer knowledge generally improve classification, achieving macro F1 scores up to 0.96 when queries are concise and knowledge is detailed; however, ground-truth quality and the choice of LLMs significantly impact results. The work demonstrates a cost-effective, scalable approach for medical text classification that could support clinicians without deploying new large-domain models, while highlighting factors that influence performance and avenues for scaling and validation ($3 \times 2 \times 3 = 18$ permutations).

Abstract

The advent of Large Language Models (LLMs) is promising and LLMs have been applied to numerous fields. However, it is not trivial to implement LLMs in the medical field, due to the high standards for precision and accuracy. Currently, the diagnosis of medical ailments must be done by hand, as it is costly to build a sufficiently broad LLM that can diagnose a wide range of diseases. Here, we explore the use of vector databases and embedding models as a means of encoding and classifying text with medical text data without the need to train a new model altogether. We used various LLMs to generate the medical data, then encoded the data with a text embedding model and stored it in a vector database. We hypothesized that higher embedding dimensions coupled with descriptive data in the vector database would lead to better classifications and designed a robustness test to test our hypothesis. By using vector databases and text embedding models to classify a clinician's notes on a patient presenting with a certain ailment, we showed that these tools can be successful at classifying medical text data. We found that a higher embedding dimension did indeed yield better results, however, querying with simple data in the database was optimal for performance. We have shown in this study the applicability of text embedding models and vector databases on a small scale, and our work lays the groundwork for applying these tools on a larger scale.

Using text embedding models as text classifiers with medical data

TL;DR

The paper addresses the need for accurate medical text classification without training large medical LLMs. It evaluates a pipeline that uses text embeddings and a vector database as a knowledge base to classify clinician notes, with ground-truth vectors generated by LLMs and queries produced by multiple LLMs. Across 8 ailments and 18 robustness permutations, higher embedding dimensionality and richer knowledge generally improve classification, achieving macro F1 scores up to 0.96 when queries are concise and knowledge is detailed; however, ground-truth quality and the choice of LLMs significantly impact results. The work demonstrates a cost-effective, scalable approach for medical text classification that could support clinicians without deploying new large-domain models, while highlighting factors that influence performance and avenues for scaling and validation ( permutations).

Abstract

The advent of Large Language Models (LLMs) is promising and LLMs have been applied to numerous fields. However, it is not trivial to implement LLMs in the medical field, due to the high standards for precision and accuracy. Currently, the diagnosis of medical ailments must be done by hand, as it is costly to build a sufficiently broad LLM that can diagnose a wide range of diseases. Here, we explore the use of vector databases and embedding models as a means of encoding and classifying text with medical text data without the need to train a new model altogether. We used various LLMs to generate the medical data, then encoded the data with a text embedding model and stored it in a vector database. We hypothesized that higher embedding dimensions coupled with descriptive data in the vector database would lead to better classifications and designed a robustness test to test our hypothesis. By using vector databases and text embedding models to classify a clinician's notes on a patient presenting with a certain ailment, we showed that these tools can be successful at classifying medical text data. We found that a higher embedding dimension did indeed yield better results, however, querying with simple data in the database was optimal for performance. We have shown in this study the applicability of text embedding models and vector databases on a small scale, and our work lays the groundwork for applying these tools on a larger scale.
Paper Structure (14 sections, 6 figures, 1 table)

This paper contains 14 sections, 6 figures, 1 table.

Table of Contents

  1. Introduction
  2. Materials and Methods
  3. Vector Database
  4. Workflow
  5. Data
  6. Prompt Engineering
  7. Robustness Test Definition
  8. Results
  9. High detail in both knowledge and query dataset
  10. Using sparse querying with detailed knowledge base
  11. Sparse data as knowledge base
  12. Same model for knowledge and query dataset
  13. Effects of vector dimensionality
  14. Discussion

Figures (6)

  • Figure 1: Flow chart of robustness test pipeline. There are 3x2x3 = 18 permutations of tests possible. Each query generation model was paired up against each text embedding model and each knowledge generation model. Embeddings from the query generation model were compared with the embeddings from the knowledge generation model using a cosine similarity.
  • Figure 2: Confusion matrix describing the results from the first robustness test where there were detailed querying and knowledge bases. Rows describe the true values, and columns describe the predicted values. The query dataset was generated by LLaMA 2 70b-chat, text-embedding-ada-002 was used as the embedding model, and the ground truth dataset in the vector database was generated by gpt-3.5-turbo.
  • Figure 3: Confusion matrix describing the results from the second robustness test where there was a sparse querying base and a detailed knowledge base. Rows describe the true values, and columns describe the predicted values. The query dataset was generated by google-flan-t5-xl, text-embedding-ada-002 was used as the embedding model, and the ground truth dataset in the vector database was generated by gpt-3.5-turbo.
  • Figure 4: Confusion matrix describing the results from the third robustness test where there was a detailed querying base and a sparse knowledge base. Rows describe the true values, and columns describe the predicted values. The query dataset was generated by gpt-3.5-turbo, text-embedding-ada-002 was used as the embedding model, and the ground truth dataset in the vector database was generated by google-flan-t5-xl.
  • Figure 5: Confusion matrix describing the results from the fourth robustness test where there were sparse querying and knowledge bases. Rows describe the true values, and columns describe the predicted values. The query dataset was generated by google-flan-t5-xl, text-embedding-ada-002 was used as the embedding model, and the ground truth dataset in the vector database was generated by google-flan-t5-xl.
  • ...and 1 more figures