English

Inverse Reinforcement Learning with Dynamic Reward Scaling for LLM Alignment

Computation and Language 2026-03-03 v6 Artificial Intelligence Machine Learning

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.

Keywords

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}
}
R2 v1 2026-06-28T22:32:49.571Z