🤖 machine learning

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

The paper proposes UQ4CT, a mixture-of-experts fine-tuning framework that integrates a calibration loss over functional space to significantly reduce Expected Calibration Error and improve the reliability of PEFT-based LLMs under both standard and shifted distributions.

Original authors: Ruijia Niu, Dongxia Wu, Rose Yu, Yi-An Ma

Published 2026-05-15

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Ruijia Niu, Dongxia Wu, Rose Yu, Yi-An Ma

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a brilliant, well-read librarian (the Large Language Model, or LLM). This librarian has read almost everything in the world, but when you ask them a specific, tricky question, they sometimes guess with 100% confidence even when they are wrong. This is called overconfidence, and it's dangerous because users might trust a wrong answer just because the librarian sounds sure.

The paper introduces a new method called UQ4CT to fix this. Think of it as a way to teach the librarian to say, "I'm pretty sure," only when they actually are sure, and "I'm not sure," when they are guessing.

Here is how it works, broken down with simple analogies:

1. The Problem: The "One-Size-Fits-All" Tuning

Usually, when we want a librarian to specialize in a new topic (like climate science), we "fine-tune" them.

The Old Way: Imagine giving the librarian a single, new pair of glasses (called LoRA) to help them see the new topic better. But if the librarian doesn't have enough practice books (data), they might put on these glasses and become too confident, even if the glasses are slightly blurry. They can't tell the difference between "I know this" and "I'm guessing."
The Issue: Existing methods try to check the librarian's confidence after they put on the glasses. But the paper argues we need to fix the glasses while the librarian is learning.

2. The Solution: The "Expert Panel" (Mixture of Experts)

Instead of giving the librarian just one pair of glasses, UQ4CT gives them a panel of 8 different experts (a "Mixture of Experts" or MoE).

The Setup: Imagine a room with 8 different specialists. When you ask a question, a "Router" (like a smart receptionist) looks at your question and decides which 2 specialists are best suited to answer it.
The Magic: Because there are multiple experts, the system can see if the experts agree or disagree. If the Router picks two experts who give very different answers, the system knows there is uncertainty. If the Router picks two experts who agree perfectly, the system knows there is confidence.

3. The "Calibration Loss": The Truth-Telling Coach

The paper's biggest innovation is a special "coach" that trains the Router during the learning process.

The Coach's Rule: The coach watches the experts.
- If the experts pick the correct answer, the coach tells the Router: "Great job! Be very confident in this choice."
- If the experts pick the wrong answer, the coach tells the Router: "You were too sure of yourself! You should have been less confident or picked different experts."
The Result: Over time, the Router learns to match its "confidence level" with the actual "truth level." It stops pretending to be an expert when it's just guessing.

4. Why This is Better Than Other Methods

No Slow-Down: Some other methods try to ask the librarian the same question 10 times to see if they get different answers (like asking a friend for advice 10 times). This is slow and expensive. UQ4CT does the job in one single pass because the "panel of experts" is built right into the system.
Better Generalization: The paper tested this on questions about common sense and specific fields like climate science. Even when the librarian was asked about topics they hadn't seen before (a "distribution shift"), UQ4CT kept its cool. It didn't get overconfident on new, weird questions; it correctly signaled that it was unsure.

The Bottom Line

The paper claims that by using this "Expert Panel" approach and training the system to align its confidence with the truth during the learning phase, they reduced the "Expected Calibration Error" (a measure of how wrong the confidence is) by over 25%.

Crucially, they did this without making the librarian slower or less accurate at answering questions. The librarian still gets the right answers, but now they are much better at knowing when they are right and when they are just guessing.

In short: UQ4CT turns a confident-but-often-wrong AI into a humble-but-accurate AI that knows the limits of its own knowledge.

Technical Summary: Functional-Level Uncertainty Quantification for Calibrated Fine-Tuning on LLMs

Problem Statement
Accurate uncertainty quantification (UQ) is critical for the reliable deployment of Large Language Models (LLMs), particularly in safety-critical domains like healthcare and scientific reasoning. However, fine-tuned LLMs frequently exhibit overconfidence when adapted to tasks with limited data, failing to reflect the true credibility of their outputs. Existing UQ methods for Parameter-Efficient Fine-Tuning (PEFT), such as those based on Low-Rank Adaptation (LoRA), are largely post hoc. They estimate uncertainty after the model has been trained by analyzing the parameter space (e.g., via ensembles or Bayesian approximations) or by perturbing prompts. These approaches often fail to address the root cause of overconfidence: the model's inability to capture fine-tuning-induced uncertainty and generalization gaps during the adaptation process itself. Furthermore, methods relying on multiple forward passes or extensive sampling introduce significant computational overhead, limiting scalability.

Methodology: UQ4CT
The authors propose Functional-Level Uncertainty Quantification for Calibrated Fine-Tuning (UQ4CT), a framework that shifts the focus from parameter-space to functional-space uncertainty. Instead of analyzing learned parameters after training, UQ4CT calibrates uncertainty over the functional space induced by prompt-dependent mixtures of LoRA experts during the fine-tuning process.

The methodology consists of three core components:

Hierarchical Functional Decomposition via MixLoRA:
To avoid the combinatorial explosion of modeling all possible expert combinations across layers, UQ4CT employs a hierarchical decomposition. It utilizes a Mixture-of-Experts (MoE) architecture where, at each layer, a sparse router dynamically selects a subset of LoRA experts based on the input hidden state. This constructs a flexible functional space where the model's output is a weighted sum of diverse basis functions (LoRA experts) conditioned on the input. The functional space is represented as a recursive composition of layer-wise mixtures, allowing the model to adaptively compose computation paths.
Functional-Level Confidence (FLC) Estimation:
The model quantifies uncertainty by analyzing the concentration of the expert selection. Specifically, the Functional-Level Confidence (FLC) is derived from the routing weights of the top- $k$ experts selected by the router at each layer. High concentration in the routing weights (i.e., the router is confident in selecting specific experts) corresponds to high functional confidence, while dispersed weights indicate higher uncertainty. This provides a principled, internal signal of the model's certainty without requiring multiple sampling passes.
Calibration Loss ( $L_{cal}$ ):
The core innovation is a novel calibration loss that aligns the functional-level confidence with predictive correctness during training. The loss is defined as:
$L_{cal} = (1\{\text{MixLoRA}(x) = y^*\} - \text{FLC}(x))^2$
where $1\{\cdot\}$ is an indicator function for correct predictions. This objective forces the FLC to approximate the probability of the model's prediction being correct. By minimizing this loss alongside the standard cross-entropy and load-balancing losses, the model learns to assign high confidence only when it is likely to be correct and higher uncertainty when it is not. This encourages the router to explore alternative expert mixtures when uncertainty is high, thereby improving generalization.

Key Contributions

Functional-Level Perspective: The paper introduces a shift from parameter-space to functional-space uncertainty quantification, explicitly calibrating the model's confidence over the input-output mappings it can realize during fine-tuning.
Calibration-Aware Fine-Tuning: A new calibration loss is proposed that dynamically aligns prompt-dependent LoRA mixture weights with predictive correctness, mitigating overconfidence without degrading accuracy.
Efficient Architecture: The method leverages a hierarchical MixLoRA architecture that captures diverse functional relationships with minimal computational overhead, requiring only a single forward pass during inference.
Theoretical Guarantee: The authors provide a proof that optimizing the calibration risk over the data distribution yields an optimal FLC that corresponds to the true probability of predictive correctness.

Experimental Results
The authors evaluated UQ4CT on four multiple-choice benchmarks (ARC-E, ARC-C, OBQA, ClimateQA) and two open-ended generative QA tasks (TruthfulQA, TriviaQA) using Llama-3.1-8B and Mistral-7B.

In-Distribution Performance: UQ4CT reduced the Expected Calibration Error (ECE) by over 25% across benchmarks compared to strong baselines (including LoRA, MC Dropout, Deep Ensembles, and Bayesian LoRA), while maintaining competitive or superior accuracy. Notably, UQ4CT achieved these results with a single forward pass, whereas methods like BLoB (N=10) required ten passes yet still yielded higher ECE.
Distribution Shift: Under distribution shifts (tested on MMLU subtasks in CS, Engineering, Law, and Health), UQ4CT maintained superior calibration and competitive accuracy, demonstrating improved robustness to out-of-distribution inputs compared to baselines.
Generative Tasks: In open-ended generation, UQ4CT achieved the lowest ECE on TruthfulQA and TriviaQA, suggesting the functional-space calibration perspective extends beyond classification-style settings.
Efficiency: Memory analysis indicates that UQ4CT's peak GPU memory usage is only ~7.5% higher than standard LoRA, making it computationally feasible for large-scale deployment.

Significance and Claims
The paper claims that UQ4CT addresses a fundamental limitation in current fine-tuning practices: the lack of uncertainty awareness during the adaptation process. By integrating uncertainty calibration directly into the fine-tuning objective via functional-level signals, the method ensures that the model's confidence better reflects its actual predictive reliability. The authors position this work as a foundational step toward trustworthy fine-tuned LLMs, particularly in data-scarce environments where overconfidence poses significant risks. They modestly note that while the approach extends to open-ended generation using F1-based proxies, scaling to long-form reasoning or multi-step generation remains a direction for future work.