fix(deepconf): Vote and calibrate as in the DeepConf reference implementation

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

optillm/deepconf/confidence.py

@@ -133,18 +133,34 @@ 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.
 
         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]:
         """

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 = {

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