Probabilistic Calibration Is a Trainable Capability in Language Models
summary
The gist
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
In short
This study investigates whether fine-tuning language models can improve their ability to follow specific probability distributions when generating text, even for implicit constraints. Researchers compared two fine-tuning methods—soft-target and hard-target—using synthetic mathematical prompts. Both methods significantly improved structured sampling fidelity on unseen data and extended gains to broader random generation tasks, suggesting probabilistic calibration is a trainable skill.
Key concepts
- Probabilistic Calibration
- This refers to a language model's ability to adjust its output probabilities so that the generated text more accurately reflects a desired underlying probability distribution. In this study, the goal was to train models to better adhere to specific randomness constraints provided by users.
- Soft-Target Method
- This fine-tuning strategy involves taking a desired output distribution and converting it into targets for each individual next token. It treats the problem as supervised learning by assigning probabilities based on enumerating valid outputs, helping the model learn a smoother transition to the target distribution.
- Hard-Target Method
- This strategy trains the model using samples directly drawn from the target distribution. Instead of predicting specific tokens, it is trained on completions that are already known to belong to the desired set of answers, focusing on high-fidelity sampling from that specific set.
Terminology used across episodes
This episode discusses
- Probabilistic Calibration Is a Trainable Capability in Language Models · Paper Radio
- gpt-oss-120b & gpt-oss-20b Model Card
- Beyond Single-Point Judgment: Distribution Alignment for LLM-as-a-Judge
- The Llama 3 Herd of Models · Paper Radio
- The Illusion of Stochasticity in LLMs
- Generative AI and Creativity: A Systematic Literature Review and Meta-Analysis
- Evaluating the Quality of Randomness and Entropy in Tasks Supported by Large Language Models
- Understanding the Effects of RLHF on LLM Generalisation and Diversity
- Your Pre-trained LLM is Secretly an Unsupervised Confidence Calibrator
- Inducing Sustained Creativity and Diversity in Large Language Models
- String Seed of Thought: Prompting LLMs for Distribution-Faithful and Diverse Generation
- Creativity Has Left the Chat: The Price of Debiasing Language Models
- TabKD: Tabular Knowledge Distillation through Interaction Diversity of Learned Feature Bins
- tinyBenchmarks: evaluating LLMs with fewer examples
- Gemma 3 Technical Report
- What and Whose Knowledge? Measuring Epistemic Diversity in Large Language Models · Paper Radio
- Flipping Against All Odds: Reducing LLM Coin Flip Bias via Verbalized Rejection Sampling
- Qwen3 Technical Report
- NoveltyBench: Evaluating Language Models for Humanlike Diversity
- Large Language Models Are Bad Dice Players: LLMs Struggle to Generate Random Numbers from Statistical Distributions
The paper
Probabilistic Calibration Is a Trainable Capability in Language Models · Read on arXiv
Davide Baldelli, Sruthi Kuriakose, Maryam Hashemzadeh, Amal Zouaq, Sarath Chandar
Chandar Research Lab · Mila – Quebec AI Institute
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.
Transcript
Introduction to the show: ident: AI Radio. Generated commentary on the latest Artificial Intelligence papers.
Tom: I'm Tom, and with me are Jane, Lu, senior AI researcher at Tsinghua, Meng, lead engineer at a mysterious AI startup and Lalam, the in-house Large Language Model.
Jane: Today's paper: "Probabilistic Calibration Is a Trainable Capability in Language Models".
Tom: 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.
Jane: First, who's behind it and why it matters.
Title and authors: Tom: Let’s start with the title and who wrote this, "Probabilistic Calibration Is a Trainable Capability in Language Models," by Davide Baldelli, Sruthi Kuriakose, Maryam Hashemzadeh, Amal Zouaq, and Sarath Chandar.
Jane: It sounds very technical at first glance because of all the names and the concept of probabilistic calibration.
Lu: The core question they are asking is whether we can improve this capability directly through fine-tuning language models on synthetic prompts that require sampling from mathematical distributions.
Meng: So, it’s not just about adjusting weights randomly; it involves training on specific scenarios where the required output distribution is already defined.
Lalam: I see them comparing two different ways to do this training, which is what makes this paper more than just one idea.
Tom: Exactly. They compare a soft-target method and a hard-target method to see which one works better for teaching the model how to calibrate its outputs correctly.
The paper's summary: Tom: The summary explains that they study fine-tuning language models on synthetic prompts that demand sampling from mathematical distributions and compare two strategies: soft-target and hard-target.
Jane: The soft-target method involves enumerating all the valid numeric outputs for a prompt, assigning a target probability to each one, and then converting that into supervised probabilities for the next token at each prefix.
Lu: And the hard-target method is different; it draws many numeric answers from that target distribution and trains on them using standard next-token cross-entropy loss.
Meng: The main finding they present is that both methods substantially improve structured-sampling fidelity on held-out distribution families and unseen parameter settings across twelve models spanning four families.
Lalam: So, it’s not just a small win; both approaches show real improvement when tested on data and model settings the model hasn't seen before.
The paper's improvements: Tom: The paper highlights that the two methods show 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.
Jane: That means if you just want perfect numeric sampling, hard-target might be your best bet right now.
Lu: But the soft-target method actually shows better performance on things like open-ended random generation, multiple choice answer position balancing, and NoveltyBench.
Meng: The evaluation protocol checks six axes: structured distribution sampling, open-ended random generation, NoveltyBench utility, capability retention—how well the model keeps its general smarts—and PALOMA perplexity.
Lalam: And they found that soft-target fine-tuning gives the best aggregate PALOMA score across all Qwen and Gemma checkpoints for held-out text likelihood.
Conclusion: Tom: So to wrap up, this paper shows that language models can be trained to behave more like controlled probabilistic samplers, with gains that extend beyond the synthetic distributions they used for supervision.
Jane: That means we are not just flattening the probability distribution; we are actually teaching the model a new way to sample based on our constraints.
Lu: The implication is that this calibration is a trainable capability in language models, and these gains transfer beyond those specific mathematical prompts into more natural random-generation settings.
Meng: For practical engineering, it suggests we can select the right fine-tuning variant based on the task; hard-target for exact numeric sampling, soft-target for broader stochastic generation.
Lalam: I think this opens up a path where we can fine-tune models to be more reliable when used in applications that require precise control over randomness.
More episodes
- 2603.17917-Weight Clustering Reveals the Importance of Relative Rank in Large Language Models
- 2604.08003-Rethinking Entropy Allocation in LLM-based ASR: Understanding the Dynamics between Speech Encoders and LLMs
- 2604.08104-Quantum Vision Theory Applied to Audio Classification for Deepfake Speech Detection
- 2604.17633-Copy First, Translate Later: Interpreting Translation Dynamics in Multilingual Pretraining
- 2604.22503-How AI Writing Assistance Distorts Perceptions of AI Users and Their Opinions
- 2605.10268-MemReread: Enhancing Agentic Long-Context Reasoning via Memory-Guided Rereading
- 2606.05405-Agents' Last Exam
- 2606.15419-Let Large Language Models Judge Each Other: Multi-Agent Peer-Reviewed Reasoning for Medical Question Answering
- 2606.16215-PACT: Privileged Trace Co-Training for Multi-Turn Tool-Use Agents
- 2607.04986-The syntax of wh-agreement in Yemeni Ibbi Arabic