🤖 AI

Learning to Watermark in the Latent Space of Generative Models

This paper introduces DistSeal, a unified approach that trains post-hoc watermarking models in the latent space of generative models and distills them into the generator or decoder, achieving robust, imperceptible, and up to 20x faster watermarking compared to traditional pixel-space methods.

Original authors: Sylvestre-Alvise Rebuffi, Tuan Tran, Valeriu Lacatusu, Pierre Fernandez, Tomáš Souček, Nikola Jovanović, Tom Sander, Hady Elsahar, Alexandre Mourachko

Published 2026-01-23

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Sylvestre-Alvise Rebuffi, Tuan Tran, Valeriu Lacatusu, Pierre Fernandez, Tomáš Souček, Nikola Jovanović, Tom Sander, Hady Elsahar, Alexandre Mourachko

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you have a magical artist (a Generative AI) that can paint incredibly realistic pictures from scratch. The problem is, once these pictures are out in the world, it's hard to tell if they were painted by a human or a machine. This makes it easy for bad actors to create fake news or deepfakes.

To solve this, we need to put a "hidden signature" on these AI pictures, like a secret watermark. This paper introduces a new, smarter way to do that called DISTSEAL.

Here is how it works, broken down into simple concepts:

1. The Old Way: Painting the Signature on the Canvas

Previously, people tried to add watermarks after the AI finished painting. Imagine the AI paints a 512x512 pixel image (a huge canvas). Then, a separate tool has to go in and carefully scribble a tiny, invisible code onto every single pixel of that finished canvas.

The Problem: This is slow and clunky. It's like hiring a second person to walk over a finished painting and add a signature. It also leaves tiny "artifacts" (glitches) that can sometimes be seen or removed easily.

2. The New Way: The "Latent Space" Kitchen

The authors realized that before the AI paints the final picture, it actually works in a compressed, hidden "kitchen" called Latent Space. Think of this as a rough sketch or a set of instructions that is much smaller and simpler than the final painting.

The Innovation: Instead of waiting for the painting to finish, DISTSEAL puts the watermark inside this hidden kitchen while the AI is still cooking.
The Benefit: Because the "kitchen" is so much smaller than the final canvas, the process is 20 times faster. It's like writing the secret code on the recipe card instead of trying to write it on the finished cake.

3. Two Ways to Use the Secret Code

The paper shows two main ways to use this new method:

A. The "Post-It Note" Method (Post-Hoc)
You can still use the system after the AI generates the image, but instead of touching the big canvas, you touch the small recipe card (the latent space) before it becomes a picture.

Result: It's incredibly fast and the pictures look just as good as the originals, but the hidden code is much harder to remove.

B. The "Baked-In" Method (In-Model)
This is the really cool part. The authors figured out how to "teach" the AI artist itself to always include the secret code while it paints.

How: They take the "Post-It Note" method and "distill" it (like transferring knowledge) directly into the AI's brain (its weights).
Result: Now, every time the AI paints, the watermark is baked in automatically. You don't need a separate tool anymore. If someone tries to copy the AI's code (open-source), the watermark stays there because it's part of the artist's DNA, not just a tool they can delete.

4. Why This is a Big Deal

Speed: It's up to 20 times faster than previous methods because it works on the small "recipe" instead of the huge "canvas."
Stealth: The watermarks are invisible to the human eye and don't ruin the quality of the art.
Durability: The paper tested these watermarks against things like cropping, resizing, changing colors, and compressing the image. The "Latent Space" watermarks survived these attacks much better than the old "Pixel Space" methods.
Better than Pixel Distillation: The authors tried to teach the AI using the old "Pixel" method, but it failed (the AI couldn't learn the complex patterns). However, when they taught the AI using the new "Latent" method, the AI learned perfectly.

Summary Analogy

Old Method: A security guard stands at the exit of a factory, stamping every single finished car with a hidden serial number. It's slow, and the stamp might rub off.
DISTSEAL: The factory workers are trained to stamp the hidden serial number onto the car's blueprint while they are building it. The stamp is part of the car's structure, it's done instantly, and it's nearly impossible to remove without destroying the car itself.

The paper concludes that this method works for both major types of AI image generators (Diffusion models and Autoregressive models) and provides a robust, fast, and efficient way to prove that an image came from an AI.

Technical Summary: Learning to Watermark in the Latent Space of Generative Models

Problem Statement

The rapid advancement of generative models (diffusion and autoregressive) has enabled the creation of highly realistic synthetic content, raising concerns regarding misuse, deepfakes, and intellectual property violations. While invisible watermarking offers a solution for provenance, existing approaches face significant limitations:

Post-hoc Pixel-Space Methods: These modify already-generated images in pixel space. While flexible and model-agnostic, they incur substantial computational overhead due to high-resolution processing (e.g., $512 \times 512 \times 3$ ) and introduce latency. Furthermore, they offer weak security for open-source deployments, as users can easily bypass them by removing a single line of code.
In-Model Methods: While more secure, previous attempts to integrate watermarks directly into model weights often lack robustness or require complex modifications to the sampling process.

The core challenge is to develop a unified watermarking framework that operates directly within the latent space of generative models, enabling in-model watermarking (embedding the watermark into the model weights) that is both computationally efficient and robust against removal, without introducing latency penalties.

Methodology: DISTSEAL

The authors introduce DISTSEAL, a unified framework for latent space watermarking applicable to both diffusion and autoregressive models. The approach consists of two primary stages:

1. Post-hoc Latent Watermarking

The first stage involves training a watermark embedder-extractor pair ( $W_\theta, E_\theta$ ) that operates in the latent space of the generative model rather than the pixel space.

Embedding: The embedder takes a latent representation $z$ $z$ and a binary message $m$ $m$ , outputting a watermarked latent $z_w = z + \epsilon W_\theta(z, m)$ $z_{w} = z + ϵ W_{θ} (z, m)$ .
- For diffusion models (continuous latents), the perturbation is applied directly to $z$ .
- For autoregressive models (discrete tokens), the embedder can operate either before quantization (modifying continuous latents) or after quantization (modifying token embeddings). The authors utilize straight-through estimation to backpropagate through the quantization step.
Training Objective: The model is trained using a combination of watermark extraction loss ( $L_w$ ) and a discriminator loss ( $L_{disc}$ ) applied in pixel space to ensure imperceptibility. Unlike some prior works, the authors avoid reconstruction or perceptual losses for the embedder, relying instead on the watermark strength $\epsilon$ and the discriminator to constrain the output.
Augmentation: Training involves various value-metric, geometric, and compression augmentations to simulate user edits and ensure robustness.

2. Distillation for In-Model Watermarking

The second stage distills the trained post-hoc latent watermarker into the generative model itself or its latent decoder, eliminating the need for a separate post-processing step.

Distillation into the Generative Model: The pre-trained generative model (diffusion transformer or autoregressive transformer) is fine-tuned to reconstruct or predict the watermarked latents $z_w$ instead of the original clean latents $z$ . This forces the model to inherently generate watermarked outputs.
Distillation into the Latent Decoder: The latent decoder $D$ is fine-tuned to reconstruct watermarked images from non-watermarked latents. The loss function combines a reconstruction loss ( $L_{rec}$ ) and a watermark extraction loss ( $L_w$ ) to guide the decoder to produce outputs that contain the watermark signal.

Key Contributions

Unified Framework (DISTSEAL): A single approach that handles both continuous latent spaces (diffusion models) and discrete token sequences (autoregressive models).
Efficiency and Robustness: The method achieves up to a 20× speedup in inference compared to pixel-space baselines by operating on compressed latent representations (e.g., $8 \times 8 \times 128$ vs. $512 \times 512 \times 3$ ) while maintaining competitive robustness against various attacks.
Effective Distillation: The authors demonstrate that latent watermarkers can be effectively distilled into model weights. Crucially, they show that distilling latent watermarkers outperforms distilling pixel-space watermarkers. Pixel-space watermarkers often fail to distill into latent decoders due to the difficulty of reconstructing high-frequency perturbations, whereas latent perturbations are more structured and easier for the decoder to learn.
Comprehensive Analysis: The paper provides a detailed study of trade-offs between post-hoc and in-model methods, multi-watermarking compatibility, and watermark forgetting (susceptibility to LoRA fine-tuning).

Experimental Results

The authors evaluated DISTSEAL on DCAE (a diffusion model) and RAR-XL (an autoregressive model).

Post-hoc Performance: Latent watermarking achieved average bit accuracy of 95.18% (DCAE) and 93.96% (RAR-XL) under combined attacks, comparable to pixel-space baselines (97.78% and 97.27% respectively). Visual quality (FID) remained high, with minimal degradation.
In-Model Performance (Distillation):
- Latent Decoder Distillation: Distilling the latent watermarker into the decoder achieved 96.77% (DCAE) and 90.57% (RAR-XL) average bit accuracy, often outperforming the teacher post-hoc model. In contrast, distilling pixel watermarkers resulted in significantly lower robustness (e.g., 65.69% for DCAE) and higher FID.
- Generative Model Distillation: Fine-tuning the generative model directly yielded robust results (94.78% for DCAE, 89.13% for RAR-XL) with slightly better visual quality than decoder distillation.
Speed: The latent approach provided a 20× speedup for DCAE and a 2.4× speedup for RAR-XL compared to pixel-space methods.
Limitations Observed:
- Watermark Forgetting: Distilled generative models are susceptible to forgetting the watermark when fine-tuned on non-watermarked data (e.g., via LoRA), whereas distilled decoders remain robust to such changes.
- Strong Attacks: The method is not yet fully robust against very strong geometric transformations, a limitation the authors acknowledge.

Significance and Claims

The paper claims that DISTSEAL provides a practical and efficient solution for embedding robust watermarks directly within generative models. By moving watermarking from the pixel space to the latent space, the authors achieve:

Reduced Overhead: Significant computational savings and latency reduction, making real-time or large-scale deployment more feasible.
Enhanced Security for Open Source: By integrating the watermark into the model weights (via distillation), the watermark cannot be trivially disabled by removing a post-processing script, addressing a critical security gap in open-source model releases.
Superior Distillation: The finding that latent watermarkers distill more effectively than pixel-space ones suggests a new paradigm for creating "watermarked-by-design" generative models that maintain high fidelity and robustness.

The authors position this work as a step toward reliable provenance mechanisms for AI-generated content, balancing the need for accountability with the preservation of content quality and generation speed. They note that while the method is robust, it is not immune to removal, and future work will explore combining latent watermarking with post-hoc synchronization methods to handle stronger geometric attacks.