Can Watermarking Techniques Help Prevent LLM Model Stealing?

Technical Summary: Watermarking Techniques for Preventing LLM Model Stealing

Problem Statement
Commercial Large Language Models (LLMs) are typically accessed via black-box APIs, keeping their internal architectures, training data, and parameters proprietary. However, recent work by [Carlini et al., 2024b] demonstrated a scalable "model stealing" attack capable of extracting precise architectural information, specifically the hidden layer dimension and the final projection layer, from these black-box models. The attack operates by querying the model with a set of controllable prompts to collect next-token prediction logit vectors. By constructing a matrix from these logits and applying Principal Component Analysis (PCA), an adversary can estimate the numerical rank of the matrix, which corresponds to the model's hidden dimension. This vulnerability is critical because the hidden dimension serves as a foundation for stealing the model's final layer parameters. The attack is versatile, as it does not rely on specific activation functions or network structures, and it remains effective even when the defender attempts to mitigate it via simple noise injection or averaging techniques.

Methodology
The authors propose a defense mechanism inspired by LLM watermarking techniques [Kirchenbauer et al., 2023], specifically designed to disrupt rank-estimation attacks while preserving model utility. The core strategy involves structurally perturbing the logits layer of the model before token sampling.

1. Prompt-Dependent Deterministic Noise: Unlike naive defenses that add independent random noise (which can be averaged out by an adversary), the proposed method generates noise that is deterministic for a given prompt but varies across different prompts. This is achieved by computing a seed using a cryptographic Pseudorandom Function (PRF) applied to the model's last hidden state ($h$) for a specific prompt $p$. The PRF ensures that identical prompts yield identical perturbations (preventing averaging attacks) while different prompts yield distinct, seemingly random perturbations.
2. Perturbation Mechanisms: The paper evaluates several noise structures:
   • Additive Noise: Direct addition of discrete or Gaussian noise vectors. While effective against basic PCA, these are vulnerable to sophisticated averaging or robust PCA attacks.
   • Multiplicative Noise: Scaling logits by random factors. Experiments show this significantly degrades model quality and fails to prevent rank recovery.
   • Softplus-Then-Perturb (STP): The authors propose a novel mechanism where the logit vector $\ell$ is first transformed via a softplus function ($\text{softplus}(\alpha\ell)$) to skew the values and preserve the relative ordering of high-confidence tokens. Random perturbations are then added to this transformed vector. This approach aims to maintain token ordering (crucial for quality) while introducing sufficient structural complexity to break the low-rank subspace required for rank estimation.

Key Contributions

• Defense Against Rank Estimation: The paper introduces a structured perturbation framework that specifically targets the mathematical vulnerability exploited by [Carlini et al., 2024b], preventing the recovery of the hidden dimension via PCA or Robust PCA.
• Theoretical Analysis of Noise: The authors analyze the stability of rank recovery under various noise conditions, utilizing Weyl's inequality and matrix perturbation theory. They demonstrate that simple additive noise is insufficient against averaging attacks and that the dependence structure between the noise and the prompt is critical.
• STP Mechanism: The proposal of the Softplus-Then-Perturb mechanism, which balances robustness against extraction attacks with the preservation of model performance.
• Empirical Validation: Extensive experiments on Mistral-7B (and GPT-2 in the appendix) demonstrate the trade-offs between defense robustness and model quality.

Results

• Attack Prevention: The proposed defenses, particularly the STP mechanism combined with Gaussian noise, successfully prevent the extraction of the hidden dimension. In experiments with 10,000 prompts, the attack failed to identify the correct rank (4,096 for Mistral-7B), with the predicted dimension showing no identifiable anomaly or significant jump in singular values.
• Robustness to Advanced Attacks: The STP mechanism with Gaussian noise withstood attempts at "softplus inversion" attacks and Robust PCA (RPCA) extensions. In contrast, pure softplus without perturbation was vulnerable to inversion, and simple additive noise was vulnerable to averaging.
• Model Quality Preservation:
  • Perplexity (PPL): Additive noise caused significant degradation in perplexity with higher magnitudes. However, STP-based perturbations maintained perplexity close to the original model (e.g., 3.37 vs. 3.40 for the original).
  • MMLU Accuracy: While additive noise reduced MMLU scores significantly (dropping to ~44-55%), the STP mechanism preserved high accuracy (56-57%), nearly matching the original model (57.78%).
  • Distributional Shift: Although STP induced high distributional divergence (high KL and JS divergence), it maintained a high Top-50 token agreement (≥0.92). This indicates that while the probability distribution shifts, the relative ranking of the most confident tokens remains intact, preserving the model's functional utility.

Significance
The paper claims that its approach offers a practical and effective defense against a specific class of model stealing attacks that threaten the intellectual property of commercial LLMs. By leveraging watermarking-inspired techniques to perturb logits in a prompt-dependent manner, the authors demonstrate that it is possible to obscure the model's internal dimensionality without rendering the model unusable. The work highlights that preserving the relative ordering of logits is more critical for maintaining model utility than minimizing raw distributional shifts, a finding that informs the design of the STP mechanism. The authors conclude that their method bolsters the security and integrity of commercial language models against rank-based extraction, providing a viable path forward for protecting proprietary model architectures.