SWaRL: Safeguard Code Watermarking via Reinforcement Learning

TL;DR

SWaRL tackles IP protection for code LLMs by embedding detectable watermarks without sacrificing code quality. It introduces a reinforcement-learning-based co-training framework (GRPO) that jointly optimizes watermark detectability and functional correctness, aided by LoRA adapters for lightweight transfer across model updates. A trainable watermark detector provides the reward signal, while periodic detector updates keep alignment with the evolving policy. Experimental results show SWaRL achieving high watermark detectability with minimal impact on code quality and strong resilience to refactoring attacks, enabling practical provenance verification at deployment.

Abstract

We present SWaRL, a robust and fidelity-preserving watermarking framework designed to protect the intellectual property of code LLM owners by embedding unique and verifiable signatures in the generated output. Existing approaches rely on manually crafted transformation rules to preserve watermarked code functionality or manipulate token-generation probabilities at inference time, which are prone to compilation errors. To address these challenges, SWaRL employs a reinforcement learning-based co-training framework that uses compiler feedback for functional correctness and a jointly trained confidential verifier as a reward signal to maintain watermark detectability. Furthermore, SWaRL employs low-rank adaptation (LoRA) during fine-tuning, allowing the learned watermark information to be transferable across model updates. Extensive experiments show that SWaRL achieves higher watermark detection accuracy compared to prior methods while fully maintaining watermarked code functionality. The LoRA-based signature embedding steers the base model to generate and solve code in a watermark-specific manner without significant computational overhead. Moreover, SWaRL exhibits strong resilience against refactoring and adversarial transformation attacks.

Figures (4)

• Figure 1: A code LLM is fine-tune with SWaRL to generate functional and watermarked code. During deployment, user sends a prompt to cloud code LLM API, which would generate code with signatures detectable by a watermark detector. Owners can leverage the watermark detector to check if a suspect code is watermarked or not.
• Figure 2: SWaRL overview. The actor model generates candidate code, which is then evaluated both by a watermark detector (encouraging embedded watermark signals) and by execution tests (ensuring functional correctness). The combined reward drives a Group Relative Policy Optimization (GRPO) update, adjusting the model via lightweight LoRA. Periodically, the detector is retrained on new outputs so it stays aligned with the evolving actor policy.
• Figure 3: Pass@1 comparison across watermarking methods (EXP-edit, WLLM, SWEET, and SWaRL) on HumanEval, MBPP, HumanEval+, and MBPP+. This highlights the substantial accuracy degradation introduced by prior watermarking methods relative to the SFT baseline versus the consistent performance gains achieved by SWaRL across all benchmarks.
• Figure 4: AUROC comparison across watermarking methods (EXP-edit, WLLM, SWEET, and SWaRL) on HumanEval, MBPP, HumanEval+, and MBPP+. The plots highlight the characteristics of each watermarking approach: WLLM and SWEET achieve high AUROC but at the cost of substantial drops in code-generation quality, while EXP-edit maintains moderate detectability. In contrast, SWaRL delivers competitive watermark detection across all benchmarks while preserving or improving code correctness