Adaptive Code Watermarking Through Reinforcement Learning
Abstract
As LLMs increasingly generate production code, protecting intellectual property demands watermarking techniques that respect code's strict syntactic constraints. In this work, we introduce CodeTracer, an innovative adaptive code watermarking framework underpinned by a 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. To enable gradient-based optimization of these discrete watermarking decisions, we employ Gumbel Top-k reparameterization. Extensive comparative evaluations demonstrate that CodeTracer outperforms state-of-the-art baselines across multiple benchmarks in both watermark detectability and code functionality. Our code is available at https://github.com/TimeLovercc/CodeTracer.
Lay Summary
AI tools can now write useful software, but it is often difficult to tell where generated code came from. This creates challenges for intellectual property, accountability, and academic integrity, especially because code is hard to watermark: even small changes can make a program fail. We introduce CodeTracer, a method that learns where and how to place subtle, detectable signals in AI-generated code while preserving its functionality. CodeTracer trains a small companion model using feedback about both code correctness and watermark detectability. Unlike some prior approaches, the watermark can be detected later from the code itself, without needing the original prompt or the full model used to generate it. Our experiments show that CodeTracer improves traceability while maintaining code correctness better than existing watermarking baselines. This can help make code-generating AI systems easier to audit and use responsibly.