💻 computer science

PuzzleMark: Implicit Jigsaw Learning for Robust Code Dataset Watermarking in Neural Code Completion Models

PuzzleMark is a robust and imperceptible watermarking framework for neural code completion models that leverages code complexity-based carrier selection and a novel variable name concatenation pattern to protect dataset intellectual property while achieving perfect verification accuracy and evading existing detection methods.

Original authors: Haocheng Huang, Yuchen Chen, Weisong Sun, Peizhuo Lv, Yuan Xiao, Chunrong Fang, Yang Liu, Xiaofang Zhang

Published 2026-05-01

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Haocheng Huang, Yuchen Chen, Weisong Sun, Peizhuo Lv, Yuan Xiao, Chunrong Fang, Yang Liu, Xiaofang Zhang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: Stealing the "Secret Sauce"

Imagine you are a chef who has spent years collecting the world's best recipes, testing them, and organizing them into a massive, perfect cookbook. This cookbook is your most valuable asset.

Now, imagine someone wants to open a restaurant using your cookbook without paying you. They can't just photocopy the book (that's too obvious), so they hire a robot to read your book and learn how to cook. Once the robot learns, they fire you and open their own restaurant. You have no way to prove the robot learned from your book, because the robot's cooking looks just like normal cooking.

In the world of computer code, this "cookbook" is a code dataset, and the "robot" is an AI coding assistant (like GitHub Copilot). Creating these datasets is expensive and hard work. But right now, there is no good way to prove an AI was trained on your specific dataset if someone steals it.

The Old Solution: The "Sticky Note" (And Why It Failed)

Previous researchers tried to solve this by hiding "watermarks" in the code. Think of this like slipping a sticky note into every 10th page of the cookbook that says, "This book belongs to Chef X."

The Flaw: These old sticky notes were too obvious. If a thief looked at the book, they could easily spot the weird notes, rip them out, and throw away the pages. Or, they could just scan the book, find the weird patterns, and remove them before training their robot. The "sticky notes" were too easy to detect and remove.

The New Solution: PuzzleMark

The authors of this paper propose a new method called PuzzleMark. Instead of a sticky note, they turn the code itself into a subtle, invisible puzzle.

Here is how it works, broken down into three simple steps:

1. Picking the Right "Hiding Spots" (Carrier Selection)

Imagine you are trying to hide a secret message in a library.

The Mistake: If you hide the message in a book that is already weird, torn, or written in a strange language, people will immediately suspect something is up.
PuzzleMark's Fix: The system scans the entire library and uses a "complexity meter" to find books that are perfectly normal, well-written, and common. It refuses to hide messages in books that are already messy or suspicious. By choosing only the most "natural" looking code snippets to carry the secret, the watermark becomes much harder to spot.

2. The "Jigsaw Puzzle" Trick (Concatenation Pattern)

Old watermarks relied on a simple pattern: "If you see the word key, the next word must be value." This is like a code saying, "If you see a red car, there must be a blue car next to it." Thieves can easily spot this pattern and break it.

PuzzleMark uses a "Jigsaw Puzzle" approach:

Instead of a fixed rule, it takes two existing variable names in the code (like key and iterator) and mashes them together to create a new, slightly unusual name (like key_iterator).
It's like taking two normal puzzle pieces and gluing them together to make a new piece that looks like it belongs, but only if you know the secret glue.
Because the code looks mostly normal, and the "glue" changes depending on the specific code snippet, it is incredibly hard for a thief to find and remove without breaking the code itself.

3. The "Black Box" Test (Verification)

How do you prove the AI learned from your book?

You don't need to see inside the AI's brain. You just give it a specific "trigger" (a specific code snippet with the hidden puzzle pieces).
If the AI was trained on your watermarked dataset, it will instinctively complete the code with the "glued" name you created.
If it wasn't trained on your data, it will just guess a normal name.
The authors use a statistical test (Fisher's exact test) to see if the AI is guessing correctly too often to be a coincidence. If the math says "yes," you have proof.

Why Is This Better? (The Results)

The paper tested PuzzleMark against the old methods and found:

It's Invisible: Human developers and automated tools couldn't tell the difference between watermarked code and normal code. It was as invisible as a ghost.
It's Tough: Even if thieves tried to "clean" the dataset by removing suspicious-looking code, the watermark survived. Even if they tried to "dilute" it by mixing in millions of other clean code files, the watermark was still detectable.
It Doesn't Break Anything: The AI models trained on this watermarked code performed just as well as models trained on clean code. The "glue" didn't ruin the recipe.

The Bottom Line

PuzzleMark is a new, robust way to protect the intellectual property of code datasets. Instead of leaving obvious "sticky notes" that thieves can easily rip out, it hides the proof of ownership inside the natural flow of the code itself, like a secret recipe ingredient that only reveals itself when you try to cook the dish. It ensures that if someone steals your data to train an AI, you can prove it in court, even if they try to hide their tracks.

1. Problem Statement

High-quality code datasets are valuable intellectual property (IP) requiring significant resources to curate. However, they face severe risks of unauthorized use by attackers training Neural Code Completion Models (NCCMs). While digital watermarking offers a post-hoc mechanism for ownership verification, existing methods for code datasets suffer from two critical flaws:

Fragile Carrier Selection: Previous works (e.g., CoProtector, CodeMark) often select watermark carriers (code snippets) randomly or based on rigid constraints (e.g., Semantics-Preserving Transformations). This ignores the inherent "naturalness" of the code. Complex or unnatural code snippets selected as carriers are easily flagged by automated detection tools (like DeCoMa or KillBadCode) even before watermarking, leading to high exposure risks.
Vulnerable Watermark Patterns: Existing methods rely on co-occurrence patterns (fixed trigger-target pairs). Recent research (DeCoMa) has demonstrated that these patterns create high-frequency outliers that are easily detected and removed, resulting in near 100% detection rates and negligible robustness against removal attacks.

2. Methodology: PuzzleMark

PuzzleMark introduces a robust watermarking framework consisting of three main phases: Carrier Selection, Watermark Embedding, and Model Validation.

A. Watermark Carrier Selection (Avoiding Exposure)

Instead of random selection, PuzzleMark evaluates the suitability of code snippets to serve as carriers based on their likelihood of being detected as "suspicious."

Feature Extraction: It extracts 10 complexity features (5 structural: e.g., cyclomatic complexity, NLOC; 5 semantic: e.g., unique function calls, expression diversity) from the Abstract Syntax Tree (AST).
Feature Projection: Using a mechanism inspired by Principal Component Analysis (PCA), it projects these features onto a principal component direction to generate a suitability score.
Filtering: Code snippets with low suitability scores (high complexity/naturalness issues) are filtered out. Only snippets with high suitability scores (low risk of detection) are selected as candidates for embedding.

B. Watermark Embedding (The "Jigsaw" Pattern)

PuzzleMark replaces the traditional co-occurrence pattern with a novel concatenation pattern, functioning like a "jigsaw puzzle" where pieces fit together dynamically.

Mechanism: It embeds watermarks by concatenating two variable names within the code (e.g., prefix + suffix $\rightarrow$ prefix_suffix).
Two Strategies:
1. Fixed-Trigger Strategy: Uses a user-defined variable name (e.g., key) as a fixed prefix. The suffix is adaptively selected from the code, and the target is the concatenated result.
2. Universal Watermark Strategy: To prevent attackers from identifying a fixed trigger, the prefix is also adaptive (e.g., the first variable in the code). This makes the trigger dynamic and harder to detect via regex or pattern matching.
Adaptive Perturbation: The method prioritizes embedding in code where the trigger prefix already exists naturally to minimize code perturbation. If the prefix is absent, it introduces it, but only when necessary to meet embedding rate thresholds.

C. Suspicious Model Validation

To verify if a model was trained on a watermarked dataset (Black-Box setting):

Input Construction: Two sets of inputs are generated: one with the trigger suffix replaced by a semantically meaningless token (e.g., unknown_token) and one without.
Statistical Testing: The model's outputs are analyzed to see if it predicts the specific concatenated target variable. Fisher's Exact Test is used (instead of t-tests) to determine the statistical significance of the association between the trigger and the target output, providing a $p$ -value for verification.

3. Key Contributions

Suitability Evaluation: First to introduce a quantitative mechanism for evaluating and selecting watermark carriers based on code complexity to minimize exposure risk.
Concatenation Pattern: First to propose and apply a concatenation pattern (replacing co-occurrence) for code dataset watermarking, significantly enhancing robustness against detection.
Dual Strategies: Implementation of both Fixed-Trigger and Universal Watermark strategies via variable name concatenation, balancing design simplicity with high imperceptibility.
Comprehensive Evaluation: Extensive experiments proving the method's effectiveness, harmlessness, imperceptibility, and robustness.

4. Experimental Results

The authors evaluated PuzzleMark on Python and Java datasets (CodeSearchNet) using three NCCMs (Qwen2.5, StarCoderBase, DeepSeek-Coder).

Effectiveness: Achieved a 100% verification success rate and 0% false positive rate across all models and languages.
Harmlessness: Negligible impact on model performance. BLEU scores showed no statistically significant difference between watermarked and bare models ( $p > 0.05$ , Cohen's $d < 0.2$ ).
Imperceptibility:
- Human: Average suspicious rate $\le 0.24$ (comparable to CodeMark).
- Machine: Significantly outperformed baselines against state-of-the-art detectors (SS, AC, KillBadCode, DeCoMa). Average recall by DeCoMa dropped to $\le 30.41\%$ (compared to ~100% for baselines).
Robustness:
- Removal Attacks: Maintained verifiability ( $p < 0.05$ ) even after removing samples detected by four different automated tools.
- Dilution Attacks: Remained effective even when the watermark embedding rate was diluted to 0.1% (though universal watermarks required slightly higher rates).
Scalability: Successfully protected project-level IP, though performance on very small projects (<100 functions) was limited due to insufficient signal density.

5. Significance

PuzzleMark addresses the critical vulnerability of current code dataset watermarking techniques. By shifting from static co-occurrence patterns to dynamic concatenation and rigorously selecting "stealthy" carriers, it creates a watermark that is:

Hard to Detect: Evades both human inspection and advanced automated detection tools.
Hard to Remove: Resists removal attacks because the watermark is distributed adaptively and does not rely on high-frequency outliers.
Practical: Offers a viable solution for protecting the intellectual property of high-value code datasets used to train commercial AI coding assistants.

The paper concludes that PuzzleMark sets a new standard for robust, imperceptible, and effective dataset watermarking in the era of large language models.