# Distributionally Robust Token Optimization in RLHF

arXiv: https://arxiv.org/abs/2604.08577 · HTML: https://mlsanity.com/paper/2604.08577/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Placed in Reward modeling and RLHF, one of 10 directions in LLM alignment / RLHF / preference optimisation.
- Builds on 27 of its 36 references that are in the map.
- First posted on arXiv 27 Mar 2026.

## Placed in

- LLM alignment / RLHF / preference optimisation / [Reward modeling and RLHF](https://mlsanity.com/alignment/reward-modeling-and-rlhf/)

## Builds on (most cited first, 27 of 27)

- [Training language models to follow instructions with human feedback](https://mlsanity.com/paper/2203.02155/) (2022)
- [Proximal Policy Optimization Algorithms](https://mlsanity.com/paper/1707.06347/) (2017)
- [Evaluating Large Language Models Trained on Code](https://mlsanity.com/paper/2107.03374/) (2021)
- [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://mlsanity.com/paper/2305.18290/) (2023)
- [Program Synthesis with Large Language Models](https://mlsanity.com/paper/2108.07732/) (2021)
- [DAPO: An Open-Source LLM Reinforcement Learning System at Scale](https://mlsanity.com/paper/2503.14476/) (2025)
- [Convex optimization](https://mlsanity.com/paper/2106.01946/) (2021)
- [High-Dimensional Continuous Control Using Generalized Advantage Estimation](https://mlsanity.com/paper/1506.02438/) (2015)
- [LiveCodeBench: Holistic and Contamination Free Evaluation of Large Language Models for Code](https://mlsanity.com/paper/2403.07974/) (2024)
- [Is Your Code Generated by ChatGPT Really Correct? Rigorous Evaluation of Large Language Models for Code Generation](https://mlsanity.com/paper/2305.01210/) (2023)
- [SmoothLLM: Defending Large Language Models Against Jailbreaking Attacks](https://mlsanity.com/paper/2310.03684/) (2023)
- [Advancing LLM Reasoning Generalists with Preference Trees](https://mlsanity.com/paper/2404.02078/) (2024)
- [Statistics of Robust Optimization: A Generalized Empirical Likelihood Approach](https://mlsanity.com/paper/1610.03425/) (2016)
- [Robust Prompt Optimization for Defending Language Models Against Jailbreaking Attacks](https://mlsanity.com/paper/2401.17263/) (2024)
- [Variance-based regularization with convex objectives](https://mlsanity.com/paper/1610.02581/) (2016)
- [OpenMathInstruct-1: A 1.8 Million Math Instruction Tuning Dataset](https://mlsanity.com/paper/2402.10176/) (2024)
- [Token-level Direct Preference Optimization](https://mlsanity.com/paper/2404.11999/) (2024)
- [DPO Meets PPO: Reinforced Token Optimization for RLHF](https://mlsanity.com/paper/2404.18922/) (2024)
- [GradSafe: Detecting Jailbreak Prompts for LLMs via Safety-Critical Gradient Analysis](https://mlsanity.com/paper/2402.13494/) (2024)
- [Group Robust Preference Optimization in Reward-free RLHF](https://mlsanity.com/paper/2405.20304/) (2024)
- [Gradient Cuff: Detecting Jailbreak Attacks on Large Language Models by Exploring Refusal Loss Landscapes](https://mlsanity.com/paper/2403.00867/) (2024)
- [Towards Robust Alignment of Language Models: Distributionally Robustifying Direct Preference Optimization](https://mlsanity.com/paper/2407.07880/) (2024)
- [Distributionally Robust Learning](https://mlsanity.com/paper/2108.08993/) (2021)
- [FlowRL: Matching Reward Distributions for LLM Reasoning](https://mlsanity.com/paper/2509.15207/) (2025)
- [Segmenting Text and Learning Their Rewards for Improved RLHF in Language Model](https://mlsanity.com/paper/2501.02790/) (2025)
- [CCFC: Core & Core-Full-Core Dual-Track Defense for LLM Jailbreak Protection](https://mlsanity.com/paper/2508.14128/) (2025)
- [Adaptive Segment-level Reward: Bridging the Gap Between Action and Reward Space in Alignment](https://mlsanity.com/paper/2411.00809/) (2024)

## Abstract (arXiv)

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.