Heuristic:Iamhankai Forest of Thought Self Correction Confidence Threshold
| Knowledge Sources | |
|---|---|
| Domains | Optimization, LLMs, Inference |
| Last Updated | 2026-02-14 03:30 GMT |
Overview
Dynamic self-correction strategy that re-generates LLM responses when average log-probability confidence drops below a threshold (default -0.5), keeping the better of the two attempts.
Description
The Forest-of-Thought framework implements a self-correction mechanism in the model pipeline. After initial generation, the system computes the average log-probability across all generated tokens as a confidence score. If this score falls below a configurable threshold (default: correct_threshold = -0.5), the model is prompted to "refine the answer according to Reflection or Feedback". The system then compares the confidence of the original and refined outputs, keeping whichever has higher average confidence.
This is a form of test-time compute scaling: instead of accepting the first output unconditionally, the model gets a second chance to improve low-confidence responses while preserving high-confidence ones.
Usage
Use this heuristic when running MCTS or CoT reasoning with the `--dynamic_self_correction` flag enabled. The threshold of -0.5 is the default; more negative thresholds trigger correction less frequently (accepting lower-confidence outputs), while less negative thresholds trigger correction more aggressively (at the cost of roughly doubling inference time for corrected samples).
The Insight (Rule of Thumb)
- Action: Enable `--dynamic_self_correction` and set `--correct_threshold -0.5` to activate confidence-based re-generation.
- Value: Default threshold is -0.5 (average log-probability). Correction triggers when confidence < threshold.
- Trade-off: Each correction roughly doubles inference time for that sample (two forward passes + one additional generation). The benefit is higher-quality answers for uncertain responses. High-confidence responses (above threshold) pass through unchanged with zero overhead.
- Decision Rule: The corrected answer replaces the original only if `new_average_confidence > average_confidence`.
Reasoning
LLM log-probabilities provide a cheap, model-intrinsic uncertainty signal. When a model is uncertain about its response (low average log-prob), the output is more likely to contain errors. By prompting the model to "refine" its answer with explicit instructions to verify and correct, the second attempt often benefits from the reasoning already present in the context. The confidence comparison ensures that correction never makes things worse — if the refined answer is less confident, the original is kept.
The threshold of -0.5 was chosen as a default that balances correction frequency against computational cost. Most well-formed LLM responses have average log-probabilities above -0.5, so correction is triggered selectively on the most uncertain outputs.
Code Evidence
Self-correction trigger from `models/load_local_model.py:L221-268`:
def self_correction(self, messages, generated_ids):
input_ids = generated_ids[:, :-1]
with torch.no_grad():
outputs = self.model(input_ids)
logits = outputs.logits
log_probabilities = torch.nn.functional.log_softmax(logits, dim=-1)
generated_log_probs = torch.gather(log_probabilities, 2,
generated_ids[:, 1:].unsqueeze(-1)).squeeze(-1)
average_confidence = generated_log_probs.mean().item()
if self.correction and average_confidence < self.correct_threshold:
# Re-prompt with reflection instructions
query = f'Please refine the your answer according to your Reflection or Feedback...'
# ... re-generate and compare confidence ...
if new_average_confidence > average_confidence:
average_confidence = new_average_confidence
generated_ids = new_generated_ids
return generated_ids, average_confidence
CLI argument for threshold from `run_with_mcf_stop_noearly.py:L613-614`:
args.add_argument("--dynamic_self_correction", action="store_true")
args.add_argument("--correct_threshold", type=float, default=-0.5, help="")