TraceMark-LDM: Robust Watermarking in LDMs

• TraceMark-LDM is an authenticatable watermarking framework for latent diffusion models that embeds watermarks via binary-guided rearrangement of Gaussian random variables.
• It integrates encoder fine-tuning and structured partitioning of latent variables to embed a 256-bit identifier while preserving the natural Gaussian distribution.
• The framework achieves superior robustness with >95% bit accuracy under various attack scenarios, all without detectable degradation in image quality metrics.

TraceMark-LDM is an authenticatable watermarking framework for Latent Diffusion Models (LDMs) that integrates binary-guided rearrangement of Gaussian random variables to achieve forensic attribution of AI-generated images, while maintaining non-destructive image quality. By encoding multi-bit watermarks directly in the stochastic generation process and coupling this with fine-tuning of LDM encoders, TraceMark-LDM attains robust, high-capacity, and near-invisible watermarking superior to existing state-of-the-art methods, even under intensive content distortion or re-generation scenarios (Luo et al., 30 Mar 2025).

1. Latent Diffusion Models and the Watermarking Problem

LDMs, such as Stable Diffusion v2.1, utilize a VAE encoder $E$ to compress images $I$ into a low-dimensional latent $z_0$. A forward diffusion process transforms $z_0$ into $z_T\sim\mathcal{N}(0,I)$, and a reverse denoising chain (DDPM/DDIM) reconstructs $z_0$, which is then decoded by $D$ to output pixels. Conventional watermarking approaches—embedding identifiers into pixels or intermediate noise—perturb latent Gaussian priors, causing measurable declines in image fidelity (increased FID, decreased CLIP-Score) and vulnerability to post-processing. The central technical challenge is to invisibly encode a $k$-bit identifier within $z_T$'s sampling, preserving the marginal $\mathcal{N}(0,I)$ and thereby retaining image quality and reliable recovery post-attack.

2. TraceMark-LDM Embedding Pipeline

TraceMark-LDM encodes a $I$0-bit watermark $I$1 via structured rearrangement of latent variables during image generation, involving partitioning, rearrangement based on bit values, and postprocessing to conceal the watermark.

2.1 Sampling and Partitioning

The algorithm samples $I$2 with latent dimension $I$3, and partitions $I$4 into negative ($I$5) and non-negative ($I$6) subsets.

2.2 Binary-Guided Rearrangement

Quartile partitions $I$7, $I$8 are extracted as the largest (absolute value) elements. For each watermark bit $I$9, a “large-element” sequence $z_0$0 is formed by cycling through bits and selecting an element from $z_0$1 if $z_0$2, $z_0$3 if $z_0$4, repeating $z_0$5 until $z_0$6 is assembled. The rearrangement operator $z_0$7 selects unused elements accordingly.

2.3 Group Rearrangement of Small Elements

Remaining elements $z_0$8 are sorted; most negative/positive halves ($z_0$9, $z_0$0) are split into $z_0$1 disjoint groups $z_0$2 and $z_0$3. Each group's sum signals the bit: $z_0$4 chosen from $z_0$5 encodes $z_0$6 $z_0$7, from $z_0$8 for $z_0$9 $z_T\sim\mathcal{N}(0,I)$0. Concatenation yields $z_T\sim\mathcal{N}(0,I)$1.

2.4 Interleaving, Permutation, and Generation

The sequences $z_T\sim\mathcal{N}(0,I)$2 and $z_T\sim\mathcal{N}(0,I)$3 are interleaved to yield $z_T\sim\mathcal{N}(0,I)$4; a secret key-dependent permutation $z_T\sim\mathcal{N}(0,I)$5 is applied, producing $z_T\sim\mathcal{N}(0,I)$6. This watermarked noise then passes through the LDM denoising chain to reconstruct $z_T\sim\mathcal{N}(0,I)$7, which is decoded to the final image $z_T\sim\mathcal{N}(0,I)$8.

2.5 Encoder Fine-Tuning

DDIM inversion and VAE encoding introduce extraction errors. Fine-tuning the encoder $z_T\sim\mathcal{N}(0,I)$9 (decoder $z_0$0 fixed) is performed: generating $z_0$1, applying random distortions $z_0$2 to $z_0$3, and optimizing

$z_0$4

$z_0$5

with $z_0$6, $z_0$7, for 100 epochs ($z_0$8200 images, distortions: median, JPEG, blur, noise, resize). This approach reduces bit-flip rates during extraction under attack.

3. Watermark Extraction and Authentication

Upon receiving a possibly attacked image $z_0$9:

1. Encode $D$0;
2. Apply DDIM inversion to recover $D$1;
3. Unshuffle via $D$2 to reconstruct $D$3;
4. De-interleave into $D$4 and $D$5;
5. Decode $D$6: $D$7 if $D$8, $D$9 otherwise;
6. Decode $k$0: $k$1;
7. Merge streams and repeat voting over repetitions to obtain final recovered $k$2;
8. Authenticate: compare $k$3 against user signatures. Attribution accepted if Hamming similarity $k$4 (threshold for FPR=$k$5).

4. Experimental Results and Benchmarking

The backbone is Stable Diffusion v2.1 ($k$6 images, latent $k$7). Sampling uses DPM-Solver ($k$8 steps, guidance $k$9), inversion by DDIM ($z_T$0 steps, null prompt, guidance $z_T$1). Attacks simulated include median filter ($z_T$2–$z_T$3), JPEG ($z_T$4–$z_T$5), Gaussian blur ($z_T$6–$z_T$7), Gaussian noise ($z_T$8–$z_T$9), salt-&-pepper ($\mathcal{N}(0,I)$0–$\mathcal{N}(0,I)$1), resize ($\mathcal{N}(0,I)$2–$\mathcal{N}(0,I)$3), VAE regen (quality $\mathcal{N}(0,I)$4–$\mathcal{N}(0,I)$5), and diffusion regen ($\mathcal{N}(0,I)$6–$\mathcal{N}(0,I)$7 DDPM steps). Watermark length is $\mathcal{N}(0,I)$8 bits, repeated $\mathcal{N}(0,I)$9 times.

Metric	Baseline	TraceMark-LDM	Statistical Test
FID	24.90	24.96	$I$00 ($I$01)
CLIP Score	0.3647	0.3649	$I$02 ($I$03)
Attribution (benign)	—	100% bit acc.	TPR@$I$04 ≈ 1.0

TraceMark-LDM demonstrates no statistical degradation of image quality. Robustness is sustained at $I$05 bit accuracy across distortions, $I$06 accuracy under salt-and-pepper $I$07, $I$08 under VAE regen $I$09, $I$10 under diffusion regen with $I$11 steps—the latter remains highest among compared methods.

5. Comparative Analysis with State-of-the-Art Methods

TraceMark-LDM is contrasted with prominent LDM watermarking approaches:

• Posterior-image methods (DwtDct, RivaGAN): degrade FID ($I$12) and lack robustness to distortions beyond mild JPEG.
• In-generation methods (Stable Signature, Latent Watermark): require costly U-Net fine-tuning or suffer quality loss for high-capacity embedding ($I$13 bits).
• Initial-noise methods (Tree-Rings, Gaussian Shading): cause distributional distortions or operational inefficiency (e.g., slow ChaCha20 encryption).

TraceMark-LDM achieves performance-lossless watermark embedding (no observable FID/CLIP drop), supports multi-bit capacity ($I$14 bits), maintains $I$15 bit accuracy under all attacks, exceeds $I$16 under extreme re-generation, and entails only moderate overhead (100-epoch encoder fine-tune, no per-image encryption or U-Net retraining). These properties yield superior bit-accuracy versus robustness trade-off (cf. Table I, Table III in source).

6. Context and Implications

TraceMark-LDM addresses the forensic attribution requirement for AI-generated content—a concern of increasing societal and legal significance. By integrating watermarking into the generative sampling step and developing resilience to post-processing and re-generation, the framework suggests a paradigm wherein provenance can be guaranteed with minimal operational disruption and strong resistance to adversarial attacks. A plausible implication is that similar binary-guided rearrangement and encoder fine-tuning methodologies could be extended to other generative architectures reliant on latent Gaussian sampling, offering broad utility in AIGC attribution and copyright protection domains (Luo et al., 30 Mar 2025).