Technical Summary: Beyond Binary Rewards: A Comparative Study of Reward Design for Reinforcement Unlearning

Problem Statement

Machine Unlearning (MU) aims to selectively remove specific knowledge (e.g., personal data, copyrighted material) from trained Large Language Models (LLMs) without the prohibitive cost of full retraining. This capability is increasingly mandated by regulations such as the GDPR and the EU AI Act.

Recent approaches have reformulated unlearning as a Reinforcement Learning with Verifiable Rewards (RLVR) problem, specifically utilizing Group Relative Policy Optimization (GRPO). In this framework, models are optimized against rewards computed directly from their outputs to verify the absence of forbidden concepts. However, existing methods, such as PURGE, rely on sparse binary rewards. These rewards provide a signal of 1 only if a completion contains no forbidden items and 0 otherwise. This sparsity creates a noisy training signal, offering the model no gradient information regarding how close a partial response is to successful unlearning. This limitation acts as a bottleneck for scaling Reinforcement Unlearning (RUL) to more complex tasks, resulting in slow convergence.

Methodology

The authors propose a principled framework to decouple verifiability from sparsity. They argue that a reward can be fully verifiable (computable without access to original training data) while still providing dense, informative signals.

1. Reward Decomposition Framework

The paper introduces a generic reward structure defined as $\phi(y; F, \psi) = \Gamma(I(y; F, \psi); H)$, where:

• Information Component ($I$): Extracts data from the forbidden-match operator (counting occurrences of forbidden items) and auxiliary side information (e.g., semantic graphs).
• Transformation Function ($\Gamma$): Maps the extracted information to a reward range $[0, 1]$.
• Hypothesis: By varying the information component and the transformation, one can create denser learning signals without sacrificing verifiability.

2. Proposed Reward Functions

The authors instantiate this framework with two new reward functions, positioned higher in an "Information Hierarchy" than the standard binary reward:

• Exponential Reward:

  • Mechanism: Uses the total count of forbidden-concept occurrences ($m$) as the information component.
  • Formula: $\phi_{exp}(y) = \exp(-m/\tau)$, where $\tau$ is a decaying constant.
  • Effect: Provides graded penalties. Instead of a binary 0/1, the reward decreases smoothly as the number of violations increases, offering a denser signal than the binary baseline.

• PageRank Inspired Reward:

  • Mechanism: Exploits the semantic relationships among forbidden concepts. It constructs a semantic graph where nodes are forget items and edges represent semantic similarity (via cosine similarity of embeddings).
  • Formula: Applies a personalized PageRank algorithm biased toward the primary target entity to compute importance weights ($w_j$) for each item. The reward penalizes violations based on these weights.
  • Effect: Assigns higher penalties to violations of semantically central concepts (e.g., the target entity itself) compared to peripheral associates. This creates a signal that guides optimization toward the most informative violations first.

3. Theoretical Analysis

The paper provides a theoretical analysis of convergence speed under GRPO. It posits that denser rewards maintain more stable advantage estimates (lower variance in the advantage term $\hat{A}$) compared to binary rewards.

• Binary Rewards: Suffer from "advantage collapse" as the success probability approaches 0 or 1, leading to vanishing gradients.
• Denser Rewards: Maintain non-zero variance in the reward distribution throughout training, enabling larger effective step sizes and faster convergence.

Experimental Setup

• Benchmark: Real World Knowledge Unlearning (RWKU), which evaluates unlearning on 100 real-world famous people.
• Model: Phi-3-Mini-4K-Instruct (3.8B parameters).
• Baselines:
  • Base: The original model before unlearning.
  • Binary (PURGE): The standard sparse binary reward.
• Metrics:
  • Forget: ROUGE-L recall on Fill-in-the-Blank (FB), Question-Answering (QA), and Adversarial Attacks (AA). Lower is better.
  • Neighbor: Retention of knowledge adjacent to the target. Higher is better.
  • MIA: Membership Inference Attack scores (privacy).
  • Utility: General capabilities (MMLU, Reasoning, Truthfulness, etc.). Higher is better.

Key Results

Experiments were conducted over 1,500 training steps with 20 unlearning targets.

1. Convergence Speed: Both proposed rewards significantly outperform the binary baseline in terms of training efficiency.

   • The PageRank Softmax variant achieves the same level of forgetting performance as the binary baseline in approximately 3× fewer training steps.
   • At step 500, PageRank Softmax reaches a 37.1% improvement over the base, a threshold the binary baseline does not reach until step 1,500.

2. Unlearning Performance:

   • PageRank Softmax achieves the strongest forgetting scores across all metrics (FB, QA, AA), consistently outperforming both the Binary and Exponential Decay baselines.
   • Exponential Decay also improves upon the binary baseline but is generally less effective than PageRank.

3. Utility and Neighbor Preservation:

   • Crucially, the improved efficiency does not come at the cost of general model utility or neighbor knowledge retention. Utility metrics (GA, RA, TRU, FAC, FLU) remain stable across all reward designs, shifting by at most 0.014 relative to the Base model.
   • Neighbor scores show uniform degradation across all methods, indicating that the trade-off is driven by the optimization objective itself rather than the specific reward design.

4. Ablation Studies:

   • Exponential Decay: The decay constant $\tau$ is critical. $\tau=0.5$ was found to be optimal, balancing gradient informativeness and optimization pressure. Too small ($\tau=0.1$) reverts to near-binary sparsity; too large ($\tau \ge 1.0$) becomes too lenient.
   • PageRank Variants: The Softmax variant (which applies temperature-scaled softmax to raw PageRank scores) outperforms the standard PageRank and the Linear variant. Softmax effectively flattens the power-law distribution of weights, spreading penalty mass more evenly while preserving semantic ordering, whereas Linear rescaling introduces distortion.

Significance and Claims

The paper claims that reward design is a key driver of unlearning efficiency within the Reinforcement Unlearning framework.

• Theoretical Contribution: It formalizes the trade-off between reward inductive bias and optimization speed, demonstrating that verifiability does not necessitate sparsity.
• Practical Impact: By introducing the Exponential and PageRank rewards, the authors provide a practical path toward scalable and efficient machine unlearning. The results show that carefully designed rewards can accelerate the unlearning process by up to 3× while maintaining strong forgetting performance and preserving general model utility.
• Future Direction: The work opens directions for adaptive rewards that remain verifiable without access to original training data, suggesting that the "information hierarchy" of rewards can be further explored.

The authors conclude that moving beyond binary rewards is not merely an implementation tweak but a fundamental improvement in the learning signal structure, enabling RUL to scale to harder unlearning tasks.