Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Learning from Failures: Understanding LLM Alignment through Failure-Aware Inverse RL

Nyal Patel, Matthieu Bou, Arjun Jagota, Satyapriya Krishna, Sonali Parbhoo

TL;DR

This work addresses the opacity of latent reward signals learned during RLHF by introducing Failure-Aware IRL (FA-IRL), which treats misclassified and near-tie preference pairs as high-information constraints. FA-IRL uses a dual-path reward model with a correction component trained on identified failures and a curriculum that prioritizes easier errors before subtler ones, enabling sharper and more faithful reward recovery. The authors prove theoretical identifiability benefits and demonstrate empirically that FA-IRL yields better toxicity reduction, finer alignment signals, and more effective downstream re-alignment than standard IRL, approaching the performance of ground-truth supervision. The approach provides a scalable auditing mechanism for LLM alignment and has practical impact for safer, more reliable detoxification and preference-guided fine-tuning.

Abstract

Reinforcement Learning from Human Feedback (RLHF) aligns Large Language Models (LLMs) with human preferences, yet the underlying reward signals they internalize remain hidden, posing a critical challenge for interpretability and safety. Existing approaches attempt to extract these latent incentives using Inverse Reinforcement Learning (IRL), but treat all preference pairs equally, often overlooking the most informative signals: those examples the extracted reward model misclassifies or assigns nearly equal scores, which we term \emph{failures}. We introduce a novel \emph{failure-aware} IRL algorithm that focuses on misclassified or difficult examples to recover the latent rewards defining model behaviors. By learning from these failures, our failure-aware IRL extracts reward functions that better reflect the true objectives behind RLHF. We demonstrate that failure-aware IRL outperforms existing IRL baselines across multiple metrics when applied to LLM detoxification, without requiring external classifiers or supervision. Crucially, failure-aware IRL yields rewards that better capture the true incentives learned during RLHF, enabling more effective re-RLHF training than standard IRL. This establishes failure-aware IRL as a robust, scalable method for auditing model alignment and reducing ambiguity in the IRL process.

Learning from Failures: Understanding LLM Alignment through Failure-Aware Inverse RL

TL;DR

This work addresses the opacity of latent reward signals learned during RLHF by introducing Failure-Aware IRL (FA-IRL), which treats misclassified and near-tie preference pairs as high-information constraints. FA-IRL uses a dual-path reward model with a correction component trained on identified failures and a curriculum that prioritizes easier errors before subtler ones, enabling sharper and more faithful reward recovery. The authors prove theoretical identifiability benefits and demonstrate empirically that FA-IRL yields better toxicity reduction, finer alignment signals, and more effective downstream re-alignment than standard IRL, approaching the performance of ground-truth supervision. The approach provides a scalable auditing mechanism for LLM alignment and has practical impact for safer, more reliable detoxification and preference-guided fine-tuning.

Abstract

Reinforcement Learning from Human Feedback (RLHF) aligns Large Language Models (LLMs) with human preferences, yet the underlying reward signals they internalize remain hidden, posing a critical challenge for interpretability and safety. Existing approaches attempt to extract these latent incentives using Inverse Reinforcement Learning (IRL), but treat all preference pairs equally, often overlooking the most informative signals: those examples the extracted reward model misclassifies or assigns nearly equal scores, which we term \emph{failures}. We introduce a novel \emph{failure-aware} IRL algorithm that focuses on misclassified or difficult examples to recover the latent rewards defining model behaviors. By learning from these failures, our failure-aware IRL extracts reward functions that better reflect the true objectives behind RLHF. We demonstrate that failure-aware IRL outperforms existing IRL baselines across multiple metrics when applied to LLM detoxification, without requiring external classifiers or supervision. Crucially, failure-aware IRL yields rewards that better capture the true incentives learned during RLHF, enabling more effective re-RLHF training than standard IRL. This establishes failure-aware IRL as a robust, scalable method for auditing model alignment and reducing ambiguity in the IRL process.

Paper Structure

This paper contains 34 sections, 2 theorems, 25 equations, 6 figures, 3 tables, 1 algorithm.

Table of Contents

  1. Introduction
  2. Related Work
  3. Preliminaries
  4. LLM behaviour as a contextual bandit.
  5. Preference data.
  6. Reward parameterization.
  7. Limitations of standard IRL.
  8. Failure-Aware Inverse Reinforcement Learning (FA-IRL)
  9. Problem Setup.
  10. Reward Model
  11. Dual-path decomposition.
  12. Base IRL Objectives
  13. Max-margin IRL.
  14. Identifying and Learning from Failures with Failure-Aware IRL
  15. Theoretical Guarantees of Failure-Aware IRL
  16. ...and 19 more sections

Key Result

Theorem 1

For any collection of preference constraints $\mathcal{C}$ and any non-empty set of failure constraints $\mathcal{C}_f$, the FA-IRL feasible set satisfies $\mathcal{F}_{\text{FA}} \subsetneq \mathcal{F}$ (see Figure fig:theorem_toy where failures act like support vectors, pruning away spurious solut

Figures (6)

  • Figure 1: The Failure-Aware IRL (FA-IRL) workflow, illustrated with a toy detoxification example. Preference pairs are generated from a base LLM and an aligned Detoxified LLM. FA-IRL analyzes these pairs to extract a reward model, which is then evaluated against a ground-truth reward (using STARC) and can be used to further retrain the LLM via RLHF.
  • Figure 2: 2D illustration of Theorem 1. Standard IRL constraints (blue outline) define a wide cone of feasible rewards (both grey and green dots). Adding a failure constraint $w^\top d_f \ge M$ (red line) prunes spurious solutions, leaving the smaller feasible set $\mathcal{F}_{\text{FA}}$ (green dots). Failures thus act as high-information constraints that reduce reward ambiguity, consistent with Corollary 1.
  • Figure 2: FA-IRL shows a clear advantage on difficult examples where standard IRL fails. This disagreement analysis on the Jigsaw test set ($N=10{,}000$) highlights a significant asymmetry in performance: FA-IRL is uniquely correct on 1001 examples, more than double the 457 cases where only the standard IRL baseline succeeds.
  • Figure 3: Left: Cross-subtype generalization test accuracy when training on a single subtype (rows) and evaluating on each subtype (columns). High off-diagonals (e.g., Identity Attack $\to$ Obscene) indicate shared cues; Threat is most distinct. Middle: Distribution of toxicity subtypes within the toxic subset of train/test (each 3,619 T$\to$NT pairs) labelled by Unitary Toxic-BERT Detoxify. Right: Toxicity reduction during Re-RLHF fine-tuning on SmolLM2-360M over 100 epochs; FA-IRL rewards approaches the ground-truth model and beats standard IRL.
  • Figure 4: Pair–mix sensitivity of reward fidelity. Training STARC (lower is better) vs. fraction of T$\to$NT pairs. Line shows the mean over 5 independent seeds; shaded band is the min–max range. Scarce T$\to$NT pairs degrade all methods and can disadvantage FA-IRL; as learnable pairs increase, STARC drops overall, with FA-IRL improving fastest.
  • ...and 1 more figures

Theorems & Definitions (3)

  • Theorem 1: Failure constraints shrink feasible reward sets
  • Corollary 1: Failure constraints reduce reward non-identifiability
  • proof