A Distribution-Matching Perspective on LLM Training Algorithms

Yanxing Huang; Bingzhang Yan; Zihan Tang; Wenzhe Li; Peng Li; Yang Liu

cs.LGOct 24, 2026

A Distribution-Matching Perspective on LLM Training Algorithms

Authors: Yanxing Huang, Bingzhang Yan, Zihan Tang, Wenzhe Li, Peng Li, Yang Liu

Organizations: Qiuzhen College, Tsinghua University, Beijing, China · Institute for AI Industry Research (AIR), Tsinghua University, Beijing, China · Weiyang College, Tsinghua University, Beijing, China · Department of Mathematical Science, Tsinghua University, Beijing, China ·

Abstract

In this work, we propose a unified framework for training algorithms across different stages of modern large language model development, based on the perspective of distribution matching. We show that a broad class of language model training algorithms can be interpreted as minimizing a certain divergence between the current model and a target distribution, and that characterizing the properties of this target distribution provides substantial insight into the behavior of the corresponding training algorithms. By analyzing specific forms of the target distribution, we establish that On-Policy Distillation (OPD) enjoys an advantage over classical Knowledge Distillation (KD) when the current model distribution is close to the target distribution, while exhibiting a disadvantage when the two distributions are far apart. Furthermore, when the target distribution itself depends on the model being optimized, we show that the accumulation of estimation variance can induce a deterministic degradation of the target distribution. These results provide a theoretical explanation, at least in part, for the long-term behavior of existing Reinforcement Learning (RL) algorithms. With this intuition we further designed Conditional KL (CKL) regularization for mitigating the degradation of the target distribution. CKL can effectively preserve mode diversity while relaxing the constraint imposed by classical KL regularization on reward improvements. We believe this theoretical perspective can bring some new understanding of language model training algorithms and provide useful guidance for future research.

Figures & tables

Algorithm	Objective function	Target distribution	Divergence
Noniterative algorithms
Pre-training	$\mathbb{E}_{x\sim p_{\mathcal{D}}}[\log\pi_{\theta}(x)]$	$p_{\mathcal{D}}(x)$	Forward KL
SFT	$\mathbb{E}_{y\sim p_{\mathcal{D}}(\cdot\mid x)}[\log\pi_{\theta}(y\mid x)]$	$p_{\mathcal{D}}(y\mid x)$	Forward KL
SFT-based KD	$\mathbb{E}_{y\sim\pi_{t}(\cdot\mid x)}[\log\pi_{\theta}(y\mid x)]$	$\pi_{t}(y\mid x)$	Forward KL
OPD	$\displaystyle\mathbb{E}_{y\sim\pi_{\theta}(\cdot\mid x)}\!\left[\log\frac{\pi_{t}(y\mid x)}{\pi_{\theta}(y\mid x)}\right]$	$\pi_{t}(y\mid x)$	Reverse KL
Iterative algorithms

Table 1: Training algorithms viewed as distribution matching. Parameter-independent constants are omitted. Conditional objectives are averaged over prompts

x\sim\mathcal{D}

. Forward KL places the target in the first argument; Reverse KL places it in the second. RL and MaxRL use binary rewards and

\rho_{\theta}(x)>0

, with the policy-dependent target

\pi_{\theta}^{+}

. Algorithms are iterative if their target distribution depends on the current model, and noniterative otherwise;

\pi_{\theta}

is highlighted in iterative targets. † RL is its first-order approximation.

Figure 1: The comparision of the entropy dynamics of distillation algorithms. The nonzero variance and iteration are two necessasy conditions for variance accumulation to occur. RL is where both conditions are satisfied.

Figure 2: Ablation study on the impact of advantage estimation methods. We can see that under the same RL setting, advantage estimation methods constantly lowers the entropy deduction. Although they eventually achieved similar reward on the training set, to some extent, lower entropy still reflects reduced diversity and weaker generalization ability.

Figure 3: The comparison of KL and Conditional KL (CKL) regularization in RL. We trained Qwen3-1.7B-Base ( Yang et al., 2025 ) on MATH ( Hendrycks et al., 2021 ) dataset with GRPO ( Shao et al., 2024 ) algorithm and different KL penalities. Two conditional entropies are measured on the sequence level. We can see that strong KL control strength severely limits the reward increment in RL, while CKL regularization allows the model to obtain higher reward while maintaining the diversity of true samples.

Figure 4: Extended experiments on KL and CKL regularization. These experiments follows the same setting as that in Figure 3 . We can see that CKL-1 and KL-0.3 exhibits similar regularization ability on conditional distribution, while CKL-1 constantly obtains higher accuracy. Its entropy will thus be slightly lower according to the accuracy-weighted decomposition.

Figure 5: The comparison of pass@k performance on MATH500 ( Lightman et al., 2023 ) , Apex Shortlist ( Balunović et al., 2025 ) and Apex2025 ( Balunović et al., 2025 ) . Models are trained for 20 epochs on MATH dataset. We can see that models trained without KL has higher accuracy on datasets similar to the training set, while severe collapse occurs on harder datasets like Apex2025. CKL constantly achieves higher generalizability over KL regularization across these experiments.

Appendix figures & tables7 assets

Supplementary material from the paper’s appendix.

Appendix

Algorithm	Objective function	Target distribution	Divergence
Iterative algorithms
RL	$\mathbb{E}_{\pi_{\theta}}[s(x,Y)]=\rho_{\theta}^{s}(x)$	$\begin{gathered}Q_{\theta}(y,1\mid x)\\ =\dfrac{s(x,y){\color[rgb]{0.5586,0.2695,0.5234}\pi_{\theta}}(y\mid x)}{\rho_{\theta}^{s}(x)}\end{gathered}$	Forward KL †
MaxRL (extended)	$\begin{gathered}\log\mathbb{E}_{\pi_{\theta}}[s(x,Y)]\\ =-D_{\mathrm{KL}}(Q_{\theta}\\|P_{\theta})\end{gathered}$	$\begin{gathered}Q_{\theta}(y,1\mid x)\\ =\dfrac{s(x,y){\color[rgb]{0.5586,0.2695,0.5234}\pi_{\theta}}(y\mid x)}{\rho_{\theta}^{s}(x)}\end{gathered}$	Forward KL

Appendix

Table 2: Continuation of Table 1 : bounded nonbinary rewards. Objectives are averaged over prompts

x\sim\mathcal{D}

. The model distribution is

P_{\theta}

; the target

Q_{\theta}

conditions it on acceptance and has

Q_{\theta}(y,0\mid x)=0

. † RL is its first-order approximation.

Figure 6: Reward curve of distal OPD with the JustRL teacher and two initial students. OPD rollout correctness increases for R1-Distill-1.5B but decreases and remains low for Qwen2.5-Math-1.5B. Note that since we do not have student’s on policy samples in SFT, we can not draw a similar reward curve for SFT. Their evaluation results are reported in Table 3 .

Student	Teacher	Model	AIME-24	AMC23	MATH500	Mean
Proximal Distillation
Qwen3-1.7B-Base	GRPO-50	Student	0.62	7.50	16.19	8.10
		Teacher	1.04	12.19	25.62	12.95
		SFT	1.88	15.78	32.69	16.78
		OPD	1.25	16.25	35.80	17.77
Qwen3-1.7B-Base	GRPO-100	Student	0.62	7.50	16.19	8.10

Appendix

Table 3: Evaluation results of OPD and SFT with 16 sampled responses per problem. We can see that OPD typically outperforms SFT in proximal distillation, while it may experience severe instability issue and performance degradation in distal distillation. Proximal teachers are obtained by training the base model with GRPO for 50, 100, or 200 steps.

Figure 7: Additional experiments on entropy dynamics of iterative self-distillation and reversed iterative self-distillation. In the legend “N” denotes the number of samples per request, “B” denotes batchsize and “lr” denotes the learning rate of SGD optimizer.

Figure 8: The ablation study on the advantage estimation on SmolLM2-360M-Instruct ( Allal et al., 2025 ) and GSM8K ( Cobbe et al., 2021 ) dataset. with batchsize 128, nsamples 4 and AdamW optimizer with learning rate 1e-5.

Figure 9: Additional metrics on the comparison of KL and CKL restrictions.

Figure 10: Pass@k performance of models trained for 4 epochs with KL or CKL regularization on MATH500 ( Lightman et al., 2023 ) , Apex2025 ( Balunović et al., 2025 ) , arxivmath ( Balunović et al., 2025 ) , and AIME ( Balunović et al., 2025 ) datasets. CKL regularization also preserves more learning potential in the sense of pass@k performance compared with KL regularization with the same strength. The evaluation settings is the same as that used in AReaL training pipeline.

LLM Training

Large language model training studies the data, objectives, optimizers, and systems used to build and adapt language models.

See all 353 papers

MaD-RL: Matching Distributions for Calibrating LLMs with Reinforcement Learning
Sep 12, 2026Sourabh Kulkarni, Ksheeraj Sai Vepuri, Basar Demir +8Distribution Matching RL for Language Models
Meta Superintelligence Labs
Do We Really Need KL Divergence for On-Policy Distillation of Large Language Models?
Sep 27, 2026Wenze Lin, Jiyuan Long, Jiale Zhao +11Data-Free Knowledge Distillation Language Model Distillation
LeapLab, Tsinghua University · Qiuzhen College, Tsinghua University · Beihang University +3
Unbiased Top- $k$ Estimation for On-Policy Distillation
Sep 28, 2026Linjian Meng, Siyuan Gan, YuHan Li +5Language Model Distillation On-Policy Distillation
Shanghai Artificial Intelligence Laboratory, Shanghai, China · State Key Laboratory of Novel Software Technology, Nanjing University, Nanjing, China

Sep 12, 2026J/K move · Enter open · S save

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

Sourabh Kulkarni, Ksheeraj Sai Vepuri, Basar Demir, Jason Bohrer+7

Meta Superintelligence Labs

Distribution Matching RL for Language Models Language Model Calibration RL for Language Model Reasoning

Open paper AskPDF 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 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

LLM Training

Explore similar work

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

Do We Really Need KL Divergence for On-Policy Distillation of Large Language Models?

Unbiased Top-kkk Estimation for On-Policy Distillation

Unbiased Top- $k$ Estimation for On-Policy Distillation