An End-to-End Model for Logits-Based Large Language Models Watermarking

Introduction

As AI-generated content (AIGC) becomes increasingly pervasive, the ability to distinguish between human and machine-written text has moved from a technical curiosity to a critical necessity. Robust watermarking is essential for protecting copyright, ensuring academic integrity, and mitigating the spread of misinformation. However, existing "active" watermarking techniques, such as the widely-used KGW algorithm, suffer from two primary weaknesses: fragility and quality trade-offs. These methods often fail when text is paraphrased or partially copied, and they frequently struggle to balance detectability with natural language quality.

To address these challenges, we present an end-to-end (E2E) logits perturbation framework. This model moves beyond fixed rules, employing jointly optimized neural components to create a watermark that is both highly resilient to editing and virtually imperceptible to readers.

The E2E Framework

The core of this approach is the joint optimization of two lightweight neural networks: an Encoder ($E$) and a Decoder ($D$). Instead of relying on static hashing functions, the Encoder learns to apply context-aware shifts to the LLM's output logits. Simultaneously, the Decoder is trained to identify these specific patterns within the generated text.

Expand

Source, p. 3

This joint training allows the system to discover watermark patterns that are inherently resistant to common distortions. To maintain differentiability throughout the discrete word-selection process, the framework utilizes Gumbel-Softmax Sampling, allowing gradients to flow through the sampling step and enabling standard backpropagation.

Technical Innovations: Online Prompting and Cross-LLM Compatibility

A major hurdle in training E2E watermarking models is the need to simulate real-world "attacks" (like paraphrasing) and measure semantic preservation in a differentiable manner. This research introduces Online Prompting, a novel technique where the LLM being watermarked acts as its own surrogate for training.

Expand

Source, p. 4

By prompting the LLM to paraphrase its own watermarked output during the training loop, the Encoder learns to embed signals that specifically survive that model's rewriting style. Furthermore, the LLM is prompted to evaluate the semantic similarity between original and watermarked text, ensuring that the watermark does not degrade the meaning of the output.

To ensure the system is LLM-agnostic, the authors developed a Cross-LLM Converter. This component maps embeddings between different tokenizers, allowing the watermarking tool to be applied to various models (such as Llama-2, OPT, or Qwen) without requiring a full retraining of the encoder and decoder for every new architecture.

Evaluation and Key Results

The E2E model was rigorously tested against state-of-the-art methods like KGW, Unigram, and DiPmark using benchmarks such as WikiText-103, C4, and HumanEval. The results demonstrate a significant leap in performance:

• Superior Robustness: The model achieved a 37–39% improvement in detection accuracy under paraphrasing attacks compared to previous methods. Across a suite of common edits (including synonym swaps and copy-pasting), it showed an average improvement of 17.2%.
• Maintained Text Quality: Perplexity scores remained competitive with existing methods, and the model achieved the highest BLEU scores in machine translation tasks, indicating minimal impact on linguistic utility.
• High-Efficiency Detection: Because the decoder is a lightweight LSTM-based classifier, detection is extremely fast—taking only 0.005 seconds per sample. This is roughly 16,000x faster than sampling-based methods like EXP-Edit.

Conclusion and Future Directions

This end-to-end framework represents a significant advancement in LLM watermarking, successfully bridging the gap between robustness and text quality. By leveraging the LLM itself as a training surrogate through online prompting, the system creates a context-aware watermark that survives sophisticated editing.

While the current implementation introduces some generation overhead (approximately 60% for a Llama2-7B model), the gains in security and detection speed are substantial. Future work will focus on optimizing generation efficiency and hardening the watermark against potential "spoofing" or watermark-stealing attacks through key-driven biasing.

@inproceedings{wong2025end,
  author = {Kahim Wong and Jicheng Zhou and Jiantao Zhou and Yain-Whar Si},
  title = {{An End-to-End Model for Logits-Based Large Language Models Watermarking}},
  booktitle = {International Conference on Machine Learning},
  year = {2025},
  url = {https://openreview.net/forum?id=9sNiCqi2RD},
}