From 736941b2dbd45f4d19b223312e8ed0c856201c25 Mon Sep 17 00:00:00 2001 From: flyingdoubleG <53119724+flyingdoubleG@users.noreply.github.com> Date: Sat, 24 Oct 2026 08:31:44 +0800 Subject: [PATCH] fix(deepconf): Vote and calibrate as in the DeepConf reference implementation - Score each answer by the sum of its traces' confidences, DeepConf's confidence-weighted majority vote. The 0.7 * mean confidence + 0.3 * frequency mix let one trace outvote many agreeing ones, since token confidences are several units while the frequency term is at most 1. - Leave early-terminated traces out of the vote, and let warmup traces vote only if they pass the threshold. - Treat a trace shorter than the window as one group of all its tokens, instead of reporting 0 for its lowest and bottom-10% group confidence. - Default confidence_metric to lowest_group, the trace confidence the online threshold is calibrated on. --- optillm/deepconf/README.md | 2 +- optillm/deepconf/confidence.py | 27 ++++++++++++--- optillm/deepconf/processor.py | 36 ++++++++++++-------- tests/test_deepconf.py | 60 +++++++++++++++++++++++++++++++++- 4 files changed, 104 insertions(+), 21 deletions(-) diff --git a/optillm/deepconf/README.md b/optillm/deepconf/README.md index 2a98e396..1317129f 100644 --- a/optillm/deepconf/README.md +++ b/optillm/deepconf/README.md @@ -70,7 +70,7 @@ print(response.choices[0].message.content) | `top_k` | `5` | Number of top tokens for confidence calculation | | `min_trace_length` | `100` | Minimum tokens before allowing early termination | | `max_tokens_per_trace` | `4096` | Maximum tokens per individual trace | -| `confidence_metric` | `"average_confidence"` | Metric used for threshold calculation | +| `confidence_metric` | `"lowest_group"` | Trace confidence used for the threshold and as voting weight (`"lowest_group"`, `"bottom_10_percent"` or `"average_confidence"`) | | `include_stats` | `false` | Include processing statistics in response | ### Advanced Usage diff --git a/optillm/deepconf/confidence.py b/optillm/deepconf/confidence.py index 1f7be1fd..0f793ffe 100644 --- a/optillm/deepconf/confidence.py +++ b/optillm/deepconf/confidence.py @@ -133,6 +133,21 @@ def get_average_trace_confidence(self) -> float: return 0.0 return np.mean(self.token_confidences) + def _trace_group_confidences(self) -> List[float]: + """ + Group confidences of the whole trace. A trace shorter than the window + forms a single group of all its tokens, as in the DeepConf reference + implementation, rather than having no groups at all. + + Returns: + List of group confidences (empty only for an empty trace) + """ + if self.group_confidences: + return self.group_confidences + if self.token_confidences: + return [float(np.mean(self.token_confidences))] + return [] + def get_bottom_10_percent_confidence(self) -> float: """ Calculate average confidence of bottom 10% groups. @@ -140,11 +155,12 @@ def get_bottom_10_percent_confidence(self) -> float: Returns: Bottom 10% group confidence """ - if not self.group_confidences: + group_confidences = self._trace_group_confidences() + if not group_confidences: return 0.0 - num_bottom = max(1, len(self.group_confidences) // 10) - sorted_confidences = sorted(self.group_confidences) + num_bottom = max(1, len(group_confidences) // 10) + sorted_confidences = sorted(group_confidences) bottom_confidences = sorted_confidences[:num_bottom] return np.mean(bottom_confidences) @@ -156,9 +172,10 @@ def get_lowest_group_confidence(self) -> float: Returns: Lowest group confidence """ - if not self.group_confidences: + group_confidences = self._trace_group_confidences() + if not group_confidences: return 0.0 - return min(self.group_confidences) + return min(group_confidences) def get_trace_statistics(self) -> Dict[str, float]: """ diff --git a/optillm/deepconf/processor.py b/optillm/deepconf/processor.py index 14fc31fb..5c93a135 100644 --- a/optillm/deepconf/processor.py +++ b/optillm/deepconf/processor.py @@ -31,7 +31,7 @@ "min_trace_length": 100, # Minimum tokens before termination allowed "max_tokens_per_trace": 4096, # Maximum tokens per individual trace "temperature": 0.7, # Generation temperature - "confidence_metric": "average_confidence", # Metric for threshold calculation + "confidence_metric": "lowest_group", # Metric for threshold and voting weights "include_stats": False, # Include debugging statistics in response } @@ -242,6 +242,9 @@ def weighted_majority_vote(self, traces: List[TraceResult]) -> Tuple[str, Dict[s if not traces: return "", {} + # Early-terminated traces were cut off mid-reasoning, so they don't vote + traces = [trace for trace in traces if not trace.terminated_early] or traces + # Group traces by answer and calculate weighted scores answer_groups = defaultdict(list) for trace in traces: @@ -251,17 +254,14 @@ def weighted_majority_vote(self, traces: List[TraceResult]) -> Tuple[str, Dict[s answer = trace.text.strip()[-50:].strip() answer_groups[answer].append(trace) - # Calculate weighted scores for each answer - answer_scores = {} - for answer, group_traces in answer_groups.items(): - # Weight by average confidence - confidences = [trace.confidence_stats['average_confidence'] for trace in group_traces] - weighted_score = sum(confidences) / len(confidences) # Average confidence - count_weight = len(group_traces) / len(traces) # Frequency weight - - # Combine confidence and frequency - final_score = weighted_score * 0.7 + count_weight * 0.3 - answer_scores[answer] = final_score + # Each answer scores the sum of its traces' confidences (DeepConf's + # confidence-weighted majority vote), so agreement and confidence + # both count + metric = self.config["confidence_metric"] + answer_scores = { + answer: sum(trace.confidence_stats[metric] for trace in group_traces) + for answer, group_traces in answer_groups.items() + } # Select best answer best_answer = max(answer_scores.keys(), key=lambda x: answer_scores[x]) @@ -315,8 +315,16 @@ def process_online(self, messages: List[Dict[str, str]]) -> Tuple[str, Dict[str, f"(ratio: {consensus_ratio:.3f})") break - # Phase 3: Final answer selection - final_answer, voting_stats = self.weighted_majority_vote(all_traces) + # Phase 3: Final answer selection. As in the DeepConf reference + # implementation, warmup traces vote only if they pass the threshold. + metric = self.config["confidence_metric"] + voting_traces = [ + trace for trace in self.warmup_traces + if trace.confidence_stats[metric] >= self.confidence_threshold + ] + self.online_traces + final_answer, voting_stats = self.weighted_majority_vote( + voting_traces or all_traces + ) # Compile processing statistics processing_stats = { diff --git a/tests/test_deepconf.py b/tests/test_deepconf.py index cd96f12d..818b74ef 100644 --- a/tests/test_deepconf.py +++ b/tests/test_deepconf.py @@ -161,6 +161,61 @@ def test_info_function(): logger.error(f"✗ Info function test failed: {e}") return False +def _trace(text, confidence, terminated_early=False): + from optillm.deepconf.processor import TraceResult + + stats = { + "average_confidence": confidence, + "bottom_10_percent": confidence, + "lowest_group": confidence, + } + trace = TraceResult([], text, stats) + trace.terminated_early = terminated_early + return trace + +def _processor(): + from unittest.mock import Mock + from optillm.deepconf.processor import DeepConfProcessor + + return DeepConfProcessor(Mock(), Mock()) + +def test_short_trace_group_confidence(): + """A trace shorter than the window is one group of all its tokens.""" + import torch + from optillm.deepconf.confidence import ConfidenceCalculator + + calculator = ConfidenceCalculator(window_size=2048, top_k=5) + for _ in range(50): + calculator.add_token_confidence(torch.randn(1000)) + stats = calculator.get_trace_statistics() + + assert stats["num_groups"] == 0 + assert stats["lowest_group"] == stats["average_confidence"] > 0 + assert stats["bottom_10_percent"] == stats["average_confidence"] + return True + +def test_weighted_vote_sums_confidences(): + """Each answer scores the sum of its traces' confidences, so many agreeing + traces outvote one slightly more confident trace.""" + traces = [_trace("The answer is 42", 9.0) for _ in range(15)] + traces.append(_trace("The answer is 7", 9.5)) + + answer, stats = _processor().weighted_majority_vote(traces) + + assert answer == "The answer is 42" + assert abs(stats["best_score"] - 15 * 9.0) < 1e-9 + return True + +def test_weighted_vote_skips_early_terminated_traces(): + """Traces cut off by early termination are unfinished and don't vote.""" + traces = [_trace("The answer is 42", 9.0) for _ in range(3)] + traces += [_trace("The answer is 7 because", 9.0, True) for _ in range(4)] + + answer, _ = _processor().weighted_majority_vote(traces) + + assert answer == "The answer is 42" + return True + def main(): """Run all tests.""" logger.info("Starting DeepConf test suite...") @@ -170,7 +225,10 @@ def main(): test_confidence_calculator, test_threshold_calibrator, test_config_validation, - test_info_function + test_info_function, + test_short_trace_group_confidence, + test_weighted_vote_sums_confidences, + test_weighted_vote_skips_early_terminated_traces, ] passed = 0