Probabilistic Calibration Is a Trainable Capability in Language Models

The gist Probabilistic calibration is a trainable capability in language models; this work studies whether fine-tuning can improve language models to better calibrate their probabilistic outputs to user-specified randomness constraints.

Calibration Fine-Tuning Strategies

The study compares two Calibration Fine-Tuning variants: a soft-target method and a hard-target method, both applied to synthetic prompts requiring sampling from mathematical distributions. The first strategy is soft-target fine-tuning, which involves enumerating valid numeric outputs for a prompt, assigning each its target probability, and converting this to supervised probabilities for the next token at each prefix. The second strategy is hard-target fine-tuning, which draws many numeric answers from the target distribution and trains on them with standard next-token cross-entropy.

Empirical Performance Across Variants

Across 12 models spanning four families, both methods substantially improve structured-sampling fidelity on held-out distribution families and unseen parameter settings. The two methods exhibit different empirical profiles: hard-target fine-tuning is often strongest on structured numeric sampling itself, while soft-target fine-tuning performs better on broader stochastic generation benchmarks, including open-ended random generation, multiple-choice answer-position balancing, and NoveltyBench.

Evaluation Metrics and Transfer

The evaluation protocol focuses on six axes: structured distribution sampling, open-ended random generation, NoveltyBench, MCQ generation with uniform answer positions, capability retention, and PALOMA perplexity. Structured distribution sampling shows that Calibration Fine-Tuning strongly improves structured distribution sampling on both held-out families and unseen parameter settings. The gains transfer beyond the synthetic training domain, with soft-target fine-tuning often performing better on broader stochastic generation benchmarks such as open-ended random generation and NoveltyBench.

Model and Configuration Tradeoffs

The results show that the selected soft-target configuration gives the most reliable language-space transfer, while the selected hard-target configuration remains strongest on exact numeric sampling. Capability retention is mixed, with the clearest systematic cost on GSM8K and aggregate gp-IRT often still favoring the base checkpoint.

Conclusion

Overall, our results show that language models can be trained to behave more like controlled probabilistic samplers, with gains that extend beyond the synthetic distributions used for supervision. Future work should study the mechanisms learned by each objective and develop retention-aware variants of Calibration Fine-Tuning. The overall findings demonstrate that probabilistic calibration is a trainable capability in language models >.

--- Page 43 ---

The gist Probabilistic calibration is a trainable capability in language models; this work studies whether fine-tuning can improve language models to better calibrate their probabilistic outputs to user-specified randomness constraints.

Conclusion

Overall, our results show that language models can be trained to behave more like controlled probabilistic samplers, with gains that extend beyond the synthetic distributions used for supervision.

1. textbf Real-time Probabilistic Calibration for Structured Sampling (Hard-Target Method): This improved system can be trained to teach substantially better stochastic fidelity than the base checkpoint, specifically being strongest on structured numeric sampling itself as seen in Figure 2 and Table 6, resulting in a Family-median normalized W1 drops sharply.

2. textbf Broader Stochastic Transfer via Soft-Target Calibration: The system can be configured with the soft-target method to broaden stochastic support, which is often better on broader stochastic generation settings and more often yields favorable perplexity diagnostics, making it superior for open-ended random generation and NoveltyBench utility, as noted in Table 10.

3. textbf Model Family-Specific Calibration Selection: The system can dynamically select the fine-tuning variant based on the task; specifically, hard-target fine-tuning is often strongest on structured numeric sampling itself, while soft-target fine-tuning performs better on broader stochastic generation settings, allowing for optimized performance across different use cases.

4. textbf Robustness to Unseen Parameters: The system demonstrates improved generalization by showing that both methods substantially improve structured-sampling fidelity on held-out distribution families and unseen parameter settings, with the hard-target method being particularly strong in this regard, as evidenced by Table 7 and Figure 2.

5. textbf Preserved Language Model Fit via PALOMA Metrics: The system can achieve improved language model fit on held-out text likelihood, as soft-target fine-tuning gives the best aggregate PALOMA score for all Qwen and Gemma checkpoints, showing that calibration is not merely a flattening of token probabilities.

Language models are increasingly used in settings where outputs should follow an intended distribution, yet their generation probabilities are often poorly calibrated to it. In many such requests the distribution is implicit and the valid outputs cannot be listed in advance, so the task cannot be passed to an external random number generator. We study whether this capability can be improved directly through fine-tuning. Concretely, we fine-tune language models on synthetic prompts that require sampling from mathematical distributions, and compare two Calibration Fine-Tuning variants: a soft-target method that converts the desired output distribution into trie-derived next-token targets, and a hard-target method that trains on sampled completions from the same target distribution. Across 12 models spanning four families, both methods substantially improve structured-sampling fidelity on held-out distribution families and unseen parameter settings, showing that training on cheap synthetic mathematical distributions improves sampling on unseen distributions, even in natural language. The gains also transfer to broader stochastic generation benchmarks, including open-ended random generation, where fine-tuned models broaden their first-token support by one to two orders of magnitude, and NoveltyBench. The gains sometimes reduce downstream capability, especially arithmetic reasoning, with costs varying by model. Overall, our results show that probabilistic calibration can be improved through fine-tuning with either objective, and that calibration learned on mathematical distributions transfers to natural-language stochastic generation. Code is available at https://github.com/chandar-lab/calibration-finetuning.