Functional-level Uncertainty Quantification for Calibrated Fine-tuning on LLMs

TL;DR

The paper introduces UQ4CT, a novel method for calibrating uncertainty in fine-tuned LLMs by focusing on functional space, leading to better confidence estimates and robustness under distribution shifts.

Contribution

It proposes a functional-level uncertainty calibration approach using a mixture-of-experts framework, improving calibration and generalization of fine-tuned LLMs.

Findings

UQ4CT reduces Expected Calibration Error by over 25% on multiple benchmarks.

UQ4CT maintains high accuracy while improving calibration under distribution shifts.

The method outperforms existing post hoc uncertainty estimation techniques.

Abstract

Accurate uncertainty quantification in large language models (LLMs) is essential for reliable confidence estimation, yet fine-tuned LLMs often become overconfident under limited adaptation data. Existing uncertainty methods for PEFT-based LLMs are largely post hoc, estimating uncertainty after fine-tuning rather than improving how adapters specialize to task-specific input-output relationships. We propose Functional-Level Uncertainty Quantification for Calibrated Fine-Tuning (UQ4CT), which calibrates uncertainty over the functional space induced by prompt-dependent mixtures of LoRA experts. UQ4CT implements this perspective through a mixture-of-experts fine-tuning framework, where a calibration loss aligns functional-level confidence with predictive correctness during training. Across four multiple-choice benchmarks and two open-ended generative QA tasks, UQ4CT reduces Expected Calibration Error (ECE) by over $25\%$ while preserving high accuracy. Under distribution shift, UQ4CT maintains superior calibration and competitive accuracy, demonstrating improved reliability and generalization for fine-tuned LLMs.