Simonlee711/Clinical_UltraBERT

Clinical UltraBERT

Clinical UltraBERT is a biomedical and clinical text embedding model. It maps text to a 768-dimensional vector for retrieval, semantic similarity, clustering and classification. It is UltraBERT-150M fine-tuned with a contrastive loss on about 100K biomedical and clinical text pairs.

Usage

from sentence_transformers import SentenceTransformer

model = SentenceTransformer("Simonlee711/Clinical-UltraBERT", trust_remote_code=True)

queries = ["What are the first-line treatments for community-acquired pneumonia?"]
documents = [
    "Empiric therapy for outpatient community-acquired pneumonia includes amoxicillin or doxycycline.",
    "Metformin is the preferred initial pharmacologic agent for type 2 diabetes.",
]
q = model.encode(queries)
d = model.encode(documents)
print(model.similarity(q, d))

trust_remote_code=True is needed because UltraBERT ships its own modelling code (modeling_ultrabert.py); from transformers >= 5.20 the architecture is supported natively.

Training

Contrastive fine-tuning with InfoNCE (CachedMultipleNegativesRankingLoss, scale 20) using in-batch negatives plus one mined hard negative where available. Every batch is drawn from a single source, so in-batch negatives are on-topic. Ten epochs, batch size 64, learning rate 5e-5 with 5% warm-up and linear decay, AdamW (β2 0.98, weight decay 0.01), sequence length 128, mean pooling, trained in fp32 on a single Apple M4 Max GPU.

10K pairs from each of these ten sources (100,000 pairs in all):

Source	Pairs
PubMedQA, BioASQ (F2LLM-v2)	question → PubMed abstract
Medical flashcards, MedMCQA, MedInstruct (F2LLM-v2)	medical question → answer
MSciNLI, SNLI, MNLI (F2LLM-v2)	premise → entailed hypothesis
PubMed (MedRAG/pubmed)	article title → abstract
Augmented clinical notes	doctor–patient conversation → clinical note

F2LLM-v2 subsets built from benchmark test data (MedQuAD, SciFact, MedQA, medRxiv clustering) and HealthCareMagic doctor answers (the ChatDoctor retrieval corpus) were left out of training.

Evaluation

A 20-task medical suite: 10 retrieval tasks, 7 classification tasks, BIOSSES and 2 clustering tasks, from MTEB plus five clinical classification datasets. Inputs are cut to 512 tokens for every model. The baselines (BERT, Bio_ClinicalBERT, ModernBERT, Clinical ModernBERT, BioClinical ModernBERT, UltraBERT) are masked-LM checkpoints used as frozen mean-pooled encoders, so they are expected to trail a contrastively trained model on retrieval. Classification scores come from a logistic-regression probe on frozen embeddings.

Scores ×100: retrieval nDCG@10; SciRepEval accuracy; MTSamples, Symptom2Diagnosis, Medical Abstracts and PubMed RCT macro-F1; HoC micro-F1; BIOSSES Spearman; clustering V-measure.

Clinical UltraBERT has the best score on 15 of 20 tasks.

Type	Task	Best baseline	Clinical UltraBERT
Retrieval	SciFact	33.2 (UltraBERT-150M)	63.5
Retrieval	NFCorpus	6.7 (UltraBERT-150M)	23.1
Retrieval	TREC-COVID	33.3 (UltraBERT-150M)	61.1
Retrieval	MedicalQA	17.4 (UltraBERT-150M)	65.1
Retrieval	PublicHealthQA	55.2 (UltraBERT-150M)	75.9
Retrieval	CovidQA	50.5 (UltraBERT-150M)	70.7
Retrieval	ChatDoctor	7.7 (UltraBERT-150M)	48.1
Retrieval	R2MED PMC-Clinical	6.6 (BioClinical-ModernBERT)	10.0
Retrieval	R2MED MedQA-Diag	1.1 (UltraBERT-150M)	4.1
Retrieval	BIRCO Clinical-Trial	9.5 (UltraBERT-150M)	16.2
Classification	SciRepEval DRSM	69.1 (BioClinical-ModernBERT)	64.7
Classification	SciRepEval MeSH	47.7 (UltraBERT-150M)	61.4
STS	BIOSSES	75.1 (UltraBERT-150M)	84.2
Clustering	MedrxivClusteringS2S	34.5 (UltraBERT-150M)	32.1
Clustering	ClusTREC-Covid	78.2 (Bio_ClinicalBERT)	80.0
Classification	MTSamples	67.9 (UltraBERT-150M)	67.5
Classification	Symptom2Diagnosis	95.3 (UltraBERT-150M)	95.7
Classification	Medical Abstracts	55.4 (BioClinical-ModernBERT)	56.9
Classification	PubMed RCT 20k	72.7 (BioClinical-ModernBERT)	71.0
Classification	HoC	63.8 (UltraBERT-150M)	61.8

Limitations

• Most training and evaluation text is biomedical literature or consumer health text. Real clinical notes appear only through MTSamples transcriptions and synthetic doctor–patient conversations; no credentialed EHR data (e.g. MIMIC) was used for training.
• Classification of document structure (study type, sentence role) is weaker than the best masked-LM encoders, and fine-tuning lowered a few probe scores relative to UltraBERT-150M (SciRepEval DRSM, HoC, MTSamples, MedrxivClusteringS2S).
• On MIMIC-IV prediction from serialized EHR records (mortality, length of stay, readmission) the model does not improve over its base model.
• The model is not a medical device and its outputs must not be used for clinical decisions without expert review.

License

Apache 2.0, like the base model. The training data keeps its own licenses: F2LLM-v2 (Apache 2.0), augmented clinical notes (MIT) and PubMed abstracts (NLM terms of use).