Learning Dynamics in RL Post-Training for Language Models

TL;DR

This work analyzes RL post-training through an empirical neural tangent kernel lens, revealing that limited variability in feature representations can cause RL updates to inflate model confidence and reduce output diversity. By decomposing the NTK into a Representation component and a Gradient component, the authors show that rapid, classifier-driven shaping of the Gradient component is crucial for effective learning. They propose classifier-first reinforcement learning (CF-RL), a two-stage method that prioritizes updating the classifier before standard RL optimization, and demonstrate via experiments that CF-RL accelerates RL convergence and increases early reward signals. The findings formalize the learning dynamics of RL post-training, explain the mechanism behind diversity loss, and distinguish CF-RL from LP-FT approaches in supervised settings, offering a principled direction for improving RL post-training.

Abstract

Reinforcement learning (RL) post-training is a critical stage in modern language model development, playing a key role in improving alignment and reasoning ability. However, several phenomena remain poorly understood, including the reduction in output diversity. To gain a broader understanding of RL post-training, we analyze the learning dynamics of RL post-training from a perspective that has been studied in supervised learning but remains underexplored in RL. We adopt an empirical neural tangent kernel (NTK) framework and decompose the NTK into two components to characterize how RL updates propagate across training samples. Our analysis reveals that limited variability in feature representations can cause RL updates to systematically increase model confidence, providing an explanation for the commonly observed reduction in output diversity after RL post-training. Furthermore, we show that effective learning in this regime depends on rapidly shaping the classifier, which directly affects the gradient component of the NTK. Motivated by these insights, we propose classifier-first reinforcement learning (CF-RL), a simple two-stage training strategy that prioritizes classifier updates before standard RL optimization. Experimental results validate our theoretical analysis by demonstrating increased model confidence and accelerated optimization under CF-RL. Additional analysis shows that the mechanism underlying CF-RL differs from that of linear-probing-then-fine-tuning in supervised learning. Overall, our study formalizes the learning dynamics of RL post-training and motivates further analysis and improvement.

Figures (15)

• Figure 1: The classifier learns faster than the other parameters. We partition the model parameters into transformer layers (layers $0$--$31$), token embeddings (Embeddings), the final layer normalization (Final LayerNorm), and the classifier. For each group, we plot the L2 norm of the parameter difference from the SFT model during RL training, scaled so that the norm equals $1$ at the end of RL.
• Figure 2: Example first-token distribution (SFT vs. GRPO)
• Figure 3: Entropy reduction during RL. Epoch 0 corresponds to the SFT model. Lower entropy indicates a more concentrated distribution.
• Figure 5: Best-of-$N$ rewards for SFT and GRPO.
• Figure 6: Semantic and style diversity over RL epochs. Higher values indicate greater diversity. Shaded regions denote standard deviation.

Theorems & Definitions (4)

• Proposition 1
• Proposition 2
• proof
• proof