Inverse Reinforcement Learning with Dynamic Reward Scaling for LLM Alignment
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.
Cite
@article{arxiv.2503.18991,
title = {Inverse Reinforcement Learning with Dynamic Reward Scaling for LLM Alignment},
author = {Ruoxi Cheng and Haoxuan Ma and Weixin Wang and Ranjie Duan and Jiexi Liu and Xiaoshuang Jia and Simeng Qin and Xiaochun Cao and Yang Liu and Xiaojun Jia},
journal= {arXiv preprint arXiv:2503.18991},
year = {2026}
}