ML Security PapersPuzzleMark: Implicit Jigsaw Learning for Robust Code Dataset Watermarking in Neural Code Completion Models
Haocheng Huang 1, Yuchen Chen 2, Weisong Sun 3, Peizhuo Lv 3, Yuan Xiao 2, Chunrong Fang 2, Yang Liu 3, Xiaofang Zhang 1
Published on arXiv
2604.27677
Output Integrity Attack
OWASP ML Top 10 — ML09
Key Finding
Achieves 100% verification success rate and 0% false positive rate with average suspicious rate ≤0.24 and detection recall ≤30.41% against SOTA watermark detectors
PuzzleMark
Novel technique introduced
Constructing and curating high-quality code datasets requires significant resources, making them valuable intellectual property. Unfortunately, these datasets currently face severe risks of unauthorized use. Although digital watermarking offers a post hoc mechanism for copyright authentication, existing methods are predominantly based on the co-occurrence pattern, which is not robust and is susceptible to watermark detection and removal attacks. In this paper, we propose PuzzleMark, a robust watermarking method for code datasets. To reduce the risk of watermark exposure, PuzzleMark introduces a carrier selection strategy that leverages code complexity to evaluate the suitability of code snippets as watermark carriers, and selects those with high suitability for watermarking. To enhance the robustness of the watermark, PuzzleMark proposes a novel concatenation pattern to replace the traditional co-occurrence pattern, and implements two watermarking strategies through variable name concatenation. PuzzleMark adaptively embeds watermarks based on the inherent characteristics of the code, making it more stealthy while maintaining design simplicity. For watermark verification, PuzzleMark employs Fisher's exact test to verify suspicious models under a black-box setting. Experimental results demonstrate that PuzzleMark achieves a 100% verification success rate and a 0% false positive rate, with negligible impact on model performance. Both our human study and our evaluation using four state-of-the-art watermark detection methods show that PuzzleMark exhibits strong imperceptibility, with an average suspicious rate $\leq$ 0.24 and an average recall $\leq$ 30.41%, respectively. As a practical digital watermarking method, PuzzleMark provides strong protection for the intellectual property of code datasets and offers new insights for future research.
Key Contributions
- Carrier selection strategy using code complexity to choose suitable watermark carriers
- Novel concatenation pattern for watermarking via variable name concatenation, replacing traditional co-occurrence patterns
- Black-box watermark verification using Fisher's exact test to detect models trained on watermarked datasets
🛡️ Threat Analysis
Watermarks TRAINING DATA (code datasets) to prove dataset ownership and detect unauthorized use — this is content/data provenance protection, not model IP protection. The watermark is embedded in the dataset itself, and verified by checking if a model was trained on the watermarked data.