Overcoming BERT’s limitations in uncertainty: A novel two-stage solution for multi-class medical text classification
Prabhashrini Dhanushika Manage, Yutong Wu, Jinglan Zhang, Thennakoon Mudiyanselage Anupama Udayangani Gunathilaka, Yuefeng Li · Knowledge-Based Systems · 2026
• Introduced a principled definition of an uncertain decision boundary for BERT-based classifiers using a discriminative confidence margin, enabling systematic identification of ambiguous predictions. • Integrated word–class probability embeddings (WCPE) to enrich document representations with explicit class-discriminative signals, particularly effective for uncertain and noisy medical texts. • Proposed a two-stage, uncertainty-aware decision framework that routes confident samples to a baseline BERT model and uncertain samples to a WCPE-enhanced variant, improving robustness under ambiguity. • Developed a data-driven strategy for optimizing the uncertainty threshold (UncertainT), balancing coverage of ambiguous samples and overall classification accuracy. • Established the general applicability of the proposed uncertainty-aware decision framework across multiple datasets and pretrained backbones, while revealing that statistically significant improvements emerge selectively depending on dataset ambiguity and backbone pretraining. Bidirectional Encoder Representations from Transformers (BERT) has achieved state-of-the-art performance in Natural Language Processing (NLP) tasks but struggles with uncertainty management in multi-class medical text classification, where overlapping categories and domain-specific critical terms pose challenges. BERT’s attention mechanism may distribute focus across irrelevant contextual patterns, reducing its ability to prioritize medically significant terms. To address these limitations, we propose a two-stage decision-making framework incorporating an “uncertain boundary” mechanism to separate high-confidence (“certain”) and low-confidence (“uncertain”) cases based on an optimized uncertainty threshold. A baseline BERT model processes high-confidence cases, while uncertain cases are handled by an enhanced BERT model integrating Word-Class Probabilistic Embedding (WCPE) to improve domain-specific representation learning. We evaluate our framework across three datasets-Ohsumed, PubMed, and Drug Review and further examine its robustness under multiple pretrained language model settings. Using standard BERT as the primary backbone, results show consistent accuracy improvements from 0.770 to 0.773 on Ohsumed ( p < 0.05), 0.973 to 0.975 on PubMed ( p < 0.1), and 0.596 to 0.609 on Drug Review ( p < 0.05) with stronger gains on uncertain subsets. These findings show that confidence-based partitioning and model specialization enhance BERT’s classification performance while effectively handling uncertainty in medical NLP.