Learning to Watermark in the Latent Space of Generative Models

AI summary

Overview

Research area: Content provenance and invisible watermarking for AI-generated images, specifically watermarking in the latent space of diffusion and autoregressive generative models.

Technical level: Intermediate to Advanced. The paper assumes familiarity with latent diffusion models, VAEs, vector quantization, autoregressive image generation, and encoder-decoder watermarking architectures.

Scope: The paper introduces and evaluates DistSeal, a framework that trains watermark embedders in the latent space of generative models and then distills them into either the generative model or its latent decoder for in-model watermarking.

What This Paper Is About

Most deployed image watermarking methods work post-hoc in pixel space, which adds computational overhead and can be trivially removed from open-source code. This paper asks whether watermarks can instead be embedded in the compressed latent representations that diffusion and autoregressive models already use internally, and whether those latent watermarks can be baked permanently into model weights. The goal is a single watermarking approach that is faster than pixel-space methods, comparably robust and imperceptible, and impossible for an open-source user to disable by deleting one line of code.

Key Contributions

1. DistSeal framework: A unified latent watermarking approach that handles both continuous latents of diffusion models and discrete token sequences of autoregressive models, covering post-hoc embedding before or after quantization.
2. Latent-to-model distillation: A demonstration that latent post-hoc watermarkers can be distilled into either the generative model itself or its latent decoder, yielding state-of-the-art in-model watermarking for open-source generative models.
3. Latent beats pixel for distillation: Evidence that distilling latent watermarkers is more effective than distilling pixel-space watermarkers, with the latent versions retaining robustness that pixel teachers lose during distillation.
4. Deployment study: An analysis of trade-offs between post-hoc and in-model watermarking, multi-watermarking compatibility, watermark formation during diffusion, and watermark forgetting, framed as guidance for real-world deployment.

Main Findings

• Latent post-hoc watermarking is competitive, not superior, on quality and robustness: On DCAE-generated images, the latent watermarker reached 95.18% average bit accuracy versus 97.78% for the pixel watermarker. The gap was concentrated in combined attacks (84.28% vs. 97.29%), while compression held at 99.23% and valuemetric at 98.08%.

• Large speedup on strongly compressed latents: The DCAE latent watermarker ran at 3 ms per image on CPU versus 63 ms for the pixel watermarker, a 20× speedup, operating on 8×8×128 latents instead of 512×512×3 pixels.

• Smaller speedup for less compressed latents: The RAR-XL latent watermarker ran at 17 ms versus 41 ms per image on CPU, a 2.4× speedup, with latents of 16×16×256 versus 256×256×3 pixels. The authors attribute the smaller gain to the less compressed latent space.

• PSNR is misleading for latent watermarking: The latent watermarker's PSNR was 31.06 dB versus 43.48 dB for pixel, but FID was 11.42 versus 10.84, and IS 98.11 versus 99.53. Against a non-watermarked FID baseline of 10.66, the increases were +0.76 and +0.18 respectively.

• Coverless latent watermarking is weaker: A coverless variant of the latent watermarker, which ignores the input, achieved 91.81% average bit accuracy versus 95.18% for the input-dependent version.

• Quantization placement matters for autoregressive models: Watermarking after quantization achieved 93.96% average bit accuracy versus 90.60% before quantization on RAR-XL, compared to 97.27% for the pixel watermarker.

• Only pre-quantization watermarking can be distilled into an autoregressive model: Post-quantization watermarking modifies token embeddings and produces no new token sequence, so it cannot fine-tune the generative model, though it can still be distilled into the latent decoder.

• Decoder distillation works well for latent teachers: Distilling the DCAE latent watermarker into the decoder with a 0 extractor weight reached 94.41% average accuracy at FID 11.34; adding an extractor weight of 0.1 reached 96.77% at FID 11.48, exceeding the teacher's 95.18% and improving combined-attack accuracy from 84.28% to 91.34%.

• Pixel teachers fail to distill cleanly: Distilling the pixel watermarker into the DCAE decoder with reconstruction only produced 51.38% average accuracy; with a strong extractor weight of 1.0 it reached 92.95% but at FID 17.08, far above the teacher's 10.84. The authors suggest high-frequency pixel perturbations are hard for the latent decoder to reconstruct.

• Same pattern in the autoregressive decoder: The RAR latent teacher distilled to 81.40% (extractor weight 0), 90.57% (0.1), and 92.33% (0.5) against a teacher of 93.96%. The pixel teacher reached only 50.15%, 50.81%, and 71.25% across weights 0, 0.1, and 1.0, the last at FID 5.05.

• Generative-model distillation is slightly less robust but higher quality: Distilling into the DCAE diffusion model gave 94.78% average accuracy at FID 10.90 versus 95.18% at FID 11.42 for the teacher. For RAR, 89.13% at FID 3.32 versus 90.60% at FID 3.56.

• Generative-model distillation generates new watermarked content: The distilled autoregressive model produces entirely new watermarked token sequences with the same robustness as the post-hoc teacher, and the distilled diffusion model produces new watermarked latents with the teacher's robustness.

• The watermark strengthens during denoising: Bit accuracy increases progressively across diffusion steps. A hybrid setup using the original model initially and the distilled model only for the final N steps still reached high accuracy, with watermark patterns visualizable when switching for the final 5 steps.

• Distillation choice is a deployment trade-off: Distilling into the generative model needs no hyperparameter tuning beyond fine-tuning on watermarked latents and is robust to latent decoder fine-tuning, but is susceptible to forgetting with LoRA. Distilling into the latent decoder needs tuning of the extractor weight and LPIPS weight and is vulnerable to decoder replacement, but is unaffected by changes to the generative model and offers stronger persistence against fine-tuning attacks.

• Training setup: All post-hoc models used 64-bit messages, 600k steps on ImageNet, batch size 128, watermark strengths ranging from 0.5 to 3.0 depending on model and placement (1.5 to 0.5 for DCAE, 2.0 to 0.5 for RAR after quantization, 3.0 to 1.5 before quantization, 0.2 to 0.02 for pixel).

Methodology in Plain English

The researchers took an existing pixel-space watermarking architecture, based on VideoSeal with a UNet embedder and ConvNeXt-tiny extractor, and adapted the embedder to operate on low-resolution latents rather than high-resolution images. They removed all downsampling and upsampling layers to match the smaller spatial dimensions while keeping the middle ResNet blocks, and left the extractor unchanged so that pixel and latent versions could be compared fairly, sharing objectives, schedules, augmentations, and hyperparameters, with only the watermark strength adjusted for the different value ranges.

The embedder adds a scaled perturbation to the latent, which is then decoded into an image. For autoregressive models with quantized latents, the perturbation is added either before quantization, creating a new token sequence, or after quantization, modifying token embeddings. Training combines a watermark extraction loss with a discriminator loss in pixel space; no reconstruction or perceptual losses are used for the embedder because latent perturbations can cause semantic changes.

To bake the watermark into the model, the researchers fine-tune the generative model on watermarked latents instead of original ones, so it learns to produce watermarks during generation. Alternatively, they fine-tune the latent decoder to reconstruct watermarked images from non-watermarked latents, combining a reconstruction loss (ℓ1 plus LPIPS) with the extractor loss. They then measured bit accuracy across identity, valuemetric, geometric, compression, and combined transformations, plus PSNR, FID, and IS on 50k generated images against the ImageNet validation set.

Why This Matters

Impact on research: The paper reframes watermarking as a latent-space problem and shows that distillation success depends on whether teacher perturbations are structured (latent) or high-frequency (pixel). It also provides methods and negative results for in-model watermarking across two distinct generative paradigms, diffusion and autoregressive.

Real-world applications:

• Embedding provenance signals into open-source image generators so watermarking survives even when users modify the training or inference code.
• Reducing inference cost for platforms that watermark large volumes of generated images, since latent watermarking costs 3 ms versus 63 ms per image on CPU for DCAE.
• Distinguishing outputs from different generative models that share the same autoencoder by using different binary messages.
• Authenticating AI-generated media for intellectual property enforcement and deepfake accountability.

Industry relevance: The authors note that post-hoc pixel-space watermarking is the method of choice for industry deployment because it is flexible and model-agnostic. DistSeal targets the two weaknesses that matter most in practice for open-source releases: latency, addressed by the reported speedups, and bypassability, addressed by in-model distillation, which removes the need for a removable post-processing step.

Future Directions

• Improving robustness against strong geometric transformations, potentially by combining DistSeal with post-hoc synchronization methods.
• Extending latent watermarking to video generative models, where watermarking every frame in pixel space is expensive and latent watermarking could yield large speedups.
• Analyzing watermark forgetting in more depth, since the trade-off table flags generative-model distillation as susceptible to LoRA-based forgetting, and the paper points to appendix coverage of this.
• Investigating multi-watermarking compatibility and watermark formation further, which the paper states are analyzed in appendices rather than the main text.

Target Audience

Researchers and engineers working on generative model safety, content provenance, and media forensics; practitioners deploying diffusion or autoregressive image generators who need watermarking without inference overhead or bypass risk; and students with intermediate background in generative modeling who want a concrete example of distilling a post-hoc signal into model weights.

Authors’ abstract

Existing approaches for watermarking AI-generated images often rely on post-hoc methods applied in pixel space, introducing computational overhead and potential visual artifacts. In this work, we explore latent space watermarking and introduce DistSeal, a unified approach for latent watermarking that works across both diffusion and autoregressive models. Our approach works by training post-hoc watermarking models in the latent space of generative models. We demonstrate that these latent watermarkers can be effectively distilled either into the generative model itself or into the latent decoder, enabling in-model watermarking. The resulting latent watermarks achieve competitive robustness while offering similar imperceptibility and up to 20x speedup compared to pixel-space baselines. Our experiments further reveal that distilling latent watermarkers outperforms distilling pixel-space ones, providing a solution that is both more efficient and more robust.