Distributionally Robust Token Optimization in RLHF

Yeping Jin; Jiaming Hu; Ioannis Ch. Paschalidis

cs.LGMar 27, 2026

Distributionally Robust Token Optimization in RLHF

Authors: Yeping Jin, Jiaming Hu, Ioannis Ch. Paschalidis

Organizations: Department of System Engineering Boston University, Boston, MA 02215, USA · Department of Math & Statistics Boston University, Boston, MA 02215, USA

Abstract

Large Language Models (LLMs) tend to respond correctly to prompts that align well with the data they were trained and fine-tuned on. Yet, small shifts in wording, format, or language can trigger surprisingly large failures, especially on multi-step reasoning problems. To address this problem, we propose a Distributionally Robust Token Optimization (DRTO) approach, which combines token-level Reinforcement Learning from Human Feedback (RLHF) with Distributionally Robust Optimization (DRO). DRTO constructs f-divergence ambiguity sets over span-level actor losses, providing a principled way to emphasize difficult response segments during policy optimization. Empirically, DRTO enhances consistency under distribution shifts in multiple reasoning benchmarks among different tasks, achieving $+4.4$ percentage points on MATH-500 and $+2.7$ percentage points on LiveCodeBench over standard RTO.

Figures & tables

Figure 1: Performance of five RLHF fine-tuning methods on math and coding benchmarks under distribution shift. DRTO improves over standard RTO on most benchmarks.

Figure 2: Overview of the DRTO pipeline. Each rollout is partitioned into spans at low-confidence cut points (where $p^{\mathrm{old}}_{i,h}<p_{\mathrm{cut}}$ ); per-token PPO losses driven by the RTO token reward are aggregated into span losses, which are then reweighted by a KL- or $\chi^{2}$ -divergence DRO adversary to form the final actor loss.

Figure 3

Models	Accuracy
	GSM8K	GSM (CoT)	GSM (ES)	MATH-500	GPQA
RTO	73.2	74.4	66.0	18.0	26.3
SFT	61.4 -11.8	64.8 -9.6	63.6 -2.4	17.4 -0.6	24.8 -1.5
PPO	73.2 0.0	75.2 +0.8	65.6 -0.4	17.0 -1.0	25.8 -0.5
DPO	73.9 +0.7	76.4 +2.0	66.4 +0.4	20.1 +2.1	20.7 -5.6
FlowRL	75.9 +2.7	76.2 +1.8	66.0 0.0	19.2 +1.2	28.8 +2.5

Table 1: Accuracy of Llama3-8B across five benchmarks. Superscripts report absolute change (percentage points) compared to standard RTO. Best results are in bold .

Table 5

Table 6

Appendix figures & tables2 assets

Supplementary material from the paper’s appendix.

Appendix

Panel of Figure 3 — Figure 3: Training dynamics comparison across methods.

Question. Terry eats 2 yogurts a day. They are currently on sale at 4 yogurts for $5.00. How much does he spend on yogurt over 30 days?
RTO. Terry eats 2 yogurts a day, so over 30 days he eats 230=<<230=60>>60 yogurts. They are currently on sale at 4 yogurts for $5.00, so Terry spends 60/4=$ <<60/4=15>>15.00 on yogurt over 30 days. #### 15
$\chi^{2}$ -DRTO ( $\rho=0.1$ ). Terry eats 2 yogurts a day, so over 30 days he eats 230=<<230=60>>60 yogurts. They are currently on sale at 4 yogurts for $5.00, so Terry spends$ 5.00 for every 4 yogurts. To find out how much Terry spends on 60 yogurts, we can divide 60 by 4 to find out how many sets of 4 yogurts he needs: 60/4=<<60/4=15>>15 sets of 4 yogurts. Since each set of 4 yogurts costs $5.00, Terry spends 15$ 5.00=<<155=75>>$75.00 on yogurt over 30 days. #### 75
KL-DRTO ( $\rho=0.01$ ). Terry eats 2 yogurts a day, so over 30 days he eats 230=<<230=60>>60 yogurts. The sale is 4 yogurts for $5.00, so the price per yogurt is$ 5.00/4 = $<<5/4=1.25>>1.25. Therefore, Terry spends 60 $ 1.25 = $<<601.25=75>>75 on yogurt over 30 days. #### 75

Appendix

Table 6: Example responses on the same prompt. RTO misses the pack price multiplication, while both DRTO variants return the correct total cost.

Language Model Robustness

Language model robustness studies how reliably language models preserve correct, safe behavior under input changes, distribution shifts, and adversarial pressure.

See all 493 papers

Distributionally Robust Reinforcement Learning with Human Feedback
Mar 1, 2025Debmalya Mandal, Paulius Sasnauskas, Goran RadanovicDirect Preference Optimization LLM Fine-Tuning
Dept. of Computer Science University of Warwick, UK · Department of Computing Science University of Alberta, Edmonton, Canada · Max-Planck Institute for Software Systems, Germany
Beyond Entropy: Learning from Token-Level Distributional Deviations for LLM Reasoning
Jun 18, 2026Xuanzhi Feng, Zhengyang Li, Zeyu Liu +6Reinforcement Learning RL for Language Model Reasoning
The Hong Kong University of Science and Technology, Hong Kong, China · Sichuan University, Chengdu, China · The Hong Kong Polytechnic University, Hong Kong, China
STAPO: Stabilizing Reinforcement Learning for LLMs by Silencing Rare Spurious Tokens
Feb 17, 2026Shiqi Liu, Zeyu He, Guojian Zhan +10Reinforcement Learning RL for Language Model Reasoning
School of Vehicle and Mobility & College of AI, Tsinghua University · Didi Voyager Labs, DiDi Autonomous Driving

Mar 1, 2025J/K move · Enter open · S save

Distributionally Robust Reinforcement Learning with Human Feedback

Debmalya Mandal, Paulius Sasnauskas, Goran Radanovic

Dept. of Computer Science University of Warwick, UK · Department of Computing Science University of Alberta, Edmonton, Canada · Max-Planck Institute for Software Systems, Germany

Direct Preference Optimization LLM Fine-Tuning Distributionally Robust RL Robust RL RL from Human Feedback

Open paper AskPDF arXiv

Reinforcement learning from human feedback (RLHF) has evolved to be one of the main methods for fine-tuning large language models (LLMs). However, existing RLHF methods are non-robust, and their performance deteriorates if the downstream task differs significantly from the preference dataset used in fine-tuning. In order to mitigate this problem, we introduce a distributionally robust RLHF for fine-tuning LLMs. In particular, our goal is to ensure that a fine-tuned model retains its performance even when the distribution of prompts significantly differs from the distribution encountered during fine-tuning. We formulate distributionally robust optimization (DRO) version of two popular fine-tuning methods -- (1) reward-based RLHF and (2) reward-free DPO (direct preference optimization). We propose a minibatch gradient descent based algorithms for both of them, and theoretically prove convergence guarantees for the algorithms. Subsequently, we evaluate our algorithms on an out-of-distribution (OOD) task by first training the model on the Unified-Feedback dataset and evaluating its performance on two different datasets. The experimental results show that our robust training improves the accuracy of the learned reward models on average, and markedly on some tasks, such as reasoning. Furthermore, we show that the robust versions of policy optimization methods, similarly improve performance on OOD tasks.

Abstract

Language Model Robustness

Explore similar work

Distributionally Robust Reinforcement Learning with Human Feedback

Beyond Entropy: Learning from Token-Level Distributional Deviations for LLM Reasoning

STAPO: Stabilizing Reinforcement Learning for LLMs by Silencing Rare Spurious Tokens