让大模型在语言变化下仍稳定推理,提升复杂任务表现
Distributionally Robust Token Optimization in RLHF

- 用分布鲁棒优化增强令牌级强化学习,聚焦难处理的响应片段
- 在MATH-500和LiveCodeBench上分别提升4.4和2.7个百分点
- 适合需要高鲁棒性推理的场景,如代码生成与数学解题
大语言模型在输入与训练数据高度一致时表现良好,但微小的语言或格式变化可能导致严重错误,尤其在多步推理任务中。为此,我们提出分布鲁棒令牌优化(DRTO),将令牌级人类反馈强化学习(RLHF)与分布鲁棒优化(DRO)结合。DRTO在跨度级策略损失上构建f-散度模糊集,提供一种合理方法,在策略优化中强调困难响应段。实验表明,DRTO在多个推理基准上均提升了分布偏移下的一致性,在MATH-500上比标准RTO提升4.4个百分点,在LiveCodeBench上提升2.7个百分点。
原文摘要 · Abstract (English)
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.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。