T ext2 T ree: Aligning Text Representation to the Label Tree Hierarchy for Imbalanced Medical Classification (2023)

Abstract: Deep learning approaches exhibit promising performances on various text tasks. However, they are still struggling on medical text classification since samples are often extremely imbalanced and scarce. Different from existing mainstream approaches that focus on supplementary semantics with external medical information, this paper aims to rethink the data challenges in medical texts and present a novel framework-agnostic algorithm called Text2Tree that only utilizes internal label hierarchy in training deep learning models. We embed the ICD code tree structure of labels into cascade attention modules for learning hierarchy-aware label representations. Two new learning schemes, Similarity Surrogate Learning (SSL) and Dissimilarity Mixup Learning (DML), are devised to boost text classification by reusing and distinguishing samples of other labels following the label representation hierarchy, respectively. Experiments on authoritative public datasets and real-world medical records show that our approach stably achieves superior performances over classical and advanced imbalanced classification methods. Our code is available at https://github.com/jyansir/Text2Tree.

## introduction medical text classification is widely recognized as an urgent yet challenging problem due to its extremely imbalanced data distribution, large variety of rare labels @xcite , and complicated label relationship @xcite . various downstream clinical tasks have been derived from this problem, including icd coding @xcite @xcite and automated diagnosis @xcite , showcasing its potential values in modern clinical practice with machine learning approaches @xcite . there have been various classical strategies for general imbalanced classification and long-tailed multi-label classification @xcite in machine learning studies. such studies probably overly focused on data in rare categories (e.g., by re-sampling @xcite , re-weighting @xcite @xcite @xcite , ensemble learning @xcite , data augmentation @xcite @xcite ) to alleviate distribution bias, or employed sophisticated learning paradigms such as contrastive learning @xcite , transfer learning @xcite , and prompttuning @xcite ) to learn fine-grained representations of hard cases. but, in medical text classification, many rare diseases lack sufficient data in representation learning and the incidence rates vary greatly for different diseases, which cannot be completely solved by the general imbalanced classification approaches since they overlook the underlying dependency of medical terminologies. although the current studies of medical text classification either leveraged the label descriptions @xcite @xcite or label dependency @xcite for precise pairwise sample-label matching @xcite , or incorporated external knowledge sources (e.g., wikipedia) to enrich semantic information @xcite @xcite , they did not explicitly cope with the data imbalance and scarcity issues. therefore, imbalanced medical text classification is still an open challenge. targeting this problem, this paper makes the first effort on learning medical text representations to resolve the data imbalance and scarcity issues with the support of disease label dependency. as shown in fig. 1 , the right part demonstrates an example of icd-10-cm codes (a clinical modi-fication of icd-10 codes) about pneumonia and hypertension-related disease dependency in a tree structure. among all the types of pneumonia, fungal pneumonia stands out as an exceptionally rare and life-threatening variant @xcite that poses significant challenges to early detection @xcite . consequently, diagnosed cases of fungal pneumonia are notably scarce, and we aim to supplement our analysis with existing samples from the other diseases. an intuitive assumption is that two diseases with some common clinical presentations (e.g., symptom, disease course, treatment) can be used to compensate for each other. naturally, we can associate fungal pneumonia (j16.8) with covid-19 (j12.82), since both of them are respiratory system diseases (j00-j99) and pneumonia (j09-j18) while the diagnosed cases of covid-19 are plentiful. similarly, chlamydial pneumonia (j16.0) shares a closer parent node (j16) and can be a better supplementary source, while some more dissimilar ones (e.g., circulatory system diseases, i00-i99) are not so beneficial. utilizing the tree hierarchy, we are able to access prior knowledge of the similarity between disease types, which is helpful in dealing with the data imbalance and scarcity issues. motivated by this, we present a new frameworkagnostic algorithm that aligns text representation to label tree hierarchy (called text2tree), which is tailored for better classification by incorporating an icd code embedding tree to guide medical text representations. text2tree has three major components: (i) a hierarchy-aware label representation learning module (hlr), (ii) similarity surrogate learning (ssl), and (iii) dissimilarity mixup learning (dml) approaches. unlike general imbalanced classification methods that compensate a sample by reusing samples with the same disease labels or itself, (ii) and (iii) are designed to reuse samples from other labels under the specification of the label hierarchy learned by (i). in this paper, the used "codes" refer to disease labels with underlying hierarchy, and we do not distinguish "code", "label", and "node". the text2tree training procedure proceeds as follows. first, cascade attention modules are built based on the prior structure of the medical icd code tree. only embeddings of labels on adjacent tree layers are interacted in each module, and label representations are computed layer by layer. hence, a strict interaction fashion is kept to constrain in-formation flow along tree edges. next, pairwise label similarity is calculated based on the label representations. based on the similarity, samples from other labels will be treated as surrogate positive anchors to provide extra contrastive signal by ssl, or apply mixup to give new samples by dml. the ssl branch learns to gather more information from more similar labels to explicitly reuse texts in representation learning following the label hierarchy, while dml generates and classifies hard cases, adversarially preventing excessive manifold distortion resulted from ssl. our main contributions are as follows: text2tree algorithm stably outperforms advanced framework-agnostic imbalanced classification algorithms, without any external medical resource. ## methodology in this section, we separately describe the three key components of our proposed text2tree algorithm: hierarchy-aware label representation (hlr) learning module, similarity surrogate learning (ssl), and dissimilarity mixup learning (dml). ## conclusions in this paper, we proposed text2tree, a new framework-agnostic learning algorithm that aligns text representations to disease label hierarchy with two novel learning methods of contrastively reusing data, so as to resolve the data imbalance and scarcity issues in medical text classification. compared with various state-of-the-art general imbalanced text classification methods, the superiority of text2tree was verified on 5 real-world datasets. we believe that text2tree will serve as a strong baseline in imbalanced medical text classification and could be extended to other domains when corresponding prior label dependency is provided.