Symptoms2Vec DoctorFinder

Problem Statement

Patients worldwide face a critical challenge: symptom interpretation and medical pathway navigation. When experiencing symptoms, individuals often resort to unreliable web searches, generic symptom checkers, or costly unnecessary hospital visits. Meanwhile, medical professionals and patients struggle with:

• Information Asymmetry: Patients lack structured access to verified medical case histories
• Diagnostic Uncertainty: Self-diagnosis is error-prone; symptoms map to multiple diseases with varying probabilities
• Care Routing Inefficiency: Patients struggle to identify which medical specialists actually handle their condition

Why AI is the Right Solution

This problem is fundamentally a semantic matching and retrieval task — precisely where large language models excel. While rule-based systems and simple keyword matching fail due to symptom variability and linguistic diversity, transformer-based models can:

• Learn nuanced relationships between symptoms and diseases
• Scale to diverse medical domains and patient populations
• Continuously improve with additional training data

Magnitude & Impact Potential

• Global Scale: Billions of people annually seek medical information online
• Economic Impact: Reducing unnecessary ED visits saves healthcare systems significantly

Overall Solution

• Application: A Gradio-based web application that helps users find doctors based on their symptoms.
• Model: Represent a given text in natural language describing symptoms by vectors.

┌─────────────────────────────────────────────────────────────┐
│          User-Facing Layer: Gradio Web Interface            │
│  (Symptom Input, Privacy Controls, Result Visualization)    │
└────────────────────┬────────────────────────────────────────┘
                     │
┌────────────────────▼────────────────────────────────────────┐
│         Query Engine: Vector Similarity Search              │
│         (Text → Embedding → FAISS Search → Ranking)         │
└────────────────────┬────────────────────────────────────────┘
                     │
┌────────────────────▼────────────────────────────────────────┐
│    Embedding Layer: MedGemma + Metric Learning Head         │
│    (Symptom Encoding to 1024D normalized vectors)           │
└────────────────────┬────────────────────────────────────────┘
                     │
┌────────────────────▼────────────────────────────────────────┐
│     Index Layer: FAISS (Fast Approximate Nearest Neighbor)  │
│     (Pre-computed disease-symptom embeddings)               │
└─────────────────────────────────────────────────────────────┘

Base model: medgemma-1.5-4b-it

• The model already contains extensive medical knowledge, providing a strong starting point for fine-tuning symptom understanding.
• The image understanding module was removed, and only the text module is used, due to the difficulty of finding high-quality image datasets for fine-tuning.
• Honest comparison: Some similarly sized LLMs perform comparably for this solution. However, medgemma-1.5-4b-it could offer a significant advantage if high-quality image datasets (e.g., dermatology images, radiology scans) become available.
• Alternatively, medgemma-27b-text-it could be trained, but this was not feasible on a single RTX 3090.
• Generated model: medgemma-1.5-4b-text-emb, an embedding model for medical text similarity and retrieval.

Technical Details

Dataset

Training Data: Disease and symptoms dataset 2023

• 773 unique diseases and 377 symptoms with binary values
• Preprocessing: Filtered out statistically insignificant data and converted symptom lists into text format (e.g., “nausea, stomach pain, chest pain…”)

Production and Evaluation Data: Symptom2Disease

• Contains 24 diseases and 50 multi-sentence symptom descriptions
• Likely LLM-generated
• Used as the evaluation dataset because it differs significantly from the training set, ensuring model generalization
• Demo examples are created from this dataset

Privacy & Safety

• Embeddings are computed locally.
• Users may share their symptoms after successful treatment to help others.
• Users can upload embeddings (not raw text) for convenient server-side search or download their data and perform local searches.
• No direct medical advice is provided, only doctor recommendations based on symptom similarity.
• The retrieval-based approach prevents hallucinated recommendations.

Model Architecture and Training Scheme

We convert given text into prompts:

Instruct: Diagnose the disease given symptoms.

Symptoms: {text}

Diagnosis: 

Rather than using MedGemma for generation, we exploit its learned medical representations as a foundation for semantic embeddings:

Input Symptoms → MedGemma Language Model → Last Hidden State
                                          ↓
                                    Dual-Head
                                    ├─ Last Token → Head (512D) → Diagnosis Vector
                                    └─ Mean Pooled Tokens → Head (512D) → Context Vector
                                          ↓
                                    Concatenation + Normalization
                                          ↓
                                    Final Embedding (1024D normalized)

Rather than standard mean pooling, we combine:

• Last token representation (captures final diagnostic conclusion)
• Masked mean embedding (preserves symptom context from earlier tokens)

We concatenate them for inference. For training, we calculate cosine similarity between last token and mean token as an auxiliary loss. This dual-head approach captures both diagnostic intent and symptom context, improving matching accuracy and speeding up convergence.

Arc-Face Loss for Metric Learning: We implemented ArcMarginProduct loss, typically used in face recognition, adapted for medical concept clustering:

• Encourages embeddings from the same disease to cluster together
• Pushes different diseases apart in embedding space

LoRA Fine-tuning:

• Target modules: q_proj, k_proj, v_proj (attention heads)
• Rank = 4, Alpha = 4, Dropout = 0.05
• Freezes base MedGemma weights (0.4% trainable parameters)
• Enables adaptation to symptom-disease relationships without catastrophic forgetting

Training Configuration:

Evaluation

• Training and validation datasets come from different sources to ensure generalization.
• We compute Top-5 Pairing Accuracy (pa@5) on the validation set, measuring whether the top matches belong to the same disease class.
• Additional qualitative evaluation was necessary because the validation dataset contains near-duplicates. Several handcrafted examples were manually inspected.

Reproducible Code

All source code is provided:

• training.py — Full training pipeline with LoRA + Arc-Face
• models.py — VecModel and TrainingModel architectures
• inference.py — Batch inference and FAISS querying
• demo/app.py — Gradio interface
• demo/query_faiss.py — Search engine implementation
• demo/mock_data.py — Mock symptoms and doctors dataset for the demo

How to reproduce the demo: First download the datasets described in Dataset Section and put them in data folder.

# Install dependencies
pip install -r requirements.txt

# Train model
PYTHONPATH=. python training.py

# Index mock doctors and symptoms data
PYTHONPATH=. python demo/mock_data.py

# Run web interface
PYTHONPATH=. python demo/app.py

This project is licensed under the Creative Commons Attribution 4.0 International License (CC BY 4.0).