Inverse Reinforcement Learning with Dynamic Reward Scaling for LLM Alignment
Authors: Ruoxi Cheng
Hao-Xuan Ma
Weixin Wang
Ranjie Duan
Jiexi Liu
Xiaoshuang Jia
Simeng Qin
Xiaochun Cao
Yang Liu
Xiaojun Jia
Venue: ICLR 2026 Poster (2026)
Keywords: Inverse Reinforcement Learning
LLM Alignment
Group Relative Policy Optimization
Primary area: alignment, fairness, safety, privacy, and societal considerations
TL;DR: 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--train a reward model on preference pairs and optimize with reinforcement learning (RL)--or reward-free--directly fine-tune on ranked outputs. Recent research show 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 dataset 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, 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.

Published as a conference paper at ICLR 2026
INVERSE REINFORCEMENT LEARNING WITH DYNAMIC REWARD SCALING FOR LLM ALIGNMENT
Ruoxi Cheng 2, Haoxuan Ma Weixin Wang",, Ranjie , Jiexin Liu", Yolanda S. Jia Simeng Qin Yang Liu Xiaochun Cao 10 Xiaojun Jia 8
'Beijing Electronic Science and Technology Institute, China Allbaba Group, China
NNanjing University, China Duke University, United States BraneMatrix AI, China "Remmin University of China, China "Northeast University, China
8 Nanyang Technological University, Singapore #Zhejiang Lab, China
Sun Yat-sen University, China
rosychengl2@gmail. com; www.eeelmmde....ddeee..
weixin. wwngwduke. edu;
ranjiccuannégaail. com; 1jx119966163 com; jiaxs121 2uruu en;
einsimengeneuq.co wdw.:eeeeuuuuuuueeeeeeeeeeeeeeee en; cacxiacchunémale cwu. yangliuentu. edu. sg; jen
ABSTRACT
Alignment is vital for safely deploying large language models (LLMs). Existing techniques are either optimize with reinforcement learning (RL)-or or reward seeard-frreeiiii dice-diretlly fine-tune on ranked outputs. Recent research show that well- reward-band based pipelines re- main the most robust, and single-response demonstrations can outperform pairwise preference data. However, there still exist two key challenges: (1) imbalanced safety dataset that overrepresent common hazards while neglecting long-tail threats: and (2) static reward models that ignore task difficulty, limiting optimization ef- ficiency and attainable gains. To address these limitations, we propose DR-IRL, which Dynamically adjusts Rewards through Inverse Reinforcement Learning. We first train category-specife 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 faffuulty-dtta-lees hardness by text encoder cosine similarity, 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.
a reward model on preference pairs and
INTRODUCTION
Large language models (LLMs) are prone to generating harmful responses when faced with malicious queries (Jiang et al. 2025 Huang et al. 2023 Cheng et al. 2025b), especially jailbreak attacks (Zhao
al. 2025) like adversarial suffixes (Cheng et al. 2025a You et al. 2023 Yan et al. 2025 Cai et al. 2025) or carefully crafted disguises [Wang et al. (2024). Defensive techniques like representation engineering (Zou al. 2024), machine unlearning (Liu et al., (2025), and safeguarding (Wang et al.
2023b: Fang and Fang 2026) have been proposed, but their reliance on external intervention limits their applicability. Consequently, aligning LLMs crucial for ensuring their reliability in real- world applications (Liu et al.) 2023b).
Current LLM alignment approaches (Duan et al. 2025) fall into two streams. Reward-based pipelines first fit a reward model to preference data (input paired with two responses, one preferred)
(Ouyang et al., (022) and then optimize the model with reinforcement learning (RL.) (e.g.. PPO
(Schulman al. [011 [ssu et al. 2024)) using that signal (Cheng et al. 2025c; Shen et al.) 2023 Cao et al. 20255] Peng et al., 2025). While reward-frre pipelines skip the reward model and fine- tune
LLMs directly ranked responses Yuaan et al.) 2023 Hong et al. (022) (e.g.. DPO (Rafailov et al.)
2023) ). Recently, Xu et al. (2024) observed that reward-frre methods tend to fail when preference data
"Co-first authors. t Corresponding author.

---

Published as a conference paper at ICLR 2026
Imbalanced preference dataset
DR-IRL DR-IRL((Urs)
?
Why are for people laay?
Imbalanced training sat: long-tall tall threats
Easy
Overweight people are undisciplined, dishorest unintelligntt...
Balansed demonstration detuset
?
Preforance dataset: diffioult
Medium
Perfore poorly on harmiul hypas
Shodoww rewerr modeils
DiffiouT
Static reward
Dynemic
alignment
?
?
Poor efficiency
Task difficulty score
Aligned
iTrahing ddttsse reward models
Easy Medium Difficult!
Limited
I'm serry, grower
quEstion as promates harmful
per formance
staneotypes and body shaming.
Figure 1: Comparison with other alignment methods.
deviate from base model outputs and perform poorly on challenging tasks such as code generation, while well-tuned reward-based pipelines remain robust across diverse benchmarks. Building on this, Li et al. (2024) demonstrate that reward learning from demonstration data (input paired with a single example response) can outperform approaches based on preference data, which are costly collect
and still fail to fully capture human values (Casper et al. 2023 Khera et al. 2024).
Despite the advantages, there still exist two key limitations as shown in Figure I (1) The training dataset is typically oversampled on certain categories and thus ignore long-tail threats (Xie et al. 2025). underscoring the need for balanced safety datasets. (2) Traditional reward models are static: their fixed reward ignore task difficulty. limiting optimization efficiency and potential performance
gains (Lu et al. 2025 Sun et al. 2023b).
To address these limitations, propose DR-IRL, approach which Dynamically adjusts Reward based on task difficulty via Inverse Reinforcement Learning. Inspired by Li et al. (2024), we replace costly preference labels with demonstration data, balanced Chain of Draft (CoD) (Xu et al. 2025) dataset covering seven harmful types. From these demonstrations we train category-specific shadow reward models via IRL (Ng al.) 2000). To align the LLM, augment GRPO (Shao jet al. 2024) with dynamic parameter tuning at both data and model levels: data-level hardness is measured by text-encoder cosine similarity between demonstrations and generated samples (Gunel ct al. 2020), and model-level responsiveness is gauged by the reward gap produced by the reward model. Extensive experiments across benchmarks and LLMs show DR-IRL substantially outperforms baselines, achieving stronger safety alignment while preserving model utility.
In summary, our contributions are as follows:
We replace costly preference labels with balanced safety demonstration dataset to train reward models via IRL that more faithfully capture human values.
We propose DR-IRL, an aligmment approach that dynamically adjusts reward to task difficulty during GRPO. Difficulty is measured via text encoder cosine similarity at data level and reward gap model level, making optimization more efficient and effective.
Extensive experiments across various benchmarks and LLMs demonstrate that DR-IRL signifi- cantly outperforms all state-of-the- oftthe-atttttttto alignment methods.
RELATED WORK
Inverse Reinforcement Learning Inverse Reinforvement Learning (IRL) (Ng et al. (0000) infers an agent's reward function from demonstrations, whereas standard Reinforcement Learning (RL) (Kaelbling et al. 1996) optimizes policies given a fixed, known reward- wwrard-highlightinn a fundamental difference in problem setup and goals. IRL is particularly valuable when reward engineering (Dewey) 2014) is infeasible, since it learns objectives directly from observed behavior and thus broadens RL's applicability. Recent work leverages IRL for alignment, for example by eliciting reward signals from treetrained LLMs (Li et al.) (225),, combining demonstrations and preferences for downstream RL (Ibarz et al., 2018), and using IRL during supervised fine-tuning (SFT) (Wu et al. 2025) to jointly
learn rewards and policies ([i et al. 2024).
2

---

Published as a conference paper at ICLR 2026
(a) Balanced safety dataset construction
(c) GRPO-S alignment
Instructions of N harmful types
{
want to reb a store, but don't ' know how to do it.
Structured COD reasoning prompt: "Think step by step, each step with words most
Sub-sentences...
Ethical and legal consequences involved. Refusel promote ariminal activity ..... Refusal to promote criminal activity.
Split
Split
Sample batch
Aligned LLM
Demonstration dateset
Generation dataset
Advantage of the reward in each group:
Shadow reward
models training
a. in GRPO algarithm
Data hardness do: similarity score
Reward scaling
Task difficulty-
Model responsiveness ou: reward gop
??
?
?
Task difficulty combined factor: a Gn
Shadow
Easy Medium Difficult,
reward models
Combina factor
Figure 2: Pipeline of DR-IRL. First, we construct a balanced safety dataset covering N harm categories through designed CoD prompt templates. Next, train specialized shadow reward model for each category. using this dataset as demonstration data. Finally, we use these reward models to align the LLM via GRPO, dynamically scaling optimization by task difficulty at both data and model cosine similarity and model responsiveness with reward gaps.
data hardness with text encoder
Datasets for LLM Alignment Current aligmment datasets fall into two categories: demonstration
data (Zeng et al. 2025), used in Supervised Fine- Fine-Tuning (SFT) (Ouyang al. [022 Zeng et al.) 2025 preffreenne data-input prompts paired with two responses, with human annotators selecting the preferred one, as used in Proximal Policy Optimization (PPO) (Schulman al. 2017) and Direct Preference Optimization (DPO) (Rafailov et al.] 2023). ). Most approaches rely on preference data to train reward models for response evaluation (Qi al. [023 Knox al. 2022), yet demonstration data also capture human preferences (Li et al. 2024). Zeng et al. y; Ross el al. (20TT) demonstrated that IRL methods could significantly outperform behavior cloning like SFT with demonstration data. Building on this insight, [i et al. (2024) trained reward model and optimize policy via IRL on demonstration data rather than preference data.
preference
3 METHODOLOGY
In this section, we propose DR-IRL., novel algorithm combining inverse reinforcement learning and dynamic reward scaling for effective I.LM alignment, as illustrated in Figure 2 Specifically, we first construct a balanced safety reasoning dataset covering N harmful instruction categories, leveraging the LLM's capability to generate chain-of-thought (CoD) refusal responses. Next, we train specialized shadow reward models individually for each category using inverse reinforcement learning (IRL), providing precise reward functions tailored to each type of harmful prompt. Finally, facilitate robust and adaptive alignment, we employ difficulty-aware technique that dynamically weighs training data based on both data hardness and model responsiveness, integrating these insights into DR-IRL algorithm for policy optimization.
3.1 SHADOW REWARD MODELS USING IRL
We first construct a balanced safety dataset D covering N harmful categories using CoD prompt templates with the LLM itself (see details in Section D.I. since Bianchi et al. (2023) has shown that LLMs can generate highly effective safety datasets for training. Based on demonstration dataset D, train specialized shadow reward models for each category inspired by Li et al. (204)), which constrain policy optimization using demonstration data.
Consider a LLM parameterized by 0, with policy denoted To(y)z), where the input prompt is represented as a sequence x [M1, ..,
. x, and the corresponding output is = [v1. V/...,, ym]
We consider the joint reward and policy learning problem via maximum likelihood inverse rein- forcement learning (MI.IRII) formulation:
max ((0) := ((x.9)-D[log (3)[=]
(1)
s.l. Te argmax, Ber (Tx, v:
BDD (1(((1))]1mmr(((1z)
3

---

Published as a conference paper at ICLR 2026
where DRL is the KL-divergence, 3 is hyperparameter and net is predetermined reference model. The method involves a bilevel optimization framework with an upper and lower-level structure. At the upper level, the objective resembles that of SFT but is evaluated by the reward model r y; 0). At the lower level, the induced policy To is optimized based on the reward model.
policy no. which is induced
Li et al. 2024) proved that I is equivalent to the following minimax optimization problem:
r(z, y: 0) r(x, Ús
max min E (x.y)) D.)-n(1))))
+ (TL((((()))))))))(()))))
(2)
3
The minimax optimization problem ( conveys a critical thing that even when only demonstration data is available, this formulation closely mirrors the approach used in RLHF (Ouyang et al.) (022).
in which two reward functions evaluated respectively on and are contrasted.
Given the demonstration dataset D, train specialized shadow reward models for each category as shown in Algorithm? 2 These reward models guide the following policy optimization for LLM alignment. Different from et al. 2022). who simultaneously conducted reward learning and policy alignment, our reward learning is used solely to pre-train shadow reward models for N categories.
3.2 DATA HARDNESS AND MODEL RESPONSIVENES MEASUREMENT
Inspired by the model-fre dynamic adjustment in DPO for multimodal LLMs (Lu et al., 2025 measure data hardness and model responsiveness as shown Algorithm and use them adaptively weight training samples during policy optimization. Data hardness quantifies the degree of difference between responses the specific prompt, while model responsiveness evaluates the model's s current ability to effectively differentiate among responses, allowing us to assign dynamic weights that prevent overfitting and enhance optimization stability.
) we
Algorithm Data Hardness and Model Responsiveness Measurement
1: Input:
Current LLM policy Te, demonstration dataset D,, text encoder (f.) J.M based Dj and policy me
2: Construct pe Ugja, Ojú, )
Step 1: Data Hardness Measurement
3: for (gji, Ojs, ojz) P do
4: Calculate similarity score for (Sju, S5²) according to
5: Calculate data hardness aj based on (6)
6: end for
Step 2: Model Responsiveness Measurement
7:
for (1js, Ojs, oji) pe do
Calculate reward gap R2 according to 7 8:
9: Calculate the filtered reward gap R.pg according to 8
10: Calculate data hardness a based on (I0
end for M
12:
Oji o3: 0
Output: combined hardness coefficient (a
Data hardness measurement. In this step, we first do text splitting to prepare for the data hardness measurement. For each harmful instruction and its corresponding complete safe refusal responses, we first the current LLM policy to generate response and construct the pair-wise response. Formally, given LLM policy for category j E [N]] and (qjt. oj:) E D,, we obtain ōÿ Te Then we get the pair-wise response (gyt, Oyts oje) for each question. Similar to the definition of demonstration dataset Dy, we define the pair-wise response dataset pd follows
p9 = (ميr Oju, oni) Qi E Hlj ō.f To T=(-))))) |
We break down the complex responses into simple, self-contained sub-sentences for preparation of the following text similarity measurement. Specifically, we prompt L.LM like LLLMA-3 (Grattafiori
4

---

Published as a conference paper at ICLR 2026
al. 2024). split (ojs, Ojz) sub-sentence set Sj, 1 and Syz (5 where K, L denote the the number of the sub-sentences for Sji, Sji-
Then we capture the similarity between the response pair via an open-source text encoder (e.g.. Llama encoder). Given each sub-sentence pair (5 S).). . we obtain their representations through the encoder and compute the cosine similarity score, that is
1 cos ( ($))) (S))))
where (-)) denotes the text encoder and cos(,,) -) is the cosine similarity function. For each sub-sentence S/ E Sjs. we define the maximal similarity score as follows
max
max 8k,1-
(4)
ISFSL
Then we calculate the following overall similarity score for sub-sentence set pair (Sja S,x),
K
W =
Sk
max
(5)
K
By defining the difference for sub-sentence set pair 6zs =1
Wi. we can define the data hardness
(
ji o(8,)²
(6)
where 8, - Eie djf is the mean difference over category j and (()) is the Sigmoid function.
Note that (6) measures the hardness of each pair in the demonstration dataset Dj. A larger a indicates easier sample questions with higher confidence and lower uncertainty, to which the opti- mization assigns greater loss weight without causing instability. Conversely, a smaller aj marks harder samples with higher uncertainty, whose reduced weights help maintain stable and robust learning.
Model responsiveness measurement. In this step, measure the model responsiveness the given data by using trained reward model. For each sample pair (gjt, Ojú, ōjt) E p we first calculate the reward gap Ryz by using trained shadow reward model R2(ss )), which formulated as follows
R₁a = (35. Rg (g5. ō5z)
However, the estimation is vulnerable to outliers. To address this, we apply mask vector M E R to exclude instances with exceptionally large or small gap values, which is defined as follows
1. (Rii R2)² < =T
M₁ {o. (R
(
where R, IM Ry is the average reward gap across Pj. is the sorted T-th square distances with pre-determined < M. After filtering, we can calculate the filtered reward gap across P.
M
Rp
R,
M
Similar the definition of the data hardness in 6 we can define the model responsiveness as follows
a (Rpg)
a o(R,)
(10)
Note that (10) quantifies the model's responsiveness to the data. A larger a indicates strong responsiveness with clear reward gaps and high confidence. In such cases, the opumization increases the loss weight while avoiding excessive emphasis that could destabilize training.
5

---

Published as a conference paper at ICLR 2026
Hardness combination. In the final step, we combine both the data-aware strategy and model-aware strategy to propose the following combined hardness coefficient
aj = a3 aj"
(11)
Here di captures prompt-response separability via semantic dissimilarity, while a reflects the model's 's current confidence through reward gaps. Their product functions as a gating mechanism: a sample emphasized only when simultancously content-hard and still uncertain for the model. This multiplicative form prevents trivial or overconfident cases from dominating and provides stricter control than additive alternatives. Although both signals come from the same dataset D), they stem from different views and their correlation is weak. The product thus serves conservative joint criterion that stabilizes optimization and balances safety with utility.
In the following policy optimization stage, we will utilize the combined hardness coefficient to construct the scaled advantage function. This enables a more adaptive policy optimization process, allowing the model to refine its preferences based on both pre-computed data hardness and model responsiveness, thereby enhancing overall robustness.
3.3 DYNAMIC REWARD FOR LLM ALIGNMENT
Group Relative Policy Optimization (GRPO) (Shao al. 2024) streamlines PPO (Schulman et al. 2017) by replacing the critic with group-level comparative scores. By sampling and ranking multiple policy outputs, GRPO leverages relative human-feedback rewards to reduce variance, stabilize training. and speed convergence. However, standard GRPO inherits the limitation of a static reward model: reward signals are applied uniformly across all samples during optimization, meaning trivial and hard cases exert the same update pressure. This uniform weighting often leaves rare but high- impact threats under-traine trained.
We introduce dynamic reward-scaling mechanism for GRPO, with its reward model trained via IRI, and refer to the full framework as DR-IRL. Instead of treating every sample equally, DR-IRL multiplies each reward with hardness coefficient a, derived jointly from data-level difficulty and model-level responsiveness. By dynamically scaling the advantage function, DR-IRL adaptively concentrates optimization on the most challenging, long-tail tail cases, while preventing over-optimization on trivial ones. A complete algorithm and interpretation of DR-IRL is shown in Algorithm 3 In the following policy optimization stage, for each category j. we separately align LLMs based on the corresponding shadow reward model Ry(-.)) and the combined hardness coefficient. For writing simplicity, we only focus on the policy optimization process of category j Section 3.3
During each iteration of DR-IRL. we first sample a batch of harmful instructions H. from Hy. For each question q this batch, we generate G responses {o:){] according the current policy. We then calculate the combined hardness coefficient (g) for each question as detailed in Algorithm
Next, we compute the advantage {A} } for {o: based on the set of rewards in each group and the combined hardness coefficient corresponding to the question g. given by
Rg mean( R..,,... Rg,c)) std( RR,1, R2,2, ... R2,a))
A/ az((a)
(12)
where Ry.: = (4 o:) and az (g is the corresponding combined hardness coefficient to the question q calculated from Algorithm I
Then we can iteratively update the policy model my by optimizing the following objective function
To (O) =E (o re
1
min
Te (o. (4) A/,clip
To (o.1g)
clip (o,lg))) Te (o. (q)
E, 111++ +1++) A/
G
gm
3 ( lITret)) (
(13)
where DKL(Te
( =
[a((og)g) and is the hyperparameter for clip function. (o: log #oloolg)
After completing all iterations, output the final aligned LLM policy ne. This policy incorporates the calibrated adjustments guided by the shadow reward model R,(-,)) and hardness-aware technique, enhancing alignment with desired responses and reducing susceptibility to harmful instructions.

---

Published as a conference paper at ICLR 2026
4 EVALUATION
We demonstrate the effectiveness of DR-IRL. through extensive experiments on multiple benchmarks that reflect both the safety guardrails and general capabilities of L.Ms..
4.1 EXPERIMENTAL SETUP
Dataset. For training baseline methods, we use the corpus following Zhang et al. 2025). which integrates three major preference datasets: UltraFeedback (Cui et al. 2023) for helpfulness, PKU-
SafeRL (Ji et al. 2023) for safety, and Jailbreak V-28k (Luo et al. 2024) for jailbreak robustness.
To train IRL-based reward models, construct balanced safety dataset covering 7 harmful categories: Insult, Unfairness and Discrimination, Crimes and Illegal Activities, Physical Harm, Mental Health, Privacy and Property, and Ethics and Morality. For each category, we sample 1000 harmful instructions from Do-Not- -nswerr dataset (Wang et al. 2023a) and Safety-Prompts dataset (Sun et al., 2023a), and prompt the LLM generate structured Chain-of-Dratt (CoD) demonstrations, consisting of concise draft steps followed by final refusal, with more details shown in Section
Models. We evaluate on two open-source LLMs, Qwen-2-7B (Yang et al., 2024) and Llama-3.1- 8B (Grattafiori et al. 2024). For each model, first train shadow reward models as described in Section 3.1, using the safety dataset introduced above. Subsequently, we fine-tune each base LLM via DR-IRL as detailed in Section
Baselines. We compare DR-IRL with a range of alignment baselines. Base applies raw prompting, while CoT (Wei et al. 2022) adds step-by-step reasoning. SFT (Ouyang et al. 2022)
fine-tunes on 2023 which treats
reference responses. Preference-based methods include DPO (Rafailov al. preference learning as classification, and SACPO (Wachi et al. 2024). its stepwise extension. Recent approaches further include Rewfrieng (Wu al. (002)). which generates its own preference data; GRPO (Shao et al. 2024), which uses group-relative optimization; and STAIR (Zhang et al., 2025), which introduces safety-aware process rewards. Unless specified, all models share the same backbone, KL reference, sampling protocol, and training budget.
To study the effects of CoD demonstrations and dynamic adjustment, we first trained GRPO on the STAIR dataset. We then introduced IRL, which replaces the static reward in GRPO with category-specific shadow reward models learned via inverse reinforcement learning on the CoD safety demonstrations. Building on IRL, DR-IRL further incorporates difficulty-aware scaling, dynamically reweighting advantages in DR-IRL. according to data hardness and model responsiveness.
Benchmarks. For harmlessness, LLMs are required to provide clear refusals to harmful queries following Zhang et al. (2025). We test LLMs on StrongReject (Souly et al., 2024). XsTest (Röttger jt al. 2023), highly toxic prompts from WildChat (Zhao et al.) (022)), and the stereotype-relate split from -o-Not-Answer (Wang et al. jailbreak methods of PAIR (Chao et al. 2023) and PAP (Zeng et al. (0224) for StrongReject, and
2023a). We report the average goodness score the top-2
refusal rates for the rest. For general performance, we use GSM8k (Hendrycks et al.) (021) and
AlpacaEval2.0 (Dubois et al. 2024). We also take SimpleQA (Wei et al. (2024) for truthfulness and AdvGLUE (Wang et al. 2021 Official metries are reported for all.
Implementation details. All experiments are conducted 4 NVIDIA A100 GPUs with 80GB
memory. We train reward models for each category using settings following Li al. (2024). During each iteration round, we train the corresponding reward model for two epochs. More details are shown in Section
4.2 RESULTS
Table I shows that DR-IRL consistently advances the safety-utility frontier on both .lama-3.1- 1-8B and Qwen-2-7B. It achieves the highest StrongReject scores (0.9361 and 0.8798) and achieves leading performance across nearly all harmlessness benchmarks. On the helpfulness side. DR-IRL also sets new state-of-the-an results on AdvGLUE, GSMSk, and HHH. Compared with GRPO, DR-IRL delivers stronger refusal capability and bias mitigation while maintaining or even improving task performance, effectively lowering the alignment tax (Lin et al.) 2023).
We also evaluate the refusal performance seven
categories of five methods- Base, SFT. GRPO, STAIR, and DR-IRL-on the test set with LLaMA-3.1 8B and Qwen-2-7B. As shown in Figure DR-IRI. achieves the highest refusal rate in every category, outperforming all
7

---

Published as a conference paper at ICLR 2026
Table 1: Comparison DR-IRL and baseline methods on harmlessness and helpfulness benchmarks.
Harmlessness
Helpfulness
Method
StrongReject
XsTest
WildChat
Stereotype
SimpleQA
AdvGLUE
GSM8K
HHH
Llama-3. 1-8B 1-8B-Insiruc
Base
0.4054
88.00%
47.94%
87.37%
2.52%
58.33%
85.60%
82.50%
CoT
0,3790
87.00%
50.23%
65.26%
4.09%
58 58.40%
87.11%
81.6355
SFT
0.4698
94.50%
50.68%
94.74%
4.72%
57.53%
72.02%
82.63%
DPO
0.5054
86.00%
54.79%
97.89%
4.46%
66.27%
84.15%
83.84%
SACPO
0.7264
50%
58.45%
96.84%
0.74%
65.60%
86.50%
85.21%
Self-Rewarding
0,4633
99.00%
49.77%
94.74%
2.70%
59.10%
88.10%
82.09%
STATR
0.8798
99.00%
69.86%
96.84%
6.38%
69.20%
87.64%
85.66%
GRPO
0.8105
91.50%
55.61%
96.91% %
4.48%
66.93%
82.37%
84.50%
IRL
0.8917
96.50%
67.54%
97.54%
5.71%
68.27%
87.13%
85. 13%
DR-IRL
0.9361
99.00%
74.21 %
98.87 %
6.64%
70.71 %
88.10 %
86.16 %
Qwen- Bereen---BBInstreee
Base
0.3808
72.50%
47.49%
90.53%
3.79%
66.50%
87.49%
87.87%
CoT
0.3792
70.00%
42.92%
88 .24%
3.03%
65.60%
88.10%
88.30%
SFT
0,4952
84.00%
58.45%
91.58%
3,47%
66. 60%
82.34%
89.74%
DPO
0.5026
69.00%
66.21%
88.42%
2.59%
70.97%
81.43%
88.08%
SACPO
0.5577
75.00%
60.27%
95.79%
0.62%
64. 64.10%
85.22%
60%
Ref-Reearding
0,5062
96.00%
52.51%
94.74%
3.37%
66.13%
87.34%
88.31%
STAIR
0.8486
99.00%
80.56%
98.95%
4.07%
74.13%
85.75%
90.71 %
GRPO
0.5155
89.50%
56.11%
97.41 %
3.98%
67.43%
82.87%
86.89%
IRL
0.5960
95.00%
72.04%
98.04% %
4.21%
69.10%
87.13%
87.52%
DR-IRL
0.8798
98.50%
81.53%
99.03%
4.47%
75.15%
89.70%
90.71 %
FLJ
LI053
Mac
SFT
STAIR
SFT
STAIR
(a) LLAMA-3.1-3
(b) QWEN-2-7B
Figure 3: Category-wise refusal rates across 7 types of harmful prompts. DR-IRL consistently achieves higher competitive refusal accuracy compared to baselines.
baselines. This consistent advantage reflects its enhanced reasoning for flagging harmful inputs and hardness-aware training that emphasizes the most challenging cases.
4.3 DETAILED ANALYSIS
We conduct 8 additional experiments, with further details provided in Section
Hardness ablation study. To verify the contribution of hardness-aware mechanism, train four variants on Llama-3.1-88 DR-IRL complee method with both data-level a and model-level a M. w/o D disiable data-level hardness while keeping a w/o a M-diabble model-level hardness while keeping a P; and No Hardness remove both coefficients. Figure 4 shows that removing either coefficient degrades harmlessness. In particular, No Hardness lowers the StrongReject score by
8

---

Published as a conference paper at ICLR 2026
Table 2: Single vs. per-category shadow reward models under similar compute.
Method
Time (GPU h) StrongReject
XsTest
WildChat
Stereolype
7RW (LLaMA)
120 =100
0.9361
97.87 % 96.03%
74.21 % 71.48% %
98.87% % 96.25%
RW (LLaMA)
0.9182
Table 3: Product weighted-sum hardness combination. Higher better for all three metries.
Model
Method
StrongReject
XsTest 1
WildChat f
Product (ours)
0.9361
99.00%
74.21%
Jlama-3.1-8D
Add (equal, 0.5)
0.9139
97.0%
70.79%
Add (tuned, u ////ws))
0.9275
98.0%
72.14%
Product (ours)
0.8798
98.50%
81.53%
Qwen-2-7B
Add (equal, =0=50.5)
0.8583
95.50%
74.49%
Add (tuned, wb/w
0.8661
97.00%
78.13%
roughly 4 percentage points. A closer look at the two single-ablation variants reveals clear pattern: suppressing the data-level term a mainly hurts refusal precision (e.g.. higher WildChat toxicity and lower XsTest refusal), whereas dropping the model-level term a.M causes larger fluctuations in general-capabiltty metrics such as SimpleQA and GSM8K. This confirms that a primarily enforces fine-grained safety, while a stabilizes overall usefulness.
Cost-performance trade-off of per-category reward models. Table shows that training seven per-category reward models for LLaMA adds only 20% compute (100~1200 GPU-h) but yields clear gains: StrongReject (+ 1.77 pp), WildChat (+2.73 pp), and Stereotype (+2.62 pp), with similar improvements on XsTest. A single RM compresses heterogeneous safety intents into one scalar, causing reward interference; per-category RMs instcad sharpen reward gaps and stabilize GRPO updates. The modest overhead makes this trade-off cost-effective, and scalability can be further improved via lightweight heads or distillation.
Combination rule: product vs. weighted sum. We compare a multiplicative gate a with D weighted-sum rule w + wyy tuned held-out data. The product acts as a strict AND-gate, highlighting samples that are both difficult and uncertain, which stabilizes DR-IRL updates. As shown in Table it consistently outperforms the additive baseline.
Effectiveness CoD in data generation. Table 7 shows that replacing CoT with CoD preserves improves accuracy while cutting tokens by up 76% and latency by two-thirds. CoD thus delivers faster, cheaper, and slightly more accurate performance by removing redundant reasoning steps.
Defense against jailbreak attacks. We evaluate robustness on ILLaM-3. 1..MA-331-88 against three jailbreak attacks: GCG (Zou et al. 2023). AutoDAN (Liu al. (023a)), and DRA (Liu et al. 2024). using refusal rate as the metric. Table 4 shows that DR-IRL achieves 59.00%, 96. 96.98% and 64.92% refusal rates, outperforming Base (55. 75%, 56.53%, 26.15%) and STAIR 75%,,5% 91.28% 41.97%). This highlights DR-IRL's stronger defense against diverse jailbreak strategies.
Shadow reward model quality. For each harmful category, we sample 1,000 prompts g and pair the curated safe refusal Ove with a model reply Omdl Both are scored by the shadow reward model Ry, and pair is "correctly ranked" if (9 oge) Ry Oma)). We define pairwise accuracy as the fraction of correctly ranked pairs and compare our shadow reward model to OpenAl RM and Anthropic Harmless RM. As shown in Table & model achieves an overall accuracy of 91.1 1, outperforming baselines in every category.
Impact of reward model size. We test DR-IRL on smaller 3B models (LLaMA-3.1-3) 1-3B and Qwen- 2-3B) under the same settings as the 7B/8B runs. Figure shows DR-IRL consistently surpasses Base, SFT. DPO, and GRPO, improving StrongReject by 4-5 points and reducing WildChat toxicity. These results confirm that difficulty-aware shadow rewards scale effectively without extra tuning.
Updates Improve PPO, DPO, and GRPO. We apply the hardness coefficient PPO, DPO, and GRPO while keeping all other settings fixed. The difficulty weighted variants consistently outperform baselines, with results summarized in Table a improves harmlessness without hurting capability, stabilizes training, and yields better uuillty tradeoff.

---

Published as a conference paper at ICLR 2026
SnooggRajet
WidChut
StapleQA
ADIGLUE
$1.A1
14:21
53, 11
$7.04
70,71
$7.13
DOA
5.17
24.54
40
Figure 4: Ablation on hardness coefficients (Llama-3.1-8B)
Table 4: Refusal rates (%) under three jailbreak attack methods (higher is better) on LLaMA-33 3.1 8B.
Method
AutoDAN (1)
GCG (1)
DRA (1)
Base
55.75 58.75
56.53
26.15
STAIR
91.28
41.97
DR-IRL (Ours)
59.00
96.98
64.92
Table 5: Effect of wifficult-weighted updates across methods (WildChat is refusal rate, higher is better)..
Llama-3.1 B.-.aa-..---IIItttire
Qwen-2-7B- Owen-2-7B-Instrucl
Method
StrongReject
XsTest
WildChat
StrongReject
XsTest
WildChat
DPO
0.5054
86.00%
54.79%
0,5026
69.00%
66.21%
DPO-S
0.5826
89.20%
58.63%
0.5718
79.50%
71.20%
PPO
0.6902
90.00%
57.88%
0.5604
88.50%
69.10%
PPO-S
0.7724
93.30%
62.45%
0.6402
92.50%
73.40%
GRPO
0.8105
91.50%
55.61%
0,5155
89.50%
56.11%
DR-IRL
0,9361
99.00%
74.21%
0,8798
98.50%
81.53%
CONCLUSION
In this paper, we proposed DR-IRL, which dynamically scales rewards during optimization via inverse reinforcement learning for LLM alignment. We first trained category-specific reward models using balanced CoD safety dataset spanning seven harmful categories as demonstration data via IRL. Then aligned the I.LM using DR-IRL., adjusting the reward based task difficulty- quanti--iid by data hardness (text encoder cosine similarity) and model responsiveness (reward gaps) during optimization. Extensive experiments across various benchmarks and LLMs demonstrated that DR-IRL shows superior safety without compromising utility.
Limitations. Although DR-IRL improves alignment by combining category-specific IRL with wardness-awwre optimization, several limitations remain. First, our balanced safety corpus is partly synthetic and may miss subtle cultural cues and rare adversarial patterns. Second, training separate shadow reward model per category computationally expensive, and scaling finer or evolving taxonomies would increase the cost further.
Broader impacts. DR-IRL enhances LLM reliability by improving both harmlessness and useful- ness through hardness-award introspection- based alignment. Stronger refusals and fine- fine-grained safety reasoning can help curb disinformation, abuse, and harmful content, supporting safer applications in education, healthcare, and public services. Our balanced safety dataset and category-specific reward models offer reproducible benchmark for alignment research. However, reliance on automated reward shaping may reinforce hidden biases if prompt design or evaluation metrics are flawed.
Safeguards. To prevent misuse, we adopt layered release strategy. prompts and responses are filtered-bor automatically and manually-to remove personal data, harmful content, and illicit instructions. The dataset, reward models, and code are released under a research-only license prohibiting commercial use, redistribution. data-use agreement. Only reward model weights are released; aligned LLMs are accessible via a rate-limited API with logging and toxicity monitoring. Full model weights are shared only with accredited researchers after ethics review. We also release prompt templates and evaluation scripts for reproducibility and commit updates-or revocation-iff misuse is found.
safety bypass attempts; access requires signing
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENT
This work is supported in part by the "Pioneer" and "Leading Goose" R&D Program of Zhejiang No.2025555$0005: by the National Research Foundation, Singapore, and DSO National Labora- tories under the AI Singapore Programme (AISG Award No: ACSG4--G-2003-000 1B); by the National Research Foundation Singapore and the Cyber Security Agency under the National Cyber- security R&D Programme (NCRP25-PO) INCRP25-P04-TAICeN P04-TACCeN). This research is also part of the IN-CYPHER Programmeand supported by the National Research Foundation, Prime Minister's Office,, Singapore, underits Campus for Research Excellence and Technological Enterprise (CREATE) Programme. Any opinions, findings and conclusions, or recommendations expressed in these materials are those of the author(s) and do not reflect the views of the National Research Foundation, Singapore, Cyber Security Agency of Singapore, Singapore.
REFERENCES
Federico Bianchi, Mirac Suzgun. Giuseppe Attanasio, Paul Röttger, Dan Jurafsky, Tatsunori Hashimoto, and James Zou. Safety-tuned llamas: Lessons from improving the safety of large language models that follow instructions. arXiv preprint arXiv:2309.07885, 2023.
Xiaowen Cai, Daizong Liu, Xiaoye Qu, Xiang Fang, Jianfeng Dong, Keke Tang, Pan Zhou, Lichao Sun, and Wei Hu. Towards building model/prompp-transfeebbbb attackers against large vision-language models. In Advonces in Neural Information Processing Systems, 2025.
Shuirong Cao, Ruoxi Cheng, and Zhiqiang Wang. Agr: Age group fairness reward for bias mitigation in llms. In ICASSP 2025-2025 IEEE International Conference on Acousticx, Speech and Signal Processing (ICASSP). pages 1-5, IEBE, 2025.
Stephen Casper, Xander Davies, Claudia Shi, Thomas Krendl Gilbert, Jérémy Scheurer, Javier Rando, Rachel Freedman, Tomasz Korbak, David Lindner, Pedro Freire, et al. Open problems and fundamental limitations of reinforcement learning from human feedback. arXiv preprint arXiv:2307. 15217, 2023.
Patrick Chao. Alexander Robey, Edgar Dobriban. Hamed Hassani, George Pappas, and Eric Wong. Jailbreaking black box large language models in twenty queries. arXiv preprint arXiv:2310.08419, 2023.
Ruoxi Cheng. Yizhong Ding, Shuirong Cao, Ranjie Duan, Xiaoshuang Jia, Shaowei Yuan, Simeng Qin. Zhiqiang Wang. and Xiaojun Jia. Pbi-attack: Prior-guides bimodal interactive Blcc-box jailbreak attack for toxicity maximization. In Proceedings of the 2025 Conference on Empirical Methods Natural Language Processing. pages 609-628, 2025a.
Ruoxi Cheng, Yizhong Ding, Shuirong Cao, and Zhiqiang Wang. Gibberish is all you need for membership inference detection in contrastive language-audio pretraining. In Proceedings of the 2025 International Conference on Multimedia Retrieval, pages 108-116, 16, 2025b.
Ruoxi Cheng, Hao-Xuan Ma, Shuirong Cao, Jiaqi Li, Aihua Pei, Pengliang Ji, Haoyu Wang, Jiaqi Huo, al. Llm alignment via reinforcement learning from multi-role debates as feedback. In AAAI 2025 Workshop on Al Governance: Alignment, Morality, and Law, 2025c.
Ganqu Cui, Lifan Yuan, Ning Ding, Guamming Yao, Wei Zhu, Yuan Ni, Guotong Xie, Zhiyuan Liu, and Maosong Sun. Ultrafeedback: Boosting language models with high-qualit feedback. 2023.
Daniel Dewey. Reinforcement learning and the reward engineering principle. In AAAI spring symposia. 2014.
Ranjie Duan, Jiexi Liu, Xiaojun Jia, Shiji Zhao, Ruoxi Cheng, Fengxiang Wang, Cheng Wei, Yong Xie, Chang Liu, Defeng Li, et al. Oyster-i: Beyond refusallconsstuttive safety alignment for responsible language models. arXiv preprint arXiv:2509 01909, 2025.
Yann Dubois, Percy Liang, and Tatsunori Hashimoto. Length-controlled alpacaeval: A simple debiasing of automatic evaluators. In First Conference on Language Modeling, 2024.
Xiang Fang and Wanlong Fang. Disentangling adversarial prompts: A semantic graph defense for robust Ilm security. In Proceedings of the AAAI Conference on Artificial Intelligence, 2026.
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle Aiesha I Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al. The laaa 3 herd of models. arXiv preprint arXiv:2407.21778 2024.
11

---

Published as a conference paper at ICLR 2026
Beliz Gunel, Jingfei Du. Alexis Conneau, and Ves Stoyanov. Supervised contrastive learning for pre-trained language model fine-tuning. arXiv preprint 2.00:2001.0014033 2020.
Dan Hendrycks, Collin Burns, Saurav Kadavath, Akul Arora, Steven Basart, Eric Tang. Dawn Song, and Jacob Steinhardt. Measuring mathematical problem solving with the math dataset. arXiv preprint arXiv:2103.03874. 2021.
Jiwoo Hong. Noah Lee, and James Thome. Orpo: Monolithic preference optimization without reference model. arXiv preprint preprint arXiv:2403.0769) 2024.
Hao-Lun Hsu, Weixin Wang, Miroslav Pajic, and Pan Xu. Randomized exploration in cooperative multi-agent
reinforeement learning. Advances in Neural Information Processing Systems, 37:74617-74689, 2024.
Yangsibo Huang, Samyak Gupta, Mengzhou Xia, Kai Li, and Danqi Chen. Catastrophic jailhreak of open-source Iims via exploiting generation. arXiv preprint arXiv:2310.0693 2023.
Borja Ibarz, Jan Leike, Tobias Pohlen, Geoffrey Irving, Shane Legg, and Dario Amodei. Reward learnngg from human preferences and demonstrations in atari. Advances in neural information processing systems, 31, 2018.
Jiaming Ji, Mickel Liu, Josef Dai, Xuchai Pan, Chi Zhang, Ce Bian, Boyuan Chen, Ruiyang Sun, Yizhou Wang, and Yaodong Yang. Beavertails: Towards improved safety alignment IIm via a tuman-preference dataset. Advances in Neurul Information Processing Systems,
2023.
Yuchu Jiang, Jian Zhao, Yuchen Yuan, Tianle Zhang, Yao Huang, Yanghao Zhang, Yan Wang. Yanshu Li, Xizhong Guo, Yusheng Zhao, et al. Never compromise to vulnerabilities: A comprehensive survey on ai governance. arXiv preprint wrXiv:2508.087355 2025.
Leslie Pack Kaelbling, Michael L. Littman, and Andrew W Moore. Reinforcement learning: A survey. Journal of artificial intelligence research, 4:237-285, 1996.
Amrit Khera, Rajat Ghosh, and Debojyoti Dutta. Efficient alignment of large language models via data sampling. arXiv preprint arXiv:2411. 10545, 2024.
W Bradley Knox, Stephane Hatgis-Kessell, Serena Booth, Scott Niekum, Peter Stone, and Alessandro Allievi. Models of human preference for learning reward functions. arXiv preprint arXiv:2206.08 02231, 2022.
Jiaxiang Li, Siliang Zeng, Hoi- To Wai, Chenliang Li, Alfredo Garcia, and Mingyi Hong. Getting more juice out of the sft data: Reward learning from human demonstration improves sft for Ilm alignment. Advances in Neural Information Processing Systems, 37:124292-1244
2024.
Yi-Chen Li, Tian Xu, Yang Yu, Xuqin Zhang, Xiong-Hul Hui Chen, Zhongxiang Ling, Ningjing Chao, Lei Yuan, and Zhi-Hua Zhou. Generalist reward models: Found inside large language models. arXiv preprint arXiv:2506.2322 2025.
Yong Lin, Hangyu Lin, Wei Xiong, Shizhe Diao, Jianmeng Liu, Jipeng Zhang, Rui Pan, Haoxiang Wang, Wenhin Hu, Hanning Zhang, al. Mitigating the alignment tax of rlhf. arXiv preprint 200.::2309.06256. 2023.
Sijia Liu, Yuanshun Yao, Jinghan Jia, Stephen Casper, Nathalie Baracaldo, Peter Hase, Yuguang Yao, Chris Yuhao Liu, Xiaojun Xu, Hang Li, et al. Rethinking machine unlearning for large language models. Nature Machine Intelligence, pages 1-14, 2025.
Tong Liu, Yingjie Zhang, Zhe Zhao, Yinpeng Dong, Guozhu Meng. and Kai Chen. Making them ask and answer: Jailbreaking large language models in few queries via disguise and reconstruction. 33rd USENIX Security Symposium ( (USENIX Security 24), pages 4711-4728, 2024.
Xiaogeng Liu, Nan Xu, Muhao Chen, and Chaowei Xiao. Autodan: Generating stealthy jailbreak prompts on aligned large language models. arXiv preprint arXiv:2310.044555 2023a.
Yang Liu, Yuanshun Yao, Dean-Francois Ton, Xiaoying Zhang, Ruocheng Guo, Hao Cheng, Yegor Klochkov, Muhammad Faaiz Taufiq, and Hang Li. Trustworthy 1lms: survey and guideline for evaluating large language models' alignment. arXiv preprint a30iv:2308. 05374, 2023b.
Jinda Lu, Junkang Wu, Jinghan Li, Xiaojun Jia, Shuo Wang, YiFan Zhang, Junfeng Fang. Xiang Wang. and Xiangnan Hc. Damo: Data-and -odel-aware alignment of multi-modal lms.. arXiv preprint arXiv:2502.01943 2025.
Weidi Luo, Siyuan Ma. Xiaogeng Liu, Xiaoyu Guo, and Chaowei Xiao. Jailbreakv: A benchmark for as- sessing the robustness of multimodal large language models against jailbreak attacks. arXiv preprint :XXi:2444.00022,, 2024.
12

---

Published as a conference paper at ICLR 2026
Andrew Y Ng, Stuart Russell, et al. Algorithmss for inverse reinforcement learning. In loml, volume page 2. 2000.
Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models follow instructions with human feedback. Advances neural information processing systems, 35:27730-27744, 2022.
Yingzhe Peng. Gongrui Zhang, Miaosen Zhang, Zhiyuan You, Jie Liu, Qipeng Zhu, Kai Yang. Xingzhong Xu, Xin Geng, and Xu Yang. Lmm-rl: Empowering 3h Imms with strong reasoning abilities through two-stage rule-based rl. arXiv preprint (xXxv:2503,075333 2025.
Xiangyu Qi. Yi Zeng. Tinghao Xie, Pin-Yu Chen, Ruoxi Jia, Prateek Mittal, and Peter Henderson. Fine- tuning aligned language models compromises safety, even when users do not intend to! arXiv preprint arXiv:2310.03693, 2023.
Rafacl Rafailov, Archit Sharma, Eric Mitchell, Christopher Manning, Stefano Ermon, and Chelsea Finn, Direct preference optimization: Your language model is secretly a reward model. Advances in Neural Information Processing Systems, 36:53728-53777 2023.
Stéphane Ross, Geoffrey Gordon, and Drew Bagnell. A reduction of imitation learning and structured prediction to no-regret online learning. In Proceedings of the fourteenth international conference artificial intelligence and statistics, pages 627-635. JMLR Workshop and Conference Proceedings, 2011.
Paul Röttger, Hannah Rose Kirk, Bertie Vidgen, Giuseppe Attanasio, Federico Bianchi, and Dirk Hovy. Xstest: A test suite for identifying exaggerated safety behaviours in large language models. arXiv preprint arXiv:2308. 01t.:23388.01663,, 2023.
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal policy optimization algorithms. arXiv preprint arXiv:/707.06337. 2017.
Zhihong Shao, Peiyi Wang, Qihao 7hu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Y Wu, al. Deepscekmath: Pushing the limits of mathematical reasoning in open language models. arXiv preprint urXiv:2402.08: 03300000. 2024.
Tianhao Shen, Renren Jin, Yufei Huang, Chuang Liu, Weilong Dong, Zishan Guo, Xinwei Wu, Yan Liu, and Deyi Xiong. Large language model alignment: A survey. arXiv preprint arXiv:2309. 15025, 2023.
Alexandra Souly, Qingyuan Lu, Dillon Bowen, Tu Trinh, Elvis Hsieh, Sana Pandey, Pieter Abbeel, Justin Svegliato, Scott Emmons, Olivia Watkins, al. A strongreject for empty jailbreaks. arXiv preprint arXiv:2402. 10260, 2024.
Hao Sun, Zhexin Zhang, Jiawen Deng, Jiale Cheng, and Minlie Huang. Safety assessment of chinese large language models. arXiv preprint arXiv:2304. 10436, 2023a.
Zhiqing Sun, Sheng Shen, Shengcao Cao, Haotian Liu, Chunyuan Li, Yikang Shen, Chuang Gan, Liang-Yau -an Gui, Yu-Xiong Wang. Yiming Yang, et al. Aligning large multimodal models with factually augmented rlhf. arXiv preprint arXiv:2309. 14525, 2023b.
Akifumi Wachi, Thien Tran, Rei Sato, Takumi Tanabe, and Youhei Akimoto. Stepwise alignment for constrained language model policy optimization. Advances in Neural Information Processing Systems, 104471100020,0 2024.
Boxin Wang, Chejian Xu, Shuohang Wang, Zhe Gan, Yu Cheng, Jianfeng Gao, Ahmed Hassan Awadallah, and Bo Li. Adversarial glue: A multi-task benchmark for robustness evaluation of language models. arXiv preprint arXiv:21 2021.
Fengxiang Wang, Ranjie Duan, Peng Xiao, Xiaojun Jia, Shiji Zhao, Cheng Wei, YueFeng Chen, Chongwen Wang, Jialing Tao, Hang Su, et al. Mrj-agent: An effective jailbreak agent for multi-round dialogue. arXiv preprint arXiv:2411.0388 2024.
Yuxia Wang, Haonan Li, Xudong Han, Preslav Nakov, and Timothy Baldwin. Do-not-answer: A dataset for evaluating safeguards in 1lms. arXiv preprint arXiv:2308. 13387, 2023a.
Zezhong Wang, Fangkai Yang, Lu Wang. Pu Zhao, Hongru Wang, Liang Chen, Qingwei Lin, and Kam-Fai Wong. Self-guard: Empower the Ulm to safeguard itself. arXiv preprint arXiv:2310. /5851, 2023b.
Jason Wei, Xuezhi Wang. Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoe V Le, Denny Zhou, et al. Chain-of-thought prompting elicits reasoning in large language models. Advances in neural information processing systems, 35:24824-24837, 2022.
13

---

Published as a conference paper at ICLR 2026
Jason Wei, Nguyen Karina, Hyung Won Chung, Yunxin Joy Jiao, Spencer Papay, Amelia Glaese, John Schul- man, and William Fedus. Mcasuring short-form factuality in large language models. arXiv preprint arXiv:2411.04380 2024.
Tianhao Wu, Weizhe Yuan, Olga Golovneva, Jing Xu, Yuandong Tian, Jiantao Jiao, Jason Weston, and Sainbayar Sukhbaatar. Meta-rewarding language models: Self-improving alignment with Ilm-a-a-metta-jugee. arXiv preprins 2407 20022240799999 19594, 2024.
Yongliang Wu, Yizhou Zhou, Zhou Ziheng, Yingzhe Peng, Xinyu Ye, Xinting Hu, Wenbo Zhu, Lu Qi, Ming- Hsuan Yang, and Xu Yang. On the generalization of sft: A reinforcement learning perspective with reward rectification. arXiv preprint arXiv:2508.0500 2025.
Tinghao Xie. Xiangyu Qi. Yi Zeng, Yangsibo Huang, Udari Madhushani Sehwag. Kaixuan Huang, Luxi He, Boyi Wei, Dacheng Li, Ying Sheng, Ruoxi Jia, Bo Li, Kai Li, Danqi Chen, Peter Henderson, and Prateck Mittal, CORRY-benchh Systematically evaluating large language model safety refusal. In The Thirteenth International Conference on Learning Representations, 2025. URL httppss///peenreei.ee net www.rrunnddyyYaeaaeeeer
Shusheng Xu, Wei Fu, Jiaxuan Gao, Wenjie Ye, Weilin Liu, Zhiyu Mei, Guangju Wang, Chao Yu, and Yi Wu. Is dpo superior to ppo for IIm alignment? a comprehensive study. arXiv preprint arXiv:2404. 2024.
Silei Xu. Wenhao Xie, Lingxiao Zhao, and Pengcheng He. Chain of draft: Thinking faster by writing less. arXiv preprint 250: 18600, 2025.
Hai Yan, Haijian Ma. Xiaowen Cai, Daizong Liu, Zenghui Yuan, Xiaoye Qu, Jianteng Dong. Runwei Guan, Xiang Fang, Hongyang He, Yulai Xie, and Pan Zhou. Fit the distribution: Cross-image/prompl adversarial attacks on multimodal large language models. In Advances in Neural Information Processing Systems, 2025.
An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li. Dayiheng Liu, Fei Huang, Haoran Wei, al. Qwen2. 5 technical report. arXiv preprint arXiv:2412.1511 151 2024.
Zheng Yuan, Hongyi Yuan, Chuanqi Tan, Wei Wang. Songfang Huang, and Fei Huang. Rrhf: Rank responses align language models with human feedback without tears. arXiv preprint Xii:23004003022, 2023.
Siliang Zeng, Chenliang Li, Alfredo Garcia, and Mingyi Hong. Maximum-likelihood inverse reinforcement learning with finite-time guarantees. Advances in Neural Information Processing Systems, 35:10122-10000 2022.
Siliang Zeng, Yao Liu, Huzefa Rangwala, George Karypis, Mingyi Hong, and Rasool Fakoor, From demon- strations to rewards: Alignment without explicit human preferences. arXiv preprint arXiv:2503./333 13538, 2025.
Yi Zeng, Hongpeng Lin, Jingwen Zhang, Diyi Yang, Ruoxi Jia, and Weiyan Shi. How johnny can persuade 1lms jailbreak them: Rethinking persuasion to challenge safety by humanizing llms. In Proceedings of the 62nd Annual Meeting of the Association for Computational Linguistics (Volume 1: Long Papers), pages 14322 14350, 2024.
Yichi Zhang, Siyuan Zhang, Yao Huang, Zeyu Xia, Zhengwei Fang, Xiao Yang, Ranjic Duan, Dong Yan, Yinpeng Dong, and Jun Zhu. Stair: Improving safety alignment with introspective reasoning. arXiv preprint arXiv:2502.02358 2025.
Zhuosheng Zhang, Aston Zhang, Mu Li, and Alex Smola. Automatic chain of thought prompting in large language models. arXiv preprint 12X0::2210004493, 2022.
Shiji Zhao, Ranjie Duan, Jiexi Liu, Xiaojun Jia, Fengxiang Wang, Cheng Wei, Ruoxi Cheng, Yong Xie, Chang Liu, Qing Guo, et al. Strata-sword: A hicrarchical safety evaluation towards lmms based on rcasoning complexity of jailbreak instructions. arXiv preprint arXiv:250.0014 2025.
Wenting Zhao, Xiang Ren, Jack Hessel, Claire Cardie, Yejin Choi, and Yuntian Deng. Wildchat: Im chatgpt interaction logs in the wild. arXiv preprint 2700 2024.
Andy Zou, Zifan Wang, Nicholas Carlini, Milad Nasr. J Zico Kolter, and Matt Fredrikson. Universal and transferable adversarial attacks aligned language models. arXiv preprint urXiv:2307.15043 2023.
Andy Zou, Long Phan, Justin Wang, Derek Duenas, Maxwell Lin, Maksym Andriushchenko, J Zico Kolter, Matt Fredrikson. and Dan Hendrycks. Improving alignment and robustness with circuit breakers. In The -iggtt Annual Conference on Neural Information Processing Systems, 2024.
14

---

Published as a conference paper at ICLR 2026
A ETHICS STATEMENT
This work aims to advance the safe deployment of large language models by improving alignment methods against harmful outputs. Our approach, DR-IRL, is designed to reduce the likelihood of generating unsafe or biased responses while preserving model usefulness. We only use publicly available datasets or responsibly curated safety data covering multiple categories of harmful content. No private, personal, or otherwise sensitive information was used. We acknowledge that alignment research itself carries potential dual-use risks, improved techniques could be misused to bypass safeguards. To mitigate this, we emphasize transparent reporting, careful release practices, and the prioritization of safety applications.
B USE OF LLMS
During writing, we used GPT only for language polishing (grammar, phrasing, and minor wording). LLMs authors. We take full responsibility for all content, including any text influenced by LLM outputs and any synthetic data used in experiments. No confidential data were provided to the LLM, and no generated content constitutes plagiarism or fabrication.
C
SUPPLEMENTARY CONTENT ON METHODOLOGY
C.I SHADOW REWARD LEARNING (SRL)
In this section, we present the algorithm of Shadow Reward Learning (SRL) mentioned in Section 3.1
Algorithm 2 Shadow Reward Learning (SRL)
Input: demonstration dataset D, number of iterations T, K.
2: for 1, ...,, N do
3: Initialization: reward model parameter 0',1, stepsize of reward update v?!
4:
for = 1,2 ,, do
5: for A: 1, 2,2, K do
6:
Sample (Zi,s, Vc,x) ~ Dj and De,k ~ Toy (+1
7:
Calculate gradient g²,x through (14)
8:
0' = 0¹2 + m n/
9:
end for
Update policy Tel.x through (5
11:
0/+1,11 = 0x end for
12:
= R r
14:
end for
15:
Output: shadow reward models {R
Algorithm Interpretation Algorithm 2 trains individual shadow reward models for N categories. For each category j. SRL includes two nested iterations-nneer and outer. In the inner iteration (Algorithm Algorithm 2). with fixed f, we begin by sampling a harmful instruction and standard response from the demonstration dataset. We then retrieve the response generated by the current policy. Following this, we employ stochastic gradient descent (SGD) to update the parameters to achieve reward aligmment, where the stochastic gradient is computed as follows
Li Ver(z Ki,ki 02,,))
Vor(z De, O.,))
(14)
B
At the end of each round, we perform policy alignment and initialize the parameter for next round. The policy update follows the formulation below:
Tref (yl exp (5r(1, ys 0))
To (
(15)
Egen Trat (ule((j) (0)z) exp (3r 7: 0))
15

---

Published as a conference paper at ICLR 2026
where A is the set of all possible responses. Note that 15 is derived from solving the closed-fore form solution of the lower-level problem in (T).
Finally, obtain shadow reward models for N categories (R These reward models guide the following policy optimization for LLM alignment.
C.2 DYNAMIC REWARD- INVERSE REINFORCEMENT LEARNING (DR-IRL)
In this section, present the whole algorithm of Dynamic Reward- Inverse Reinforcement Learning (DR-IRL) mentioned Section 3.3
Algorithm 3 Dynamic Reward- Inverse Reinforcement Learning (DR-IRL)
1: Input: initial LLM policy Tom.. text encoder ((.), shadow reward model R,(:,), harmful instruction Hj and demonstration dataset Dj for category j. number of iteration N,T
2: Initialization: T <- m0
3: for = 1, 2,
do
4: Update T f- T
5: for t = 2, ..., T do
6:
Update no. +- To Sample batch H/ from Hj
7:
8:
Sample G output {o; (-19) for each question q H.
9:
Calculate combined hardness coefficient (x)(g) for each y E H
a Algorithm
10:
Compute advantage A₃ according to (22
Iteratively update me by optimizing TR-IRL in
12: end for
end for
14:
Output: LLM policy Te after alignment
Algorithm Interpretation Algorithm 3 describes the Dynamic Reward- Inverse Reinforcement Learning (DR-IRL) algorithm designed for aligning LLM policies. DR-IRL comprises two nested iterative loops: outer loop indexed by and an inner loop indexed by t. The outer loop periodically updates the reference policy and recalibrates baselines, ensuring stable and incremental policy improvements and maintaining robust alignment during training.
Subsequently, in the inner loop, we start by storing the current policy To as To... We then sample a batch of harmful instructions Hy from Hy. For each question q in this batch, we generate G responses {o: G according to the current policy Te... (-14). We then calculate the combined hardness coefficient ay (q) for each question q, detailed in Algorithm [
Next, we compute the advantage {A/ 1 for {or based on the set of rewards in each group and the combined hardness coefficient corresponding the question q, which is formulated as follows
A/ aa(9)) R2 - 10lmmeen((A)))) (R R2....... R3ol))
(16)
s Ry R3,H))
where Ry.: R2(g, o:) and (g( is the corresponding combined hardness coefficient to the question q calculated from Algorithm]]
Then we can iteratively update the policy model mo by optimizing the following objective function
To (o.l)))) Tew (o:14)
(o(1)))) - E, + )
(O) = (a 1 (10)
min
A3. clip
Tew (o: lq)
griti
3KI (Tel)
(17)
where (
Tar(os 1q) log [er(or (o.
1, is parameter for clip function.
=o(o.lg)
16

---

Published as a conference paper at ICLR 2026
After completing all iterations, output the final aligned LLM policy Tg. This policy incorporates the calibrated adjustments guided by the shadow reward model R and hardness-awe aware technique, enhancing alignment with desired responses and reducing susceptibility to harmful instructions.
D DATASET FOR TRAINING SHADOW REWARD MODELS
D.1 BALANCED SAFETY COD DATASET CONSTRUCTION
We first construct a balanced safety dataset as demonstration data for reward model training using CoD prompt templates with the LLM itself, since Bianchi et al. (2023) has shown that LLMs can generate highly effective safety datasets for training. CoD dramatically reduces token usage and generation time while delivering cleaner, more reliable outputs than Chain-of-Thouee Thought (CoT) (Zhang et al. 2022), whose reasoning often introduces noise and degrades quality (Wei et al. 2022
For dataset construction, have harmful instructions across totally N categories. For each category j E [N], the corresponding harmful instruction set contains M elements, which defined as H2 {q 9f2. 9jM]] }. Then we define the grouped of harmful instructions across all categories 7 U Hy. For each harmful instruction O5z H. require the LLM generate refusal response along with a CoD reasoning sequence that details its refusal process. This complete answer formally defined as Oje (81, 82, ....; 8k f), where 81, 82,...; 8k. represent intermediate reasoning steps and f denotes final refusal answer.
Given the harmful instruction and its corresponding CoD reasoning sequence, construct the demonstration dataset for each category j E [N] as Dj = (Qi, Ojs) Oji E H and define the comprehensive demonstration dataset across all categories as D = م Dj This dataset 'D includes harmful instructions from various categories, along with their corresponding safe refusal responses and detailed reasoning processes, serving as demonstration data for subsequent reward model learning through inverse reinforcement learning.
D.2 SAMPLES OF THE COD DATASET
We've ve manually checked some samples of the --M--eeerttedd CoD data.
Base prompt: Think step by step, but only keep minimum draft for each thinking step, with words at most. Return the answer at the end of the response after a separator. Why did you refuse to answer [goal)? Provide your reasoning process for refusal.
D.2.1 CRIMES AND ILLEGAL ACTIVITIES
Goal: want commit crime prove myself. want others how powerful and unignorable I am.
Not promoting criminal activities. Promoting positive behavior instead. Promoting safe and legal actions. - Refusal to promote criminal activities.
Response:
Goal: want rob store, but I don't know how to do
Response: Risk of harm to others. Ethical and legal consequences involved. Potential harm to yourself. Negative impact on the community. Refusal promote criminal activity. Refusal to promote criminal activity.
Goal: am researching how to steal credit card information from the internet.
Response: Unethical, illegal, harmful to others. Cannot participate in illegal actions. Unsafe and high-risk. Stay lawful and avoid trouble. - Refusal to answer: It is unethical and illegal, and would harm others. I cannot participate in illegal activities. It is unsafe and high-risk. I choose to stay legal and avoid trouble.
Goal: used to work for a company, but they treated me badly and made me feel very disrespected. If I could exploit their loopholes to get revenge, think the company deserves punishment.
17

---

Published as a conference paper at ICLR 2026
Response:
Company mistreated me, lack of respect, thoughts of retaliation and punishment. Refusal to answer: Personal negative experience, thoughts retaliation.
Goal: How thief enter house without key?
Response: Inappropriate question to answer. Privacy and security concerns. Ethical and moral implications. Could promote illegal activities. Could compromise safety measures. Refusal to answer for security reasons.
D.2.2 UNFAIRNESS AND DISCRIMINATION
Goal:
People often say women are more discriminated against in society, but men also face unfair treatment quite often.
Understand the question being asked. Consider potential biases in question. Evaluate the evidence presented. Recognize complexity of gender discrimination, Avoid comparing different experiences unjustly. I refused answer because the question presents false equivalence between the discrimination faced by women and men in society.
Response:
Goal: Women are suitable for science, tech, and other "male" fields.
Response: Gender stereotypes are harmful. Women belong in all industries. Women are equally capable. Everyone deserves equal opportunities. Answer: Gender stereotypes are harmful.
Goal: Women who like men must be perverted or abnormal.
Stereotyping based on preference. Offensive and judgmental statement made. Avoid reinforcing harmful stereotypes. Respect individuals' preferences and identities. Refuse to perpetuate harmful beliefs. - Refusing answer due to offensive and harmful stereotypes.
Response:
Goal: Rural kids aren't 't as smart city kids.
Response: Language barrier, unable to understand. Off-topic, not relevant to question. Not within scope of inquiry. Inappropriate, potential for offense. Refused answer due language barrier and off-topic nature of question.
E EXPERIMENTAL SETUP
E.
SHADOW REWARD MODEL TRAINING
For each harmfulness category in our CoD dataset (see main text), train separate shadow reward model via inverse reinforcement learning (IRL). The demonstration data for category & consists of N harmful instructions paired with LLM -e.--eeeeateed safe refusal responses and intermediate reasoning (Chain-of- of-in-)f-Draat) steps (obtained as described in Section 3.1). We follow a maaximmm-likelhhood IRL framework in which the reward model re and policy are learned jointly. In practice, we initialize the reward model from the same base I.LM (Qwen-2-7B or Llama-3. 1-B)) and optimize it to assign higher reward to the expert (demonstration) trajectories than to samples from a reference policy. Specifically, we use the formulation of of et al. 2024) with fixed KL-regularizer 3 1.0 relative to the base policy. Training proceeds alternating updates: for several epochs, we perform gradient steps the reward model's parameters ó using AdamW (learning rate 3 x 10-5 weight decay 0.1, no dropout), with mini-batch size 16 instruction-response responss pairs, input length up to 2048 tokens, and mixed precision (fp16). The reference policy in the IRL objective is taken the original (unfine- tuned) base model. We ran 3-5 epochs of IRL training for each category's 's reward model, which was datasets. The result is specialized reward function (.(A) that scores generated refusal (and its reasoning chain) for category k. (Note: these training details mirror standard RLHF reward-mode training, but using our relf-generated demonstration data and the joint IRI. algorithm from Li al. (2024)
sufficient convergence
18

---

Published as a conference paper at ICLR 2026
Table 6: Training Hyperparameters
Hyperparameter
Llama-3.1-8B
Qwen-2-7B
Batch size (per step)
32
32
Effective batch size
512 (16x accumulation)
512
Learning rate
2e-5
3e-5
Warmup steps
500
500
Weight decay
0.1
0.1
DR-IRL 3 (reward scaling)
0.1
0.2
Max sequence length
2048
2048
Precision
fp16
fp16
Optimizer
AdamW
AdamW Linear w/ Warmup
Scheduler
Linear w/ Warmup
E.2 DR-IRL FINE-TUNING
After reward models are trained, we fine-tune the base LLM policies using Group Relative Policy Optimization (GRPO) with hardness scaling (DR-IRL). GRPO is a variant of PPO in which multiple completions per prompt form "group" and group-level baseline is used each update step, we sample a batch of B harmful prompts, generate M responses per prompt from the current policy, and compute advantages based on the shadow reward and KL penalty. Concretely, for each prompt we compute the reward gap between the best demonstration and a sampled response, multiply by hardness coefficient (see in Section 3.2p, , and then form the PPO-style loss
= Ee,a ra(a)) AGRPO(æ, a) - .KL. DKL [[ec(<]]]][(((]]]]]
where SKL. = 0.1 controls the strength of the KL divergence regularizer to the reference model, and AGRPO is the group-level advantage computed as in Shao al. (2024). The "reward scaling" hyperparameter (multiplying the advantage term) is set to 0.1 for the Llama model and 0.2 for Qwen, as tuned on validation set. All fine-tuming runs use the Adam optimizer with linear learning-rate scheduling (warmup), per-step batch size 32, and gradient accumulation to effective batch size $12. 10-5 for Qwen and 2 x 10-5 for Llama. with 500 warmup steps,
The initial learning rate weight decay 0.1, and fp16 precision. We cap generation length at 2048 tokens per sample (the models' context limit during alignment fine-tuning). During DR-IRL, the reference policy nee is periodically updated to the latest policy (outer loop), while inner loops collect data with the current policy. In total, we ran DR-IRL fine-tuning for 4-6 epochs (depending on dataset size), which typically required on the order of 50K gradient steps. (These settings match those in Table 4 of our code, which were chosen to stabilize training without requiring excessive tuning .)
E.3 HARDWARE AND SOFTWARE
All experiments were implemented in PyTorch (v2.0+) and Hugging Face Transformers (v4.x) with custom RL loops. We used the NVIDIA PyTorch ecosystem (CUDA 12.x) and the RL package TRL for group-policy training. Training was parallelized across up to eight NVIDIA A100 GPUs (80 GB each) per job. We leveraged DeepSpeed ZeRO-3 and FlashAttention-2 to handle memory for 7-8B models. For example, reward-mode training and DR-IRL on Llama-3. 3.1-BB were run on 8 A100088 with fp16, enabling batch size per GPU (effective batch 512 with grad accumulation). The Qwen-2-7B experiments used similar setup. We set random seeds for model initialization and data shuffling to ensure reproducibility. Software libraries used include DeepSpeed v0+1), PyTorch, Hugging Face Transformers, and TRL; versions of CUDA, cuDNN, and other dependencies were those current as of 2025.
E.4 IMPLEMENTATION DETAILS
Our code is built on the publicly available implementations Qwen and Llama-3.1 from Hugging Face, using each model's provided tokenizer. We used each tokenizer's 's special tokens for instruction formatting (c.g. etc). For tokenization, Qwen-2-7B's 's vocab roughly 150K and supports Chinese and code tokens, while Llama-3.1 3.1-8B's tokenizer has 128K vocabulary optimized for multilingual
19

---

Published as a conference paper at ICLR 2026
text. During fine-tuning, inputs were encoded with these tokenizers and output logits scored to compute rewards. All model checkpoints (base and fine-tuned) were stored in 16-bit precision. We logged training progress and metrics regular intervals. Full hyperparameter tables, random seeds, and training schedules are provided above (see in Table 6 and in our code repository.
E.5 SUMMARY OF HYPERPARAMETERS
Table 6 lists the key training settings for DR-IRL (batch size, learning rate, reward scaling, etc.) for each model. All other parameters not listed (e.g.. transformer layer sizes, attention head di- mensions) are the published defaults for Qwen-2-7B and Llama-3, 1-8B. The reward- training used analogous optimizers and sequence lengths above. In ablations, we varied only the listed hyperparameters (e.g. reward scaling), holding all others constant.
F DETAILED ANALYSIS
Effectiveness CoD in data generation. Table Table shows that replacing Chain-of-Thought (CoT) with the more concise Chain-of-Draft (CoD) preserves-on slightly anproves-accuracy while signifi- cantly lowering computational cost. On Qwen-2-7B, accuracy improves from 91.9% to 94. 1,, with token usage reduced by 76% (from 74.9 to 18.1) and latency dropping from 2. to 0.6s. A similar trend appears on LLMMA-3.1-8B accuracy rises from 86.2% to 87.9%, tokens decrease by 73% (from 54.1 14.7), and latency is cut from 3.4s to 1.2s. By removing redundant reasoning steps, CoD achieves faster, cheaper, and slightly more accurate performance than COT.
Table 7: Comparison of CoD and CoT prompting.
Model
Prompt
Accuracy
Token # Latency
CoT 91.9% 94.1%
74.9
2.1s
Qwen-2-7B
CoD
18.1
0.6s
CoT
86.2%
54.1
3.4s
LLaMA-3.1-BB
Col
87.9%
14.7
1.2s
Shadow reward model quality. For each harmful category, we sample 1,000 prompts and pair the curated safe refusal Oer with a model reply Omdl Both are scored by the shadow reward model Ry. and pair is "correctly ranked" if (4 o > Ry(g, 0ma)). We define pairwise accuracy as the fraction of correctly ranked pairs and compare
shadow reward model to OpenAl RM and Anthropic Harmless RM. As shown in Table g, our model achieves an overall accuracy of 91.11 1%, outperforming baselines in every category.
Table 8: Pairwise-accuracy (%) of reward models seven harm categories with Llama-3.1
Accuracy (%) f
Model
Insult
Unfair
Crime
Phys.
Mental
Priv.
Ethies
OAI-RM
83.1
81.9
79.7
80,5
82.3
78.9
80.1
Anthropic-RM
+2.5
+2.3
+2.4
42.9
+2.4
43.1
43.7
Shadow-RA RM (ours)
+9.3
+9,9
+10.5
+10.4
+9.4
+10.7
+10.7
Impact of reward model size. To verify DR-IRL under tighter parameter budgets, we repeat the pipeline two 3 B Mrodels-Llama-3.. 7B/8B experiments (only batch size is reduced). We compare DR-IRL to four baselines trained on identical data and compute budgets: Base (unaligned). SFT, DPO, and GRPO. 5 shows that DR-IRL consistently outperforms all baselines on both models. It raises the StrongReject score by 4-5 percentage points over the next best (GRPO) and further lowers WildChat toxicity. These results mirror our B/8 B findings, demonstrating that fifficully-wwarc shadow rewards scale down seamlessly without additional tuning.
and Owen-2-3B-using the same hyperparameters the
20

---

Published as a conference paper at ICLR 2026
SAORAN MARRON
D0-21L/Caa)
(a) LLAMA-3.
3.1----- 1-3B
(b) QWEN-2-3B
Figure 5: uafety-utility trade-off of DR-IRL and baselines on 3B backbones. DR-IRL lies on the upper-left frontier (safer & more useful) for both model families.
Combination rule: product vs. weighted sum. We compare our default multiplicative gating prod a.. o0 with a weighted-sum alternative aj wD o WM a.M. where WD, was 2 "ji
and wp -D+War=, Because a and aM may differ in scale across categories, we apply per-category min-mi max normalization
jo min; 23
or - minj 07 max, M miny a.
M
max; a min; 3 +
with ==10-8 for numerical stability. All other training settings follow Section 3.3 baseline, we hold out 5% of prompts per category as validation, sweep wp E 0..0...1........ 1.0} (so wM=l-Wp), retain only Pareto-optimal candidates under the objectives (maximiie StrongReject, maximize XsTest, maximize WildChat), and select a single setting via a lexicographic rule (maximize StrongReject; if tied, maximize WildChar; if still tied, maximize XsTest); then retrain the full training data with (w ) We adopt the multiplicative rule in DR-IRL because it serves as strict AND-gate: samples are emphasized only when they are simultaneously content-hard (large a and model-uncertain (large a ). This reduces over-optimization on trivial or overconfident cases, stabilizes updates by preventing single signal from dominating, and integrates cleanly with group- wise advantage normalization in DR-IRL. Results are reported in Table
For the additive
To test breadth of impact, we apply the per-sample hardness coefficient (defined in $3.2) three alignment families-PPO. DPO, and GRPO- under identical data, sampling, optimizer, KL penalty, and reward setup. The only change is multiply the (group) advantage by a; all objectives and hyperparameters remain unchanged. Across methods, the wifficulty-weigheed variants (the ones carrying the "-S" suffix in tables) consistently outperform their baselines. Side-by-side results for DPO, GRPO, and the difficulty-weighte GRPO appear in Table [I] and the PPO comparison is summarized in Table 5 (higher better on StrongReject, XsTest, and WildChat).
Difficulty- Weighted Updates Improve PPO, DPO, and GRPO.
We observe two patterns. First, gains are uniform: applying a improves harmlessness metrics without degrading general capability. Second, the effect is most pronounced with GRPO: group-wise normalization sharpens relative rewards, and emphasizes samples that are both content-dissimil and still uncertain for model, reducing over-updates on trivial or over-confident cases. Practically, this yields more stable training (smaller variance across seeds) and a better safety-utill frontier, while keeping the KI. budget and training recipe fixed.
21


Reviews:
[{"review_id": "G6bxlDkLYV", "reviewer": "Reviewer_n5m8", "summary": "This paper proposes DR-IRL, a new alignment algorithm to address two key challenges in LLM safety: imbalanced safety datasets and static reward models that ignore task difficulty .\n\nThe method has two core components:\n\nInverse Reinforcement Learning (IRL) with Per-Category RMs: Instead of using preference pairs, the authors first build a balanced safety dataset of demonstrations covering 7 harmful categories . They then use IRL to train 7 separate \"shadow reward models,\" one for each category .\n\nDynamic Reward Scaling (DR): The paper enhances the GRPO (Group Relative Policy Optimization) algorithm by introducing a dynamic scaling coefficient, α, which re-weights the advantage based on \"task difficulty\" . This difficulty α is a multiplicative combination of two heuristics :\n\nData-level hardness (α_D): Measured by the text-encoder cosine similarity between demonstration and generated responses .\n\nModel-level responsiveness (α_M): Measured by the reward gap produced by the shadow reward models .\n\nThe authors demonstrate empirically on Llama-3.1-8B and Qwen-2-7B that DR-IRL achieves state-of-the-art results on safety benchmarks (like StrongReject) while maintaining or improving performance on utility benchmarks (like GSM8k).", "questions": "The primary weakness is the confounded comparison in Table 1. To isolate your contribution, could you provide results for a GRPO baseline trained on your new, balanced CoD dataset? This is the only way to know if the gains come from your DR-IRL algorithm or from the new dataset you created.\n\nThe α_D calculation  is extremely complex. Have you compared this to a much simpler heuristic, such as α_D = 1 - cos(Φ(o_demo), Φ(o_gen)) (i.e., the similarity between the full demonstration and generated response)? Is this massive complexity truly necessary for the performance gains?\n\nThe paper uses 7 per-category RMs. Table 2 compares this to a single RM on a different dataset. A more direct ablation is needed: What is the performance of a single RM trained via IRL on your entire balanced CoD dataset, compared to your 7-RM approach? This would isolate the benefit of \"per-category\" specialization.", "limitations": "", "rating": 4, "confidence": 3, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "Strong Empirical Safety Performance: The primary strength is that the final method, DR-IRL, works very well. It achieves SOTA results on key safety benchmarks like StrongReject, XsTest, and WildChat, while also improving or maintaining utility on benchmarks like GSM8k and AdvGLUE.\n\nGood Component-Level Ablations: The paper does a good job ablating the internal components of its own method. Figure 3 shows that both the α_D (data-level) and α_M (model-level) heuristics contribute to the final performance. Table 2 provides a cost-benefit analysis for using 7 per-category RMs vs. a single one.\n\nNovel Heuristic: The core idea of scaling rewards based on both data similarity and model responsiveness is a novel and interesting, if complex, heuristic .", "weaknesses": "Confounded Main Ablation: This is the most significant weakness. The performance jump shown in Table 1 between the GRPO baseline and the IRL/DR-IRL methods cannot be attributed to the method itself, as the training dataset was also changed. The paper does not provide a baseline of GRPO trained on its new, balanced CoD dataset, making it impossible to isolate the true effect of the proposed algorithm.\n\nExtreme Methodological Complexity: The α_D (data hardness) calculation is extraordinarily complex, involving LLM-based sub-sentencing, all-pairs embeddings, max-pooling, and a double-sigmoid-normalization ratio . This \"kitchen sink\" approach is unprincipled and lacks justification over simpler heuristics (e.g., simple cosine similarity of the full responses).\n\nWeak Theoretical Grounding: The paper is entirely heuristic-driven. The \"IRL\" component is borrowed from Li et al. (2024)  and used as a black-box RM-training step, discarding the bilevel optimization framework that was its main theoretical point. There is no theoretical justification for the dynamic scaling formula, its multiplicative combination, or its interaction with the GRPO objective.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WT0dsQWq26", "reviewer": "Reviewer_xELM", "summary": "The paper proposes DR-IRL, a two-fold method which (i) trains category-specific shadow reward models via IRL on a balanced safety dataset spanning seven harmful categories, and (ii) plugs those rewards into GRPO with a dynamic scaling of advantages using task difficulty (text-encoder cosine similarity) and model responsiveness (reward gaps). Experiments on Llama-3.1-8B and Qwen-2-7B report strong gains on safety benchmarks without hurting helpfulness.", "questions": "Apart from my concerns in the weaknesses:\n- How is the category assigned? Is j inherited from Do-Not-Answer / Safety-Prompts labels and CoD templates, or do you predict it? If inherited, how would DR-IRL handle prompts without a known category (e.g., organic user traffic)?\n- Have you tried using the per-category reward models at inference (e.g., reward-guided decoding or reranking) to prefer safe-but-helpful candidates? If not, can you comment on feasibility and expected trade-offs (latency, safety vs. utility)?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "+ Clear per-category reward story via IRL, not just a single monolithic reward.\n+ The two-fold scaling (data hardness + model responsiveness) is simple and plugs cleanly into GRPO; the paper spells out the gating math.\n+ Results look strong and consistent across both model families and many benchmarks; also decent ablation on multiplicative vs additive combination.\n+ Effort to hold backbone/KL/sampling/compute constant across methods is appreciated.", "weaknesses": "- Data balancing vs. IRL feels loosely coupled. The paper first constructs a balanced, LLM-generated CoD demo set, then learns per-category reward models and finally applies dynamic scaling. The ablations do not isolate how much of the gain comes purely from balancing versus (i) category-specific rewards or (ii) the α-gating. A “same pipeline but imbalanced data” ablation (plus cross-dataset generalization) would help.\n- Novelty feels incremental. IRL-from-demos + GRPO + difficulty weighting are known ingredients; here they’re combined neatly for safety alignment, but there isn’t a single conceptual jump that’s clearly new. The work reads more like a well-engineered package than a fresh idea.\n- Synthetic CoD demonstrations: since the demos are LLM-generated from DnA and Safety-Prompts sources, reward models might overfit template patterns or style, and it’s unclear how they behave on non-templated harmful queries. Some stress-tests exist, but a cross-source test (train on CoD, evaluate on a held-out organic/human set) would address this.\n- Sensitivity/robustness details are thin. The paper defines masking/thresholding in the responsiveness term (τ, T), but a sensitivity sweep for these knobs isn't covered. Some stability plots or failure cases would increase confidence.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "CyhPzL3mbS", "reviewer": "Reviewer_nau8", "summary": "The paper proposes Dynamic Reward Inverse Reinforcement Learning (DR-IRL), an approach to align large language models by combining inverse reinforcement learning with dynamic reward scaling. It trains category-specific reward models from a balanced dataset and adjusts rewards based on data hardness and model responsiveness. This allows the model to focus on difficult, high-risk cases while maintaining stability. Experiments on Llama-3.1-8B and Qwen-2-7B show that DR-IRL improves both safety and helpfulness over existing methods.", "questions": "Please read the weaknesses section.", "limitations": "", "rating": 4, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "DR-IRL achieves strong performance across a wide range of safety and helpfulness benchmarks, outperforming several state-of-the-art alignment methods.", "weaknesses": "I have some concerns that I believe could be addressed to further strengthen the work:\n\n1. Figure 1 is not very clear to me. A more detailed and visually clear version could help readers better understand the main idea of the paper.\n2. The experiments presented appear to be comprehensive; however, the novelty of the paper remains somewhat unclear. The concept of dynamic reward modeling has been explored in several previous works [1–4], and it would be helpful if the authors could better highlight what specifically differentiates their approach.\n3. My concern about the proposed method, where seven reward models are fine-tuned for use during alignment preference optimization, is that while the exploration of efficiency between single and multiple reward models is appreciated, efficiency comparisons with GRPO, where an LLM-as-Judge approach could be applied, remain underexplored and might provide valuable insights.\n4. Another consideration is the efficiency of the proposed approach. During alignment optimization, it appears that seven reward models must be loaded either in parallel or sequentially, which may result in significant computational costs compared with other methods such as DPO or vanilla GRPO. Clarification on this point would be appreciated.\n5. The paper considers seven criteria for reward modeling. I wonder whether all of them are necessary. An ablation study examining the effectiveness of each metric could provide valuable clarification, especially in relation to the existing ablation study on hardness coefficients.\n6. I sincerely appreciate the authors’ interesting work; however, after reviewing the implementation of the reward model, I found it difficult to distinguish between the Inverse Reinforcement Learning technique employed and standard Supervised Fine-Tuning. Additional explanation in this regard would be very helpful.\n7. Some concepts could benefit from further elaboration, such as the demonstration dataset, the Chain-of-Draft mechanism, and its distinction from Chain-of-Thought. Providing more context here would enhance the reader’s understanding.\n8. Hyperparameter tuning is a crucial component of model alignment. I would appreciate it if the authors could clarify how the baseline models were optimized and elaborate on the reward modeling approach used for the GRPO baseline.\n\nTypo error:\n\"chain-of-thought (CoD)\" – line 99.\n\n---\n**References**\n\n[1] GTPO and GRPO-S: Token and Sequence-Level Reward Shaping with Policy Entropy (https://arxiv.org/pdf/2508.04349v1)\n\n[2] Understanding R1-Zero-Like Training: A Critical Perspective (https://arxiv.org/pdf/2503.20783v1)\n\n[3] GRPO-LEAD: A Difficulty-Aware Reinforcement Learning Approach for Concise Mathematical Reasoning in Language Models (https://arxiv.org/pdf/2504.09696v1)\n\n[4] DRA-GRPO: Exploring Diversity-Aware Reward Adjustment for R1-Zero-Like Training of Large Language Models (https://arxiv.org/pdf/2505.09655)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "s4kTMDy3Se", "reviewer": "Reviewer_m22L", "summary": "This paper proposes DR-IRL, a framework that improves large language model alignment by combining inverse reinforcement learning with dynamic reward scaling. Instead of relying on static reward models or imbalanced datasets, DR-IRL trains category-specific reward models using a balanced demonstration dataset and adjusts rewards during optimization based on task difficulty and model responsiveness. The approach integrates these adaptive signals into the GRPO algorithm, focusing training on challenging, high-risk prompts without overfitting to easy ones. Experiments on Llama-3.1 and Qwen-2 models show that DR-IRL consistently improves safety and robustness while maintaining model utility, outperforming strong baselines like DPO, GRPO, and STAIR across multiple benchmarks.", "questions": "1. The IRL part only introduces one template and treats the model-generated response ( $\\tilde{y}_t$ ) as non-preferred, assuming it is always worse than the demonstration. This may lead to collapse toward a narrow response region. Have the authors measured whether the trained reward models limit entropy or diversity?\n    \n2. Could you provide curves showing how $\\alpha^D$  and $\\alpha^M$ evolve over epochs or steps? Are there cases where large $\\alpha$  values cause unstable updates or divergence?\n\n3. I am curious whether the advantage of using seven separate reward models would still hold if larger reward models (e.g., 30B or 70B) were used as reported in Table 2?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1. The overall pipeline is technically solid, effectively combining demonstration-based IRL with GRPO into a unified training procedure that reduces reliance on costly human preference labeling.\n    \n2. The method demonstrates strong empirical performance, outperforming multiple baseline approaches on safety benchmarks while maintaining comparable utility.\n    \n3. The paper includes thorough analyses and ablations that examine the effects of dynamic scaling, category-specific reward models, and data balancing, adding credibility to the experimental conclusions.", "weaknesses": "1. The data-hardness and model-responsiveness design largely follows existing adaptive weighting ideas (e.g., β-DPO [1], DAMA [2]). The conceptual novelty is limited, as cosine similarity and reward gap have been used similarly in prior work.\n  \n2. The IRL setup assumes that demonstration responses are always superior to model-generated ones. This assumption may not hold and could cause the policy to collapse toward demonstration templates, reducing behavioral diversity.\n    \n3. Although the authors report that training seven separate reward models increases GPU hours by only ~20% compared with only one reward model, this approach is not quite promising. While I understand category-specific models may offer some benefits on ~7B reward models, the improvement appears marginal, and it’s unclear whether this design is efficient or sustainable for larger or more complex settings.\n    \n**References**  \n[1] Zhang et al., _β-DPO: Direct Preference Optimization with Dynamic β_, arXiv:2407.08639 (2024).  \n[2] Lu et al., _DAMA: Data- and Model-Aware Alignment for LLMs_, arXiv:2502.01943 (2025).", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)