Optimizing Token Choice for Code Watermarking: An RL Approach

TL;DR

This paper presents CodeTracer, a reinforcement learning-based framework for embedding watermarks in code generated by large language models, ensuring detectability while preserving code functionality.

Contribution

It introduces a novel RL training paradigm with a policy-driven approach and Gumbel reparameterization for effective, subtle code watermarking.

Findings

CodeTracer outperforms existing methods in watermark detectability.

It maintains code functionality with minimal disruption.

The framework effectively balances watermark strength and code correctness.

Abstract

Protecting intellectual property on LLM-generated code necessitates effective watermarking systems that can operate within code's highly structured, syntactically constrained nature. In this work, we introduce CodeTracer, an innovative adaptive code watermarking framework underpinned by a novel reinforcement learning training paradigm. At its core, CodeTracer features a policy-driven approach that utilizes a parameterized model to intelligently bias token choices during next-token prediction. This strategy ensures that embedded watermarks maintain code functionality while exhibiting subtle yet statistically detectable deviations from typical token distributions. To facilitate policy learning, we devise a comprehensive reward system that seamlessly integrates execution feedback with watermark embedding signals, balancing process-level and outcome-level rewards. Additionally, we employ Gumbel Top-k reparameterization to enable gradient-based optimization of discrete watermarking decisions. Extensive comparative evaluations demonstrate CodeTracer's significant superiority over state-of-the-art baselines in both watermark detectability and the preservation of generated code's functionality.