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

Technical Summary: Robust Zero Knowledge Verifiable Watermarking of Code LLMs with ML/Crypto Co-Design

Problem Statement
The proliferation of AI-generated code by Large Language Models (LLMs) raises significant intellectual property (IP) concerns, including plagiarism, vulnerability injection, and unauthorized commercial use. Watermarking is proposed as a solution to identify LLM-generated code. However, existing methods face a critical tri-objective conflict: detectability (ability to find the watermark), fidelity (preserving code functionality), and robustness (resisting removal).

A specific bottleneck in code watermarking is the low-entropy nature of code. Unlike natural language, code has a limited space for high-quality signatures. Current verification methods require LLM owners to disclose their secret signatures to third-party arbitrators for legal validation. This disclosure creates a security paradox: once a signature is revealed, it must be re-encoded for future artifacts to prevent forgery. In low-entropy code, re-encoding often degrades code functionality or reduces watermark resilience, limiting the reusability of the watermarking scheme. Furthermore, existing neural-based methods often rely on shallow networks trained from scratch, limiting feature extraction capabilities, while inference-based methods manipulate tokens in ways that can break syntax.

Methodology: RoSeMary Framework
The paper introduces RoSeMary, a framework utilizing ML/Crypto co-design to address these challenges. It consists of two primary components: a robust ML-based watermarking architecture and a customized Zero-Knowledge Proof (ZKP) verification protocol.

1. ML-Based Watermarking Architecture:

   • Backbone: RoSeMary utilizes CodeT5, a pre-trained sequence-to-sequence model on millions of code snippets, as the backbone for feature extraction. This leverages pre-trained knowledge to better understand code semantics compared to models trained from scratch.
   • Insertion: The encoder fuses the original code features with a watermark message. Two decoders ($S_{d1}$ and $S_{d2}$) predict probabilities for syntactic transformations (e.g., loop types, operator substitution) and variable renaming. The watermarked code is generated by applying the most probable transformations.
   • Adversarial Training: To ensure robustness, the system generates adversarial variants during training by perturbing the transformation probability distributions. The extraction module is trained to recover the message from both the standard watermarked code and these adversarially transformed versions.
   • Training Objectives: The system is trained end-to-end to minimize three losses:
     • Functionality Loss ($L_f$): Ensures the watermarked code compiles and executes identically to the original (using a feature approximation network $R_f$).
     • Detectability Loss ($L_d$): Minimizes the error between the encoded and extracted messages.
     • Robustness Loss ($L_r$): Minimizes the error between the encoded message and the message extracted from adversarial examples.

2. Zero-Knowledge Verification (ZKP):

   • Goal: Enable legal verification of ownership without the LLM owner revealing the secret signature to the arbitrator, thereby preserving the signature for reuse.
   • Mechanism: RoSeMary employs non-interactive Zero-Knowledge Proofs (NIZK) based on the Halo2 zk-SNARK system.
   • Process:
     • The LLM owner (Prover) generates a proof that a specific watermarked code snippet contains a valid signature matching their private key, without revealing the signature itself.
     • The arbitrator (Verifier) checks the proof using a public verification key.
   • Optimizations: To make ZKP feasible for neural network inference, the authors implement four custom optimizations:
     1. Polynomial Approximation: Replacing ReLU activation with a polynomial ($x^2 + x$) to avoid expensive range-check lookups.
     2. Quantization: Using Bfloat16 for parameters to reduce memory overhead.
     3. Custom zkBER Circuit: A dedicated circuit to compute Bit Error Rate (BER) within the zero-knowledge proof, binding the extracted message to a binary signature.
     4. Composability: Fusing the feed-forward and BER circuits into a single computational graph to reduce proof size and verification time.

Key Contributions

• RoSeMary Framework: A novel code LLM watermarking system that jointly optimizes detectability, fidelity, and robustness using a pre-trained CodeT5 backbone.
• Confidential Verification: The first application of customized ZKPs to code watermarking, allowing ownership verification without disclosing the signature, thus solving the reusability issue in low-entropy code.
• Adversarial Training: An end-to-end training strategy that explicitly models and defends against variable renaming, refactoring, and re-watermarking attacks.
• System-Level Co-Design: Integration of ML and Crypto components with specific optimizations (quantization, polynomial approximations) to achieve efficient proof generation and verification.

Experimental Results
Evaluations were conducted on HumanEval, MBPP (Python, C++, Java), and EvalPlus benchmarks:

• Detectability & Fidelity: RoSeMary achieved an AUROC of 0.97 across benchmarks while maintaining a Pass Rate of ~95-97%, significantly outperforming inference-based baselines (KGW, SWEET) in fidelity and neural baselines (SrcMarker) in detectability.
• Robustness: Under variable renaming (up to 50% of variables) and LLM-based refactoring attacks, RoSeMary maintained an AUROC above 0.73, outperforming SrcMarker and SWEET. It also showed resilience against re-watermarking attacks.
• ZKP Efficiency:
  • Proof Generation: Takes approximately 7.15 seconds (one-time cost by the owner).
  • Verification: Takes approximately 118.6 ms by the arbitrator.
  • Communication: Proof size is 17.12 KB, and the verification key is 577.3 KB.
  • Scalability: Verification time remains statistically invariant across different programming languages (Python, Java, C++) and watermark lengths (4 to 16 bits), though it increases with decoder width.

Significance and Claims
The paper claims that RoSeMary provides a balanced solution for practical code watermarking systems. By co-designing the ML and Crypto components, it achieves:

1. High-Quality Watermarks: Preserving code functionality while ensuring high detection rates and robustness against common attacks.
2. Signature Reusability: Eliminating the need to re-encode watermarks after legal verification, which is critical for low-entropy code where the signature space is limited.
3. Privacy-Preserving Verification: Allowing third-party arbitrators to verify ownership legally without accessing the proprietary signature, thereby protecting the IP owner's trade secrets.

The authors emphasize that while the ZKP generation is computationally heavier than standard inference, it is a one-time cost per artifact, making the overall system efficient for deployment. The framework is presented as a viable tool for regulating LLM-generated content and protecting intellectual property in the software development lifecycle.