aiqbench analysis
This research introduces a reinforcement learning framework for distribution matching in large language models, addressing the limitations of current post-training methods that often reduce output diversity. By optimizing for specific divergence metrics like KL and Jensen-Shannon rather than just expected reward, the approach enables better control over synthetic data generation and fairness constraints without relying solely on token-space entropy regularization. The method demonstrates improved performance in mathematical reasoning and programming tasks by effectively managing latent categorical attributes across model generations.
MaD-RL: Matching Distributions for Calibrating LLMs with Reinforcement Learning
arXiv:2609.31644v1 Announce Type: new 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 \textit{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.
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 \textit{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.
Machine Learning (cs.LG); Artificial Intelligence (cs.AI)
arXiv:2609.31644 [cs.LG]
(or arXiv:2609.31644v1 [cs.LG] for this version)
https://doi.org/10.48550/arXiv.2609.31644 Focus to learn more arXiv-issued DOI via DataCite (pending registration)
arXivLabs is a framework that allows collaborators to develop and share new arXiv features directly on our website.
Both individuals and organizations that work with arXivLabs have embraced and accepted our values of openness, community, excellence, and user data privacy. arXiv is committed to these values and only works with partners that adhere to them.
Have an idea for a project that will add value for arXiv's community? Learn more about arXivLabs.
Full text under the publisher's open licence · 288 words
Related coverage
[AINews] Opus 5.5 is good at explainer videos
Latent Space · Sep 29, 2026
[AINews] The Future of Latent Space
Latent Space · Sep 25, 2026
Granite 4.1 LLMs: How They’re Built
Hugging Face · Apr 29, 2026
Towards Full Pipeline FP8 Reinforcement Learning for LLMs
arXiv · Sep 22, 2026