TLDR: Code Language Models (CLMs) often inadvertently memorize sensitive data like API keys and passwords from their training. This paper introduces CodeEraser, a novel machine unlearning method that effectively and efficiently removes specific sensitive information from trained CLMs without requiring costly full retraining. CodeEraser selectively targets and erases sensitive code segments while preserving the model’s overall performance and functional correctness, addressing critical privacy concerns and compliance with regulations like GDPR.
Code Language Models (CLMs) have become incredibly powerful tools in software engineering, assisting with tasks like generating code and summarizing programs. However, recent studies have uncovered a significant privacy concern: these models can unintentionally memorize sensitive information from their training data. This means they might reproduce confidential details, such as emails, passwords, API keys, and cryptographic secrets, if prompted in a specific way.
This issue directly clashes with global privacy regulations like the European Union’s General Data Protection Regulation (GDPR) and the California Consumer Privacy Act (CCPA), which enshrine the “Right to Be Forgotten.” These regulations empower individuals to request the deletion of their personal data. Traditional approaches to mitigate this, such as de-duplicating training data or using differential privacy, often require completely retraining the model. For already deployed CLMs, which are massive and computationally expensive, this is a time-consuming and costly process, making it impractical for dynamic user requests.
To tackle this challenge, a new research paper titled Scrub It Out! Erasing Sensitive Memorization in Code Language Models via Machine Unlearning by Zhaoyang Chu, Yao Wan, Zhikun Zhang, Di Wang, Zhou Yang, Hongyu Zhang, Pan Zhou, Xuanhua Shi, Hai Jin, and David Lo, introduces a pioneering investigation into machine unlearning for CLMs. Machine unlearning is a post-hoc modification method that aims to remove specific information from trained models without the need for full retraining.
Introducing CodeEraser: Selective Unlearning for CLMs
The researchers developed an advanced variant called CodeEraser. Unlike previous unlearning methods that might indiscriminately forget entire code instances, CodeEraser is designed to selectively target and erase only the sensitive memorized segments within the code. This approach ensures that the structural integrity and functional correctness of the surrounding, non-sensitive code are preserved.
CodeEraser works by applying a technique called gradient ascent specifically to the sensitive segments, actively diminishing their memorization. Simultaneously, it uses gradient descent on the non-sensitive parts to maintain their integrity. Additionally, it incorporates a Kullback-Leibler (KL) divergence-based constraint to ensure that the model’s overall utility and performance are not significantly degraded after the unlearning process.
Quantifying Memorization and Experimental Setup
Before developing CodeEraser, the team conducted a preliminary study to quantify the extent of sensitive memorization in CLMs. They found that approximately 7% of training samples in models like CodeParrot-small and CodeGen-350M-Mono contained sensitive data that was memorized beyond safe thresholds. Based on this, they curated a dataset of 50,000 highly memorized sensitive samples to serve as targets for unlearning experiments.
The effectiveness and efficiency of CodeEraser were rigorously tested against two widely used gradient ascent-based unlearning approaches: the vanilla method and the constraint-based method. Experiments were conducted on three families of CLMs: CodeParrot, CodeGen-Mono, and Qwen2.5-Coder. The evaluation focused on three key aspects:
- Effectiveness: How well the methods removed targeted sensitive information.
- Efficiency: The computational resources (GPU time and memory) required.
- Model Utility: The impact on the CLM’s general code generation performance using the HumanEval benchmark.
Key Findings and Impact
The results demonstrated that CodeEraser is both effective and efficient. For instance, with the Qwen2.5-Coder-7B model, CodeEraser achieved an impressive 93.89% reduction in memorization of targeted sensitive data, successfully bringing scores below the established memorization thresholds. Crucially, it did so efficiently, completing the unlearning process for a group of 32 samples in approximately 1500 seconds (averaging 46.88 seconds per sample), a significantly lower cost compared to full model retraining.
Furthermore, CodeEraser had only a minor impact on the CLM’s overall utility. For the Qwen2.5-Coder-7B model, it preserved an outstanding 99.99% of the model’s code generation performance. This highlights the success of its selective unlearning mechanism in maintaining valuable code knowledge while erasing sensitive information.
The study also explored how characteristics of the forgotten data, such as the number of samples, their duplication frequency in training, and the type of sensitive segments, influenced unlearning performance. Interestingly, removing API/SSH keys sometimes led to an improvement in model performance, suggesting that these atypical patterns might act as outliers that negatively impact generalization. The learning rate was found to be a critical hyperparameter, while other regularization factors offered more flexibility.
Also Read:
- Enhancing AI in Education: The Role of Machine Unlearning for Responsible and Adaptive Systems
- Targeted Forgetting: A New Method to Safely Unlearn Harmful Knowledge in AI Models
Conclusion
This research marks a significant step forward in addressing privacy vulnerabilities in Code Language Models. By pioneering the use of machine unlearning and introducing CodeEraser, the authors provide a practical and efficient technique to actively mitigate the risks of sensitive data memorization. This work not only enhances data privacy in CLMs but also offers a pathway for these powerful AI tools to comply with evolving privacy regulations, ensuring they can be deployed responsibly and securely.
