Boosting Long-Tail Data Classification with Sparse Prototypical Networks

Alexei Figueroa, Jens-Michalis Papaioannou, Conor Fallon, Alexandra Bekiaridou, Keno K. Bressem, Stavros Zanos, Felix A. Gers, Wolfgang Nejdl, Alexander Löser

Abstract: Clinical Decision Support Systems (CDSS) have become ubiquitous in healthcare facilities, leveraging the increasing presence of Electronic Health Records (EHR). Predicting clinical outcomes from clinical text, such as identifying diagnoses based on the admission state of patients, is among the core tasks that a CDSS must address. The state-of-the-art for this task has been set by transformer encoder models, recently superseded by encoders enhanced with a prototypical network. This task remains a significant challenge due to the substantial imbalance of the outcome labels, which is characterized by a long-tailed distribution where the majority of diagnoses are under-represented. Motivated by recent biologically inspired findings in deep learning, we propose S-Proto, a novel, efficient, and sparse prototypical layer. Our method achieves state-of-the-art performance in outcome diagnosis prediction, without compromising on the explainability characteristics of prototypical encoders. Quantitative results demonstrate that our approach is robust to the challenges presented by clinical notes, and transfers successfully to a second, unseen dataset. Qualitative evaluation with medical doctors shows that S-Proto is capable of disaggregating the representations of a disease that manifests differently in patient cohorts.