T2SMark: Balancing Robustness and Diversity in Noise-as-Watermark for Diffusion Models

AI summary

Overview

Research area: Image watermarking for diffusion-model-generated images, specifically the Noise-as-Watermark (NaW) paradigm, sitting at the intersection of generative AI, intellectual property protection, and content provenance.

Technical level: Intermediate. The core idea is intuitive, but the paper assumes familiarity with diffusion sampling, DDIM inversion, latent spaces, and Gaussian truncation.

Scope (one sentence): The paper proposes T2SMark, a training-free, two-stage Noise-as-Watermark scheme built on Tail-Truncated Sampling, and evaluates it against post-processing, fine-tuning, and inversion-based watermarking baselines on Stable Diffusion v2.1 (U-Net) and Stable Diffusion v3.5 Medium (DiT) for robustness, generation diversity, undetectability, and visual quality.

What This Paper Is About

Diffusion models can be watermarked by encoding information into the initial Gaussian noise vector used for generation and recovering it later by inverting the diffusion process. The problem is that existing Noise-as-Watermark methods cannot simultaneously keep the watermark robust to distortions and keep generation diverse — some (like Gaussian Shading) gain robustness by heavily constraining the initial noise, which harms user experience, while others (like PRC-Watermark) preserve diversity but break under common inversion errors. T2SMark's goal is to achieve both at once by embedding bits only in the statistically reliable tails of the Gaussian distribution and by using a two-stage, session-key-based encryption pipeline to restore randomness.

Key Contributions

1. Tail-Truncated Sampling (TTS). Instead of mapping each bit to a positive or negative value, TTS partitions the Gaussian distribution into a bit-0 region, a bit-1 region, and an undecided central zone. Bits are embedded only in the tail regions (at least a distance τ from the boundary), and the central zone is sampled randomly to compensate for diversity while preserving the overall distribution.

2. A two-stage key hierarchy. A static master key encrypts a randomly generated session key in the first stage; that session key then encrypts the actual watermark bits in the second stage. The session key plays a dual role — embedded payload in the first segment and key for the second — randomizing the full noise vector to support generation diversity.

3. Projection-based decoding and detection. Each bit is recovered from the sign of the reconstructed noise's projection onto pseudorandom hyperplane normals, while detection uses the L1 norm of the first-stage projection vector as a confidence statistic compared to a threshold calibrated to a target false-positive rate.

4. A theoretical bit-error-probability analysis under an AWGN model showing that initiating truncation (a small τ > 0) strictly decreases the bit-error probability, predicting an optimal threshold that is determined empirically.

Main Findings

• Robustness and traceability: On SD v2.1, T2SMark reports TPR of 1.000 (clean) / 0.998 (adversarial) and bit accuracy of 1.0000 / 0.9754, matching GS's detection (1.000 / 0.998) while exceeding GS's adversarial bit accuracy (1.0000 / 0.9548). PRCW performs worst in traceability (0.6494 / 0.5024) and drops below 30% TPR under adversarial conditions (1.000 / 0.294).

• Diversity: Measured by LPIPS, PRCW scores highest at 0.7074, with T2SMark close behind at 0.7069 — a gap the authors describe as less than 10⁻³ and negligible. GS has the lowest diversity at 0.6446; TRW reaches 0.6943 and Stable Signature 0.6917. The unwatermarked SD v2.1 baseline is 0.7072.

• Visual quality: Only T2SMark and PRCW consistently satisfy the no-degradation criterion across both CLIP score and FID tests. T2SMark reports CLIP 0.3227 ± 0.0008 (t = 0.5081) and FID 56.9317 ± 0.42 (t = 0.5490) against an SD v2.1 baseline of 0.3224 ± 0.0010 and 56.8132 ± 0.48. GS shows a competitive CLIP score (0.3242 ± 0.0027) but its FID deviates (58.1377 ± 1.19) and has a notably larger standard deviation.

• Undetectability: Using a ResNet-18 classifier trained on 8,000 watermarked and 8,000 clean samples with 500 test samples each, detection accuracy is 0.971 for TRW, 0.994 for GS, 0.532 for PRCW, and 0.578 for T2SMark. T2SMark is second-best and described as difficult to detect.

• Generalization to a DiT backbone: On SD v3.5 Medium, T2SMark achieves TPR 1.000 / 0.985, bit accuracy 1.0000 / 0.9768, diversity 0.6102, and detection accuracy 0.518, close to PRCW's 0.516. TRW's robustness declines markedly versus SD v2.1 (0.878 / 0.318) and GS's diversity falls to 0.5176.

• TTS ablation: Removing TTS lowers adversarial TPR from 0.998 to 0.996, adversarial bit accuracy from 0.9754 to 0.9307, and diversity only from 0.6746 to 0.6743 — a substantial robustness gain at negligible diversity cost.

• Hyperparameter behavior: Bit accuracy degrades as capacity rises — 1.0000/0.9754 at 256 bits, 1.0000/0.9595 at 384, 1.0000/0.9437 at 512, 0.9992/0.9145 at 768, and 0.9968/0.8789 at 1024. Robustness is stable across 5, 10, 25, 50, and 100 inversion steps, so inversion cost can be reduced. Longer session keys hurt: adversarial bit accuracy falls from 0.9776 (8 bits) to 0.9754 (16), 0.9687 (24), and 0.9481 (32), making 16 bits the chosen trade-off. Performance degrades only marginally across guidance scales spanning 2–18.

• Key weakness: T2SMark is highly vulnerable to Gaussian noise — its performance degrades dramatically even at a noise standard deviation of 0.1, a weakness also reported for other inversion-based methods.

Methodology in Plain English

The authors start from the observation that Gaussian noise samples near the origin are far more likely to flip sign under perturbation than samples far out in the distribution's tails. A sign flip, in a scheme where bits are encoded as signs, is a decoding error. So instead of assigning every dimension to a bit, they reserve only the tail regions for bits and let the central region be sampled freely at random. This raises the signal-to-noise ratio per bit even though fewer dimensions carry information, because the retained dimensions are the reliable ones.

Concretely, the noise vector's dimensions are split by a binary mask into a watermark-encoding subspace and a random-noise subspace, with the split determined by τ and controlled by a secret key via a pseudorandom number generator that produces orthogonal, non-overlapping normal vectors. Along each normal, samples are drawn at least τ away from the boundary, and the bit's sign is applied to their magnitudes. Decoding regenerates the same normal vectors from the key, projects the reconstructed noise onto them, and reads the sign of each projection.

The two-stage structure addresses diversity. The n-dimensional noise is split into two segments. The first encodes a random session key under a fixed master key; the second encodes the actual watermark bits under that session key. Because both segments depend on a random key, the full vector stays randomized, so multiple users and multiple generations do not collapse onto the same noise pattern. Detection relies on the first stage only, to avoid error propagation, and compares the L1 norm of the first-stage projection against a threshold; traceability uses full two-stage decoding and matches against a registered identity-watermark database.

Evaluation covers three baseline families: post-processing transforms (dwtDct, dwtDctSvd, RivaGAN), a fine-tuning approach (Stable Signature), and inversion-based schemes (Tree-Ring, Gaussian Shading, PRC-Watermark). The SD v2.1 setup uses guidance scale 7.5, 50 DDIM denoising steps, 512×512 output, a 16-bit session key, a 256-bit watermark, and τ = 0.674, with the session key placed in the first channel of the initial noise. Robustness uses 500 prompts from the SDP training split, 500 watermarked images, and nine distortions, reporting TPR at FPR = 10⁻⁶ and per-bit accuracy. Diversity uses 10 images per each of 1,000 COCO test prompts with 45 unique pairs per prompt. Visual quality uses 10 independent trials of 1,000 images each with two-sample t-tests. All experiments ran in PyTorch 2.4.1 on a single NVIDIA RTX A6000 GPU.

Why This Matters

Impact on research: The paper reframes the robustness–diversity trade-off in Noise-as-Watermark methods as a sampling-distribution design problem rather than a discrete-encoding problem, and backs the intuition with an AWGN bit-error analysis. It also provides a direct comparison across three watermarking paradigms on two diffusion backbones (U-Net and DiT), which is useful for benchmarking.

Real-world applications:

• Proving provider copyright over images produced and served through a closed API, since even a single recovered bit constitutes evidence of the provider's legitimate ownership.
• Tracing misuse back to a specific API account by extracting an identity watermark and matching it against an account database.
• Flagging images as synthetic in origin, which supports provenance and authenticity verification of AI-generated content.
• Attribution in disinformation investigations, such as the paper's cited example of generative AI used to fabricate fake news and videos during the 2024 U.S. presidential election.

Industry relevance: The threat model directly describes commercial generative-image services with proprietary weights and training data. Because T2SMark is training-free — no weight updates, no additional training — it can be layered onto an existing API pipeline far more cheaply than fine-tuning approaches like Stable Signature, and the finding that robustness holds across 5 to 100 inversion steps matters for the extraction-side latency that NaW pipelines pay.

Future Directions

• Defending against forgery attacks. The authors note that the same high robustness enabling watermark recovery could let an adversary use a proxy model to invert the diffusion process, recover the watermarked noise vector, and generate forged images carrying a valid watermark.
• Resisting geometric distortions. NaW methods lack explicit mechanisms against geometric attacks; the paper points to GaussMarker as a promising direction.
• Broadening sampler compatibility. T2SMark depends on invertible, ODE-based sampling (DDIM, DPM-Solver), so it cannot be applied to diffusion models that do not support reconstruction of the initial noise.
• Reconciling with controllable generation. Methods that modify core sampling logic may interfere with watermark embedding and detection, and the paper flags this as an unresolved conflict.

An additional open issue is T2SMark's unique vulnerability: embedding the session key in truncated tails concentrates energy in a way that introduces subtle distributional anomalies that can be detected. Reducing that signature without losing the two-stage randomness is an explicit gap. The paper also does not report a principled procedure for choosing τ beyond determining it empirically.

Target Audience

Researchers and practitioners working on watermarking, generative model provenance, and AI content authentication will get the most from this paper. It is also relevant to engineers building image-generation APIs who need a training-free watermarking layer, and to graduate students studying the intersection of information hiding and diffusion models. Readers without background in diffusion sampling or DDIM inversion will need to consult the cited foundational work first.

Authors’ abstract

Diffusion models have advanced rapidly in recent years, producing high-fidelity images while raising concerns about intellectual property protection and the misuse of generative AI. Image watermarking for diffusion models, particularly Noise-as-Watermark (NaW) methods, encode watermark as specific standard Gaussian noise vector for image generation, embedding the infomation seamlessly while maintaining image quality. For detection, the generation process is inverted to recover the initial noise vector containing the watermark before extraction. However, existing NaW methods struggle to balance watermark robustness with generation diversity. Some methods achieve strong robustness by heavily constraining initial noise sampling, which degrades user experience, while others preserve diversity but prove too fragile for real-world deployment. To address this issue, we propose T2SMark, a two-stage watermarking scheme based on Tail-Truncated Sampling (TTS). Unlike prior methods that simply map bits to positive or negative values, TTS enhances robustness by embedding bits exclusively in the reliable tail regions while randomly sampling the central zone to preserve the latent distribution. Our two-stage framework then ensures sampling diversity by integrating a randomly generated session key into both encryption pipelines. We evaluate T2SMark on diffusion models with both U-Net and DiT backbones. Extensive experiments show that it achieves an optimal balance between robustness and diversity. Our code is available at \href{https://github.com/0xD009/T2SMark}{https://github.com/0xD009/T2SMark}.