Technical Summary: Replay What Matters: Off-Policy Replay for Efficient LLM Reinforcement Unlearning

Problem Statement

Large Language Model (LLM) unlearning aims to remove hazardous knowledge (e.g., copyrighted content, private information) from pretrained models without the prohibitive cost of full retraining. Recent reinforcement learning (RL)-based approaches, such as RULE, reformulate unlearning as learning a refusal behavior: teaching the model to refuse queries related to a "forget set" while preserving responses on a "retain set" and a "boundary set."

However, the authors identify a critical inefficiency in the Stage 2 on-policy optimization (specifically using GRPO) employed by methods like RULE. During training, a pronounced hard–easy disparity emerges:

1. Easy Cases: Many prompts rapidly converge to high-reward states. Continuing to sample these prompts in later training stages yields little informative gradient signal (advantage $\approx 0$) but consumes significant computational resources.
2. Hard Cases: Prompts near the semantic boundary between forgetting and retaining often remain low-reward throughout training. In standard on-policy settings, the low-reward rollouts generated for these hard cases are discarded after a single use, despite containing valuable information about the model's failure modes.

This dynamic leads to suboptimal training efficiency, where computation is wasted on converged easy cases while hard cases remain under-optimized, potentially degrading the model's ability to preserve general utility (Retain Quality) while forgetting specific targets.

Methodology: ReRULE

To address this, the authors propose ReRULE (Replay-enhanced Reinforcement UnLEarning), an off-policy replay enhancement mechanism integrated into the RULE framework.

Framework Overview

ReRULE retains the two-stage structure of RULE:

• Stage 1 (Rejection Steering): Supervised Fine-Tuning (SFT) on a refusal dataset to warm-start the model's ability to refuse.
• Stage 2 (Refusal Boundary Optimization): RL-based optimization to refine the boundary between forget and retain queries.

The ReRULE Pipeline

The core innovation occurs in Stage 2, where ReRULE introduces a replay buffer and hybrid training:

1. Hard-Case Identification and Storage:
   During early GRPO training, the system monitors the mean reward of rollout groups generated for each prompt. If a group's mean reward falls below a predefined threshold $\tau$, it is classified as a "hard case" and stored in a replay buffer $\mathcal{B}$. This buffer accumulates low-reward trajectories that represent the model's current failure modes.

2. Hybrid Training (Off-Policy Updates):
   Once the buffer contains a sufficient number of hard cases ($|\mathcal{B}| \ge B_{min}$) and a warm-up period ($t \ge T_{warm}$) is completed, ReRULE activates a hybrid training regime. At each step, it performs:

   • Standard On-Policy Update: A standard GRPO update using fresh samples from the current policy $\pi_\theta$.
   • Off-Policy Replay Update: An additional gradient update using samples drawn from the replay buffer $\mathcal{B}$.

3. Importance Sampling:
   To correct for the distribution shift between the policy that generated the stored rollouts ($\pi_{\theta_{old}}$) and the current policy ($\pi_\theta$), ReRULE employs importance sampling. The off-policy gradient is computed as:
   $$\nabla_\theta J_{off}(\theta) = \mathbb{E}_{g \sim \mathcal{B}} \left[ \mathbb{E}_{y \sim \pi_{\theta_{old}}} \left[ \tilde{\rho}(y|x) \nabla_\theta \log \pi_\theta(y|x) \hat{A}(x, y) \right] \right]$$
   where $\tilde{\rho}(y|x)$ is a clipped and normalized importance weight. This allows the model to retrain on difficult cases without requiring new, expensive sampling.

Key Contributions

1. Identification of Training Inefficiency: The paper empirically demonstrates that on-policy GRPO in LLM unlearning suffers from diminishing returns on easy cases and under-utilization of hard cases, leading to a "hard/easy disparity" that hinders convergence.
2. Proposal of ReRULE: The authors introduce an off-policy replay mechanism that stores low-reward hard-case rollouts and reuses them via importance-sampled updates. This redirects computational focus toward boundary cases requiring further optimization.
3. Theoretical Analysis: The authors provide a theoretical proof showing that ReRULE yields a tighter convergence bound for hard cases compared to pure on-policy RULE. By increasing the effective number of hard-region training tokens ($N_{H}^{ReRULE} = N_{H}^{on} + \eta_\rho N_{H}^{off}$), the method reduces the exploration error in the hard region.
4. Empirical Validation: Extensive experiments across three benchmarks (RWKU, MUSE-Books, TOFU) validate the method's effectiveness.

Experimental Results

The authors evaluated ReRULE against baselines including GA, NPO, SimNPO, and the original RULE across three datasets:

• MUSE-Books (Copyrighted Unlearning): This setting, involving Harry Potter content, exhibits a high proportion of hard cases. ReRULE significantly improved Retain Quality from 46.3 (RULE) to 56.2, while maintaining top-tier Forget Quality. The training time increased by only 5–11% (from 6.8h to 7.1h).
• RWKU (Real-World Knowledge): ReRULE achieved the best Forget Quality and maintained top-2 performance on Forget Naturalness and Retain Quality, outperforming RULE in efficiency on hard cases.
• TOFU (Entity Unlearning): In this setting, the forget set is smaller and generally easier (fewer hard cases). ReRULE showed limited improvement over RULE, which the authors cite as evidence that the replay mechanism is most beneficial when a pronounced hard/easy disparity exists.
• Ablation Studies:
  • Replacing the hard-case selection with random storage caused performance to degrade to the level of standard RULE, confirming that targeting hard cases is crucial.
  • Varying the threshold $\tau$ showed that setting it too high (including easier cases) reduces efficiency, while the optimal setting (0.4 for MUSE) balances selectivity and buffer size.

Significance and Claims

The paper claims that ReRULE offers a cost-effective enhancement to RL-based unlearning. By recycling "failed" rollouts from hard cases, it improves the Retain Quality of the unlearned model—a critical metric often compromised in unlearning tasks—without requiring full retraining or significant additional compute.

The authors emphasize that the method is conditional: its benefits are most pronounced in scenarios with a significant disparity between easy and hard prompts (e.g., MUSE-Books). In simpler settings (e.g., TOFU), the gains are marginal, supporting the hypothesis that the replay mechanism specifically targets the inefficiency of under-optimized hard cases. The work positions ReRULE as a practical step toward more efficient and robust machine unlearning, particularly for preserving general model utility while removing specific hazardous knowledge.