MaD-RL: Matching Distributions for Calibrating LLMs with Reinforcement Learning

Sourabh Kulkarni; Ksheeraj Sai Vepuri; Basar Demir; Jason Bohrer; Emily Shen; Jianfa Chen; Nan Jiang; Ankit Jain

cs.LGSep 12, 2026

MaD-RL: Matching Distributions for Calibrating LLMs with Reinforcement Learning

Authors: Sourabh Kulkarni, Ksheeraj Sai Vepuri, Basar Demir, Jason Bohrer, Emily Shen, Jianfa Chen, Nan Jiang, Ankit Jain,

Organizations: Meta Superintelligence Labs

Abstract

Reinforcement learning (RL) is widely used in language-model post-training to maximize rewards assigned to individual model outputs, such as scores from binary verifiers or reward models trained on human feedback. However, applications such as synthetic-data generation, fairness-related constraint satisfaction, and policy exploration require controlling the distribution of outputs across model generations rather than only maximizing expected reward. We propose a general RL-based framework for Distribution Matching allowing matching the distribution of a latent categorical attribute of model outputs to a specified target distribution. Empirically, we demonstrate that dominant post-training recipes such as Group Relative Policy Optimization (GRPO) reduce output diversity by concentrating policy probability towards a single mode. Entropy regularization and sampling temperature can improve the spread of the distribution but have constrained effectiveness, limited to apply only in token space and toward uniform distributions. We show that prior work in this area is a specific case of Distribution Matching involving the $L_2$ divergence. We then propose reward functions for other divergences such as KL and Jensen-Shannon and motivate them with theoretical justification. Finally, we demonstrate the effectiveness of our approach on a set of experiments involving mathematical reasoning and programming.

Figures & tables

Figure 1: Schematic comparison of the category-level objectives. Eight completions sampled for one prompt, one glyph per category, under three post-training objectives and two independent runs. Correctness-only RL concentrates on a single category but does not control which one: the two runs settle on different ones. An entropy bonus can restore spread, but does not specify the desired proportions. MaD-RL instead targets a user-specified distribution including a non-uniform one (Column C).

Divergence	Objective	Reward $\texttt{R}_{D}(z,{\bm{x}})$
$L_{2}$	$\displaystyle\frac{1}{2}\sum_{d}(\mathsf{P}_{\pi}(d)-\mathsf{P}^{*}(d))^{2}$	$\displaystyle\mathsf{P}^{*}(z)-\mathsf{P}_{\pi}(z)$
$D_{\textrm{KL}}(\mathsf{P}^{*}\\|\mathsf{P}_{\pi})$	$\displaystyle\sum_{d}\mathsf{P}^{}(d)\log\frac{\mathsf{P}^{}(d)}{\mathsf{P}_{\pi}(d)}$	$\displaystyle\frac{\mathsf{P}^{*}(z)}{\mathsf{P}_{\pi}(z)}-1$
$D_{\textrm{KL}}(\mathsf{P}_{\pi}\\|\mathsf{P}^{*})$	$\displaystyle\sum_{d}\mathsf{P}_{\pi}(d)\log\frac{\mathsf{P}_{\pi}(d)}{\mathsf{P}^{*}(d)}$	$\displaystyle\log\mathsf{P}^{*}(z)-\log\mathsf{P}_{\pi}(z)$
$\textrm{JSD}(\mathsf{P}^{*},\mathsf{P}_{\pi})$	$\displaystyle\frac{1}{2}D_{\textrm{KL}}(\mathsf{P}^{*}\\|\mathsf{M})+\frac{1}{2}D_{\textrm{KL}}(\mathsf{P}_{\pi}\\|\mathsf{M})$	$\displaystyle\frac{1}{2}\bigl(\log\mathsf{M}(z)-\log\mathsf{P}_{\pi}(z)\bigr)$

Table 1: Reward functions corresponding to for the different divergence objectives. Category-independent constants are chosen so each reward vanishes at

\mathsf{P}^{*}=\mathsf{P}_{\pi}

Figure 2: Divergence across RL Steps on SYNTHETIC. The upper block lists the options in the prompt (closed-set), the lower block does not (open-set). Rows are targets, columns are evaluation metrics, curves are the training objective; every run is scored under all four metrics. Markers are means over three seeds; bands span their minimum and maximum.

Appendix figures & tables3 assets

Supplementary material from the paper’s appendix.

Appendix

options shown (main text)
		train	held-out
target	reward	JSD	JSD	off-support
uniform
	$L_{2}$	$0.0510$	$0.0901$	$0.0070$
	KL $(p\\|q)$	$0.0265$	$0.0467$	$0.0189$
	KL $(q\\|p)$ , $\varepsilon$ -sm.	$0.0263$	$0.0504$	$0.0135$

Appendix

Table 3: Closed-set results, mean over three seeds,

K{=}5

. JSD to target in nats over valid answers; lower is better. Train JSD is the mean of the last ten logged points; held-out is the final evaluation on

80

prompts from

8

topics disjoint from training. Off-support is the held-out rate of answers outside the declared set; it is not logged on the training pool. Reverse KL is

\varepsilon

-smoothed: its coefficient is unbounded where the target assigns zero mass, so on the truncated target its value is set partly by the floor

\varepsilon

. The main text reports the options-shown rows; the options-inferred rows are discussed in appendix B .

		uniform target			peaked target
Reward	Weight	GSM8K	MATH	CodeContests	GSM8K	MATH	CodeContests
l2	gated	$0.0033$	$0.0081$	$0.0442$	$0.0267$	$0.0281$	$0.1056$
	add90	$0.0190$	$0.0043$	$0.0615$	$0.0236$	$0.0300$	$0.2178$
	add80	$0.0023$	$0.0067$	$0.0303$	$0.0210$	$0.0300$	$0.0819$
	add70	$0.0027$	$0.0073$	$0.0177$	$0.0681$	$0.0280$	$0.1110$
	add55	$0.0279$	$0.0047$	$0.0285$	$0.0617$	$0.0237$	$0.0554$

Appendix

Table 4: Full reward grid, all three datasets. JSD to target in nats, mean over three seeds, at step

1200

. Rows add90 to add55 are the additive form, the suffix giving the percentage weight on correctness (

\alpha

in eq. 23 ); gated is the correctness-gated form. The main-text figure and table 2 pool the four additive rows; the gated row is shown here but excluded there, since it is a different combination rule rather than a different mixing weight. The mixing weight moves JSD by about as much as the divergence does (median spread across the four weights

0.0057

0.0444

nats against

0.0053

0.0911

across divergences), which is why no single weight is reported as representative.

Figure 4: Per-category holdout pass@1. Rows are datasets, columns are targets. Each reward bar pools the twelve additive runs behind it by summing correct completions and completions. Grey is the SFT init. A dot on the baseline marks a category that falls below the

30

-completion reporting floor; it is not a measured zero.

Distribution Matching

Distribution matching measures and aligns probability distributions to transfer data structure across models, domains, and generated outputs.

See all 80 papers

DISA: Offline Importance Sampling for Distribution-Matching LLM-RL
May 17, 2026Shaobo Wang, Yujie Chen, Yafeng Sun +9Reinforcement Learning RL for Language Model Reasoning
Shanghai Jiao Tong University · The Hong Kong University of Science and Technology (Guangzhou) · The University of Science and Technology of China +2
Reaching Beyond the Mode: RL for Distributional Reasoning in Language Models
Mar 25, 2026Isha Puri, Mehul Damani, Idan Shenfeld +3RL for Language Model Reasoning Efficient Language Model Reasoning
Massachusetts Institute of Technology
Sample What You Say: Aligning Language Models to Sample the Distributions They State
Sep 28, 2026Kasra Arabi, Virginia Smith, Chhavi YadavRL for Language Models Group Relative Policy Optimization
New York University · Carnegie Mellon University

May 17, 2026J/K move · Enter open · S save

DISA: Offline Importance Sampling for Distribution-Matching LLM-RL

Shaobo Wang, Yujie Chen, Yafeng Sun, Wenjie Qiu+8

Shanghai Jiao Tong University · The Hong Kong University of Science and Technology (Guangzhou) · The University of Science and Technology of China+2

Reinforcement Learning RL for Language Model Reasoning

Open paper AskPDF arXiv

Modern reasoning agents are increasingly evaluated on their ability to generate multiple valid solution paths, plans, or tool-use traces for a given input. Standard reward-maximizing RL tends to collapse onto the most easily reinforced high-reward mode, whereas distribution-matching RL aims to allocate probability mass across the entire reward-shaped solution set. Achieving this objective requires computing a prompt-dependent partition function over the trajectory space. Because existing distribution-matching methods learn this partition function online alongside the policy, calibration errors in the partition function directly distort policy updates and remain impossible to diagnose independently. We introduce DISA, short for Decoupled Importance-Sampled Anchoring, which moves this calibration problem outside the RL loop. DISA draws proposal trajectories offline, estimates the partition function via importance sampling, and freezes the resulting partition-function estimate before policy optimization begins. This decoupling preserves the distribution-matching objective while strictly separating partition-function estimation from policy learning in data, gradients, loss, and diagnostics. Empirically, on two open-weight backbones across six math and three code benchmarks, DISA matches or exceeds the online-coupled distribution-matching baseline FlowRL, outperforms rewardmaximization baselines GRPO and GSPO on math averages, and exceeds LoRASFT distillation by up to 13.8 Mean@8 points on the same offline trajectories. An LLM-as-judge evaluation further shows that DISA retains substantially more strategy-level diversity than reward-maximization baselines, and sensitivity studies on the proposal strength and inverse temperature follow the bias-variance pattern predicted by the analysis.

Abstract

Distribution Matching

Explore similar work

DISA: Offline Importance Sampling for Distribution-Matching LLM-RL

Reaching Beyond the Mode: RL for Distributional Reasoning in Language Models

Sample What You Say: Aligning Language Models to Sample the Distributions They State