Title: Leveraging robust optimization for llm alignment under distribution shifts
Venue: NeurIPS 2025
Accepted: yes
Keywords: preference alignment

distribution shifts

robust optimization
Abstract: Preference alignment methods are increasingly critical for steering large language models (LLMs) to generate outputs consistent with human values. While recent approaches often rely on synthetic data generated by LLMs for scalability and cost-efficiency reasons, this reliance can introduce distributional shifts that undermine the nuanced representation of human preferences needed for desirable outputs. In this paper, we propose a novel distribution-aware optimization framework that improves preference alignment despite such shifts. Our approach first leverages well-learned classifiers to assign a calibration value to each training sample, quantifying its alignment with the target human-preferred distribution. These values are then incorporated into a robust optimization objective that minimizes the worst-case loss over regions of the data space most relevant to human preferences. By explicitly focusing optimization on the target distribution, our approach mitigates the impact of distributional mismatch and improves the generation of responses that better reflect intended values.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes DoRA (Distributionally Robust Alignment), a new preference learning framework inspired by distributionally robust optimization (DRO), applied to LLM alignment. The key idea is to model preference optimization as minimizing regret over a worst-case distribution shift between the original preference dataset and the LLM’s evolving policy. To do this, the authors train classifiers to estimate the density ratio between the preferred and dispreferred responses, use those to compute importance weights, and apply a regularized regret loss to update the model. The method is evaluated primarily on the HH-RLHF dataset with LLaMA-7B, showing improved win rate over DPO and SimPO.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths: The formulation is grounded in well-studied robust optimization theory.\n\nWeaknesses:\n\nDisconnect between method and LLM context:\nAlthough the paper positions itself as a contribution to LLM alignment, the methodology appears to be a generic instance of distributionally robust optimization (DRO) applied to a preference dataset. The classifier-based density ratio estimation and reweighting procedure is standard in DRO literature, and the paper does not explain how the unique challenges of LLM alignment inform its algorithmic design. There is little adaptation to the sequential or structured nature of LLM outputs. As a result, the method feels like a transplant from a different domain rather than a carefully tailored alignment strategy.\n\nUnclear evaluation metric:\nThe primary metric reported in the experiments is “HH dataset win rate,” but the paper fails to explain how this metric is computed. Is the win rate measured via GPT-4 preferences, a learned reward model, or some internal heuristic? What exactly constitutes a \"win\" in this setting? Against which model? Are responses sampled once per prompt or averaged over multiple runs? Without a clear description, it's hard to judge the reliability or fairness of this metric. Given that the win rate is the only reported outcome, this lack of transparency significantly weakens the empirical claims.\n\nLimited empirical validation:\nThe evaluation is restricted to a single dataset (HH-RLHF) and a single win rate metric. No comparisons are made on other standard benchmarks used for LLM alignment evaluation, such as MT-Bench, AlpacaEval, Arena-Hard, or Open LLM Leaderboard metrics. Additionally, there is no breakdown of performance across prompt categories or qualitative error analysis. This narrow evaluation makes it hard to judge whether the proposed method offers any practical advantages in broader or more realistic alignment settings.\n\nUnnecessarily complex and opaque setup:\nThe proposed method requires training separate classifiers to estimate the density ratio between preferred and dispreferred responses. These are then used to reweight training samples when optimizing the policy. However, the motivations for these design choices are not clearly justified, and the pipeline introduces substantial complexity. The use of SFT initialization is briefly mentioned but not analyzed. Overall, the system feels overengineered relative to the simplicity of existing approaches like DPO, and it is unclear whether the added complexity yields meaningful or robust benefits.\n\nQUESTIONS:\nCan the authors clearly define how the HH dataset win rate is computed? Is it from GPT-4 comparisons, pairwise preferences, or external rankers?\n\nWhat specific role does the classifier play in estimating the density ratio? How is it trained and validated? More clarity and ablations are needed here.\n\nWhy is the method described as LLM-specific? Apart from the use of preference data from LLM outputs, the algorithm seems agnostic to language modeling.\n\nCan the authors provide evaluation on more diverse benchmarks (e.g., MT-Bench, Arena-Hard, AlpacaEval)? One win rate metric is insufficient to support alignment claims.\n\nGiven the added complexity, how does DoRA perform under simpler ablations? For example, removing the density ratio weighting or using a simpler robustness penalty.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 2, "REVIEWER_CONFIDENCE": 3, "CLARITY": 1, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-20"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes DoRA (Distribution-aware optimization for Robust Alignment) to address distribution mismatch in preference learning. The core issue is that maximum likelihood estimators for optimal policies require samples from the target distribution, but training data typically contains a mixture of human-generated and LLM-generated responses from various sources.\n\nDoRA uses distributionally robust optimization to handle this distribution shift. The method trains classifiers to identify the source distribution of each sample, then incorporates these predictions as importance weights in a robust optimization objective. This transforms the standard sum-of-log-exponentials loss into a log-sum-exponential form that minimizes worst-case expected loss. The approach is method-agnostic, working with various preference learning algorithms beyond DPO.\n\nExperiments on HH-RLHF and UltraFeedback datasets using Mistral-7B and Llama-8B show improvements over baselines like Robust DPO and Dr.DPO across most settings.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n- **Technically sound approach** with clear motivation addressing a practical problem in preference learning\n- **Method-agnostic approach** that extends beyond DPO to various preference learning algorithms\n- **Well-organized presentation** with focused contributions and thorough experimental evaluation\n\n## Weakness\n* **Limited practical gains vs. complexity**: Requires training separate BERT classifiers for each source distribution, adding significant implementation complexity for small improvements on ArenaHard and AlpacaEval 2.\n\nQUESTIONS:\n* (traditionally) Log-sum-exponential objectives are notoriously difficult to optimize, prone to numerical instability and vanishing gradients. While the authors add a 1/n stabilizing term, insufficient discussion of **practical** optimization considerations may hinder reproducibility\n* Consider renaming the method, as \"DoRA\" is already used for a LoRA variant in parameter-efficient training.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper introduces DoRA (Distribution-aware Optimization for Robust Alignment), a framework aiming to address the LLM alignment in the presence of distribution shifts. Taking the modern trend of using synthetic preference data as an example, the paper propose the method that model remains resilient to distribution shifts between the training data and the target distribution while still benefiting from their scalability.\nThe authors propose a two-phase approach that first uses distribution classifiers to estimate the likelihood that training samples originate from the target distribution Q₀, then incorporates these calibration scores into a KL-divergence-based Distributionally Robust Optimization objective to reweight the loss function toward samples better aligned with the target distribution.\nEmpirical evaluation across multiple alignment tasks demonstrates that DoRA consistently improves model performance and win rates compared to various baseline approache\n\nSTRENGTHS AND WEAKNESSES:\n### Strengths\n\n- The proposed framework is well-motivated, addressing a crucial challenge in alignment training in the divergence between synthetic and target distributions.\n- The expansion beyond traditional pairwise noise robustness to broader distributional concerns is a timely and necessary shift in the alignment literature.\n- The proposed method is method-agnostic and can be integrated into diverse alignment objectives, including different synthetic sources or online updates\n- The writing is generally clear and well-structured. Key concepts such as mixture response shift, KL-DRO, and the calibration mechanism are explained clearly. The proposed method is demonstrated in a both principled and practical way.\n- Experiments span pairwise and listwise settings across multiple datasets and models, showing consistent (although sometimes modest) gains in alignment quality.\n\n### Weaknesses\n- The contribution of introducing new sample-level classifier on existing method DRO offers limited novelty and add marginal contributions. \n- While the paper provides theoretical motivation, the empirical gains are sometimes marginal (e.g., +0.6% in win rate for DPO with Mistral-7B), raising questions about cost-effectiveness. The added complexity of training additional classifiers and additional inference might not be justified.\n- There is insufficient clarity on how DoRA’s calibration is different from reward models. Both attempt to score sample quality and may collapse to a combined reward models. The classifier based approach still hold the limitation of reward models that are prone to biases, potentially leading to suboptimal labelling. The authors should discuss this distinction more clearly.\n- The approach might amplify label noise rather than mitigate it, since it emphasizes Q₀ without validating whether those labels are truly correct.\n\nMinor feedback on Clarity:\n- The details about win rate (e.g. using GPT-4o) is introduced first in Table 2. Earlier context would help orient readers.\n- Figure 1 misleadingly depicts $Q_₁$ as easily separable from $Q_₀$, contradicting evidence that modern LLMs approximate human preferences well. In addition, unclear definition of $Q_1$ makes P unclear in Figure1, which prevent it from delivering the clear picture of overall method and its strength.\n- notation of $h(z)$ vs. $h^{\\sim} (z)$ is confusing, as one refers to the Radon-Nikodym derivative and the other to the calibration term.\n\nQUESTIONS:\n- Could the authors elaborate more on the conceptual (and technical) novelty of introducing a sample-level classifier within the DRO framework? In particular, how does this differ fundamentally from prior work, and what capabilities does it enable beyond existing methods?\n- How does DoRA perform under adversarial distribution shifts or highly corrupted synthetic data? An ablation study with listedwise dataset from LLMs of varying quality might be helpful.\n- Although the author mentioned about the flexibility of the proposed method with diverse distributional shift, the author only demonstrated on listwise preference setting with synthetic response. Other source such as combining different datasource with different quality, online updates would help strengthen the contribution.\n- Since DoRA emphasizes Q₀-labeled samples, how do the authors ensure that this does not amplify label noise or miscalibration? What happens when classifier calibration itself is unreliable?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-29"}

comments: SUMMARY:
The paper consider the problem of robust preference alignment in LLM. The approach is to first learn a classifier to assign a calibration value to each training sample, and then leverage it to the target human-preferred distribution. The objective is formulated as a distributionally robust optimization with respect to KL divergence.

STRENGTHS AND WEAKNESSES:
**Strengths:** The paper consider a distribution-aware optimization framework to robustly align LLM outputs with human preferences under distribution shifts. The framework relies on the mixture response shift assumption, which posits that the training data comprises a mixture of heterogeneous sub-distributions. Then the procedure is broken into two parts: 1. employ a calibration score to each training sample, estimating its alignment to human preferences, and 2. integrate the calibration scores to a distribution-aware optimization objective that minimizes the worst case loss. 

**Weaknesses:** I have some concerns about clarity of the expositions and generalizability of the experimental results. Please see my questions in the **Questions** section.

QUESTIONS:
1. Construction of $\tilde h(z)$ is not completely clear to me. How does one compute the constants $\gamma_j$ and $\beta_j$?

2. How well does this method generalize to datasets not used for training? Did the authors test the aligned model on a dataset that has not been used for training (not considering train-test split)?

3. I see that golden/reference distribution is approximated using trusted human-labeled data. Clarifying questions, is the "golden" distribution same as the target/preferred distribution? How one may obtain such a data?

LIMITATIONS:
Please see the **Strengths And Weaknesses** and **Questions** sections.
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-06-30