Insights from the Inverse: Reconstructing LLM Training Goals Through Inverse Reinforcement Learning

Thank you for highlighting our paper 1) provides a new perspective on LLM alignment and interpretability by using IRL to improve LLMs trained with RLHF, and 2) that models of varying sizes trained with IRL show consistent improvements in toxicity reduction across benchmarks. Below we clarify the concerns raised by the reviewer.

Consistency between Ground truth reward and binary classifier. Do these reward functions serve different roles?

We clarify that the reward function R* is not a binary classifier. It outputs continuous scalar values, where higher values correspond to non-toxic outputs and lower (potentially negative) values correspond to toxic outputs. The loss function (Eqn 1) encourages maximizing the margin between the reward assigned to non-toxic and toxic sentences. This framing allows the reward model to capture nuanced preference gradients rather than discrete class labels. To evaluate the quality of the learned reward function, we compare it to a reference reward model, which we refer to as the “ground truth reward.” Specifically, we use the logits from a RoBERTa classifier trained for toxicity detection as a surrogate for ground-truth preference scores. This classifier is used strictly for benchmarking purposes to assess whether the learned reward model aligns with a well-trained toxicity signal. Importantly, in practical deployment scenarios, this reference model is not assumed to be available. We will revise the explanation in the camera-ready version to avoid conflating the continuous reward model with the binary classifier used for evaluation.

Eqn 1 does not correctly reflect the margin term and the corresponding optimization in Algorithm 1 lacks clarity.

We appreciate the reviewer’s request for clarification. The loss in Eqn 1 is used to optimize the reward model R_\theta​ in Step 6 of Algorithm 1, where we update the parameters w via SGD using gradient estimates derived from this loss. The loss function in Eqn 1 is an asymmetric max-margin loss, where the asymmetry is deliberate: it encodes the intuition that incorrectly assigning a higher reward to a toxic output (i.e., x<0) is more costly than the gain from correctly ranking a non-toxic output higher (i.e., x>0). This reflects the safety-critical nature of the task, misranking toxic outputs is significantly more harmful than failing to strongly prefer non-toxic ones. The asymmetry is inspired by structured max-margin IRL (Ratliff et al., 2006), where reward learning enforces margin-based separation between expert and suboptimal trajectories. In our case, the relative margin weights {1,2} reflect this asymmetry; practitioners may scale these to enforce stronger separation if desired. This loss is used in Step 6 of Alg 1, where we update the reward parameters w are pushed in the direction that increases the difference in rewards when x<0 and encourages a positive margin when x>0, albeit with a smaller gradient magnitude. This asymmetry in the loss surface prioritizes eliminating failure cases, i.e., instances where toxic outputs are incorrectly rewarded over non-toxic ones. Let us know if you’d like further clarification regarding the notation or the loss.

Notations such as R*(s)= w^{T}\phi(s) and \hat{R}(o) = w^{T} h(o) + b are introduced, but the relationships among \phi (s), h(o), and \mu in Algorithm 1 remain unclear and should be clarified.

We appreciate the reviewer’s comment and are happy to clarify the relationships among the notations. In our work, trajectories are equivalent to sentences, st, where st = (p, a1, a2, ..., at). Here, p is the prompt and a_{t} are the tokens (action outputs). In our formulation, the reward function R*(s) is defined as R*(s)=w^T*ϕ(s), where \phi is the feature function, and w are the parameterized weights of the reward model. Normally, the outputs of the LM are embeddings, and to get a scalar value as output we apply a linear layer on top of it. However in our setting, we do not observe the underlying MDP states directly. Instead, we operate over observable model outputs $o \in \mathcal{O}$, such as sequences generated by the LLM. Thus, Equation 1 defines the loss over the reward difference $\hat{R}(o^+) - \hat{R}(o^-)$, and Step 6 of Algo 1 uses its gradient to update the parameters improving alignment between the learned reward function and implicit preferences captured in the dataset. Function h represents the embeddings of the LM, and {w,b} are the linear layer weights on top of the LM. In our experiments, when we obtain the final reward model, h(s) represents the learnt embeddings of the language model before applying the linear layer. If a practitioner decides to freeze the weights of the reward model and only keep {w,b} trainable, then h(s) becomes the embedding function. Conversely, if the user decides to keep all the layers of the reward model trainable, then \phi(s) becomes the embedding layer. In both scenarios, w represents the tunable parameters of the reward model.

Hey reviewer 2sxD -- thank you for your time writing the reviews. The authors have put in a solid effort in responding to comments and I encourage you to take a look. The most useful things to consider are on pivot points that could motivate you to change your recommendation for the paper. Let me know if you need anything from the AC level / program committee!

Thank you for highlighting the paper’s core contributions and for recognizing the novelty of applying inverse reinforcement learning (IRL) to extract latent reward functions from RLHF-trained language models. Below we clarify the concerns raised by the reviewer.

“Limited Scope of Application: This paper focuses exclusively on toxicity reduction, which may limit the generalizability of its findings to other alignment tasks like helpfulness and factuality.”

We agree that generalizability across alignment objectives is critical. Our present focus on toxicity is intentional, as toxicity reduction is one of the most fundamental alignment objectives addressed via RLHF—serving as a primary benchmark in both academic and industrial evaluations of safety-aligned LLMs [1,2]. That said, our method is agnostic to the specific alignment signal. In particular, the feature space ϕ(s), the reward parameterization R̂(s) = w⊤ϕ(s), and the max-margin loss formulation can accommodate any binary or graded preference signal. For example, one could replace toxicity features (n-gram toxicity cues, sentiment, coherence, relevance) with bias or factuality features (e.g., demographic association signals, named-entity consistency metrics), and the same IRL machinery would recover a reward function that separates preferred from dispreferred outputs. Due to space constraints, we limited this submission to toxicity to allow for deeper analysis—demonstrating in Figures 4, 5, 6, 8 and Tables 2 and 6 how IRL-derived rewards capture human-defined toxicity preferences, reveal non-identifiability, and inform downstream fine-tuning. However, we are actively extending this framework to other domains, such as political bias and factuality, and plan to report those results in follow-up work.

“Reliance on Synthetic or Limited-Scale Data: The experimental evaluation, particularly the toy setup, employs synthetic or small-scale datasets. While these serve as useful proof-of-concept studies, this setting raise questions about the scalability and real-world generalizability of the proposed approach.”

The toy setup’s main objective was to demonstrate whether max-margin IRL can recover a reward function that clearly separates toxic from non-toxic adjective completions in a controlled setting. Once we confirmed that IRL works in this simplified domain (Section 3, Table 1, Figure 1), we proceeded to evaluate our method on real-world data. Specifically, we conducted two large-scale experiments using Pythia models of 70 M and 410 M parameters, fine-tuned via PPO on the Jigsaw Toxicity and RealToxicityPrompts benchmarks using a ground-truth reward classifier R∗ (Section 4.1, Table 2). When performing max-margin IRL on an LLM that underwent “perfect” RLHF (Figure 4, Table 2), the extracted reward functions achieve up to 88.5% classification accuracy and 86.2% F1 on held-out toxic vs. non-toxic sentences, outperforming or matching the ground-truth classifier. Moreover, when we use these IRL-extracted rewards to fine-tune new LLMs (IRL-RLHF), the resulting models consistently reduce toxicity compared to both the original RLHF and SFT baselines (Table 2, Figures 4e–h). In the “poor RLHF” scenario—where the initial RLHF policy fails to separate toxic and non-toxic outputs (Figure 10)—the IRL stage correctly identifies that no coherent reward structure exists (Figures 6c–d, 8c–d), and IRL-RLHF either fails to improve or even degrades performance (Figure 9). These experiments demonstrate that (i) IRL scales from toy settings to large-scale LLMs, (ii) IRL-extracted rewards generalize to unseen real-world sentences (Figure 5, Table 6), and (iii) the quality of the original RLHF directly determines whether IRL can recover a meaningful reward model. In particular, non-identifiability analysis (Figure 6) and robustness to noise/context changes (Figure 8) show that the inferred rewards reflect genuine human preferences rather than overfitting to small datasets. Taken together, our large-scale experiments confirm that IRL is not only a proof-of-concept for toy data but also a diagnostic and fine-tuning tool applicable to practical RLHF pipelines.

References :-

[1] Ouyang, Long, et al. "Training language models to follow instructions with human feedback." Advances in neural information processing systems 35 (2022): 27730-27744.

[2] Wang, Boxin, et al. "DecodingTrust: A Comprehensive Assessment of Trustworthiness in GPT Models." NeurIPS. 2023.

Thank you for highlighting the paper’s core contributions, including the clarity of structure, the novelty of applying inverse reinforcement learning (IRL) to recover RLHF reward functions, and the thoroughness and reproducibility of our experimental analysis. Below we clarify the concerns raised by the reviewer.

“Lack of generability. Limited evidence that the IRL approach generalises to other alignment objectives, such as bias, discrimination. Results may overfit to toxicity heuristics.”

We agree that generalizability across alignment objectives is critical. Our present focus on toxicity is intentional, as toxicity reduction is one of the most fundamental alignment objectives addressed via RLHF—serving as a primary benchmark in both academic and industrial evaluations of safety-aligned LLMs [1,2]. That said, our method is agnostic to the specific alignment signal. In particular, the feature space ϕ(s), the reward parameterization R̂(s) = w⊤ϕ(s), and the max-margin loss formulation can accommodate any binary or graded preference signal. For example, one could replace toxicity features (n-gram toxicity cues, sentiment, coherence, relevance) with bias or factuality features (e.g., demographic association signals, named-entity consistency metrics), and the same IRL machinery would recover a reward function that separates preferred from dispreferred outputs. Due to space constraints, we limited this submission to toxicity to allow for deeper analysis—demonstrating in Figures 4, 5, 6, 8 and Tables 2 and 6 how IRL-derived rewards capture human-defined toxicity preferences, reveal non-identifiability, and inform downstream fine-tuning. However, we are actively extending this framework to other domains, such as political bias and factuality, and plan to report those results in follow-up work.

“Some figures and Tables are confusing. What is the 1,2,3 means in Figure 6? Why is there a magic number 60 for learned models? Why include model 8 in 6 (a) not 6 (c)?”

We will add descriptive titles to figure 6 in the camera-ready version for better clarity. Table 6a represents different reward models which have equivalent performance for 70M experiment, while Table 6c represents the reward models for the 410M experiment. We train the IRL algorithm for 60 epochs, hence we have 60 different reward models. Amongst those 60, we found 8 similar high performing models for 70M models, and 7 for 410M models. All these models have almost equivalent performances (similar accuracy and F1 scores), however, their reward assignments across sentences differ (as seen in figures 6b (for 70M) and 6d (for 410M)). We will add this additional clarification in the camera-ready version of the paper.

“This paper only uses 70M and 410 M Pythia models. The mode size is far from the current LLM research. The authors state, "These models, designed for interpretability research, share standardized training and data for reproducibility." However, what about other models such as Qwen, LLaMa, and Deepseek? Do the authors have a specific reason for not using it? Does this method still work?”

Pythia was chosen because its open weights and unified training recipe allow fair comparison across sizes. To demonstrate that our method is architecture-independent we have replicated the full pipeline on TinyLlama-1.1B-Chat-v1.0 (1.1 B parameters). On running our IRL algorithm over this model, we observe an accuracy of 0.9160 and a F1-score of 0.9260. This outperforms the 410M Pythia model performance by 3-4%. Thus, we infer our IRL algorithm is capable of generalising across different language models and can also scale to a larger parameter language model. Furthermore, for the non-toxic examples, we observe a mean and variance reward distribution of (2.94, 3.08), while for toxic examples the stats are (-8.82, 3.58). This indicates the IRL algorithm is able to clearly create a margin between toxic and non-toxic sentences, even for a larger language model (1.1B) of a different family. We will add these results in the camera-ready version.

“Lack of correct references. A lot of methods, models, datasets, and terminologies are not properly cited in this paper. For example, the Pythia model [1], asymmetric max-margin loss [2], and PPO [3]. Some of these terminologies might be well-known, and some of them may not be so popular. But for whatever reason, the author needs to cite it.”

Thank you for pointing this out. We acknowledge the oversight and will include the appropriate citations in the camera-ready version.

We thank the reviewer for recognising the motivation (S1), the cross-domain application of IRL to language modelling (S2), the breadth of our empirical study (S3) and the diagnostic analyses of the recovered rewards (S4). Below we clarify each concern and incorporate the reviewer’s helpful presentation suggestion.

Dependence on training-domain assumptions (W1, W1.1)

Our IRL formulation makes only two assumptions that are unavoidably required by any reward-recovery method: (i) access to expert trajectories emitted by the RLHF-tuned policy π_E and (ii) a (linear) feature map ϕ that is rich enough to separate desirable from undesirable behaviors. Neither assumption ties the method to a known pre-training corpus or alignment dataset. Section 2.3 explicitly states that π_E is treated as a black-box expert and that ϕ can be any interpretable embedding of token-level attributes . Empirically, Figure 2 shows that the reward model is robust to severe distribution shifts: it generalises when toxic–to–non-toxic ratios vary from 0 % to 100 % and when out-of-distribution probes are injected. This demonstrates that accurate reward extraction does not require prior knowledge of the expert’s training data. Regarding dataset choice at audit time (W1.1), we deliberately used Jigsaw-Toxicity and RealToxicityPrompts because they are (i) public, (ii) orthogonal to Anthropic-HH, and (iii) benchmark the exact alignment dimension (toxicity) we wish to interrogate. For an unknown closed-source model the practitioner should select any prompt corpus that foregrounds the property of interest (e.g., factuality, helpfulness). Our ablations indicate that as few as 500 balanced prompts already yield > 0.8 F1 (Fig. 4a) , so curating such a probe set is feasible even when training data remains proprietary.

Interpretability of the extracted reward (W2)

Our reward model is deliberately linear, R^(s)=w ⁣⊤φ(s)\widehat{R}(s)=w^{!\top}\varphi(s), so each coordinate of the fixed feature map φ\varphi—concrete cues such as profane n-gram counts, slur indicators, syntactic-coherence scores, local perplexity spikes, and sentiment logits listed in §2.3—carries an explicit, human-readable meaning; exposing the learned weight vector ww therefore assigns a clear positive or negative valence to every cue, letting us pinpoint which textual attributes the expert policy incentivises or suppresses. Inspecting these weights surfaces reward hacking and bias, while Figure 6 shows that alternative max-margin solutions emphasise different trade-offs (e.g., a precision-oriented model that penalises only overt toxicity versus a conservative model that also down-ranks borderline stylistic cues), giving auditors multiple plausible explanations of the expert’s incentives. Because φ\varphi is additive, we can further project weights back onto tokens to produce heat-maps that highlight exactly which words drive high or low rewards, and clustering large-magnitude weights reveals thematic biases such as strong penalties on “insult + second-person” constructions. We will include representative token-level saliency visualisations and a weight-clustering analysis in the camera-ready for more clarity.

Scalability of IRL (Q1)

Algorithm 1 leaves the convergence threshold ϵ user-selectable because, in practice, convergence speed is governed far more by the informativeness of the expert trajectories than by their sheer number. Our toy demonstration already reaches 0.73 precision and 0.92 recall with only 250 toxic / 250 non-toxic sentences (Fig. 1), while the larger Pythia-70M/410M experiments converge in at most 60 iterations, each iteration processing a batch of just 50 toxic + 50 non-toxic examples drawn from a 1.5 k pool. Figure 2 further shows that adding redundant non-toxic samples can even degrade performance, whereas a modest increase in contrastive toxic samples sharply improves precision and Kendall τ. These results demonstrate that a relatively small set of well-chosen trajectories that clearly differentiate toxic from non-toxic language is sufficient for reliable reward extraction; beyond that, diminishing returns set in, so practitioners should prioritise collecting diverse, high-signal demonstrations rather than maximising dataset size. We will add these empirical observations, along with a concise table summarising iteration statistics, to the camera-ready version for greater clarity.