Technical Summary: Task-Aligned Embeddings for Clinical Code Retrieval

Problem Statement

Clinical coding is a high-stakes, fine-grained semantic retrieval task where short clinical mentions must be mapped to specific ICD-10/CIE-10 codes. Unlike general semantic search, this task requires distinguishing between clinically similar conditions that differ only by subtle qualifiers such as severity, anatomical location, laterality, temporality, or etiology. For example, codes for "Brucellosis due to Brucella suis" and "Brucellosis due to Brucella melitensis" are semantically close but represent distinct, non-interchangeable codes.

The field faces four primary limitations:

1. Language Bias: Most public embedding models (e.g., BioBERT, ClinicalBERT) and biomedical resources are optimized for English, leaving Spanish and other clinical languages under-represented.
2. Granularity Mismatch: General-purpose encoders capture broad semantic relatedness but fail to learn the fine-grained decision boundaries required for specific code assignment.
3. Data Scarcity: Manually annotated clinical data is expensive and scarce, particularly for multilingual settings where local coding vocabularies and granularity levels vary.
4. Architectural Misalignment: Two-stage retrieval systems often suffer from misalignment; a high-recall bi-encoder may fail to retrieve the correct code, or a reranker may be unable to order clinically adjacent candidates if the initial retrieval is poor.

Methodology

The authors propose a two-stage neural retrieval architecture trained on high-quality, task-aligned synthetic supervision generated by a frontier Large Language Model (LLM).

1. Synthetic Data Generation

Instead of relying on web-scale weak supervision, the authors used Google DeepMind Gemini 2.5 Flash Pro to act as a "tutor" for a controlled label space. The LLM generated two distinct datasets grounded in the ICD-10 hierarchy:

• Dataset A (Bi-Encoder Training): 19,502 synthetic examples across six languages (English, Spanish, Catalan, Italian, Portuguese, French). Each example includes a clinical query, positive pairs (paraphrases/canonical descriptions), hard negatives (semantically adjacent codes within the same chapter), and soft negatives (unrelated chapters). This design forces the model to learn fine-grained distinctions.
• Dataset B (Cross-Encoder Training): 10,628 listwise groups anchored in real Spanish clinical mentions (from CodiESP). Each group contains one query, one positive, and up to nine hard negatives sampled by lexical and embedding adjacency.

2. Model Architecture

The system follows a retrieve-and-rerank paradigm:

• Bi-Encoder (Stage 1): Initialized from PlanTL-GOB-ES/bsc-bio-ehr-es (a Spanish biomedical BERT model). It is fine-tuned using MultipleNegativesRankingLoss on Dataset A. It projects queries and candidates into a shared vector space for efficient top-k retrieval.
• Cross-Encoder (Stage 2): Initialized from the fine-tuned bi-encoder checkpoint. It jointly processes the query and candidate pairs to produce a precise relevance score. It is trained on Dataset B using a listwise softmax cross-entropy objective with hard negatives.

3. Evaluation Setup

The system was evaluated on two held-out Spanish clinical coding datasets:

• CodiESP v4: 3,664 mention-to-code pairs (primary benchmark).
• DISTEMIST: 1,224 disease mentions (secondary benchmark for generalization).

Performance was measured against lexical baselines (BM25) and off-the-shelf sentence-transformers (MiniLM-L6-v2, BioBERT-ST, MPNet-v2) using F1, MAP@10, Precision, and Recall at exact-code, category (3-char), and chapter levels.

Key Results

The proposed system significantly outperformed all public baselines and previous state-of-the-art results on the CodiESP benchmark.

• Performance on CodiESP v4:
  • TietAI Cross-Encoder: Achieved F1 = 0.709 (exact code) and MAP@10 = 0.747.
  • Comparison: This represents a 48.4 percentage point improvement over the best public encoder (MiniLM-L6-v2, F1=0.225) and a 51.6 point improvement over the in-domain English S-BioBERT baseline (F1=0.193).
  • Hierarchy: At the category level, the system achieved F1 = 0.823 and MAP@10 = 0.851.
• Performance on DISTEMIST:
  • The system achieved F1 = 0.776 (exact code) and MAP@10 = 0.812, confirming the results generalize to a different corpus.
• Baseline Failures: Counter-intuitively, larger models (MPNet-v2) and domain-specific English pretraining (S-BioBERT) underperformed smaller, general-purpose models (MiniLM-L6-v2) and the proposed task-aligned system. This suggests that generic biomedical pretraining does not transfer well to Spanish clinical coding without specific task alignment.

Contributions and Significance

The paper makes three primary contributions:

1. Reproducible Recipe: It documents a method for training two-stage ICD retrievers using multilingual LLM-generated supervision, demonstrating that high-quality synthetic data can compensate for the lack of large-scale annotated clinical corpora.
2. Data Curation over Model Scale: The results indicate that task-aligned supervision and language coverage are more critical than model scale or generic biomedical pretraining. The authors argue that carefully curated synthetic datasets (approx. 19.5k pairs for the bi-encoder and 10.6k groups for the cross-encoder) can enable compact, effective systems without web-scale training.
3. Role of Reranking: The study quantifies the necessity of cross-encoder reranking for resolving qualifier-level distinctions (e.g., laterality, severity) common in ICD-10. While the bi-encoder provides strong recall (MAP@10 ~0.69 at category level), the cross-encoder is essential for calibrating the top-1 ordering.

Conclusion and Future Work

The authors conclude that for specialized medical retrieval, data curation is the binding constraint, not model size. Their pipeline achieves an F1 of 0.709 on CodiESP, surpassing the best results from the 2020 CodiESP shared task (F1=0.687), despite using a bi-encoder trained exclusively on out-of-sample synthetic data.

The paper suggests a production architecture consisting of a universal, curated bi-encoder for candidate generation, paired with language-specific or specialty-specific cross-encoders for reranking. Future work proposed includes:

• Multilingual evaluation to isolate language effects (testing Italian, Portuguese, and Catalan).
• Scaling the synthetic corpus by an order of magnitude to cover rare qualifier combinations.
• Replacing the cross-encoder with an LLM-based coding agent that operates on the bi-encoder's shortlist.

The authors maintain a modest stance, noting that synthetic data is not a substitute for real clinical evaluation and that the generated examples may lack the noise of real-world notes. However, they assert that their findings demonstrate the viability of small-data, task-aligned approaches for clinical semantic search.