""" Uncertainty Quantification Module for LLMs This module implements various uncertainty quantification methods for large language models, including softmax confidence, Monte Carlo dropout, ensemble disagreement, and calibration metrics. """ import numpy as np import torch from typing import List, Dict, Any, Union, Optional from scipy.special import softmax import torch.nn.functional as F class UncertaintyQuantifier: """Base class for uncertainty quantification methods.""" def __init__(self, name: str): """ Initialize the uncertainty quantifier. Args: name: Name of the uncertainty quantification method """ self.name = name def quantify(self, model_outputs: Dict[str, Any]) -> Dict[str, float]: """ Quantify uncertainty in model outputs. Args: model_outputs: Outputs from the LLM interface Returns: Dictionary of uncertainty metrics """ raise NotImplementedError("Subclasses must implement this method") class SoftmaxConfidence(UncertaintyQuantifier): """Uncertainty quantification based on softmax confidence scores.""" def __init__(self): """Initialize the softmax confidence quantifier.""" super().__init__("softmax_confidence") def quantify(self, model_outputs: Dict[str, Any]) -> Dict[str, float]: """ Quantify uncertainty using softmax confidence scores. Args: model_outputs: Outputs from the LLM interface, must include logits Returns: Dictionary of uncertainty metrics: - mean_confidence: Average confidence across tokens - min_confidence: Minimum confidence across tokens - entropy: Average entropy of token distributions """ if "logits" not in model_outputs: raise ValueError("Model outputs must include logits for softmax confidence") logits = model_outputs["logits"][0] # Use first sample's logits # Calculate softmax probabilities and confidence metrics confidences = [] entropies = [] for token_logits in logits: probs = softmax(token_logits, axis=-1) max_prob = np.max(probs) confidences.append(max_prob) # Calculate entropy of the probability distribution entropy = -np.sum(probs * np.log(probs + 1e-10)) entropies.append(entropy) return { "mean_confidence": float(np.mean(confidences)), "min_confidence": float(np.min(confidences)), "entropy": float(np.mean(entropies)) } class MonteCarloDropout(UncertaintyQuantifier): """Uncertainty quantification based on Monte Carlo dropout sampling.""" def __init__(self): """Initialize the Monte Carlo dropout quantifier.""" super().__init__("mc_dropout") def quantify(self, model_outputs: Dict[str, Any]) -> Dict[str, float]: """ Quantify uncertainty using Monte Carlo dropout sampling. Args: model_outputs: Outputs from the LLM interface, must include multiple samples Returns: Dictionary of uncertainty metrics: - sample_variance: Variance across different samples - sample_diversity: Lexical diversity across samples """ if "samples" not in model_outputs or len(model_outputs["samples"]) <= 1: raise ValueError("Model outputs must include multiple samples for MC dropout") samples = model_outputs["samples"] # Calculate sample diversity using token overlap from nltk.tokenize import word_tokenize try: tokenized_samples = [set(word_tokenize(sample.lower())) for sample in samples] except: # Fallback to simple whitespace tokenization if nltk is not available tokenized_samples = [set(sample.lower().split()) for sample in samples] # Calculate Jaccard similarity between all pairs of samples similarities = [] for i in range(len(tokenized_samples)): for j in range(i+1, len(tokenized_samples)): intersection = len(tokenized_samples[i].intersection(tokenized_samples[j])) union = len(tokenized_samples[i].union(tokenized_samples[j])) if union > 0: similarities.append(intersection / union) else: similarities.append(1.0) # Empty sets are considered identical # Convert similarity to diversity (1 - similarity) diversity = 1.0 - np.mean(similarities) if similarities else 0.0 # Calculate variance in sample lengths as another diversity metric sample_lengths = [len(sample) for sample in samples] length_variance = np.var(sample_lengths) if len(sample_lengths) > 1 else 0.0 return { "sample_diversity": float(diversity), "length_variance": float(length_variance), "num_samples": len(samples) } class EnsembleDisagreement(UncertaintyQuantifier): """Uncertainty quantification based on ensemble disagreement.""" def __init__(self): """Initialize the ensemble disagreement quantifier.""" super().__init__("ensemble_disagreement") def quantify(self, ensemble_outputs: List[Dict[str, Any]]) -> Dict[str, float]: """ Quantify uncertainty using ensemble disagreement. Args: ensemble_outputs: List of outputs from different models Returns: Dictionary of uncertainty metrics: - response_diversity: Lexical diversity across model responses - confidence_variance: Variance in confidence scores across models """ if not ensemble_outputs or len(ensemble_outputs) <= 1: raise ValueError("Ensemble outputs must include results from multiple models") # Extract primary responses from each model responses = [output["response"] for output in ensemble_outputs] # Calculate response diversity using token overlap (similar to MC dropout) from nltk.tokenize import word_tokenize try: tokenized_responses = [set(word_tokenize(response.lower())) for response in responses] except: # Fallback to simple whitespace tokenization if nltk is not available tokenized_responses = [set(response.lower().split()) for response in responses] # Calculate Jaccard similarity between all pairs of responses similarities = [] for i in range(len(tokenized_responses)): for j in range(i+1, len(tokenized_responses)): intersection = len(tokenized_responses[i].intersection(tokenized_responses[j])) union = len(tokenized_responses[i].union(tokenized_responses[j])) if union > 0: similarities.append(intersection / union) else: similarities.append(1.0) # Empty sets are considered identical # Convert similarity to diversity (1 - similarity) diversity = 1.0 - np.mean(similarities) if similarities else 0.0 # Extract confidence scores if available confidences = [] for output in ensemble_outputs: if "mean_confidence" in output: confidences.append(output["mean_confidence"]) # Calculate variance in confidence scores confidence_variance = np.var(confidences) if len(confidences) > 1 else 0.0 return { "response_diversity": float(diversity), "confidence_variance": float(confidence_variance), "num_models": len(ensemble_outputs) } class CalibrationMetrics(UncertaintyQuantifier): """Uncertainty quantification based on calibration metrics.""" def __init__(self): """Initialize the calibration metrics quantifier.""" super().__init__("calibration_metrics") def expected_calibration_error( self, confidences: List[float], accuracies: List[bool], num_bins: int = 10 ) -> float: """ Calculate Expected Calibration Error (ECE). Args: confidences: List of confidence scores accuracies: List of boolean accuracy indicators num_bins: Number of bins for binning confidences Returns: Expected Calibration Error """ if len(confidences) != len(accuracies): raise ValueError("Confidences and accuracies must have the same length") if not confidences: return 0.0 # Create bins and calculate ECE bin_indices = np.digitize(confidences, np.linspace(0, 1, num_bins)) ece = 0.0 for bin_idx in range(1, num_bins + 1): bin_mask = (bin_indices == bin_idx) if np.any(bin_mask): bin_confidences = np.array(confidences)[bin_mask] bin_accuracies = np.array(accuracies)[bin_mask] bin_confidence = np.mean(bin_confidences) bin_accuracy = np.mean(bin_accuracies) bin_size = np.sum(bin_mask) # Weighted absolute difference between confidence and accuracy ece += (bin_size / len(confidences)) * np.abs(bin_confidence - bin_accuracy) return float(ece) def maximum_calibration_error( self, confidences: List[float], accuracies: List[bool], num_bins: int = 10 ) -> float: """ Calculate Maximum Calibration Error (MCE). Args: confidences: List of confidence scores accuracies: List of boolean accuracy indicators num_bins: Number of bins for binning confidences Returns: Maximum Calibration Error """ if len(confidences) != len(accuracies): raise ValueError("Confidences and accuracies must have the same length") if not confidences: return 0.0 # Create bins and calculate MCE bin_indices = np.digitize(confidences, np.linspace(0, 1, num_bins)) max_ce = 0.0 for bin_idx in range(1, num_bins + 1): bin_mask = (bin_indices == bin_idx) if np.any(bin_mask): bin_confidences = np.array(confidences)[bin_mask] bin_accuracies = np.array(accuracies)[bin_mask] bin_confidence = np.mean(bin_confidences) bin_accuracy = np.mean(bin_accuracies) # Absolute difference between confidence and accuracy ce = np.abs(bin_confidence - bin_accuracy) max_ce = max(max_ce, ce) return float(max_ce) def quantify( self, confidences: List[float], accuracies: List[bool] ) -> Dict[str, float]: """ Quantify uncertainty using calibration metrics. Args: confidences: List of confidence scores accuracies: List of boolean accuracy indicators Returns: Dictionary of calibration metrics: - ece: Expected Calibration Error - mce: Maximum Calibration Error """ return { "ece": self.expected_calibration_error(confidences, accuracies), "mce": self.maximum_calibration_error(confidences, accuracies) } # Factory function to create uncertainty quantifiers def create_uncertainty_quantifier(method: str) -> UncertaintyQuantifier: """ Create an uncertainty quantifier based on the specified method. Args: method: Name of the uncertainty quantification method Returns: Uncertainty quantifier instance """ if method == "softmax_confidence": return SoftmaxConfidence() elif method == "mc_dropout": return MonteCarloDropout() elif method == "ensemble_disagreement": return EnsembleDisagreement() elif method == "calibration_metrics": return CalibrationMetrics() else: raise ValueError(f"Unsupported uncertainty quantification method: {method}")