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

• Pointwise rewards do not directly control output mixtures across generations; correctness-only RL can concentrate on one category, while entropy bonuses do not specify desired proportions.The figure contrasts those objectives with MaD-RL's targeting of user-specified distributions, including non-uniform ones.Source: §1 · p. 1, Figure 1 · p. 2
• MaD-RL combines task rewards with distributional feedback from sampled category frequencies, deriving GRPO-compatible rewards from L2, forward KL, reverse KL, and JS objectives.Source: §1 · p. 2
• For non-uniform targets, KL- and JS-based objectives achieve closer target adherence than the L2 objective.Source: §1 · p. 2
• MaD-RL steers category frequencies toward specified targets while measuring the associated trade-off in task accuracy.Source: §1 · p. 2

2 Matching Distributions with RL (MaDRL)

• For each prompt, MaD-RL compares sampled completion-category frequencies with a target distribution and combines distributional feedback with task correctness in GRPO.The target can vary by prompt, though the experiments use the same target for every prompt.Source: §2 · p. 2, §2.1 · p. 3

2.1 Formalization

• The formalized objective minimizes a discrepancy between target and policy category distributions, with the score-function identity converting it into rollout-level rewards.The derivation assumes outputs belong to one of the target categories; experiments penalize off-support outputs and retain the full group size as the count denominator.Source: §2.1 · pp. 2–3

2.2 GAPO and the Squared-L2 Objective

• Squared-L2 recovers GAPO’s frequency-dependent reward, and its population reward gives the negative gradient of the squared-L2 divergence.The population reward for category z is the target probability minus the policy probability.Source: §2.2 · p. 3
• The squared-L2 identity is population-level; plug-in estimates, GRPO normalization, clipping, and task-reward mixing further separate implementation from that gradient.The theorem uses population policy probabilities, while the rollout-based estimate includes each rollout in its own probability estimate.Source: §2.2 · p. 3

2.3 Forward and Reverse KL Objectives

• Forward and reverse KL yield different rewards because KL is asymmetric, and each direction is singular under a different zero-probability condition.Forward KL is singular when the target assigns positive probability but the policy assigns zero; reverse KL is singular when the policy assigns positive probability but the target assigns zero.Source: §2.3 · p. 4

2.4 Jensen–Shannon Divergence

• Jensen–Shannon divergence uses the symmetric mixture of target and policy distributions, but its policy-gradient reward requires epsilon-smoothing despite the divergence’s upper bound.MaD-RL does not need a learned discriminator for this connection because the target is known and rollout frequencies are counted.Source: §2.4 · p. 4

2.5 Finite-Group Training and Task Rewards

• At finite group size, empirical policy frequencies make nonlinear KL and JSD rewards biased, and on-policy updates cannot directly reward categories absent from the group.A sampled category’s empirical frequency is at least 1/G.Source: §2.5 · p. 4
• The experiments combine distributional rewards with binary correctness additively or through a gated heuristic, then use GRPO’s group-relative, clipped, KL-regularized update.The gated form applies distributional feedback only through correct rollouts and is affected by category-independent shifts that rescale correctness rewards.Source: §2.5 · p. 5

3.1 Experiments and Datasets

• The three experiments move from synthetic choices, where category choice does not affect correctness, to five-language mathematics and five-language programming, with uniform and exact-zero peaked targets.The peaked target excludes two categories and divides probability equally among the other three; the math and programming tasks make category choice relevant to correctness.Source: §3 · p. 5, §3.1 · pp. 5–6

3.2 Model and Training Protocol

• MaD-RL samples completion groups, estimates category frequencies, and compares them with the target; math and programming runs combine distributional reward with binary correctness using additive or correctness-gated forms.Runs use a shared supervised warm-start per dataset, KL regularization to that initialization, and no language-specific prompting at evaluation.Source: §3.2 · p. 7

3.3 Analysis and Results

• Correctness-only RL shifts the output distribution from near-uniform initialization; every uniform-target MaD-RL configuration has lower JSD than the best RLVR seed on all three datasets.Correctness-only RL concentrates outputs in Chinese on GSM8K, Spanish on MATH, and largely Python on CodeContests.Source: §3.3 · p. 7
• On synthetic held-out prompts, FKL has the lowest mean divergence in all 16 form–target–metric combinations, while L2 finishes last in 14.Against L2 training, FKL reduces final L2 divergence by roughly one-half on the uniform target and one-third on the peaked target; its correction strengthens as a positive-target category becomes rare, but cannot recover a category absent from the rollout group.Source: §3.3 · pp. 7–8
• RKL fits best in five of six dataset–target combinations, while L2 fits worst in four; neither divergence nor reward form is uniformly best across settings.Across post-training experiments, additive rewards yield lower JSD on all three peaked targets and the programming uniform target; the gated form is lower on the two mathematics uniform targets by less than the finite-sample floor.Source: Table 2 · p. 8, §3.3 · p. 9
• On GSM8K and MATH, every objective closely fits the uniform target; for peaked targets, RKL has the lowest median JSD and L2 the largest residual.Median JSD is below 0.005 on GSM8K and 0.008 on MATH for uniform targets; for peaked targets, RKL records 0.0064 and 0.0065, while L2 records 0.0218 and 0.0291, respectively.Source: §3.3 · p. 8, Figure 3 · p. 8
• In CodeContests, FKL fits the uniform target best and RKL the peaked target best, where RKL reaches 0.0059 JSD against L2’s 0.0949.Pooled median pass@1 ranges from 0.317 to 0.352 across targets, versus 0.379 for correctness-only RLVR.Source: Figure 3 · p. 8, §3.3 · p. 9
• Across post-training datasets, MaD-RL trails correctness-only RLVR by 1.2–8.5 pass@1 points; additive runs can also lose recovery signal when a category is absent from a rollout group.Pooled pass@1 does not cleanly separate mixture effects from within-category behavior, and supervised warm-start coverage does not eliminate category-collapse risk.Source: §3.3 · p. 9

Post-training and diversity

• Post-training can concentrate probability on fewer outputs (Kirk et al., 2024), while entropy regularization and reference-policy KL control do not specify probability shares among semantic categories.Source: §4 · p. 9

Controllable generation and repeated sampling

• Controllable generation methods such as CTRL (Keskar et al., 2019) and Plug and Play Language Models (Dathathri et al., 2020) act on individual generations, whereas repeated sampling and synthetic-data reuse make aggregate output proportions relevant (Wang et al., 2023; Yu et al., 2023).Source: §4 · p. 9

Distribution matching for generated text

• Prior distribution-matching work includes constrained target text distributions (Khalifa et al., 2021; Korbak et al., 2022), while groupwise methods such as GAPO (Anschel et al., 2025) promote uniform coverage rather than matching specified categorical proportions alongside task reward.Other approaches use answer-set credit, dual games over answer marginals, or Bregman divergences.Source: §4 · pp. 9–10

5 Conclusion

• MaD-RL uses rollout-group category frequencies to steer outputs toward a specified target while optimizing the original task, making the task-accuracy trade-off and divergence choice explicit.Source: §5 · p. 10
• MaD-RL fits uniform and peaked targets on all three math and programming datasets, with median holdout pass@1 from 0.2 points below to 10.4 points above SFT but 1.2–8.5 points below correctness-only RLVR.The authors conclude distributional control can preserve or improve starting-model accuracy, while trading off against correctness optimization.Source: §5 · p. 10
• L2 is the weakest distribution-matching choice: its fit is poorest in both post-training experiments, with similar task accuracy across all four rewards.In synthetic closed-set choice, L2 is strictly dominated on the fit-versus-validity frontier in two of four prompt-form × target cells and is best in none.Source: §5 · p. 10
• The best of forward KL, reverse KL, and JSD depends on task: forward KL fits synthetic closed-set targets best but has higher invalidity, while JSD best balances fit and validity.In the two correctness-trained post-training experiments, reverse KL fits the peaked target most closely on all three datasets.Source: §5 · p. 10
• The authors hypothesize that marginal accuracy improvements reflect the use of a 4B-scale model and relatively small post-training tasks; larger-scale benefits remain an expectation, not a demonstrated result.They also identify further study of the divergence metrics' different optimization dynamics as future work.Source: §5 · p. 10

A Proof and Derivation of MaD-RL

• The score-function identity expresses a category-probability gradient as an expectation of the completion's log-policy gradient when the category classifier is fixed.Source: A Proof and Derivation of MaD-RL · p. 13
• For a completion's category, the distributional reward is the signed gap between target frequency and realized group frequency.The derivation substitutes the score-function identity for category probabilities to obtain the signed-gap gradient implemented by MaD-RL.Source: A Proof and Derivation of MaD-RL · p. 13
• Additive rewards combine correctness and distributional feedback for every completion, whereas correctness-gated rewards apply distributional modulation only to correct completions.The additive form uses α as the correctness weight; the gated form multiplies correctness reward by a distributional adjustment.Source: A Proof and Derivation of MaD-RL · p. 13
• GRPO standardizes each completion's reward into a group-relative advantage using the group's reward spread, with clipping and the KL-regularized update otherwise unchanged.Source: A Proof and Derivation of MaD-RL · p. 13
• The method estimates category frequencies within each prompt's sampled group, or pools counts across a minibatch for population-level matching.The displayed algorithm samples completions, classifies them, computes rewards and group-relative advantages, then applies the standard clipped, KL-regularized GRPO update.Source: A Proof and Derivation of MaD-RL · pp. 13–14

B Synthetic Choice: Full Results

• Table 3 reports JSD over valid answers separately from off-support rate; at K=5 and G=32, the finite-sample JSD floors are 0.017 nats for uniform and 0.008 for peaked targets.These floors are common to all four rewards and do not affect the comparisons.Source: B Synthetic Choice: Full Results · p. 14
• Held-out JSD rose from 0.0648 to 0.1726 on the uniform target and from 0.1806 to 0.3920 on the peaked target when the option list was withheld.Off-support answers also rose by more than an order of magnitude, with little improvement over training.Source: B Synthetic Choice: Full Results · p. 14
• JSD lies on the fit-versus-validity frontier in all four prompt-form × target cells, while L2 is strictly dominated in two, including the open-set peaked-target comparison.There, JSD has 0.4139 against L2's 0.4294 JSD and 0.0979 against L2's 0.1065 off-support.Source: B Synthetic Choice: Full Results · p. 14
• Withholding the option list leaves the ranking across the four divergence objectives unchanged, despite the higher divergence and off-support rate.Source: B Synthetic Choice: Full Results · p. 14

C.1 Experimental setup

• GSM8K and MATH categorize solutions by language across English, Chinese, Spanish, Hindi, and French; CodeContests categorizes code across Python, C++, Java, Go, and JavaScript.Each dataset has its own SFT initialization, training pool, and held-out evaluation set.Source: §C.1 · p. 14
• The base Qwen3-4B-Instruct-2507 policy initially places categorized completions in one category, so each dataset receives a multilingual SFT warm-start to make all five categories reachable.The checkpoint is selected for category coverage on a split disjoint from RL training and evaluation, rather than for accuracy.Source: §C.1 · p. 14
• Table 3 reports closed-set results averaged over three seeds for K=5, with held-out evaluation on 80 prompts from eight topics disjoint from training.It reports JSD to target over valid answers and held-out off-support rate separately; lower JSD is better.Source: Table 3 · p. 15
• The distributional reward is a score-function coefficient for one of four divergences and combines with correctness additively or through correctness gating.The four additive settings allocate 10, 20, 30, or 45 percent of reward to the distributional term; outputs outside the declared set receive a fixed validity penalty.Source: §C.1 · p. 15
• Training runs 1,200 GRPO steps with group size 16, learning rate 5 × 10^-6, per-token KL penalty 0.04, and temperature 1.0.Each configuration uses three seeds; evaluation samples four completions for each of 200 held-out problems at temperature 1.0.Source: §C.1 · p. 15
• Finite-sample JSD is biased upward, so an on-target policy need not score zero; simulated estimator means are 0.00063 nats for uniform and 0.00032 for peaked targets.Differences within a small factor of these floors are not separable, so the uniform column is treated as a check rather than a comparison.Source: §C.1 · p. 15

C.2 Full reward grid

• 0.0057–0.0444 nats across correctness weights versus 0.0053–0.0911 across divergences; the mixing weight moves JSD by about as much as the divergence does.Values are median spreads across the four weights and across divergences, respectively.Source: Table 4 · p. 16
• Table 4 reports JSD to target in nats, averaged over three seeds at step 1200, across all three datasets.Source: Table 4 · p. 16, §C.2 · p. 16
• Rows add90 through add55 vary correctness weight in additive rewards; the gated row is a separate correctness-gated form excluded from the main-text pooling.The main-text figure and Table 2 pool the four additive rows.Source: Table 4 · p. 16

C.3 Per-category accuracy

• Category accuracy is not a clean competence measure because runs choose which problems to answer; counts are pooled across twelve additive runs and categories need 30 completions to be reported.Under the peaked target, two categories fall below that floor for most rewards.Source: §C.3 · p. 17
• Across three datasets, weaker SFT categories generally gained while stronger ones gave ground; on MATH, almost every category gained.Source: §C.3 · p. 17, Figure 4 · p. 17
• On CodeContests, Go rose from 0.151 to 0.165–0.210, while GSM8K Hindi rose from 0.605 to 0.659–0.706.JavaScript ranged from 0.252 to 0.370 against an SFT init of 0.295, gaining under seven of eight reward and target combinations; C++ fell from 0.465 to 0.352–0.438.Source: §C.3 · p. 17
• Figure 4 uses dataset rows and target columns; reward bars pool twelve additive runs, grey marks the SFT init, and baseline dots indicate fewer than 30 completions.A dot is not a measured zero.Source: Figure 4 · p. 17