Data leakage detection in machine learning code: transfer learning, active learning, or low-shot prompting?

Area of Science:

• Computer Science
• Artificial Intelligence
• Software Engineering

Background:

• Machine learning (ML) code quality is crucial for reliable model performance, with data leakage being a significant issue.
• Data leakage, where test set information influences training, inflates metrics and harms generalization.
• Existing detection methods for code-level data leakage are often manual or lack advanced automation.

Purpose of the Study:

• To explore and propose ML-based approaches for detecting code-level data leakage, particularly with limited annotated datasets.
• To address challenges in identifying quality issues within large and complex ML codebases efficiently.
• To develop automated methods for handling imbalanced code data relevant to leakage detection.

Main Methods:

• Investigated three ML-based approaches: transfer learning, active learning, and low-shot prompting.
• Developed an automated method to manage data imbalance issues specific to code data.
• Evaluated the effectiveness of these approaches in detecting code-level data leakage.

Main Results:

• Active learning demonstrated superior performance, achieving an F-2 score of 0.72.
• Active learning reduced the required number of annotated samples from 1,523 to 698.
• The proposed ML-based methods effectively addressed challenges posed by limited data availability.

Conclusions:

• ML-based approaches, especially active learning, are effective for detecting code-level data leakage in ML code.
• These methods significantly improve efficiency by reducing the need for extensive manual annotation.
• Automated detection of data leakage enhances ML code quality and model reliability.