Inverse Reinforcement Learning with Dynamic Reward Scaling for LLM Alignment
Ruoxi Cheng, Haoxuan Ma, Weixin Wang, Ranjie Duan, Jiexi Liu, Xiaoshuang Jia, Simeng Qin, Xiaochun Cao, Yang Liu
Abstract
Alignment is vital for safely deploying large language models (LLMs). Existing techniques are either reward-based (training a reward model on preference pairs and optimizing with reinforcement learning) or reward-free (directly fine-tuning on ranked outputs). Recent research shows that well-tuned reward-based pipelines remain the most robust, and single-response demonstrations can outperform pairwise preference data. However, there still exist two key challenges: (1) imbalanced safety datasets that overrepresent common hazards while neglecting long-tail threats; and (2) static reward models that ignore task difficulty, limiting optimization efficiency and attainable gains. To address these limitations, we propose DR-IRL, which Dynamically adjusts Rewards through Inverse Reinforcement Learning. We first train category-specific reward models using a balanced safety dataset of seven harmful categories as demonstration via IRL. Then we enhance Group Relative Policy Optimization (GRPO) by introducing dynamic reward scaling: adjusting rewards by task difficulty, data-level hardness by text encoder cosine similarity, and model-level responsiveness by reward gaps. Extensive experiments across various benchmarks and LLMs demonstrate that DR-IRL outperforms all baseline methods in safety alignment while maintaining usefulness.
Ask about this paper
Your agent reads all of it.
Lune indexed this paper to the last equation, along with the top-tier papers that cite it. Ask a question and the answer quotes them.
Your agent calls
Luneget_paper_fulltext
Free to start. No credit card required.
Terminal
Install the CLIlune papers fulltext 4a98c009-121a-4c04-8912-deec1863e268Cited by top-tier papers5
- VRPO: Rethinking Value Modeling for Robust RL under Noisy Supervision in LLM Post-TrainingDingwei Zhu, Shihan Dou, Zhiheng Xi, Senjie Jin et al.ACL 2026 · 9 citations
- Obscure but Effective: Classical Chinese Jailbreak Prompt Optimization via Bio-Inspired SearchXun Huang, Simeng Qin, Xiaoshuang Jia, Ranjie Duan et al.ICLR 2026 · 9 citations
- EcoAlign: An Economically Rational Framework for Efficient LVLM AlignmentRuoxi Cheng, Haoxuan Ma, Teng Ma, Hongyi ZhangCVPR 2026 · 6 citations
- Quantized but Deceptive? A Multi-Dimensional Truthfulness Evaluation of Quantized LLMsYao Fu, Xianxuan Long, Runchao Li, Haotian Yu et al.EMNLP 2025
- MESA: Improving MoE Safety Alignment via Decentralized ExpertiseYitong Sun, Yao Huang, Teng Li, Ranjie Duan et al.ICML 2026
Builds on25
- Training language models to follow instructions with human feedbackLong Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida et al.NeurIPS 2022 · 24,707 citations
- Chain-of-Thought Prompting Elicits Reasoning in Large Language ModelsJason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma et al.NeurIPS 2022 · 22,562 citations
- Direct Preference Optimization: Your Language Model is Secretly a Reward ModelRafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D. Manning et al.NeurIPS 2023 · 10,924 citations
- AutoDAN: Generating Stealthy Jailbreak Prompts on Aligned Large Language ModelsXiaogeng Liu, Nan Xu, Muhao Chen, Chaowei XiaoICLR 2024 · 722 citations
- WildChat: 1M ChatGPT Interaction Logs in the WildWenting Zhao, Xiang Ren, Jack Hessel, Claire Cardie et al.ICLR 2024 · 504 citations
Related papers
- Group Robust Preference Optimization in Reward-free RLHFShyam Sundhar Ramesh, Yifan Hu, Iason Chaimalas, Viraj Mehta et al.NeurIPS 2024 · 122 citations
- Dynamic Rewarding with Prompt Optimization Enables Tuning-free Self-Alignment of Language ModelsSomanshu Singla, Zhen Wang, Tianyang Liu, Abdullah Ashfaq et al.EMNLP 2024 · 1 citation
- SafeGRPO: Self-Rewarded Multimodal Safety Alignment via Rule-Governed Policy OptimizationXuankun Rong, Wenke Huang, Tingfeng Wang, Daiguo Zhou et al.CVPR 2026 · 13 citations
- ToolRL: Reward is All Tool Learning NeedsCheng Qian, Emre Can Acikgoz, Qi He, Hongru Wang et al.NeurIPS 2025 · 387 citations
- Reinforcement Learning for Large Language Models via Group Preference Reward ShapingHuaisheng Zhu, Siyuan Xu, Hangfan Zhang, Teng Xiao et al.EMNLP 2025
