mlsanityData to 7 Oct 2026

Distributionally Robust Token Optimization in RLHF

arXiv 2604.08577 ↗First posted 27 Mar 2026cs.LG cs.AIcited by 0

On the map: LLM alignment / RLHF / preference optimisation → Reward modeling and RLHF

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.

Key facts. No citing papers in the map yet (data through 7 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.

Builds on

27 of 36 references.

Closest ancestors · 4

Same direction, newest first

  1. 2025Segmenting Text and Learning Their Rewards for Improved RLHF in Language Model6 Jan
  2. 2024Adaptive Segment-level Reward: Bridging the Gap Between Action and Reward Space in Alignment23 Oct
  3. 2024DPO Meets PPO: Reinforced Token Optimization for RLHF29 Apr
  4. 2023Direct Preference Optimization: Your Language Model is Secretly a Reward Model29 May

Other references · 23

  1. 2022Training language models to follow instructions with human feedbackcited by 15,712
  2. 2017Proximal Policy Optimization Algorithmscited by 14,441
  3. 2021Evaluating Large Language Models Trained on Codecited by 8,053
  4. 2021Program Synthesis with Large Language Modelscited by 3,509
  5. 2025DAPO: An Open-Source LLM Reinforcement Learning System at Scalecited by 2,678
  6. 2021Convex optimizationcited by 2,479
  7. 2015High-Dimensional Continuous Control Using Generalized Advantage Estimationcited by 1,916
  8. 2024LiveCodeBench: Holistic and Contamination Free Evaluation of Large Language Models for Codecited by 1,906
  9. 2023Is Your Code Generated by ChatGPT Really Correct? Rigorous Evaluation of Large Language Models for Code Generationcited by 1,332
  10. 2023SmoothLLM: Defending Large Language Models Against Jailbreaking Attackscited by 434
  11. 2024Advancing LLM Reasoning Generalists with Preference Treescited by 191
  12. 2016Statistics of Robust Optimization: A Generalized Empirical Likelihood Approachcited by 182
11 more
  1. 2016Variance-based regularization with convex objectivescited by 153
  2. 2024Robust Prompt Optimization for Defending Language Models Against Jailbreaking Attackscited by 150
  3. 2024OpenMathInstruct-1: A 1.8 Million Math Instruction Tuning Datasetcited by 143
  4. 2024Token-level Direct Preference Optimizationcited by 135
  5. 2024Group Robust Preference Optimization in Reward-free RLHFcited by 94
  6. 2024GradSafe: Detecting Jailbreak Prompts for LLMs via Safety-Critical Gradient Analysiscited by 87
  7. 2024Gradient Cuff: Detecting Jailbreak Attacks on Large Language Models by Exploring Refusal Loss Landscapescited by 63
  8. 2024Towards Robust Alignment of Language Models: Distributionally Robustifying Direct Preference Optimizationcited by 52
  9. 2021Distributionally Robust Learningcited by 44
  10. 2025FlowRL: Matching Reward Distributions for LLM Reasoningcited by 38
  11. 2025CCFC: Core & Core-Full-Core Dual-Track Defense for LLM Jailbreak Protectioncited by 3

Built on by

Not cited yet.

Similar, not cited

Close in wording, no citation either way

  1. 2025Distributionally Robust Reinforcement Learning with Human Feedback
  2. 2025Online Distributionally Robust LLM Alignment via Regression to Relative Reward
  3. 2026STAPO: Stabilizing Reinforcement Learning for LLMs by Silencing Rare Spurious Tokens
  4. 2025Diversity or Precision? A Deep Dive into Next Token Prediction
  5. 2025Do Not Let Low-Probability Tokens Over-Dominate in RL for LLMs
  6. 2026Sparse but Critical: A Token-Level Analysis of Distributional Shifts in RLVR Fine-Tuning of LLMs
  7. 2024MA-RLHF: Reinforcement Learning from Human Feedback with Macro Actions
  8. 2025Token-Efficient RL for LLM Reasoning