Deng-Bao Wang

Author directory

2026

Token probabilities play a central role in large language models. However, these probabilities often fail to reflect either answer correctness or probability distributions explicitly specified in the context. We attribute this failure to the **Context-Logit Probability Gap**: contextual probabilities are not reliably reflected in the model’s next-token logits. We show that post-hoc temperature scaling has a limited effect on addressing this gap, because it rescales logits without correcting erroneous probability rankings. Motivated by this, we propose Bag-of-Tokens (BoT) tuning, a very simple fine-tuning approach that directly calibrates the logits with the probability distribution conveyed by synthetic contexts. More specifically, we construct BoT examples by forming a bag of candidate tokens with an explicitly assigned probability distribution. This distribution is expressed in the context and used as soft supervision to align the model’s logit probabilities via KL-divergence minimization. Experiments show that BoT tuning generalizes to LLM-based sampling tasks, such as random-number generation, and improves downstream QA accuracy and calibration. These results suggest that calibrating LLM logits with contextual probabilities provides a lightweight and effective way to enhance their uncertainty awareness.