Uncertainty-Aware Expert Allocation for Efficient Multitask Fine-Tuning of Large Language Models

Large Language Models increasingly rely on Mixture-of-Experts architectures to scale model capacity while controlling computational cost. However, most MoE systems employ static expert routing strategies that allocate identical computational resources to all tokens, regardless of their uncertainty or difficulty. This paper proposes Confidence-Based Dynamic Routing (CBDR), an uncertainty-aware expert allocation mechanism for multitask fine-tuning of large language models. CBDR estimates token-level confidence from the entropy of the gating distribution and dynamically adjusts the number of activated experts accordingly. Tokens with high confidence are processed with fewer experts, while uncertain tokens receive additional expert capacity. To mitigate expert imbalance, we introduce Confidence-Based Importance Balancing (CBIB), a training objective that encourages balanced expert utilization. Experiments on multiple biomedical question-answering datasets demonstrate that CBDR improves model performance across multiple biomedical QA datasets, achieving absolute improvements of up to 3% in F1 score on the BioASQ benchmark while maintaining efficient expert utilization. These results suggest that uncertainty-aware routing provides a practical mechanism for adaptive compute allocation in sparse large language models.