Provably avoiding over-optimization in Direct Preference Optimization without knowing the data distribution

Problem Statement

Direct Preference Optimization (DPO) has become a standard for post-training large language models (LLMs) due to its simplicity and ability to bypass explicit reward modeling. However, DPO is known to suffer from over-optimization: as training progresses, the model may maximize the training loss while degrading the actual quality of generated responses relative to the initial model. This issue is exacerbated by the fact that standard DPO and Reinforcement Learning from Human Feedback (RLHF) treat the estimated reward as a perfect ground truth, ignoring uncertainty.

Theoretical guarantees for mitigating over-optimization typically rely on pessimism, ensuring the learned policy is robust to estimation errors. Existing pessimistic variants (e.g., $\chi^2$-regularized DPO, DPO+SFT) achieve strong theoretical bounds depending only on a single-policy concentrability coefficient ($C^\star$), which implies robustness even if the dataset does not cover the entire policy space. However, these guarantees critically depend on knowing the data-generating distribution ($\pi_{data}$) or having access to a reference policy well-covered by $\pi_{data}$. In practice, this assumption is often violated, such as when fine-tuning smaller models on preference data generated by larger, proprietary systems (e.g., GPT-4) where $\pi_{data}$ is unknown and inaccessible.

The central research question addressed is: Can we design a DPO-like algorithm that is provably robust to over-optimization without knowledge of the data-generating distribution $\pi_{data}$?

Methodology: PEPO

The authors introduce PEPO (Pessimistic Ensemble based Preference Optimization), a single-step algorithm that achieves pessimism through an ensemble of policies trained on disjoint data subsets, without requiring knowledge of $\pi_{data}$.

1. Ensemble Training:
The preference dataset $D$ is partitioned into $L$ disjoint subsets. For each subset $D_\ell$, a lightweight policy $\tilde{\pi}_\ell$ is trained using a pessimistic DPO objective. This objective modifies the standard DPO loss by shifting the sigmoid function horizontally by a parameter $\lambda_\ell$, effectively reducing the margin required for a winning probability. This shift induces a conservative (pessimistic) estimate of the reward gap.

2. Pessimistic Aggregation:
The output policy $\pi_{out}$ is defined as the solution to a "virtual" RLHF problem that optimizes the worst-case reward across the ensemble. The implicit reward for the ensemble is constructed by taking the minimum over the $L$ models of their estimated reward gaps, adjusted for the data distribution in the other subsets (to ensure independence).
$$r^-(x, a) = \beta \min_{\ell \in [L]} \left( \log \frac{\tilde{\pi}_\ell(a|x)}{\pi_{ref}(a|x)} - \zeta_\ell(x) \right)$$
where $\zeta_\ell(x)$ accounts for the contribution of data outside the $\ell$-th subset.

3. Efficient Sampling:
The closed-form solution for $\pi_{out}$ involves a normalization constant over the entire response space, which is intractable. To bypass this, PEPO employs rejection sampling (Algorithm 2). Any member of the ensemble $\tilde{\pi}_\ell$ serves as a proposal distribution because the numerator of $\pi_{out}$ is bounded by $\tilde{\pi}_\ell$. For practical large-scale applications, the authors also propose a token-level approximation that applies the pessimistic aggregation at each decoding step, avoiding the need for trajectory-level rejection sampling and significantly reducing generation latency.

Key Contributions

Theoretical Contributions:

• Sample Complexity without $\pi_{data}$: In the tabular setting, PEPO achieves sample complexity guarantees depending only on the single-policy concentrability coefficient ($C^\star$). This contrasts with standard DPO, which suffers from an all-policy concentrability coefficient ($C_{all}$), and existing pessimistic DPO variants that require knowledge of $\pi_{data}$.
• Optimal Ensemble Size: The analysis characterizes the optimal ensemble size $L$ as the smallest value sufficient to guarantee pessimism with high probability. Theoretical bounds show that $L$ scales logarithmically with the state and action space sizes.
• Robustness: The bounds hold without assuming the dataset covers the entire policy class, addressing the "unknown $\pi_{data}$" gap in prior theoretical works.

Practical Contributions:

• Empirical Validation: PEPO is evaluated on 7B, 8B, and 34B parameter models (Zephyr-7B, Llama-3.1-8B, Mistral-7B, Yi-34B) using the UltraFeedback dataset.
• Mitigation of Over-optimization: Experiments demonstrate that while standard DPO exhibits a sharp drop in win rates after initial gains (over-optimization), PEPO's performance stabilizes and improves.
• Superiority over Baselines: PEPO consistently outperforms $\chi^2$PO and DPO+SFT in settings where $\pi_{data} \neq \pi_{ref}$ (e.g., UltraFeedback), where the theoretical guarantees of those baselines are invalid.
• Efficiency: The token-level implementation achieves a favorable trade-off between win rate and generation time, offering near-optimal performance with significantly lower latency than full rejection sampling.

Results

• Over-optimization: Figure 1 shows that DPO's win rate against the initial model peaks and then declines, whereas PEPO (with appropriate $L$) shows sustained improvement. Early stopping based on validation loss is shown to be unreliable for selecting the best DPO checkpoint.
• Comparison: In Figure 3, PEPO achieves higher win rates on AlpacaEval2 compared to $\chi^2$PO and SFT+DPO across all tested models. The performance gap narrows only when $\pi_{data} = \pi_{ref}$, confirming that PEPO's advantage stems from its ability to handle unknown data distributions.
• Ablation: Experiments confirm that $L$ is a critical hyperparameter; too small a value fails to induce sufficient pessimism, while too large a value reduces the effective data per model, degrading performance.

Significance and Claims

The paper claims to present the first DPO-like algorithm that provably avoids over-optimization without knowledge of the data-generating distribution. This is significant because it bridges the gap between theoretical robustness (single-policy concentrability) and practical constraints (unknown $\pi_{data}$), a common scenario in real-world LLM alignment where data is generated by proprietary or diverse models.

The authors position PEPO as a practical solution that retains the simplicity of DPO (single-step training, no explicit reward model) while incorporating the theoretical safety of pessimistic regularization. They acknowledge that their theoretical guarantees are currently established in the tabular setting, and extending these results to general function approximation classes without knowing $\pi_{data}$ remains an open challenge. Additionally, the current method outputs a non-parametric policy (via sampling), which is ideal for the final step of a training pipeline but not for intermediate checkpoints, suggesting a direction for future work on parametric outputs.