Leveraging robust optimization for llm alignment under distribution shifts

Thanks for the questions, and we hope the following responses can address the confusion or misunderstanding.

$Q_1$: "Can the authors clearly define how the HH dataset win rate is computed? Is it from GPT-4 comparisons, pairwise preferences, or external rankers?"

As noted in line 550, we use GPT-4o to compute win rates according to the provided evaluation prompt. To improve clarity, we will move this specification to the “Experimental Settings” section in the main paper.

$Q_2$: "What specific role does the classifier play in estimating the density ratio? How is it trained and validated? More clarity and ablations are needed here."

The classifier provides scores used to compute the calibration factor in Equation (6). Training procedures, model architecture, and hyperparameters are detailed in Algorithm 1 and Appendix C.3 (line 536). Concretely, we train a BERT-base classifier using a standard binary classification loss, where the "chosen" response in each original dataset entry is labeled as 1 and all other responses as 0. We use a learning rate of 2e-5 and train for 3 epochs. We appreciate the suggestion for improved clarity and will make this training pipeline more explicit in the final version of the paper.

$Q_3$: "Why is the method described as LLM-specific? Apart from the use of preference data from LLM outputs, the algorithm seems agnostic to language modeling."

Please note that we do not describe this method as LLM-specific. The framework is method-agnostic by design and is motivated as such throughout the paper. We also test it across different model architectures to demonstrate its generalization ability. We are happy to make further clarifications if needed.

$Q_4$: "Can 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."

Thanks for asking. Please see Table 3 in the main paper, where we provide results on Arena-Hard and AlpacaEval. The results suggest that DoRA consistently improves performance on diverse benchmarks.

$Q_5$: "Given the added complexity, how does DoRA perform under simpler ablations? For example, removing the density ratio weighting or using a simpler robustness penalty."

Please note that the calibration factor is one of the core contributions to the proposed method; removing this, it falls back to a standard DRO baseline. We present such ablations in Table 4 (Line 259), comparing DoRA against simpler, robustness-controlled variants such as vanilla DRO and reweighting-based optimization. These results show that DoRA achieves higher performance while maintaining stability.

Thanks again for the questions, and we hope the above responses can mitigate the concerns raised. We are happy to make further clarifications and we welcome further discussion.

Thank you for the insightful comments and helpful suggestions! We are grateful for the positive rating and hope that our responses alleviate the questions that were raised.

$Q_1$: 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.

Thank you for raising this point. We acknowledge that DoRA introduces some additional implementation overhead due to training source-specific classifiers. However, these classifiers are lightweight, fast to train, and only trained once per source. Their outputs can be pre-computed and cached, resulting in negligible overhead during policy optimization.

It’s also important to note that our current experiments are based on relatively clean datasets, which may understate DoRA’s practical utility. As a robustness control method, DoRA’s benefits become more apparent under noisy or significantly-shifted distributions. We show this in our stress tests (see the table below), including: 1. larger shifts in response (we use a model distributionally far from the target (Alpaca‑7B) to generate the synthetic responses. 2. larger label noise (preference labels are randomly corrupted in 40% of the training set).

In both cases, DoRA consistently improves baseline performance on AlpacaEval 2. We expect its value to be even greater in real-world or evolving deployment scenarios, where data often includes noisy synthetic outputs or imperfect human annotations. Overall, DoRA provides strong robustness gains at a modest and manageable implementation cost.

We are happy to include this in the discussion and provide more ablation/analysis in the updated version. Thanks again for the feedback.

$Q_2$: "insufficient discussion of practical optimization considerations"

We appreciate this thoughtful comment. To reduce the numerical instability of log-sum-exp–based objectives, we adopt a stabilizing term $\frac{1}{n}$ as an offset to the denominator (Line 149). This ensures that $\tilde{h}(z)$ is bounded in $(0,n)$, ensuring no single sub-distribution can dominate training. Besides, we also leverage the LSE trick by subtracting the max value before exponentiating to maintain numerical stability in practice. These techniques help stabilize training across all tasks in our experiments. We will include more details in the updated version, and we are happy to release the code to ensure reproducibility.

$Q_3$: "Consider renaming the method, as "DoRA" is already used for a LoRA variant in parameter-efficient training."

Thanks for the kind advice. We will consider another name to avoid confusion in the updated version.

Thanks again for the questions, and we hope the above responses can mitigate the concerns raised. We are happy to make further clarifications and we welcome further discussion.

Thanks for the insightful comments and helpful suggestions! We hope that our responses help alleviate the concerns that were raised.

$Q_1$: "The contribution of new sample-level classifier on existing DRO framework"

Thanks for asking. Our method is not merely a straightforward extension of DRO by adding a classifier. Instead, it introduces a fundamentally new formulation and training objective for robust alignment under synthetic data mixtures:

(i) Problem reframing: We step beyond traditional empirical risk minimization (ERM) approaches used in alignment, and formally define a new robustness setting: mixture response shift. This setting, though increasingly common in practice due to the prevalence of synthetic data, has been under-explored in the literature.

(ii) Robust objectives with calibration factor: Building on this, we derive the robust objective for existing alignment methods by introducing a calibration factor into the uncertainty set, where the probabilistic classifiers are leveraged to compute these factors.

(iii) Closed-form solution with over-pessimism mitigation: The resulting objective admits a closed-form dual solution, and directly mitigates the over-pessimism commonly observed in classical DRO, where the model can over-focus on outliers or noisy examples.

We also conclude some key differences with previous DRO literature in the following table.

$Q_2$: "Performance under adversarial distribution shifts/highly corrupted synthetic data?"

This is a good point. We focus on “clean” benchmark datasets (no synthetic corruptions) in the main paper, but we are happy to add the following corrupted-label experiments to serve as stress tests to confirm DoRA’s ability under adversarial shifts. Specifically, we introduce corruption to the data by randomly flipping the labels. Experiments show that incorporating DoRA consistently improves robustness under increasing label corruption, yielding consistent gains on different benchmarks.

We will add these new experimental results to the revised manuscript and highlight them in the discussion.

$Q_3$: "how DoRA’s calibration is different from reward models."

DoRA's calibration is fundamentally different from reward models, and the two operate on different axes of the alignment pipeline. First, DoRA relies on a learned probabilistic binary classifier which estimates distribution membership, not human-preference utility. While reward models directly score “how good a sample is”, DoRA’s classifier scores "how likely a sample comes from the target distribution". Therefore, DoRA’s calibration factor is a density ratio proxy, not a reward score. Its role is to modulate the DRO uncertainty set (to determine how much robustness weight to assign to each sample).

We agree that any learned model can be biased if trained on unrepresentative or mislabeled data. However, reward model bias directly impacts what the model learns to generate. If it reflects skewed or noisy preferences, it can misdirect optimization and lead to alignment failures. In contrast, our classifier only distinguishes between sources (e.g., human vs. synthetic), enabling selective robustness without steering generation itself. Moreover, the impact of classifier miscalibration is bounded. As we describe in Line 146 and Eq. (6), we introduce a stabilizing factor (e.g., a small additive term in the denominator) that prevents any individual sample from dominating the optimization. This ensures that even overconfident classifier outputs do not result in extremely unstable or biased learning.

$Q_4$:"performance on combining different datasource with different quality, online updates"

Thanks for raising this point. While the primary focus of this paper is on mixture response shift with synthetic data in offline training, we briefly mention online updates as a promising direction in “Limitation and Future Work” section.

To illustrate DoRA’s robustness in such scenarios, we conduct a small-scale experiment: we train LLama3.2‑3B on the HH dataset (Iter 1), then continue training on its own generations for two more iterations. This setup induces an evolving mixture shift as the response distribution becomes increasingly synthetic. The results show that DoRA remains effective under this shift, suggesting that its instance-level calibration also generalizes to online or self-generated data settings, where distribution drift occurs naturally.

$Q_5$: "how do the authors ensure that this does not amplify label noise or miscalibration? What happens when classifier calibration itself is unreliable?"

Thanks for asking, and we try to answer this from the following two perspectives.

1. "Label noise": We first clarify that $Q_0$ denotes the target distribution (trusted human-preferred responses from known sources/ground truth), while label noise in the alignment context typically refers to incorrect “chosen/rejected” annotations. From this perspective, if label noise exists in the original dataset, DoRA can actually mitigate its impact by downweighting those samples that are distributionally misaligned with the target (i.e., with low calibration scores).

2. "Classifier accuracy": If the concern instead refers to errors in the classifier itself, we agree that any learned model may be imperfect. However, we observe that our lightweight classifiers achieve consistently high accuracy across datasets, as shown below, indicating reliable calibration in practice:

"Minor feedback on Clarity"

1. Thanks for pointing this out and we will place the explanation of "Win Rates" earlier in the main paper.

2. "misleadingly depicts $Q_1$ as easily separable from $Q_0$, contradicting evidence that modern LLMs approximate human preferences well."

Thank you for the feedback. We clarify that Figure 1 illustrates distributional separability, not preference similarity. While it is true that modern LLMs (e.g., GPT‑4) may approximate human preferences in aggregate utility—as reflected by similar reward scores—this does not imply that their response distributions are identical. In fact, we find that $Q_1$ (synthetic responses) and $Q_0$ (human-preferred responses) often exhibit distinct generation patterns, which are learnable by a classifier with over 90% accuracy in our setting. This supports the distribution-level separability shown in Figure 1.

We will revise Figure 1 to (i) explicitly clarify the meaning of $Q_1$ and $P$ in the caption and (ii) adjust the visual to better reflect realistic overlaps while still communicating the method’s key intuition.

3. "notation of $h(z)$ vs. $\tilde{h}(z)$ is confusing..."

We clarify that $h(z)$ denotes the Radon-Nikodym derivative $\frac{d Q_0}{dQ}$, which expresses the ideal density ratio between the two distributions, but is generally incomputable in practice. The symbol $\tilde{h}(z)$ refers to our practical, learned approximation of $h(z)$, obtained via a calibrated probabilistic classifier.

Thus, the notation distinction emphasizes that $\tilde{h}(z)$ is an implementable surrogate, not the true derivative. We will clarify this in the final version by explicitly linking the two in notation and explanation.

Thanks again for the questions, and we hope the above responses can address the confusion or misunderstanding. We are happy to make further clarifications and we welcome further discussion.

Thank you for the insightful comments and helpful suggestions! We hope that our responses below help alleviate the concerns that were raised.

$Q_1$. "Construction of $\tilde{h}(z)$ is not completely clear to me. How does one compute the constants $\gamma_j$ and $\beta_j$ ?"

Thanks for the detailed review. We clarify the definitions of $\beta_j$ and $\gamma_j$ as follows:

$\beta_j$ denotes the relative propotion of sub-distribution $Q_j$ in the dataset;

$\gamma_j$ is defined as the imbalance ratio between a synthetic distribution and the golden distribution $\frac{q(y=0)}{q({y=1})}=\frac{\beta_j}{\alpha}$.

For example, suppose for a given query, the responses set includes: 1 response from golden distribution $Q_0$, 2 responses from $Q_1$, and 3 responses from $Q_2$, then $\beta_1=\frac{2}{1+2+3}=\frac{1}{3},\gamma_1$ = $\frac{\beta_1}{\alpha}=\frac{\frac{1}{3}}{\frac{1}{1+2+3}}=2$. Then $\tilde{h}(z)$ can be computed according to Eq. (6), where $c_{\phi_j}$ represents the score from the learned classifier for sub-distribution $Q_j$.

$Q_2$: "How well does this method generalize to datasets not used for training? ..."

Thank you for the insightful question. Table 3 (line214) reports results on AlpacaEval 2 and Arena-Hard using models trained on UltraFeedback, where the test queries come from a different distribution to the training data. These results demonstrate that DoRA generalizes effectively when evaluated out-of-distribution.

We further evaluate generalization capability under the following scenarios: (1) Training on HH-RLHF and testing on summarization (different tasks), and (2) Training with label noise (40% of labels randomly shuffled). As shown in the table below, DoRA improves baseline performance in all cases, demonstrating robust generalization.

$Q_3$. "Clarifying questions, is the "golden" distribution same as the target/preferred distribution? How one may obtain such a data?"

Yes, in our setting, the “golden” distribution corresponds to the positive class when training the probabilistic classifiers. It is constructed with samples drawn from the target/preferred distribution $Q_0$. In practice, we approximate this golden distribution using "chosen" responses from trusted, high-quality preference datasets, where preferences are made either by humans or by a very strong LLM such as GPT‑4. These responses serve as a reasonable proxy for the target distribution.

Thank you again for the detailed review and insightful feedback, and we will include these clarifications in the revised manuscript. We believe your input has already helped improve the paper and look forward to engaging with you further during the discussion period.