Randomized Masked Finetuning: An Efficient Way to Mitigate Memorization of PIIs in LLMs

1. Problem Statement

Large Language Models (LLMs) are known to memorize training data, often regurgitating sensitive Personally Identifiable Information (PII) when prompted. This poses severe privacy risks, particularly when training datasets contain structured PII like email addresses.

• Root Cause: Research indicates that memorization is driven by data duplication; token sequences appearing frequently in the training set are more likely to be memorized.
• Limitations of Existing Solutions:
  • Unlearning: Post-training removal of data is often inefficient and may not fully eliminate memorization.
  • Standard Deduplication: While effective at removing repeated text sequences (e.g., 50+ tokens), it operates at the document level. It fails to specifically target structured PII (like email addresses) embedded within unique contexts, often requiring the removal of entire data points to eliminate a single repeated email, which reduces dataset size and linguistic diversity.

2. Methodology: Randomized Masked Fine-Tuning (RMFT)

The authors propose Randomized Masked Fine-Tuning (RMFT), a novel technique designed to reduce PII memorization while preserving the dataset's utility and size.

• Core Mechanism:

  1. Identification: The system identifies all occurrences of a specific PII (e.g., an email address) across the dataset.
  2. Preservation: The first occurrence of the PII is left untouched to maintain the original data structure for downstream tasks.
  3. Masking: All subsequent duplicate occurrences are replaced with a randomized, structurally similar PII.
     • Implementation: The system maintains pools of first names, last names, and domains. For a duplicate email (e.g., kay.mann@enron.com), it randomly selects an "anchor" component (e.g., the last name mann) and replaces the other components with sampled values (e.g., bob.mann@yahoo.org).
     • Goal: This breaks the exact token sequence repetition that triggers memorization while maintaining the syntactic format (e.g., FIRST.LAST@DOMAIN.TLD) to preserve linguistic coherence.

• Baseline Comparison (Adapted Deduplication):
  The authors adapted standard deduplication for PII contexts. Instead of removing entire documents, they remove only the headers containing repeated email addresses while retaining the main email text. This serves as the primary baseline for comparison.

• Evaluation Framework (MaxTER):
  To assess the trade-off between privacy and utility, the authors introduce MaxTER, a Pareto-optimal evaluation framework:

  • Privacy Metrics:
    • Total Extraction Rate (TER): The percentage of original emails leaked across the entire dataset.
    • Seen Extraction Rate (SER): The percentage of leaked emails relative to the specific data seen up to a training checkpoint.
  • Utility Metric:
    • Mean Delta Perplexity (MDP): The average increase in perplexity compared to the baseline, measuring performance degradation.
  • Efficiency Metric:
    • Area Under the Response Curve (AURC): Quantifies the overall privacy-performance efficiency by integrating the trade-off curve.

3. Experimental Setup

• Dataset: Enron Email Dataset (10,000 training, 1,000 validation, 100 test samples).
• Models: GPT-2-XL (primary) and GPT-Neo-1.3B (for generalizability).
• Training: Fine-tuned for 3 epochs with checkpoints saved every 10% of training steps.
• Evaluation:
  • Memorization: Tested using 16,000 prompts from Common Crawl to extract emails via regex.
  • Performance: Measured using perplexity on the Enron test split and WikiText2.

4. Key Results

The study compares RMFT, Adapted Deduplication, and Baseline Fine-Tuning.

• Privacy Superiority: RMFT outperforms deduplication significantly, achieving an 80.81% reduction in TER compared to the baseline, versus 69.38% for deduplication.
• Utility Preservation: RMFT incurs a minimal performance penalty (+5.73% perplexity), whereas deduplication causes a much larger degradation (+23.85%).
• Pareto Efficiency (MaxTER):
  • RMFT demonstrates a feasible trade-off across all perplexity thresholds.
  • Deduplication has an "infeasible region" where no checkpoint achieves a perplexity increase below 20%.
  • AURC: RMFT achieves an AURC of 92.52, vastly outperforming deduplication's 17.09, indicating a superior privacy-utility balance.
• Generalizability: Similar results were observed on GPT-Neo-1.3B, confirming the method's architecture-agnostic nature.

5. Significance and Contributions

1. Novel Technique: Introduces RMFT, a targeted approach that reduces PII repetition without discarding valuable training data or linguistic context.
2. Evaluation Framework: Proposes MaxTER and AURC as robust metrics for jointly evaluating privacy and utility, moving beyond isolated assessments.
3. Practical Impact: Demonstrates that it is possible to significantly mitigate PII leakage (reducing extraction rates by ~80%) with negligible impact on model performance (<6% perplexity increase).
4. Scalability: While deduplication is faster in preprocessing, RMFT offers a much better trade-off for scenarios where data utility and domain-specific performance are critical (in-distribution deployment).

6. Limitations and Future Work

• Scope: Currently restricted to fine-tuning; the impact on large-scale pretraining is unexplored.
• PII Specificity: The method relies on the structured nature of PII (like emails). Extending it to arbitrary token sequences where semantic coherence is harder to maintain is challenging.
• Out-of-Distribution (OOD): RMFT excels in in-distribution scenarios (matching training and deployment data). Deduplication showed slightly better generalization on OOD data (WikiText2), suggesting RMFT is best suited for domain-specific privacy preservation.
• Model Diversity: Future work will validate the approach across a wider range of model families (e.g., Llama, Qwen).

Conclusion:
The paper establishes Randomized Masked Fine-Tuning as a highly effective strategy for privacy-preserving LLM training. By intelligently masking duplicate PII occurrences rather than deleting data, RMFT achieves superior privacy protection with minimal performance loss, offering a practical solution for deploying LLMs in sensitive domains like healthcare and finance.