Robust and Secure Code Watermarking for Large Language Models via ML/Crypto Codesign

TL;DR

RoSeMary tackles IP protection and misuse concerns for LLM-generated code by introducing a ML/Crypto codesign watermarking framework. It leverages CodeT5 as the insertion backbone and end-to-end training to preserve functionality while enhancing watermark detectability and robustness against adversarial Code transformations. A Halo2-based non-interactive zero-knowledge proof system enables secure, public verification without disclosing the watermark itself, achieving practical verification times and compact proofs. Experimental results across code benchmarks show high AUROC around 0.97–0.99, strong fidelity, and resilience to attacks, with modest insertion overhead and efficient ZKP verification, highlighting RoSeMary’s potential for scalable IP protection in AI-assisted software development.

Abstract

This paper introduces RoSeMary, the first-of-its-kind ML/Crypto codesign watermarking framework that regulates LLM-generated code to avoid intellectual property rights violations and inappropriate misuse in software development. High-quality watermarks adhering to the detectability-fidelity-robustness tri-objective are limited due to codes' low-entropy nature. Watermark verification, however, often needs to reveal the signature and requires re-encoding new ones for code reuse, which potentially compromising the system's usability. To overcome these challenges, RoSeMary obtains high-quality watermarks by training the watermark insertion and extraction modules end-to-end to ensure (i) unaltered watermarked code functionality and (ii) enhanced detectability and robustness leveraging pre-trained CodeT5 as the insertion backbone to enlarge the code syntactic and variable rename transformation search space. In the deployment, RoSeMary uses zero-knowledge proofs for secure verification without revealing the underlying signatures. Extensive evaluations demonstrated RoSeMary achieves high detection accuracy while preserving the code functionality. RoSeMary is also robust against attacks and provides efficient secure watermark verification.

Figures (6)

• Figure 1: Overview of watermark insertion and extraction. The Code LLM owner watermarks the code before distributing the snippets to end users. The third-party arbitrator leverages zero-knowledge proofs to verify the ownership without requiring the owner to reveal the encoded watermark.
• Figure 2: RoSeMary watermarking procedure. The watermark insertion takes the original code and watermark message as input and fuses their features by CodeT5's encoder $\mathbf{S}_e$. Two sets of decoders $\mathbf{S}_{d1}$ and $\mathbf{S}_{d2}$ predicts the probability over the available syntactic transformations and the renamed variable over the vocabulary. Then, the watermark extraction module decodes watermarks from the syntactic-transformed and variable-renamed watermarked code $S(T, M)$, as well as its malicious transformation $\hat{S}(T, M)$. The two parts are trained jointly to ensure (i) functionality-invariant by minimizing functionality loss $L_f$ and (ii) accuracy and robust message decoding by minimizing detectability loss $L_d$ and robustness loss $L_r$.
• Figure 3: Robustness evaluation results under Variable-rename Attack (VA) and Refactor Attack (RA).
• Figure 4: Watermarked example randomly selected from HumanEval chen2021codex. The upper code shows the original code and the lower code shows the watermarked code, where all watermarks are successfully extracted.
• Figure 5: Watermarked example from HumanEval chen2021codex and MBPP austin2021program. The left code shows the original code and the right code shows the watermarked code, where all watermarks are successfully extracted.