MCGMark: An Encodable and Robust Online Watermark for Tracing LLM-Generated Malicious Code

2024-08-02

AI Summary

To address the challenge of tracing malicious code generated by large language models (LLMs), this paper proposes the first online watermarking scheme specifically designed for malicious code. The method integrates three core components: (1) probability-anomaly-driven token selection, (2) code-structure-aware constraints on embedding positions—explicitly avoiding easily modifiable regions such as comments—and (3) robustness-optimized watermark encoding. It innovatively models the syntactic and semantic structures of malicious code and leverages token-level controllable sampling to ensure high fidelity. Evaluated on DeepSeek-Coder, the approach achieves an 88.9% watermark embedding success rate within ≤400 tokens, with negligible impact on functionality and readability. The embedded watermark demonstrates strong resilience against common adversarial manipulations—including deletion, insertion, and reformatting—achieving, for the first time, a balanced optimization of high success rate, strong robustness, and minimal quality degradation.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: Code Generation / Program Synthesis from Natural Language
• Computer Vision: Adversarial Attacks & Robustness

Application Category

• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs

Abstract

With the advent of large language models (LLMs), numerous software service providers (SSPs) are dedicated to developing LLMs customized for code generation tasks, such as CodeLlama and Copilot. However, these LLMs can be leveraged by attackers to create malicious software, which may pose potential threats to the software ecosystem. For example, they can automate the creation of advanced phishing malware. To address this issue, we first conduct an empirical study and design a prompt dataset, MCGTest, which involves approximately 400 person-hours of work and consists of 406 malicious code generation tasks. Utilizing this dataset, we propose MCGMark, the first robust, code structure-aware, and encodable watermarking approach to trace LLM-generated code. We embed encodable information by controlling the token selection and ensuring the output quality based on probabilistic outliers. Additionally, we enhance the robustness of the watermark by considering the structural features of malicious code, preventing the embedding of the watermark in easily modified positions, such as comments. We validate the effectiveness and robustness of MCGMark on the DeepSeek-Coder. MCGMark achieves an embedding success rate of 88.9% within a maximum output limit of 400 tokens. Furthermore, it also demonstrates strong robustness and has minimal impact on the quality of the output code. Our approach assists SSPs in tracing and holding responsible parties accountable for malicious code generated by LLMs.

Problem

• Traces malicious code generated by LLMs using watermarks
• Ensures watermark robustness by avoiding easily modified code positions
• Maintains output quality while embedding encodable information

Innovation

• Encodable watermark via token control
• Robust by code structure awareness
• Minimal impact on code quality