Technical Summary: Mitigating Sensitive Information Leakage in LLMs4Code through Machine Unlearning

Problem Statement

Large Language Models for Code (LLMs4Code), such as CodeLlama and AIXCoder, have demonstrated significant proficiency in code generation tasks. However, these models inherit a critical privacy vulnerability known as the "memorization problem," where they inadvertently retain and reproduce sensitive information (e.g., Personally Identifiable Information, API keys, passwords) from their training data. While recent studies have highlighted this risk, there is a lack of comprehensive solutions that can effectively remove such sensitive data without retraining the model from scratch or degrading its code generation capabilities. Traditional data cleaning during pre-processing is often inflexible and unscalable. Consequently, there is a need for a lightweight, post-training mechanism to "forget" specific sensitive data points while preserving the model's utility.

Methodology

This study presents the first comprehensive empirical investigation into applying Machine Unlearning (MU) to mitigate sensitive information leakage in LLMs4Code. The authors constructed a dedicated benchmark and evaluated three representative unlearning algorithms across multiple model architectures.

1. Benchmark Construction

The study utilized a dual-component dataset:

• Forget Set: A synthetic dataset containing 5,000 fictional character resumes with diverse sensitive attributes (e.g., names, addresses, bank balances, passwords, political views). These were transformed into code completion prompts using GPT-4 to align with LLM4Code functionalities. The set was split 9:1 for training (unlearning) and testing.
• Retain Set: 5,000 samples drawn from established code generation datasets (HumanEval, MBXP, LiveBench) to evaluate the preservation of general code generation capabilities.

2. Unlearning Techniques Evaluated

Three state-of-the-art unlearning algorithms were applied to the forget set while monitoring the retain set:

• Gradient Ascent (GA): Maximizes the loss on the forget set to encourage the model to "forget" specific patterns associated with privacy queries.
• GA + Gradient Descent (GA+GD): Simultaneously performs gradient ascent on the forget set and gradient descent on the retain set to maintain general performance.
• GA + KL Divergence (GA+KL): Combines gradient ascent on the forget set with Kullback-Leibler (KL) divergence minimization against a reference model (fine-tuned only on the retain set). This acts as a regularization term to preserve the model's generative distribution on non-sensitive inputs.

3. Models and Metrics

The study evaluated these techniques on three widely used LLMs4Code families: AIXCoder, CodeLlama, and CodeQwen, covering parameter sizes from 1.5B to 34B.

• Privacy Metric: Leak Rate, defined as the proportion of test samples where sensitive information is present in the output (evaluated via human expert review).
• Utility Metric: Pass@1, measuring the functional correctness of the generated code on the retain set.

Key Results

1. Effectiveness of Unlearning

Machine unlearning significantly reduces direct privacy leakage while maintaining code generation performance.

• Leak Rate Reduction: Across all models, unlearning techniques reduced the leak rate by more than 50% on average. For instance, in AIXCoder-7B, the leak rate dropped from 23.2% (Zero-Shot) to 10.4% (GA+KL). The GA+GD method consistently achieved the lowest leak rates, reducing CodeLlama-34B's leak rate from 32.7% to 5.0%.
• Performance Preservation: The utility of the models remained largely intact. Pass@1 scores showed negligible degradation (often within 1-2%), indicating that the models retained their ability to generate functional code. For example, CodeQwen-14B maintained a Pass@1 of 73.3% post-unlearning compared to 74.5% originally.

2. Shift in Leakage Forms

A critical finding is the transformation of leakage behavior post-unlearning:

• Protection Strategies: Models adopted diverse strategies to avoid leaking sensitive values, most commonly replacing sensitive fields with variable names or abbreviations (17.23%), returning variables without values (15.02%), or using symbolic placeholders (11.48%).
• Indirect Leakage: While Direct Privacy Leakage (explicitly requested sensitive data) decreased significantly, Indirect Privacy Leakage (unintended disclosure of related, non-requested sensitive data) increased. For example, in AIXCoder-7B, the ratio of indirect leakage rose from 51% (Zero-Shot) to 62% (GA+KL). This suggests that while models stop outputting exact memorized strings, they may still leak contextually related information (e.g., leaking a user's email when asked for their password).

3. Fine-Grained Analysis

The effectiveness of unlearning varied by data type:

• Structured Data (Emails, Passwords): Gradient-based methods (GA, GA+GD) were highly effective, reducing leakage to near-zero levels (e.g., Email leakage in CodeLlama dropped from 6.01% to 0.77%).
• Context-Dependent Data (Addresses): These showed less consistent reductions, particularly with KL regularization, highlighting the difficulty of unlearning diverse, context-heavy data patterns.
• GA+GD Superiority: The combined Gradient Ascent and Gradient Descent approach consistently outperformed others across both structured and semi-structured data types.

Significance and Contributions

The paper claims the following contributions to the field:

1. Feasibility of MU for LLMs4Code: It demonstrates that machine unlearning is a viable and effective solution for enhancing privacy in code generation models, capable of reducing direct leak rates by over 50% with negligible side effects on code generation quality.
2. Identification of Indirect Leakage: The study uncovers a previously underexplored vulnerability where unlearning shifts leakage from direct to indirect forms. This implies that future unlearning processes must address not just explicit memorization but also contextual associations to ensure robust privacy.
3. Comprehensive Benchmark: The authors provide the first dedicated benchmark for evaluating unlearning in LLMs4Code, including a synthetic forget set and a retain set, along with public access to all code and data.
4. Practical Guidance: The results suggest that GA+GD is the most robust strategy for balancing privacy and utility, offering a practical alternative to full retraining for organizations deploying LLMs4Code.

The authors conclude that while machine unlearning is a promising step toward privacy governance in LLMs4Code, the emergence of indirect leakage indicates that the field requires further research into holistic unlearning strategies that can simultaneously mitigate both direct and indirect privacy risks.