Technical Summary: SWaRL – Safeguard Code Watermarking via Reinforcement Learning

1. Problem Statement

The proliferation of Code Large Language Models (LLMs) has accelerated software development but introduced significant challenges regarding Intellectual Property (IP) protection and code accountability. High-quality code generated by these models can be redistributed or modified without attribution, obscuring provenance. Furthermore, malicious actors can exploit these models to generate obfuscated or vulnerable code, complicating forensic analysis.

Existing code watermarking approaches face a fundamental trade-off between functional correctness and watermark robustness:

• Inference-based methods (e.g., "green/red" token lists) manipulate token probabilities during generation. While they require no model updates, they often degrade code fidelity, causing syntax errors or functional failures.
• Neural-based methods rely on handcrafted transformation rules (e.g., variable renaming, syntax restructuring) to embed signatures. While they preserve functionality better, they are vulnerable to refactoring attacks that systematically restructure the code, and they lack generalizability to new syntactic patterns.

There is a need for a framework that automatically aligns code LLMs to generate functionally correct code while embedding verifiable watermarks, without relying on manual transformation rules or sacrificing code quality.

2. Methodology: SWaRL

SWaRL is a robust, fidelity-preserving watermarking framework that employs Reinforcement Learning (RL) to co-train a code generation model (Actor) and a watermark detector.

Core Components

1. Actor Model (Code LLM): A base LLM (e.g., Qwen2.5-Coder) fine-tuned to generate code that satisfies two objectives: passing functional unit tests and embedding a detectable watermark signature.
2. Watermark Detector (Reward Model): A trainable neural network (based on CodeBERT) that acts as a reward signal. It distinguishes between watermarked code generated by the Actor and non-watermarked reference code.
3. Low-Rank Adaptation (LoRA): To ensure computational efficiency and modularity, SWaRL injects watermarking behavior via small, trainable low-rank matrices into the Actor's linear layers, rather than fine-tuning all parameters.

Optimization Framework: GRPO

SWaRL utilizes Group Relative Policy Optimization (GRPO), a lightweight RL algorithm, to optimize the Actor. The process involves:

• Sampling: For a given prompt, the Actor generates a group of $G$ candidate code outputs.
• Hybrid Reward Calculation: Each candidate $\tau$ is scored using a composite reward function $R(\tau)$:
  $$R(\tau) = \lambda_{wm} R_{wm}(\tau) + \lambda_{exec} R_{exec}(\tau) - \beta D_{KL}(\pi_\theta || \pi_{ref})$$
  • Watermark Reward ($R_{wm}$): The output of the frozen detector model $D(\tau; \theta_d)$, estimating the likelihood the code is watermarked.
  • Execution Reward ($R_{exec}$): The pass rate of the code against a suite of unit tests, ensuring functional correctness.
  • KL Divergence Penalty: Constrains the policy to stay close to the reference model to prevent distributional drift and ensure stable training.
• Policy Update: The Actor is updated using a clipped policy-gradient objective based on the relative advantage of candidates within the group.
• Co-Training: The watermark detector is periodically unfrozen and fine-tuned on a mix of recent Actor outputs and reference code. This ensures the detector remains aligned with the evolving generation policy, preventing the Actor from "hacking" a static detector.

Deployment

At inference, the LoRA adapter is merged with the base model. The system operates as a standard code-generation API, producing natural, executable code. Authorized parties can use the trained detector to verify the provenance of any code snippet without needing access to the original model or retraining.

3. Key Contributions

The paper claims the following contributions:

1. SWaRL Framework: Introduction of an RL-based watermarking framework that automatically aligns code LLMs to generate functional code with embedded signatures, eliminating the need for manually crafted transformation rules.
2. Algorithmic Integration:
   • Utilization of GRPO to jointly optimize for functional correctness, watermark detectability, and robustness against refactoring attacks.
   • Integration of LoRA to provide a modular, lightweight mechanism for injecting watermarking behavior into base LLMs.
3. Performance Superiority: Demonstration that SWaRL achieves state-of-the-art watermark detectability while maintaining the highest code-generation accuracy among existing baselines. It consistently surpasses non-watermarked Supervised Fine-Tuning (SFT) baselines in pass rates and outperforms prior watermarking methods in robustness.

4. Experimental Results

Experiments were conducted on HumanEval, MBPP, and their extended variants (HumanEval+, MBPP+), comparing SWaRL against baselines: No Watermark (Base+SFT), EXP-edit, WLLM, and SWEET.

Functional Correctness

• SWaRL achieved the highest Pass@1 and Pass@10 scores across all benchmarks among watermarking methods.
• Notably, SWaRL surpassed the non-watermarked SFT baseline in several settings (e.g., HumanEval Pass@1 improved from 55.23% to 58.22%), suggesting the RL objective acts as an additional alignment stage for syntactic validity.
• In contrast, statistical methods like WLLM and SWEET significantly degraded functional correctness (e.g., WLLM dropped HumanEval Pass@1 to 30.82%).

Watermark Detectability

• SWaRL achieved strong AUROC scores (e.g., 0.908 on MBPP), balancing detectability with fidelity better than prior methods.
• While WLLM and SWEET showed high AUROC, they did so at the cost of code functionality. EXP-edit preserved functionality but offered weak detectability (near random guessing on some datasets).

Robustness to Attacks

• Refactoring Attacks: When code was refactored by external LLMs (Qwen2.5-3B and DeepSeek-Coder-7B), SWaRL exhibited the lowest AUROC degradation.
  • Under the strongest attack (DS Attack on MBPP), SWaRL showed only 2.21% degradation, compared to 14–17% for WLLM and SWEET.
• Variable Renaming: SWaRL maintained high AUROC even when 50% of identifiers were renamed, whereas token-level methods (WLLM, SWEET) suffered significant performance drops.

Efficiency

• Generation Latency: SWaRL introduced minimal overhead (0.039s/token), comparable to statistical methods.
• Detection Latency: SWaRL achieved the fastest detection time (0.0002s/token), significantly outperforming EXP-edit (0.588s/token) and statistical methods.

5. Significance and Claims

The paper positions SWaRL as a practical and scalable solution for attribution and ownership verification in AI-assisted programming. Its significance lies in resolving the traditional trade-off between code utility and security:

• Fidelity Preservation: Unlike inference-based methods that break syntax, SWaRL ensures the generated code is functionally correct and passes unit tests.
• Robustness: By learning watermark-consistent generation behaviors through RL rather than relying on static rules, SWaRL resists systematic refactoring and adversarial transformations that typically defeat rule-based or token-level watermarks.
• Efficiency: The use of LoRA and a lightweight detector ensures the framework is computationally feasible for real-world deployment without substantial inference overhead.

The authors conclude that SWaRL enables reliable provenance tracking and IP protection without altering the user workflow or degrading the quality of the generated code.

6. Limitations and Future Work

The paper acknowledges specific limitations and outlines future directions:

• Binary-Level Watermarking: Current implementation targets source code; extending to compiled binary code is a potential future extension.
• Repository-Level Scope: Evaluation is currently at the function level. The authors note that repository-level watermarking remains unexplored, presenting a larger adversarial modification space that requires further research.