Distributionally Robust Token Optimization in RLHF

• URL: http://arxiv.org/abs/2604.08577v1
• Date: Fri, 27 Mar 2026 21:36:32 GMT
• Title: Distributionally Robust Token Optimization in RLHF
• Abstract summary: Large Language Models (LLMs) tend to respond correctly to prompts that align to the data they were trained and fine-tuned on.<n>Small shifts in wording, format, or language can trigger surprisingly large failures, especially on multi-step reasoning problems.<n>We propose a Distributionally Robust Token Optimization approach, which combines token-level Reinforcement Learning from Human Feedback with Distributionally Robust Optimization.
• Score: 5.00213782531351
• License: http://arxiv.org/licenses/nonexclusive-distrib/1.0/
• Abstract: Large Language Models (LLMs) tend to respond correctly to prompts that align to 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 bounds worst case token-wise rewards by constructing an f-divergence ambiguity set over a loss minibatch, leading to a theoretical robustness. Empirically, DRTO enhances consistency under distribution shifts in mathematical reasoning benchmarks, achieving 9.17\% improvement on GSM8K and 2.49% improvement on MathQA.

Related papers

• Learning from Failures: Correction-Oriented Policy Optimization with Verifiable Rewards [73.44333771806282]
  Reinforcement Learning with Verifiable Rewards (RLVR) has emerged as an effective paradigm for improving the reasoning capabilities of large language models.<n>We propose Correction-Oriented Policy Optimization (CIPO), a simple and effective extension to RLVR that converts on-policy failed trajectories into correction-oriented supervision.<n>CIPO improves learning effectiveness while explicitly enhancing the model's ability to correct its own errors.
  arXiv  Detail & Related papers  (2026-05-14T08:22:21Z)
• Hierarchical Chain-of-Thought Prompting: Enhancing LLM Reasoning Performance and Efficiency [4.139013391483979]
  Chain-of-Thought (CoT) prompting has significantly improved the reasoning capabilities of large language models (LLMs)<n>In this work, we introduce Hierarchical Chain-of-Thought (Hi-CoT) prompting.<n>Hi-CoT decomposes the reasoning process into hierarchical substeps by alternating between instructional planning and step-by-step execution.
  arXiv  Detail & Related papers  (2026-03-31T18:35:22Z)
• Towards Efficient Large Language Reasoning Models via Extreme-Ratio Chain-of-Thought Compression [55.63153956934198]
  Chain-of-Thought (CoT) reasoning successfully enhances the reasoning capabilities of Large Language Models (LLMs)<n>Existing CoT compression methods often suffer from a critical loss of logical fidelity at high compression ratios.<n>We propose a novel EXTreme-RAtio Chain-of-Thought Compression framework, termed Extra-CoT, which aggressively reduces the token budget while preserving answer accuracy.
  arXiv  Detail & Related papers  (2026-02-09T06:57:15Z)
• Short Chains, Deep Thoughts: Balancing Reasoning Efficiency and Intra-Segment Capability via Split-Merge Optimization [68.89915707647138]
  Large Reasoning Models (LRMs) have demonstrated impressive capabilities in solving complex tasks through the generation of long reasoning chains.<n>We propose textbfCoSMo (textbfSplit-textbfMerge textbfOptimization), a framework designed to eliminate structural redundancy rather than indiscriminately restricting token volume.
  arXiv  Detail & Related papers  (2026-02-03T05:54:28Z)
• Hard Negative Sample-Augmented DPO Post-Training for Small Language Models [4.425580048633862]
  We propose a lightweight and pragmatic post-training pipeline that targets structured errors under realistic compute budgets.<n>We introduce a compact MathVerifier that decomposes a candidate solution into a six-dimensional error profile and aggregates it into interpretable wrongness and absurdity scores.<n> Experiments show that verifier-guided, weighted DPO yields more targeted improvements than vanilla SFT and unweighted DPO.
  arXiv  Detail & Related papers  (2025-12-17T06:15:52Z)
• Token Constraint Decoding Improves Robustness on Question Answering for Large Language Models [4.078176555898098]
  We introduce and evaluate Token Constraint Decoding (TCD)<n>This simple yet effective inference-time algorithm enforces alignment between token-level predictions to enhance robustness in noisy settings.<n>Our findings establish TCD as a practical, model-agnostic approach for improving reasoning stability under real-world imperfections.
  arXiv  Detail & Related papers  (2025-06-11T05:33:56Z)
• Reinforcing Video Reasoning with Focused Thinking [65.85683941058916]
  We propose TW-GRPO, a novel framework that enhances visual reasoning with focused thinking and dense reward granularity.<n>Specifically, we employ a token weighting mechanism that prioritizes tokens with high informational density.<n>We also reformulate RL training by shifting from single-choice to multi-choice QA tasks.
  arXiv  Detail & Related papers  (2025-05-30T15:42:19Z)
• Token-Efficient RL for LLM Reasoning [0.02488650627593658]
  We propose reinforcement learning strategies tailored for reasoning in large language models (LLMs) under strict memory and compute limits.<n>Building on early policy gradient methods with baseline subtraction, we design critic-free methods that operate on a small, informative subset of output tokens.<n>We show that our methods raise accuracy on the SVAMP benchmark from 46% to over 70% and show strong performance on multi-digit multiplication.
  arXiv  Detail & Related papers  (2025-04-29T14:58:43Z)
• DR-DSGD: A Distributionally Robust Decentralized Learning Algorithm over Graphs [54.08445874064361]
  We propose to solve a regularized distributionally robust learning problem in the decentralized setting. By adding a Kullback-Liebler regularization function to the robust min-max optimization problem, the learning problem can be reduced to a modified robust problem. We show that our proposed algorithm can improve the worst distribution test accuracy by up to $10%$.
  arXiv  Detail & Related papers  (2022-08-29T18:01:42Z)
• Detached Error Feedback for Distributed SGD with Random Sparsification [98.98236187442258]
  Communication bottleneck has been a critical problem in large-scale deep learning. We propose a new distributed error feedback (DEF) algorithm, which shows better convergence than error feedback for non-efficient distributed problems. We also propose DEFA to accelerate the generalization of DEF, which shows better bounds than DEF.
  arXiv  Detail & Related papers  (2020-04-11T03:50:59Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.