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

arXiv: https://arxiv.org/abs/2609.31644 · HTML: https://mlsanity.com/paper/2609.31644/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 32 of its 33 references that are in the map.
- First posted on arXiv 12 Sep 2026.

## Builds on (most cited first, 30 of 32)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Training Verifiers to Solve Math Word Problems](https://mlsanity.com/paper/2110.14168/) (2021)
- [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://mlsanity.com/paper/2402.03300/) (2024)
- [Self-Consistency Improves Chain of Thought Reasoning in Language Models](https://mlsanity.com/paper/2203.11171/) (2022)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [Soft Actor-Critic: Off-Policy Maximum Entropy Deep Reinforcement Learning with a Stochastic Actor](https://mlsanity.com/paper/1801.01290/) (2018)
- [Deep reinforcement learning from human preferences](https://mlsanity.com/paper/1706.03741/) (2017)
- [The Curious Case of Neural Text Degeneration](https://mlsanity.com/paper/1904.09751/) (2019)
- [Learning to summarize from human feedback](https://mlsanity.com/paper/2009.01325/) (2020)
- [Generative Adversarial Imitation Learning](https://mlsanity.com/paper/1606.03476/) (2016)
- [Competition-Level Code Generation with AlphaCode](https://mlsanity.com/paper/2203.07814/) (2022)
- [Tulu 3: Pushing Frontiers in Open Language Model Post-Training](https://mlsanity.com/paper/2411.15124/) (2024)
- [Large Language Monkeys: Scaling Inference Compute with Repeated Sampling](https://mlsanity.com/paper/2407.21787/) (2024)
- [Plug and Play Language Models: A Simple Approach to Controlled Text Generation](https://mlsanity.com/paper/1912.02164/) (2019)
- [CTRL: A Conditional Transformer Language Model for Controllable Generation](https://mlsanity.com/paper/1909.05858/) (2019)
- [DoReMi: Optimizing Data Mixtures Speeds Up Language Model Pretraining](https://mlsanity.com/paper/2305.10429/) (2023)
- [Understanding the Effects of RLHF on LLM Generalisation and Diversity](https://mlsanity.com/paper/2310.06452/) (2023)
- [MMD GAN: Towards Deeper Understanding of Moment Matching Network](https://mlsanity.com/paper/1705.08584/) (2017)
- [Large Language Model as Attributed Training Data Generator: A Tale of Diversity and Bias](https://mlsanity.com/paper/2306.15895/) (2023)
- [Does Writing with Language Models Reduce Content Diversity?](https://mlsanity.com/paper/2309.05196/) (2023)
- [A Distributional Approach to Controlled Text Generation](https://mlsanity.com/paper/2012.11635/) (2020)
- [Sequence Tutor: Conservative Fine-Tuning of Sequence Generation Models with KL-control](https://mlsanity.com/paper/1611.02796/) (2016)
- [On Reinforcement Learning and Distribution Matching for Fine-Tuning Language Models with no Catastrophic Forgetting](https://mlsanity.com/paper/2206.00761/) (2022)
- [SetPO: Set-Level Policy Optimization for Diversity-Preserving LLM Reasoning](https://mlsanity.com/paper/2602.01062/) (2026)
- [Where does output diversity collapse in post-training?](https://mlsanity.com/paper/2604.16027/) (2026)
- [Group-Aware Reinforcement Learning for Output Diversity in Large Language Models](https://mlsanity.com/paper/2511.12596/) (2025)
- [Uniform-Correct Policy Optimization: Breaking RLVR's Indifference to Diversity](https://mlsanity.com/paper/2605.00365/) (2026)
- [Distributional Alignment Games for Answer-Level Fine-Tuning](https://mlsanity.com/paper/2604.27166/) (2026)
- [Controlling Distributional Bias in Multi-Round LLM Generation via KL-Optimized Fine-Tuning](https://mlsanity.com/paper/2604.05756/) (2026)

## Abstract (arXiv)

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 textitDistribution 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.