REMEDI: A Benchmark for Retention and Unlearning Evaluation in Multi-label Clinical Disease Inference

Technical Summary: REMEDI – A Benchmark for Retention and Unlearning Evaluation in Multi-label Clinical Disease Inference

Problem Statement
The paper addresses the challenge of "machine unlearning" (MU) in the clinical domain, specifically regarding the removal of sensitive patient data from trained language models. While regulations like GDPR establish a "right to be forgotten," retraining models from scratch after data removal is computationally prohibitive. Existing MU methods and benchmarks are largely restricted to non-medical domains or rely on synthetically curated datasets that fail to represent the complexities of real-world clinical systems. The authors argue that current benchmarks do not adequately capture the unique constraints of medical data, such as label correlations, longitudinal structures, and strict safety requirements, leaving the effectiveness of unlearning methods in healthcare largely unclear.

Methodology and Benchmark Construction
To address these gaps, the authors introduce REMEDI (Retention and Unlearning Evaluation in Multi-label Clinical Disease Inference), a benchmark tailored for multi-label and multiclass clinical disease inference.

• Data Source: The benchmark utilizes the MIMIC-III clinical database, focusing on discharge summaries and associated diagnostic outcomes. The task involves mapping electronic health records (EHRs) to ICD-9-CM codes. To balance task complexity with clinical relevance, the authors selected the ten most frequent disease categories (e.g., Hypertension, Heart failure, Diabetes mellitus), framing the task as a multi-label classification problem where patients may have multiple concurrent diagnoses.
• Task Setting: The unlearning task is defined as instance-level removal. Given a trained model $f_\theta$, a forget set $D_f$ (sensitive records to be removed), and a retain set $D_r$ (records to be preserved), the goal is to produce an unlearned model $f_{\theta'}$ that behaves as if it were trained exclusively on $D_r$, while satisfying privacy constraints.
• Unlearning Severity Levels: REMEDI defines three distinct scenarios to test robustness:
  1. Distinct Instances (Low Severity): Simulates individual patient deletion requests with no overlap between forget and retain sets. Forget set sizes range from 64 to 512 instances.
  2. Concurrent Instances (Medium Severity): Simulates realistic scenarios where forget and retain sets contain overlapping clinical characteristics (e.g., similar demographics or admission patterns). Half of the forget instances overlap with the retain set.
  3. Large-Scale Instances (High Severity): Simulates bulk data removal requests (e.g., regulatory compliance), with forget sets comprising 3%, 4%, and 5% of the training data.
• Evaluation Metrics: The benchmark evaluates two critical dimensions:
  1. Utility: Measured via Macro-F1 scores on test, retain, and forget sets. The ideal unlearned model should match the performance of a fully retrained model on the retain and test sets while showing degraded performance on the forget set.
  2. Privacy: Assessed via Membership Inference Attacks (MIA). A binary classifier is trained to distinguish between forget set samples and holdout set samples based on model loss. A lower MIA score (closer to 0) indicates stronger privacy guarantees, with 0.5 representing random guessing (perfect unlearning).

Experimental Setup
The authors evaluate four representative unlearning methods across two biomedical language models: BioLinkBERT and BioBERT. The baselines include:

1. Gradient Ascent (GA): Maximizes loss on forget data (with a variant GA+R that incorporates retain-awareness).
2. Adversarial Unlearning (AU): Uses adversarial perturbations to diminish the influence of forget samples (with a weight-regularized variant AU+WR).
3. Bad Teacher (BT): Uses knowledge distillation with a weak teacher for forget data and a strong teacher for retain data.
4. SCRUB: Fine-tunes the model to maximize disagreement between representations of forget and retain samples.

Key Results

• Utility-Privacy Trade-off: The experiments reveal a significant trade-off between utility preservation and unlearning efficacy.
• Method Performance:
  • Gradient Ascent (GA): Caused severe performance degradation, effectively nullifying the model’s prediction capabilities. The retain-aware variant (GA+R) mitigated this but still showed limitations.
  • Adversarial Unlearning (AU): Demonstrated moderate effectiveness but struggled with stability, particularly in large-scale scenarios.
  • Bad Teacher (BT): Achieved moderate utility preservation but suffered significant performance drops in large-scale forgetting scenarios due to residual teacher influence.
  • SCRUB: Achieved the best balance between utility and privacy. It maintained performance close to the retrained model across various forget set sizes and maintained strong privacy metrics.
• Model Sensitivity: BioLinkBERT consistently outperformed BioBERT in both utility preservation and privacy maintenance, attributed to its entity-linked pretraining which captures richer clinical semantics. BioBERT exhibited higher sensitivity to unlearning interventions.
• Scalability: Most unlearning methods optimized for small-scale data removal struggled to maintain utility when faced with large-scale unlearning requests. SCRUB was the only method to consistently maintain performance close to the retrained model across all forget set sizes for BioBERT.

Significance and Contributions
The paper claims three primary contributions:

1. Realistic Medical Testbed: REMEDI enables the evaluation of unlearning methods on real-world patient data (MIMIC-III), providing a more realistic assessment than synthetic benchmarks.
2. Diverse Removal Scenarios: It introduces three unlearning severity levels (distinct, concurrent, and large-scale) that capture diverse real-world data removal needs, including challenging forget-retain overlap scenarios.
3. Complex Multi-label Task: By focusing on multi-label classification with correlated disease labels, REMEDI presents a more demanding testbed than existing single-label benchmarks. This complexity highlights the difficulty of disentangling learned representations without disrupting valid label correlations in the retain set.

The authors conclude that REMEDI reveals failure modes of unlearning methods that perform well on non-medical benchmarks, emphasizing the need for methods that can handle the unique complexities of clinical data, such as label correlations and large-scale deletion requests.