We sincerely thank Reviewer Nk1b for the thoughtful and constructive review. We are glad to hear that you found our paper to be “well-motivated by the distributionally robust optimization problem.” We are also encouraged by your recognition that “the proposed approach is clean and easy to follow” and that “the experiments cover multiple scenarios and show consistent improvement.” We hope our responses and new results address your concerns and strengthen the overall contribution of the paper.

Q1: In Table 1 there are two rows of DPO at Epoch 2 (early stopping) for LLaMA-3.2-1B but the number on BBH is different (0.33 vs 0.35). What leads to this difference?

We thank Reviewer Nk1b for catching this typographical error. Both entries refer to the same model checkpoint, and the correct BBH score is 0.35. As verification, Table 3 in the appendix reports all 39 subtask scores on the OpenLLM Leaderboard v2, and their average confirms the 0.35 value. We will correct this in the final version.

Q2: Are there any in-depth analyses and ablation studies on the effect of the parameters $\rho$ and $\tau$?

We thank the reviewer for raising this important point. As discussed in Appendix E, we performed extensive hyperparameter analysis in both the ArmoRM Multi-Objective Alignment and OpenLLM Leaderboard v2 experiments. For instance, Figure 4 and Table 3 (Appendix E) show that both WDPO and KLDPO exhibit relatively low sensitivity to their respective hyperparameters across tasks.

To further investigate, we conducted an additional ablation using the checkpoints from Figure 3, where preferences were generated from 0.5 × Ultrafeedback-Truthfulness + 0.5 × Helpsteer-Complexity. Specifically, we evaluated:

• KLDPO with $\tau$ = 0.5, 1, 5
• WDPO with $\rho$ = 0.01, 0.005, 0.001

on 500 HelpSteer2 prompts, generating 5 completions per prompt, and scoring all completions using armoRM’s 19 sub-rewards.

The results confirm our earlier findings: KLDPO and WDPO demonstrate robustness across a wide range of hyperparameter values. While tuning $\tau$ and $\rho$ can further improve robustness in specific applications, both methods remain stable and effective without exhaustive tuning.

Q3: This paper claims to target on distributional robustness of the algorithm. How is this property related to the benchmark settings in the experiment section? In other words, how are the metrics there connected with distributional robustness?

We thank the reviewer for this important question. Our experimental design specifically targets varying degrees of preference distribution shift to evaluate distributional robustness. We summarize our approach below:

1. Emotion Dataset (Toy Example) – This controlled experiment introduces explicit distributional shift by varying the convex mixing coefficient between two competing objectives (e.g., “anger” vs. “fear”). It demonstrates how traditional DPO degrades under even modest shifts.

2. ArmoRM Multi-Objective Alignment – This more realistic setup simulates partial coverage of reward space. Models are trained on a mixture of two ArmoRM subrewards and then evaluated on other subrewards unseen during training. This mimics real-world scenarios where deployed models encounter user preferences outside their training signal.

3. OpenLLM Leaderboard v2 – This is our most challenging and practical testbed. Models are trained using ArmoRM (ranked #32 on RewardBench, then evaluated on OpenLLM Leaderboard v2, which aggregates feedback from various human-preference and judge models. This highlights a genuine failure mode: a developer finetunes a model using ArmoRM but sees performance drop on broader evaluation. Our methods aim to mitigate this degradation.

Across all settings, distributional shift is either explicitly modeled or naturally induced. The metric of interest in each case is the model’s reward under out-of-training-distribution preferences. Hence, distributional robustness is defined as the ability to maintain performance across such shifts, whether synthetic or real-world (e.g., varying mixtures of user goals, cultural expectations, or evaluation benchmarks).

We believe this progression from toy to real-world tasks demonstrates the practical importance of robust preference optimization.

We thank Reviewer oGJc for taking the time to carefully review our work and for providing thoughtful and detailed feedback. We are glad to know that the reviewer found that "preference distribution shift is an important problem", "[our] methodology effectively addresses this issue", and "experiments were conducted across various distribution shift setups and evaluated on multiple models". Below, we address your specific concerns.

Q1: Comparison between SFT and DPO.

We thank Reviewer oGJc for raising a critical point regarding the observation that DPO underperforms the SFT model (LLaMA-3.2-1B-Instruct) on the OpenLLM Leaderboard v2.

To clarify upfront: when evaluated on the same reward model used to generate the preference dataset, armoRM, DPO consistently outperforms the SFT baseline, as expected. We verified this through a controlled evaluation (see supplementary tables below), using the exact model checkpoints from Table 1, and scoring model outputs across all 19 subrewards with armoRM. This confirms the experimental pipeline is sound and not flawed. As expected from classical ML, longer training improves in-distribution performance on the training reward.

We sampled 500 test prompts from the HelpSteer 2 dataset. For each model, we generated five completions per prompt. These were then scored using armoRM across 19 subrewards, aggregated across 2500 completions.

Now back to Table 1 in our paper, the underperformance of DPO stems from a well-known issue in RLHF: reward hacking and distributional overfitting. As studied in Azar et al. (2024) and other works on reward hacking (Amodei et al., 2016; Skalse et al., 2022; Eisenstein et al., 2023), methods like DPO (and RLHF PPO) can overfit to the training-time reward model, here, armoRM, when the test-time reward or preference distribution shifts, e.g., to OpenLLM Leaderboard v2.

We hope this resolves the concern and reinforces the importance of developing preference optimization methods that are robust to reward and preference distribution shifts.

Q2: ...this study lacks baseline experiments and only compares against DPO...

We thank Reviewer oGJc for highlighting the importance of comparing against other robust preference learning methods such as IPO and DRDPO. In response, we trained IPO and DRDPO using preferences generated according to 0.5 × Ultrafeedback-Truthfulness + 0.5 × Helpsteer-Complexity. Following our evaluation protocol, we sampled 500 prompts from the HelpSteer 2 test set and generated five completions per model per prompt. Our results show that both KLDPO and WDPO outperform DPO and IPO on several key sub-rewards under the armoRM evaluation. While DRDPO performs competitively on some tasks, KLDPO and WDPO consistently achieve stronger overall scores, confirming the benefit of our proposed distributionally robust methods.

Q3: While preference distribution shift is indeed a problem, it is not as extreme as the current experimental setup suggests. Switching objectives between anger and fear for emotional alignment represents an overly dramatic scenario.

We appreciate Reviewer oGJc’s perspective on the severity of preference distribution shifts in our experimental setup. While we agree that the emotion alignment experiment is a simplified toy example designed to isolate and illustrate core failure modes, our remaining experiments address far more realistic and impactful scenarios.

In particular, our second experiment uses armoRM, a strong, publicly available reward model that currently ranks #32 on RewardBench, outperforming many proprietary models. For comparison:

• google/gemini-2.5-pro scores #30
• openai/gpt-4.1-mini-2025-04-14 scores #33

Notably, Arena-Hard 2, a widely used benchmark for LLM evaluation, includes models such as Gemini-2.5, Gemini-2.5-Flash, and OpenChat 4, which are in a similar performance tier. armoRM also surpasses GPT-4 in LLM-as-a-Judge setups, making it a viable and meaningful alternative for annotation and benchmarking tasks (e.g., MT-Bench, AlpacaEval-2.0, Arena-Hard).

In our third experiment, we evaluate models on the OpenLLM Leaderboard v2, which aggregates scores across multiple axes and reflects the state-of-the-art in community-driven LLM evaluation. These setups represent genuine real-world scenarios where preference shift arises, such as switching reward models between training, tuning, and deployment phases.

Thus, our proposed methods address not only extreme preference shifts but also moderate and subtle ones that are common yet hard to detect in practice. Their value lies in mitigating performance degradation across this full spectrum of shift severity, which is central to ensuring robust alignment in deployed systems.

References

Mohammad Gheshlaghi Azar, Mark Rowland, Bilal Piot, Daniel Guo, Daniele Calandriello, Michal Valko, and Rémi Munos. A general theoretical paradigm to understand learning from human preferences. ArXiv, abs/2310.12036, 2023

Amodei, D., Olah, C., Steinhardt, J., Christiano, P., Schulman, J., and Mané, D. (2016). Concrete problems in ai safety. arXiv preprint arXiv:1606.06565.

Skalse, J., Howe, N., Krasheninnikov, D., and Krueger, D. (2022). Defining and characterizing reward gaming. Advances in Neural Information Processing Systems, 35:9460–9471.

We thank Reviewer S16R for the thoughtful and detailed review. We are encouraged to hear that you found our work to provide “practical solutions (WDPO and KLDPO) for addressing preference distribution shift,” which you identify as “a critical issue in real-world LLM deployment.” We also appreciate your positive remarks on the combination of “solid theoretical guarantees with scalable, implementable algorithms,” and your recognition that our “empirical results consistently demonstrate robustness and improvement over standard DPO.” We hope our responses further clarify the contributions and scope of our work.

Q1: Can you show the effectiveness of your methods on larger models and larger datasets?

We appreciate the reviewer’s interest in scaling our methods. Due to resource constraints common in academic settings, we were limited to training models up to the size of LLaMA 8B. While this may not match the scale of industry models, we believe it still provides valuable insights into scalability. Notably:

• Our algorithms were applied across a range of model sizes: LLaMA 1B (used in the ArmoRM multi-objective alignment task), 3B, and 8B.
• ArmoRM itself is a powerful LLaMA-8B reward model, offering a strong and challenging signal for preference optimization.
• Our successful extension of KLDPO to LLaMA 8B demonstrates practical scalability within feasible academic compute budgets.

We agree that further evaluation at larger scales is an important future direction, and we hope to pursue this as additional compute resources become available.

Q2: Can your distributional robustness framework be extended to non-DPO alignment methods such as PPO-based RLHF or Imitation Learning?

Thank you for this insightful question. Extending distributionally robust alignment beyond DPO is a compelling future direction.

One promising avenue is to robustify the reward model itself, learning a distributionally robust reward function that generalizes well under test-time distribution shifts. For instance, in preference modeling, we can consider a reward function of the form $r_\theta(s, a) = \langle \theta,\; \phi(s, a) \rangle$, where $\phi$ is a feature map, and estimate it via maximum likelihood using a preference dataset $\mathcal{D} = \{(s\_i, a^1\_i, a^2\_i, y\_i)\}\_{i=1}^N \sim P^o$. The negative log-likelihood loss of $\theta$ is

$$\mathcal{L}(\theta \mid \mathcal{D}) = - \sum\_{i=1}^N \log[\mathbf{1} (y\_i=1) \sigma( \langle {\theta,\;\phi(s_i,a_i^1)-\phi(s_i,a_i^2)}\rangle) + \mathbf{1} (y\_i=0) \sigma( \langle {\theta,\;\phi(s_i,a_i^2)-\phi(s_i,a_i^1)}\rangle)].$$

Robustifying this objective with respect to distributional shifts in $P^o$, which is used to sample $\mathcal{D}$, could yield a reward model that is resilient to preference drift.

This robust reward can in turn be used for PPO-style RLHF, as part of a broader robust alignment pipeline.

We also note that recent works have proposed distributionally robust behavior cloning and imitation learning algorithms (Bashiri et al.; Panaganti et al.), which could complement our framework in future explorations.

References:

M. A. Bashiri, B. Ziebart, and X. Zhang, “Distributionally robust imitation learning,” Advances in neural information processing systems.

Panaganti, K., Xu, Z., Kalathil, D., and Ghavamzadeh, M. (2023b). Distributionally robust behavioral cloning for robust imitation learning. In 2023 62nd IEEE Conference on Decision and Control (CDC), pages 1342-1347.vol. 34, pp. 24404–24417, 2021.

Q3: Have you considered evaluating your methods under real-world preference shift scenarios, such as regional, demographic, or language-based distributions, rather than synthetic datasets designed for simplified testing?

We thank the reviewer for this valuable suggestion. While we agree that regional, demographic, and language-based preference shifts represent important real-world challenges, public datasets with explicit annotations for such axes of variation remain limited or proprietary. As a result, conducting controlled and reproducible experiments along these lines is currently nontrivial.

That said, we would like to emphasize that our experimental setup already captures meaningful and realistic distribution shifts:

• In our ArmoRM multi-objective alignment experiment, we simulate preference shifts by training on a mixture of sub-objectives (e.g., helpfulness + correctness) and evaluating on sub-objectives held out during training (e.g., coherence or verbosity). This mimics real-world scenarios where a deployed model may face user goals not seen in training.

• In our OpenLLM Leaderboard evaluation, we train using ArmoRM, which represents a single (though strong) judge model, and test on a diverse aggregation of preferences across multiple human and automated evaluation sources. This shift from training to evaluation naturally induces real-world preference mismatch, even if not explicitly labeled by region or demographic.

We agree that future work could explore explicit real-world axes such as culture, age, or language, especially as more datasets with these attributes become publicly available. We are excited by this direction and view our current work as a foundation for robust preference optimization under broader forms of distributional variability.

Q4: Can you discuss more about the computational overhead of your methods?

We thank the reviewer for raising this important point. We provide a detailed discussion on computational overhead and implementation optimizations in Appendix F of our paper.

Briefly:

• WDPO introduces a gradient penalty term that requires differentiating the DPO loss w.r.t. model input. Since tokenized inputs are discrete, we track gradients at the embedding layer. Moreover, to reduce variance from small micro-batch sizes in distributed LLM training, we upper-bound the square-root penalty with its squared counterpart, leading to a tractable per-example regularizer.

• KLDPO introduces a reweighting scheme across the full batch. To stabilize this under distributed training, we insert an all_gather operation to synchronize losses across workers, mitigating the high variance from local micro-batch estimates.

Despite these added steps, the rest of the DPO pipeline (optimizer, learning rate schedule, model parallelism setup) remains unchanged. We found the overall overhead to be modest and manageable across both small-scale (single A100) and large-scale (8xH100 with DeepSpeed ZeRO-2) experiments.

We hope this reassures the reviewer that our methods are practical and scalable with careful engineering, and we will make these implementation strategies explicit in our updated appendix.

We sincerely thank Reviewer 4S36 for their thoughtful and encouraging review. We are especially grateful that you found our formulation of distributionally robust preference optimization to be “very novel” and “well motivated,” and that you appreciated the theoretical depth of our work, including the error bounds and SGD-style algorithms. We are also pleased that you found our empirical evidence convincing and the paper clearly written, despite the technical complexity. Below, we address your insightful questions and suggestions, particularly regarding dataset choices, robustness to noisy data, and integration into existing DPO training frameworks.

Q1: When technically very impressive and comprehensive, it is unclear how easily their variants of SGD (Algorithms 1-2) would fit into existing DPO training libraries, which could limit the impact of these new approaches. I would like to see the authors address this. A code example would be very useful to have in the appendix.

We thank the reviewer for highlighting the importance of integration with existing DPO pipelines. Our implementations of WDPO and KLDPO are designed for minimal disruption: we modify only the loss function, keeping the rest of the DPO pipeline (data, sampling, training loop, etc.) unchanged.

To illustrate this, we provide the following pseudocode sketch built directly upon the DPOLoss structure in openrlhf. The core DPO logic remains intact and is reused in both KLDPOLoss and WDPOLoss. Only the lines marked between === blocks are additions specific to each variant.

class (KL/W)DPOLoss(nn.Module):
    def __init__(self, ...) -> None:
        super().__init__()
        # DPO parameters
        ...

        # KLDPO and WDPO hyperparameters
        self.kldpo_tau = ...
        self.wdpo_rho = ...

    def forward(
        self,
        policy_chosen_logps, policy_rejected_logps,
        reference_chosen_logps, reference_rejected_logps,
        ...
    ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

        # Calculating vanilla DPO loss (common to DPO, KLDPOLoss, and WDPOLoss)
        pi_logratios = policy_chosen_logps - policy_rejected_logps
        ref_logratios = reference_chosen_logps - reference_rejected_logps
        logits = pi_logratios - ref_logratios

        losses = (
            -F.logsigmoid(self.beta * logits) * (1 - self.label_smoothing)
            - F.logsigmoid(-self.beta * logits) * self.label_smoothing
        )

        # ======== KLDPO Extension (added in KLDPOLoss) ========
        all_losses = [torch.zeros_like(losses) for _ in range(dist.get_world_size())]
        dist.all_gather(all_losses, losses.detach())
        all_losses = torch.cat(all_losses)  # Gather losses across all distributed ranks

        kl_reweights = torch.exp((losses - all_losses.mean()) / self.kldpo_tau).clamp(max=10).detach()
        # KL reweighting approximates the adversarial distribution in KL-DRO

        loss = (kl_reweights * losses).mean()  # Final KLDPO objective
        # =======================================================

        # ======== WDPO Extension (added in WDPOLoss) ========
        grads = torch.autograd.grad(
            outputs=losses,
            inputs=model_embeddings,
            grad_outputs=torch.ones_like(losses),
            create_graph=(train_eval == 'train')
        )[0]
        # Calculate WDPO gradient penalty

        grads = grads.pow(2).sum(dim=(-1, -2))
        grads = grads[: len(grads) // 2] + grads[len(grads) // 2 :]
        losses = losses + self.wdpo_rho * grads
        loss = losses.mean()  # Final WDPO objective
        # ===================================================

        chosen_rewards = self.beta * (policy_chosen_logps - reference_chosen_logps).detach()
        rejected_rewards = self.beta * (policy_rejected_logps - reference_rejected_logps).detach()

        return loss, chosen_rewards, rejected_rewards

Q2: Did you train on datasets besides HelpSteer? If so, why not?

Yes, we use two distinct prompt sources in our experiments. The distinction we study is not just prompt diversity, but preference distribution shift, so we fix the prompt pool (e.g., HelpSteer2) and vary how reward signals are assigned.

---

Why not UltraFeedback prompts?

While UltraFeedback has 60k prompts, its pairwise preference labels are fixed and cannot be flexibly recomposed. In contrast, HelpSteer2 prompts combined with ArmoRM enable structured and interpretable simulations of preference shift.

---

What is ArmoRM and why use it?

ArmoRM is not a black-box chatbot judge like GPT-4 or Claude. It is a multi-objective reward model trained with absolute ratings across human-aligned dimensions (e.g., honesty, helpfulness, verbosity). This design allows us to explicitly control which sub-objectives are used for training vs. testing, precisely modeling the kind of distribution shift our methods are built to handle.

Q3: It seems like your approach still assumes that data is of high quality. Do you think that your approach could be extended for modeling data with noise? A discussion of this in the paper would be helpful.

We thank the reviewer for raising this important point. While our methods were not originally designed to explicitly model noisy preference labels, we explored this direction in a follow-up experiment and observed promising robustness trends.

Specifically, we trained DPO, KLDPO, and WDPO on preferences generated by a synthetic reward:
0.5 × HelpSteer-Complexity + 0.5 × UltraFeedback-Truthfulness,
and then introduced 10% and 20% label noise by randomly flipping a subset of the pairwise preferences.

All models were trained on these corrupted datasets and evaluated on 500 prompts from the HelpSteer2 test set, each with 5 completions. These completions were scored using the same two ArmoRM subrewards: HelpSteer-Complexity and UltraFeedback-Truthfulness, to isolate the effect of label noise as the only varying factor.

The results below show that WDPO consistently maintains performance under increasing noise, especially compared to DPO and KLDPO:

This suggests that our distributionally robust formulation, though originally motivated by test-time preference shifts, may also confer some resilience against noisy training data.