Technical Summary: Probabilistic Calibration Is a Trainable Capability in Language Models

Problem Statement

Language models (LMs) are increasingly deployed in scenarios requiring user-specified randomness, such as generating diverse outputs, balancing choice probabilities, or sampling from specific mathematical distributions. However, standard autoregressive LMs are poor native stochastic generators. Their induced output distributions often collapse onto narrow subsets, violate support constraints, and fail to align with requested target distributions. Existing prompting workarounds (e.g., injecting reasoning or seeds) can partially mitigate these issues but introduce inference-time costs, prompt sensitivity, and computational overhead without fundamentally altering the model's sampling behavior.

The central question addressed is whether probabilistic calibration—the ability of an LM to generate samples reflecting a known target distribution—can be learned directly through fine-tuning, rather than relying on inference-time heuristics.

Methodology: Calibration Fine-Tuning

The authors propose a framework to train LMs to act as calibrated probabilistic samplers using synthetic prompts that explicitly request sampling from mathematical distributions (e.g., "Generate exactly ONE random number from a [distribution] distribution with parameters [params]").

Two distinct fine-tuning strategies are compared, both utilizing Low-Rank Adaptation (LoRA) on frozen base models:

1. Soft-Target Fine-Tuning:

   • Mechanism: The target distribution is discretized into a finite canonical output space. A prefix trie is constructed from all valid numeric outputs, where each node represents a token prefix.
   • Supervision: For any given prefix, the method computes a target next-token distribution by normalizing the remaining probability mass of all valid continuations in the trie.
   • Objective: The model minimizes the Kullback-Leibler (KL) divergence between its predicted next-token distribution and the trie-derived target distribution at every visited prefix. This provides dense, prefix-level supervision.

2. Hard-Target Fine-Tuning:

   • Mechanism: The model is trained on completions sampled directly from the target distribution.
   • Supervision: For each prompt, a numeric answer is sampled from the target distribution, tokenized, and used as the ground-truth completion.
   • Objective: The model minimizes standard masked autoregressive cross-entropy on the sampled completion. This provides sparse, path-level supervision, requiring more training steps (more samples per prompt) to achieve comparable coverage.

Experimental Setup:

• Models: 12 models across four families (Qwen3, Gemma-3, Llama-3.2, GPT-OSS) ranging from 0.6B to 27B parameters.
• Data: 30 distribution families (24 seen for training, 6 held-out for Out-of-Distribution (OOD) testing).
• Evaluation:
  • Structured Sampling: Measured via Wasserstein-1 distance (W1) and logit-level KL divergence against target distributions.
  • Transfer Benchmarks: Open-ended random generation (diversity support), MCQ answer-position balancing, and NoveltyBench (semantic diversity).
  • Retention: Capability retention on TinyBenchmarks (MMLU, GSM8K, etc.) and language modeling fit on PALOMA.

Key Contributions

1. Trainability of Calibration: The paper demonstrates that probabilistic calibration is a trainable capability. Both soft-target and hard-target fine-tuning substantially improve structured-sampling fidelity on held-out distribution families and unseen parameter settings.
2. Comparative Analysis of Strategies:
   • Hard-Target: Generally achieves superior performance on the in-domain structured numeric sampling benchmark, particularly for unseen parameters. It excels at exact numeric fidelity.
   • Soft-Target: While competitive on numeric tasks, it demonstrates superior transfer to broader stochastic generation settings. It consistently broadens stochastic support in open-ended generation, improves answer-position balancing in MCQs, and yields higher utility on NoveltyBench.
3. Trade-offs and Costs: The study highlights that gains in calibration often come with trade-offs. Specifically, arithmetic reasoning (GSM8K) tends to regress, with costs varying by model size and fine-tuning variant. However, language modeling fit (PALOMA perplexity) is often preserved or improved, suggesting the models do not merely flatten token probabilities but learn specific stochastic behaviors.

Results

• Structured Sampling: Both methods significantly reduce normalized W1 and logit KL compared to base checkpoints. Hard-target fine-tuning is typically the strongest performer on numeric benchmarks, reducing errors by an order of magnitude in many cases.
• Open-Ended Generation: Soft-target fine-tuning consistently increases the top-90% support size (the number of tokens required to cover 90% of probability mass) across all models, often by one to two orders of magnitude. This leads to more diverse and less concentrated outputs in categories like cities, animals, and names.
• MCQ and Novelty: Soft-target configurations generally yield better answer-position balance in MCQ generation and higher utility scores on NoveltyBench, indicating that the learned stochastic behavior transfers to implicit natural-language constraints.
• Retention: The impact on general capabilities is mixed. While some medium-to-large models maintain or improve aggregate scores on TinyBenchmarks, smaller models and specific tasks (notably GSM8K) show degradation. This suggests that the training objective (short, direct completions without reasoning traces) may conflict with reasoning-heavy tasks.

Significance and Claims

The paper claims that probabilistic calibration is not an inherent limitation of language models but a capability that can be explicitly trained. By framing the problem as learning both the valid output space and the probability mass allocation across it, the authors show that LMs can be transformed into controlled probabilistic samplers.

The significance lies in moving beyond inference-time prompting to training-time solutions that fundamentally alter the model's distributional behavior. The authors note that while the gains are selective and do not universally improve all downstream tasks (with arithmetic reasoning being a notable exception), the ability to tune models for specific stochastic behaviors opens new avenues for applications requiring diversity, fairness, or strict adherence to statistical constraints. The work suggests that future research should focus on mechanisms learned by these objectives and developing retention-aware variants to mitigate capability degradation.