How we approached Cohort X Task 3

Note: this blog is designed to be beginner friendly. If you are interested in the technical details, you can view the repository at: https://github.com/debug420/cohortx-task3

Introduction of Challenge

The overall goal of the challenge was to determine a "cohort" of individuals who are eligible for a trial, or as the organisers more formally define it: "extracting executable cohort definitions for medical imaging research".

There are actually three tasks within this challenge that come together to address the above problem:

1. Extracting cohort selection criteria from free texts
2. Structuring cohort eligibility criteria in the form of triples
3. Resolving medical conditions to ICD-10-CM Codes

You can read further about the challenge at:

https://csisc.github.io/CohortX-MICCAI/

Particularly, our research was aimed at step 3 of the pipeline, wherein our methodology scored 5th internationally.

So what does "Resolving medical conditions to ICD-10-CM Codes" actually mean?

First, we must understand what ICD-10-CM codes are. ICD codes simply standardise the classification of medical diagnoses and conditions across the healthcare industry. Without a standardised coding system, 2 different doctors will write different things. For example:

• The Emergency Room Doctor writes: "Traumatic blunt force trauma to the left orbit resulting in sudden vision loss and direct pupillary defect."
• The Ophthalmologist (Eye Specialist) writes: "Traumatic left optic neuropathy secondary to retrobulbar pressure."

To someone without a medical degree, these two descriptions look like entirely different medical issues. One talks about "orbit trauma and a pupillary defect," while the other calls it "optic neuropathy."

Both of the above descriptions can actually be mapped to S0401, which is the universal ICD code that refers to "Injury of optic nerve." As such, everyone instantly knows the exact condition/diagnoses that is being referred to.

Cohort X Task 3 aims to address this exact code resolution. Using the same example above, the goal is to derive S0401 when the model is fed something like "optic neuropathy." (although technically, the challenge is a bit more complicated - however, we will delve into the specific details later)

Embedding the ICD Codes

In order to capture the meaning/semantics of the ICD codes, we must convert them to vectors. This allows us to quantitively compare how similar 2 codes are to each other. For example, consider the following codes:

• S0401: Injury of optic nerve
• S0402: Injury of optic chiasm
• E06: Thyroiditis

Our model should be able to recognise that S0401 & S0402 are closely related, while E06 is a different medical concept. As a first step, we decided to use SapBERT, as it has been trained on medical papers, and therefore would be able to "understand" medical concepts better than generic text embedding models.

As you can see from the below visualisation, SapBERT is able to map codes by their "semantic" meaning based on their textual descriptions, bringing similar medical concepts closer together in the embedding space.

Below is a clearer example of how SapBERT is able to differentiate the different medical concepts.

As shown in the mapping space above, codes cluster logically by pathology: J449 (Chronic obstructive pulmonary disease) is placed near other respiratory codes and far from unrelated codes like S48012S (Amputation at shoulder). However, the model struggles to differentiate semantically adjacent concepts, such as J208 (Acute bronchitis) and S25409 (Injury of pulmonary vessels).

Having embedded the ICD codes using SapBERT, we decided to retrieve the closest n=500 seeds (out of the ~97k seeds) to see how well this general approach performs as a baseline. In addition, we experimented with a sparse encoder, SPLADE, which has been trained on PubMed medical papers. Our rationale was that a sparse encoder helps resolve dense retrieval failure modes such as disambiguating opposing clinical modifiers (e.g., hypothyroidism vs. hyperthyroidism) or acuity and laterality distinctions (e.g., acute vs. chronic heart failure, or left vs. right femur fractures); essentially where terms share nearly identical semantic embedding spaces but demand strict lexical precision. The below figure shows the mapping space of SPLADE.

In order to test how well our retrieval methods work, we calculated the recall of the different retrieval modes:

1. Purely SabBERT
2. Hybrid of SabBERT & SPLADE (0.5/0.5 weight)

As seen in the graph above, we evaluated on only five labeled conditions, two of which were clinical acronyms (URTI and SOB). Unsurprisingly, these two performed the worst. To improve retrieval recall, we implemented acronym expansion and text preprocessing to give the embedding models richer context to work with. For example "URTI" expanded to "Upper Respiratory Tract Infection".

Integrating these techniques significantly improved target code retrieval. As shown in the 2-D PCA projections below, the retrieved candidates consistently map to the semantic neighbourhood of the gold KEEP seeds, effectively concentrating predictions within the ground-truth region.

After seeing that this methodology can be used, we realised that improving this recall was not entirely feasible. Instead, we realised that if an ICD code is retrieved with a very high confidence, chances are, the children of that code are also part of the gold set. For example, if our model retrieved E11 (type 2 diabetes mellitus) confidently, then all seeds under E11, such as E112 (Type 2 diabetes mellitus with kidney complications), were most likely part of the golden set.

So, instead of banging our heads against the wall trying to get the model to catch every hyper-specific leaf code in a massive label space, we used a shortcut: hierarchical expansion. We grabbed the high-confidence parent code and simply cascaded downwards to pull in its children automatically. Leaning into the ICD tree structure this way gave our recall an instant boost without having to retrain a massive model from scratch.

Determining DIFF/ASSOC Seeds

Resolving the direct ICD codes of the input text was only half of the challenge. The other half was determining DIFFERENT, and ASSOCIATED codes.

To see why this gets tricky, take a condition like URTI (Upper Respiratory Tract Infection) from the labelled dataset.

• The KEEP bucket is straightforward: it holds direct matches like acute nasopharyngitis or common colds (codes starting with J00–J06), acute laryngitis, and influenza (J10–J11). These directly describe an upper respiratory infection.
• But then you hit the ASSOCIATION category: here, the ground truth flags conditions like Otitis Media (H65 and H66 series - middle ear infections). Pathologically, a bad head cold frequently leads to a blocked Eustachian tube and an ear infection. They aren't the exact same condition, but clinically they run in the same pack.
• Finally, the DIFF (Different) bucket contains diagnoses like viral and bacterial pneumonia (J12–J18) or acute bronchitis (J20–J21). To a naive embedding model, "pneumonia" and "severe respiratory infection" sound virtually identical in vector space. But anatomically and clinically, pneumonia is a lower respiratory tract infection, making it fundamentally different from a URTI, which is why it has to be cleanly separated.

Because models love to conflate semantically close medical terms, separating these out wasn't just a matter of calculating cosine similarity. We had to teach the pipeline the subtle boundaries between what is the condition, what tags along with it, and what is strictly something else entirely.

We determined the only way to tackle this is through some sort of "reasoning" which an LLM may be capable of. Since the challenge had strict hardware requirements, we were strictly limited in the selection of models. In the end, we settled on MedPsy for its performance to parameters ratio.

We determined 2 fundamentally different approaches we can take:

1. Use MedPsy to generate DIFF & ASSOC search terms, which are then fed into our retrieval model the same way our KEEP search term is fed in.
2. Use MedPsy to classify the retrieved seeds into the correct bucket (e.g. J00 (Acute nasopharyngitis [common cold]) has been retrieved by the model and is fed into MedPsy for it to determine whether it is a KEEP, ASSOC, or DIFF seed)

Long story short, we determined that the LLM is not really able to "reason" correctly and seemed to convince itself on whatever answer looked "right". In fact, it performed so poorly that it lowered our score further from our more traditional approach. However, do note that some of the other teams that entered this competition did have some luck with the usage of LLMs. Therefore the above conclusion is not to say that LLMs are completely useless; due to the strict time deadline of the competition, we simply did not have the resources to investigate this further.

In the limited time that we did have with LLMs, we encountered the following challenges:

1. Parse failures (model did not emit correct JSON as per request).
2. Output consistency was all over the place across the ontology tree. The model misclassified the parent code J01 (Acute sinusitis) as not being a URTI, yet somehow labeled its specific subtype J0110 (Acute frontal sinusitis) as one.
3. Inference was also too slow, putting us in a gray zone of potentially getting disqualified on a timeout.

Once again, I want to reinstate that LLMs may still have potential in this domain, however due to the time and compute limitation we were presented with, we were not able to optimise it.

---

Conclusion

Ultimately, we decided to play to our strengths. Rather than submitting noisy predictions, we left the ASSOC and DIFF buckets empty and focused purely on maximising recall and precision on KEEP. Looking back, while there are a ton of things we'd love to polish with more time, the benchmark we hit under such intense time pressure was impressive.

---

CohortX Task 3 Team

• Parsa Ehyaee
• Maximus Debski
• Sarah Huang