RLHF in an SFT Way: From Optimal Solution to Reward-Weighted Alignment

1. Problem Statement

Reinforcement Learning from Human Feedback (RLHF) is the standard for aligning Large Language Models (LLMs) with human values (helpfulness and harmlessness). However, existing approaches face significant challenges:

• Online Methods (e.g., PPO, GRPO): These require complex pipelines involving multiple models (policy, reference, critic, reward), massive computational resources, and online sampling. They suffer from training instability, high latency, and a high risk of model collapse.
• Offline Methods (e.g., DPO, KTO, A-LoL): While simpler, they often rely on contrastive losses or importance sampling that can lead to:
  • Training Instability: Caused by negative weights in the loss landscape (e.g., when penalizing rejected responses), leading to unbounded optimization.
  • Suboptimal Performance: Methods like PPO and A-LoL use "clipping" mechanisms to stabilize training, which flattens reward distinctions between responses with similar values, preventing the model from reaching the true optimal policy.
  • Data Requirements: Some methods require high-quality preference pairs (chosen/rejected), incurring data collection costs.

The core problem is to find an alignment method that matches the performance of online RLHF but retains the efficiency and stability of Supervised Fine-Tuning (SFT), without the pitfalls of negative weights or clipping.

2. Methodology: Variational Alignment with Re-weighting (VAR)

The authors propose VAR, a novel framework that reformulates RLHF as a variational inference problem over positive measures.

Theoretical Foundation

Instead of optimizing the policy $\pi_\theta$ via the standard RLHF objective (maximizing expected reward minus KL divergence), VAR minimizes the forward KL divergence between the optimal policy $\pi^*$ and the learning policy $\pi_\theta$:
$$\min_\theta KL[\pi^* \parallel \pi_\theta] = \max_\theta \mathbb{E}_{\pi^*}[\log \pi_\theta(y|x)]$$

Using the closed-form solution of the optimal RLHF policy (derived from the Bradley-Terry model):
$$\pi^*(y|x) = \frac{1}{Z(x)} \pi_{ref}(y|x) \exp\left(\frac{1}{\beta} r(x, y)\right)$$

The authors apply importance sampling to rewrite the objective in terms of the reference policy $\pi_{ref}$:
$$\mathbb{E}_{\pi^*}[\log \pi_\theta] = \mathbb{E}_{\pi_{ref}} \left[ \underbrace{\frac{\exp(r(x,y)/\beta)}{Z(x)}}_{\text{Reward-driven Weight}} \cdot \log \pi_\theta(y|x) \right]$$

This transforms the alignment task into a reward-weighted Supervised Fine-Tuning (SFT) problem.

Key Technical Innovations

1. Positive Measure Optimization: Unlike DPO or A-LoL, which may assign negative weights to rejected responses (leading to unbounded loss), VAR uses an exponential transformation of the reward. This ensures all weights are strictly non-negative, guaranteeing a well-defined, stable optimization landscape.
2. In-Batch Partition Estimation: The normalization term $Z(x)$ (partition function) is computationally expensive to calculate over the entire output space. VAR approximates $Z(x)$ using Monte Carlo estimation within a micro-batch:
   $$\hat{Z}(x_i) \approx \frac{1}{B} \sum_{j=1}^B \exp\left(\frac{1}{\beta} r(x_i, y_j)\right) \pi_{ref}(y_j|x_i)$$
   where $y_j$ are samples from the reference model within the current batch. This avoids the need for online sampling or pre-sampling large datasets.
3. No Clipping: The method does not require clipping ratios (as in PPO) or negative weights, allowing the model to naturally learn fine-grained preferences.

3. Key Contributions

• Novel Formulation: The first approach to derive RLHF alignment directly from the optimal solution via variational inference, resulting in a simple, reward-weighted SFT loss.
• Stability via Positive Measures: By eliminating negative weights and clipping, VAR solves the training instability issues prevalent in DPO and A-LoL.
• Efficiency: It removes the need for online sampling, critic models, and complex RL pipelines, achieving convergence speeds 5× faster than GRPO.
• Scalability: The in-batch normalization technique allows the method to scale efficiently across different model sizes (from 0.5B to 32B parameters) without significant computational overhead.

4. Experimental Results

The authors evaluated VAR on the Helpful and Harmless Assistant (HHA) benchmark and various generative tasks (MMLU, GSM8k, HumanEval, BBH).

• Performance vs. DPO:
  • VAR outperforms DPO across most model sizes (Llama and Qwen series).
  • On HHA benchmarks, VAR achieved an average 7.16% improvement in helpfulness and harmlessness metrics compared to DPO.
  • VAR demonstrates superior robustness when starting from Base models (pre-trained only), whereas DPO often struggles without an initial SFT step.
• Performance vs. Online Methods (GRPO):
  • VAR achieves comparable or better performance than online GRPO methods.
  • Speed: VAR is >5× faster per epoch than GRPO (e.g., 37 mins vs. 2h 54m for Qwen2.5-3B).
• Training Stability:
  • Validation reward curves show VAR increases steadily, while DPO exhibits volatility and early decline.
  • VAR maintains stable output sequence lengths, whereas DPO tends to generate excessively long sequences to game reward scores.
• Generative Benchmarks:
  • On MMLU, GSM8k, and HumanEval, VAR consistently outperforms DPO and A-LoL, achieving the highest average scores.
  • On Arena-Hard 0.1, VAR significantly outperforms DPO (10.8% win rate vs. 8.6%), indicating better performance on challenging, high-quality prompts.

5. Significance

This paper bridges the gap between the efficiency of SFT and the performance of RLHF.

• Practical Impact: It offers a "plug-and-play" alternative to complex RLHF pipelines. Practitioners can achieve state-of-the-art alignment results with standard SFT infrastructure, significantly reducing computational costs and engineering complexity.
• Theoretical Insight: It challenges the necessity of negative weights and clipping in preference alignment, demonstrating that a purely positive, reward-weighted objective is sufficient and superior for stability.
• Future Direction: The work suggests that the "optimal solution" of RLHF can be approximated directly through variational inference, potentially rendering complex online RL loops obsolete for many alignment tasks.

In summary, VAR provides a mathematically grounded, computationally efficient, and highly stable method for aligning LLMs, effectively solving the "efficiency vs. performance" trade-off that has plagued the field.