SSG: Logit-Balanced Vocabulary Partitioning for LLM Watermarking

2026-04-24

AI Summary

This work addresses the significant performance degradation of existing large language model (LLM) watermarking methods—such as KGW—in low-entropy scenarios like code generation and mathematical reasoning, which stems from random vocabulary partitioning that yields an unacceptably low theoretical lower bound on watermark strength. To overcome this limitation, the authors propose SSG, a novel approach that explicitly aligns vocabulary partitioning with the token prediction probability distribution. By sorting tokens according to their logits and constructing two logit-balanced subsets, SSG elevates the lower bound of watermark strength for each token. Theoretical analysis and empirical evaluations demonstrate that SSG substantially outperforms baseline methods like KGW in code and mathematical reasoning tasks, achieving markedly higher watermark detection accuracy.

Technology Category

• Machine Learning: Large Multimodal Models (LMMs)
• Natural Language Processing: (Large) Language Models
• Computer Vision: Large Vision Models

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Graph Algorithms and Modeling for the Web: Foundation models and LLMs for Web-related graphs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

Abstract

Watermarking has emerged as a promising technique for tracing the authorship of content generated by large language models (LLMs). Among existing approaches, the KGW scheme is particularly attractive due to its versatility, efficiency, and effectiveness in natural language generation. However, KGW's effectiveness degrades significantly under low-entropy settings such as code generation and mathematical reasoning. A crucial step in the KGW method is random vocabulary partitioning, which enables adjustments to token selection based on specific preferences. Our study revealed that the next-token probability distribution plays an critical role in determining how much, or even whether, we can modify token selection and, consequently, the effectiveness of watermarking. We refer to this characteristic, associated with the probability distribution of each token prediction, as \emph{watermark strength.} In cases of random vocabulary partitioning, the lower bound of watermark strength is dictated by the next-token probability distribution. However, we found that, by redesigning the vocabulary partitioning algorithm, we can potentially raise this lower bound. In this paper, we propose SSG (\textbf{S}ort-then-\textbf{S}plit by \textbf{G}roups), a method that partitions the vocabulary into two logit-balanced subsets. This design lifts the lower bound of watermark strength for each token prediction, thereby improving watermark detectability. Experiments on code generation and mathematical reasoning datasets demonstrate the effectiveness of SSG.

Problem

• LLM watermarking
• vocabulary partitioning
• watermark strength
• low-entropy generation
• token probability distribution

Innovation

• watermarking
• vocabulary partitioning
• logit-balanced
• watermark strength
• large language models