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 L2 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 RD(z,x)
L2
21d∑(Pπ(d)−P∗(d))2
P∗(z)−Pπ(z)
DKL(P∗∥Pπ)
d∑P∗(d)logPπ(d)P∗(d)
Pπ(z)P∗(z)−1
DKL(Pπ∥P∗)
d∑Pπ(d)logP∗(d)Pπ(d)
logP∗(z)−logPπ(z)
JSD(P∗,Pπ)
21DKL(P∗∥M)+21DKL(Pπ∥M)
21(logM(z)−logPπ(z))
Table 1: Reward functions corresponding to for the different divergence objectives. Category-independent constants are chosen so each reward vanishes at P∗=Pπ .
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.
Figure 4
Appendix figures & tables3 assets
Supplementary material from the paper’s appendix.
Appendix
train
held-out
target
reward
JSD
JSD
off-support
options shown (main text)
uniform
L2
0.0510
0.0901
0.0070
KL (p∥q)
0.0265
0.0467
0.0189
KL (q∥p) , ε -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 ε -smoothed: its coefficient is unbounded where the target assigns zero mass, so on the truncated target its value is set partly by the floor ε . 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 ( α 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 to 0.0444 nats against 0.0053 to 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.
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.
Shaobo Wang, Yujie Chen, Yafeng Sun +9
Shanghai Jiao Tong University · The Hong Kong University of Science and Technology (Guangzhou) · The University of Science and Technology of China +2
Given a question, a language model (LM) implicitly encodes a distribution over possible answers. In practice, post-training procedures for LMs often collapse this distribution onto a single dominant mode. While this is generally not a problem for benchmark-style evaluations that assume one correct answer, many real-world tasks inherently involve multiple valid answers or irreducible uncertainty. Examples include medical diagnosis, ambiguous question answering, and settings with incomplete information. In these cases, we would like LMs to generate multiple plausible hypotheses, ideally with confidence estimates for each one, and without computationally intensive repeated sampling to generate non-modal answers. This paper describes a multi-answer reinforcement learning approach for training LMs to perform distributional reasoning over multiple answers during inference. We modify the RL objective to enable models to explicitly generate multiple candidate answers in a single forward pass, internalizing aspects of inference-time search into the model's generative process. Across question-answering, medical diagnostic, and coding benchmarks, we observe improved diversity, coverage, and set-level calibration scores compared to single answer trained baselines. Models trained with our approach require fewer tokens to generate multiple answers than competing approaches. On coding tasks, they are also substantially more accurate. These results position multi-answer RL as a principled and compute-efficient alternative to inference-time scaling procedures such as best-of-k. Code and more information can be found at https://multi-answer-rl.github.io/.
Language models are increasingly used to sample from a specified distribution, for instance, to simulate survey respondents or generate synthetic data. Instruction-tuned models can state such a distribution correctly and still fail to sample from it. Prompting and changes to decoding reduce this mismatch only partly, which motivates training with policy optimization. Group relative policy optimization (GRPO) is a natural fit for this problem because it already samples a group of rollouts per prompt, and the group's empirical distribution can be compared with the target. However, scoring the group as a whole gives every rollout the same reward. Group-relative centering then sets all advantages to zero, and the model receives no learning signal. To give each rollout its own signal, we introduce the witness advantage, a per-rollout advantage derived from maximum mean discrepancy (MMD). It trains a model to match a target distribution over a finite set of outcomes. The MMD between the model's distribution and the target has a witness function that measures how over- or under-produced each outcome is. Each rollout's advantage estimates the negative witness at its outcome, so a rollout is rewarded for an outcome the group under-produces and penalized for one it over-produces. The witness advantage is computed in closed form from the group's outcome counts, and we use it as the reward in GRPO. On unseen target distributions, training with the witness advantage substantially reduces the total variation distance to the target while largely preserving the model's general capabilities.