Complementing reinforcement learning with SFT through logit averaging in the post training of LLMs

Technical Summary: Complementing Reinforcement Learning with SFT through Logit Averaging

1. Problem Statement

The post-training of Large Language Models (LLMs) typically relies on two paradigms: Supervised Fine-Tuning (SFT) and Reinforcement Learning (RL) with verifiable rewards (RLVR). While SFT provides dense per-token supervision to learn formatting and style, it is limited by the distribution of training demonstrations. Conversely, RL methods like Proximal Policy Optimization (PPO) and Group Relative Policy Optimization (GRPO) can surpass these distributional limits but suffer from sparse, noisy reward signals.

The standard approach combines these sequentially: SFT is applied first, followed by RL with a Kullback-Leibler (KL) regularization term, $\beta \text{KL}(\pi_\theta \| \pi_{\text{sft}})$, to anchor the trainable policy $\pi_\theta$ to the SFT reference $\pi_{\text{sft}}$. The authors identify a fundamental tension in this "KL-anchored" recipe: the KL penalty pulls $\pi_\theta$ toward $\pi_{\text{sft}}$, potentially constraining the policy to the limitations of the SFT checkpoint (e.g., algebraic errors, repetition loops) and preventing it from fully leveraging the reasoning capabilities developed during RL. Practitioners often mitigate this by choosing a small $\beta$, but the authors propose a structural alternative to regularization.

2. Methodology

The paper introduces a method that replaces KL regularization with logit averaging. Instead of penalizing divergence, the trainable policy $\pi_\theta$ and a frozen reference policy $\pi_{\text{ref}}$ (e.g., SFT) are coupled by averaging their logits to form a mixed policy $\pi^{\text{mix}}_{\theta, \alpha}$.

Core Mechanism:
Given a weight $\alpha \in [0, 1]$ and state $s$, the mixed logits $z^{\text{mix}}_{\theta, \alpha}(s)$ are defined as:
$$z^{\text{mix}}_{\theta, \alpha}(s) = (1 - \alpha) z_\theta(s) + \alpha z_{\text{ref}}(s)$$
The mixed policy is then $\pi^{\text{mix}}_{\theta, \alpha}(a | s) = \text{softmax}(z^{\text{mix}}_{\theta, \alpha}(s))$.

Training Objective:
The authors adapt the GRPO objective to this mixed policy:

• Rollouts: Trajectories are sampled from the mixed policy $\pi^{\text{mix}}_{\theta, \alpha}$ rather than the pure trainable policy $\pi_\theta$.
• Advantage Estimation: Group-relative advantages are computed based on rewards from these mixed rollouts.
• Optimization: The trainable parameters $\theta$ are updated to maximize the reward under the mixed policy. The reference logits $z_{\text{ref}}$ remain frozen throughout.
• Weighting Schemes:
  1. Fixed Weights: $\alpha$ is held constant (e.g., $\alpha=0.5$) throughout training.
  2. Adaptive Weights: $\alpha$ is updated dynamically based on a held-out validation set. The update rule uses a "net-gain" metric comparing the correctness of the pure policy versus the mixed policy, pushing $\alpha$ toward 1 if mixing yields a net gain and toward 0 otherwise.

Comparison to RLVR:
Unlike standard RLVR methods (PPO, GRPO, DAPO) that use a critic and KL regularization, this approach has no KL term or critic. The reference policy participates in exploration via logit averaging, and the importance ratios in the loss function are calculated between the current and old mixed policies.

3. Key Contributions and Findings

Complementary Skills (Pattern 1 & 2):
Empirical analysis across MATH, cn-k12, and MMLU datasets reveals that SFT and the RL-trained policy possess complementary skills:

• SFT excels at formatting (adhering to output templates) but may fail on complex reasoning.
• RL Policy ($\pi_\theta$) excels at reasoning (deriving correct numerical answers) but often loses formatting discipline.
• Complementarity: A significant fraction of problems (22%–38%) are solved correctly by exactly one of the two policies. Logit averaging acts as a "Product of Experts," combining the formatting strength of SFT with the reasoning strength of the RL policy.

Performance Results:
Evaluated on Qwen2.5-Instruct models (1.5B, 3B, 7B) across three datasets:

• Accuracy: The logit averaging method achieves higher or comparable answer accuracy relative to canonical KL-regularized GRPO in all nine (dataset $\times$ model size) configurations.
• Ablation: Replacing the SFT anchor with a base Instruct model results in significantly lower performance, confirming that the gains stem from mixing with a format-aware reference.
• Logit vs. Probability Averaging: Logit averaging (geometric mean of distributions) outperforms probability averaging (arithmetic mean). Logit averaging achieves higher solve rates with smaller mixing weights and peaks at $\alpha=0.5$, whereas probability averaging requires higher weights and peaks at $\alpha=0.6$.

Theoretical Connection:
The paper establishes that logit averaging is mathematically equivalent to a Product of Experts (POE). The resulting policy minimizes a weighted sum of KL divergences to the RL and reference policies. This structure allows the policy to concentrate probability on tokens favored by both experts, effectively correcting SFT's algebraic errors while retaining its formatting constraints.

4. Significance and Claims

The authors claim that this method offers a distinct advantage over the standard KL-regularized recipe by decoupling the "formatting anchor" from the "gradient constraint."

• Overcoming Limitations: While KL regularization forces $\pi_\theta$ to stay close to $\pi_{\text{sft}}$ (potentially inheriting its reasoning flaws), logit averaging allows $\pi_\theta$ to overcome the reference's limitations while maintaining formatting advantages through the structural coupling of the mixed policy.
• Implicit Trust Region: The mixing acts as an implicit trust region, preventing the distributional drift away from formatting styles that often plagues pure RL at scale.
• Dense Priors: Unlike RL which propagates signals only from the final reward token, logit averaging incorporates a dense, token-level prior from the SFT anchor at every decoding step.
• Scalability: The approach naturally extends to ensembles of Small Language Models (SLMs), allowing a system to leverage specialized experts (e.g., for code, safety, or domain terminology) without the cost of training a massive single policy.

The paper concludes that logit averaging provides a robust alternative to KL regularization, successfully leveraging the complementarity between SFT and RL to improve reasoning accuracy without sacrificing formatting quality.