The Stable Signature: Rooting Watermarks in Latent Diffusion Models

In plain words

Image generators built on latent diffusion turn a compact hidden representation into pixels with a decoder. This paper changes only that decoder, so every picture it outputs carries an invisible, fixed pattern of 48 bits. A separate, pre-trained reader checks an image for that pattern. If enough bits match, the image is flagged as coming from the model. If each customer gets a different pattern, the matching pattern also shows which customer made the image.

The generation process itself is left alone, and the decoder change is a short fine-tuning run. The pattern survives common edits such as cropping, brightening and recompression, although strong edits lower the match rate. All numbers are the authors' own.

TL;DR

• Fine-tunes only the decoder of a latent diffusion model (LDM) so that all its outputs carry a fixed 48-bit key §4.2.
• The key is read by a HiDDeN-style extractor that is pre-trained with simulated distortions and then whitened §4.1.
• Detection counts matching bits against a threshold; identification picks the best of N user keys §3.
• Author-reported: 0.99 bit accuracy on unmodified text-to-image output and 0.92 after a combined crop, brightness and JPEG edit, with FID shifts of at most 0.6 tab. 1.
• Not yet reproducible from the public code alone: the headline text-to-image experiments use an internal 2.2B-parameter model §5.

Problem

A watermark added after generation can be removed once a generator leaks or is open-sourced §1. The authors want the watermark to be part of generation itself, so that a model provider can do two things. For detection, it decides whether an image came from its model while bounding the rate of false flags on natural images §3.1. For identification, it gives each downstream user a key and attributes an image to a user, without wrongly accusing someone who did not make it §3.2.

Both tests rely on the extracted bits of an unmarked image being independent fair coin flips. The authors observed that this holds only approximately for their extractor, which is why they whiten its output §4.1 app. B.5.

Method

The method has two stages, plus a statistical test on top fig. 2.

1. Pre-train the extractor. A HiDDeN-style encoder embeds a k-bit message into training images, and an extractor recovers it after random distortions. The encoder's residual is bounded by a tanh and scaled by a factor α = 0.3. The encoder is discarded afterwards, and the extractor is whitened with a PCA-based affine layer so that its bits on vanilla images are close to i.i.d. Bernoulli(1/2) §4.1 app. A.1.
2. Fine-tune the LDM decoder for one key. Fix a key m. Copy the autoencoder's decoder D to D_m, and train only the copy for about 100 steps. The loss is the message loss between the extractor's output on the watermarked reconstruction and m, plus λ_i times a Watson-VGG perceptual distance between the watermarked reconstruction and the reconstruction of the original decoder §4.2 eq. 4. The paper reports under 500 images and about one minute on one GPU §4.2.

Generation then runs unchanged with D_m in place of D, because the diffusion process works in latent space and never touches the decoder §4.2 fig. 1.

# Pre-train the watermark encoder W_E and extractor W  [§4.1] [app. A.1]
# Input:  cover images (B, 3, H, W); messages (B, k) in {0, 1}
# Output: scalar message loss

def pretrain_step(covers, messages, encoder, extractor, attack):
    # Residual from the encoder, bounded by tanh  [§4.1]
    residual = encoder(covers, messages)
    # Scale the residual by alpha = 0.3  [§4.1] [app. A.1]
    watermarked = covers + alpha * residual
    # Random distortion from the simulation layer  [§4.1] [app. A.1]
    attacked = attack(watermarked)
    # Soft bits from the extractor, scored with BCE  [§4.1]
    soft_bits = extractor(attacked)
    return bce(sigmoid(soft_bits), messages)

# Whiten the extractor on vanilla images  [§4.1] [app. A.1]
# Input:  raw extractor outputs on about 10k vanilla images, shape (N, k)
# Output: affine map appended to the extractor

def whitening_map(vanilla_bits):
    # Mean and covariance of the raw outputs  [app. A.1]
    mean = average(vanilla_bits, axis=0)
    cov = covariance(vanilla_bits)
    # Eigendecomposition cov = U Lambda U^T  [app. A.1]
    eigenvalues, eigenvectors = eigh(cov)
    # Whitening weight and bias  [app. A.1]
    weight = diag(eigenvalues ** -0.5) @ eigenvectors.T
    bias = -weight @ mean
    return weight, bias

# Fine-tune a copy of the LDM decoder for one key  [§4.2] [eq. 4]
# Input:  key m (k bits); frozen encoder E, original decoder D, extractor W
# Output: fine-tuned decoder D_m

def finetune_decoder(key, encoder, decoder, extractor, batches):
    # Only this copy is trained  [§4.2]
    decoder_m = copy(decoder)
    # About 100 steps, batch size 4, AdamW  [§4.2]
    for images in batches:
        # Latents from the frozen encoder  [§4.2]
        latents = encoder(images)
        # Reconstruction with the watermarked and the original decoder  [§4.2]
        recon_m = decoder_m(latents)
        recon_orig = decoder(latents)
        # Message loss: BCE between the extractor output and the key  [§4.2]
        loss_m = bce(sigmoid(extractor(recon_m)), key)
        # Perceptual distance to the original reconstruction  [§4.2]
        loss_i = watson_vgg(recon_m, recon_orig)
        # Total loss with weight lambda_i  [§4.2] [eq. 4]
        loss = loss_m + lambda_i * loss_i
        update(decoder_m, loss)
    return decoder_m

# Detection: flag an image as generated by the model with key m  [§3.1]
# Input:  image x; key m (k bits); threshold tau
# Output: True if the image is flagged

def detect(image, key, tau):
    # Hard bits are the signs of the soft outputs  [§4.1]
    bits = extractor(image) > 0
    # Matching bits M(m, m')  [§3.1] [eq. 1]
    matches = count(bits == key)
    return matches >= tau

# Identification: attribute an image to one of N user keys  [§3.2]
# Input:  image x; keys (N, k); threshold tau
# Output: index of the best user, or None

def identify(image, keys, tau):
    bits = extractor(image) > 0
    scores = [count(bits == key) for key in keys]
    best = argmax(scores)
    if scores[best] < tau:
        return None
    return best

Key equations, as printed in the paper:

M(m, m') \geq \tau eq. 1

\text{FPR}(\tau) = I_{1/2}(\tau+1,\; k-\tau) eq. 2

\text{FPR}(\tau, N) = 1 - (1-\text{FPR}(\tau))^N \approx N \cdot \text{FPR}(\tau) eq. 3

\mathcal{L} = \mathcal{L}_m + \lambda_i \mathcal{L}_i eq. 4

The message loss is written as BCE between σ(m') and the key §4.1 §4.2. It is not numbered in the paper. The image perceptual loss is Watson-VGG, a variant of LPIPS §4.2.

Hyperparameters

Results

Author-reported. Setup: 512×512 images with 48-bit keys §5 tab. 1. The text-to-image rows come from the paper's internal 2.2B-parameter diffusion model §5; the other rows use the base models named in the table. Robustness is bit accuracy averaged over 10 models and 1k images per task §6.3 tab. 1.

Image quality and robustness per task (Stable Signature rows of Table 1 tab. 1):

"Crop 0.1" keeps 10% of the image area app. A.2. The combined edit is crop 50%, brightness 1.5 and JPEG 80 app. A.2.

Comparison with post-hoc watermarks on text-generated images, 512×512, 48 bits, 5k images. These are the authors' own implementations §6.4 app. A.4 tab. 1:

The authors say HiDDeN's bits are not i.i.d., so its false-positive guarantees do not hold §6.4.

Detection §5.1 fig. 3:

• Setup: one fine-tuned model per key, 1000 generated images per key, 10 random keys. TPR is measured. FPR is inferred from eq. 2 and checked empirically down to 1e-7 app. B.6.
• Unmodified images: TPR 0.99 at FPR 1e-9, with one vanilla image in 1e9 flagged.
• At FPR 1e-9: 0.84 TPR after a crop that keeps 10% of the image; 0.65 TPR after the combined edit.
• The abstract reports 90%+ accuracy for a 10% crop at FPR below 1e-6 [abstract]. The body does not state that figure directly, so check it against fig. 3.
• A passive forensic baseline (Corvi et al.) is said to need FPRs orders of magnitude higher §5.1.

Identification §5.2 fig. 4:

• Setup: N′ = 1000 fine-tuned models, each generating 100 images, so 100k watermarked images. τ is chosen for FPR 1e-6 §5.2.
• N = 10^5 users: 98% identification accuracy on unmodified images; 40% after the combined edit §5.2.
• The text says a user who generates 3 images is identified 80% of the time, without naming the edit §5.2.
• Zero false accusations at that scale, because τ is high §5.2.

Quality trade-off, weight λ_i §6.5 tab. 3:

Robustness per edit on text-to-image images, averaged over images and keys §6.3 tab. 2: none 0.99; crop 0.1 0.95; JPEG 50 0.88; combined 0.92; brightness 2.0 0.97; contrast 2.0 0.98; saturation 2.0 0.99; sharpness 2.0 0.99; median filter k=7 0.94; resize 0.7 0.91; text overlay 0.99. Appendix tables cover more tasks and edits app. B.2. The authors say bit accuracy is above 0.9 for most cases, while resizing to 0.7 and JPEG 50 sometimes fall below that app. B.2.

Robustness without edit-aware training §6.6 tab. 4. Pre-training at 128×128, 16 bits, 50 epochs:

The authors note that robustness to crops and brightness comes from the extractor, not from simulated edits during LDM fine-tuning §6.3. Adding the transformations improves rotation and JPEG but makes training harder §6.6.

Attacks (author-reported; threat models are removal and collusion) §7:

• Removal with neural auto-encoders: to push bit accuracy near 50%, distortion must be strong (PSNR below 26 dB). If the attacker knows the generator's own auto-encoder, the watermark is erased at PSNR above 29 dB §7.1 fig. 7.
• White-box removal with the extractor available: an adversarial image edit removes the watermark at a lower PSNR budget §7.1 fig. 7.
• Model purification (fine-tuning the decoder to drop the message loss): bit accuracy is hard to reduce without visible artifacts §7.2 fig. 8.
• Collusion by weight averaging: the averaged model's bit is 0 or 1 when both colluders agree on it and random when they disagree. The paper reports the same observation for image-level collusion §7.2.
• Network compression on the decoder: 8-bit and 4-bit quantization keep bit accuracy at 0.99; L1 pruning at 30% keeps 0.99 and at 60% gives 0.95 app. B.3.

Paper vs code

Code: official repository facebookresearch/stable_signature at c91217c06e953a4eaf270ef9de4a3939bcea3366, linked from the paper's conclusion §8. The repository is licensed CC BY-NC 4.0 (non-commercial) LICENSE. This card cites it by path and line for facts only and quotes no code.

• Activation in the watermark blocks. The paper describes Conv-BN-ReLU blocks app. A.1. The code uses GELU in ConvBNRelu models.py.
• Fine-tuning message loss. The paper uses BCE between σ(m') and the key, with no temperature §4.2. The code applies the BCE to the extractor output multiplied by 10 finetune_ldm_decoder.py.
• λ_i default. The paper sets λ_i = 2.0 by default §4.2. The code default is 0.2 finetune_ldm_decoder.py. The code's perceptual loss is a sum over the batch divided by batch size finetune_ldm_decoder.py. The paper does not say how its perceptual loss is normalized, so the same λ_i value may act differently. The paper's trade-off table reaches about 30 dB between λ_i = 0.2 (29.7 dB) and λ_i = 0.4 (30.6 dB) tab. 3, so it is unclear which λ_i produced the 30 dB in Table 1 tab. 1.
• Fine-tuning learning rate. The paper's schedule ends at 1e-4 §4.2. The code default is AdamW with lr 5e-4, and the cosine schedule with 20-step warmup and decay to 1e-6 uses that value as its peak finetune_ldm_decoder.py utils.py.
• Whitening scale. The paper gives weight Λ^(-1/2) U^T and bias −Λ^(-1/2) U^T μ app. A.1. The code multiplies both by √k finetune_ldm_decoder.py. Decoding by sign is unaffected, but the scale changes the logits that enter the BCE loss.
• Pre-training settings. The paper uses Lamb, batch size 64 per GPU and peak learning rate 1e-2 app. A.1. The repository's reproduction command uses Lamb with lr 2e-2, batch size 16 and --scaling_w 0.3 (matching α = 0.3) README.md. The README warns that the learning rate is rescaled with the world size README.md. The script's own default optimizer is Adam with 400 epochs main.py.
• Simulation layer. The paper describes crop, resize or identity, plus JPEG with probability 0.5 app. A.1. The code's combined augmentation also includes horizontal flip, Gaussian blur, rotations of ±10° and ±90°, and colour jitter data_augmentation.py. The reproduction command sets JPEG probability to 1.0 and turns off blur, rotation and colour jitter README.md. The README says JPEG robustness is somewhat worse than the paper's because the augmentation implementation differs README.md.
• Input normalization for the extractor. The paper does not say how images are normalized before the extractor. The code maps the autoencoder's [-1, 1] output to [0, 1], then applies ImageNet mean and standard deviation finetune_ldm_decoder.py utils_img.py.
• Detection rule versus FPR formula. The rule flags an image when M ≥ τ §3.1 eq. 1. The FPR formula uses P(M > τ) §3.1 eq. 2. The released notebook computes the FPR as the binomial CDF at n − τ, which is P(M ≥ τ) decoding.ipynb. The paper's own rule and formula disagree by one matching bit.
• Center crop meaning. The paper's crop parameter is an area ratio app. A.2. The code takes the square root of the area ratio for each side utils_img.py, which agrees with the paper.
• Decoder depth. The paper says the extractor has 7 blocks plus a block with k filters app. A.1. The code's default decoder depth is 8 and produces one first conv block, seven more, and a final block with num_bits filters models.py. Whether the paper's "7" counts the first block is not stated.

Pitfalls

• The headline numbers are not reproducible from the public code alone. The text-to-image experiments use an internal 2.2B-parameter model trained on 330M licensed image-text pairs §5. The repository expects the user to build Stable Diffusion configs and checkpoints from external sources README.md. Its released material is an extractor and an example decoder, plus a command for the extractor README.md. The paper's attack section uses the KL-f8 auto-encoder of LDM §7.1 app. A.5.
• The paper's own FPR rule is off by one. Use the rule consistently. Either matches >= tau with P(M >= tau), or matches > tau with eq. 2. Note that at τ = 0 the rule flags every image, while eq. 2 gives 1 − 2^−k, which is almost 1 but not exactly.
• Bit-independence is approximate. The FPR is computed under an i.i.d. Bernoulli(1/2) assumption. Whitening makes the bits closer to that, but the covariance still differs from a simulated Bernoulli one. Empirically the FPR is slightly higher than the empirical one below 1e-6, which is conservative app. B.5 app. B.6.
• Identification thresholds must be recomputed as N grows. The global FPR is about N times the single-key FPR eq. 3. The paper's τ moves from 41 to 44 when N goes from 1 to 1000 at FPR 1e-6 §5.2. Accuracy falls with N because τ must rise §5.2.
• Robustness comes from the extractor. The paper says robustness to crops and brightness appears even when the LDM is fine-tuned without any augmentation §6.3. Robustness to rotation and JPEG needs simulated edits during extractor pre-training §6.6.
• The quality–robustness trade-off is set by λ_i. Higher λ_i gives images closer to the original decoder but weaker bit accuracy under the combined edit §6.5 tab. 3. The paper's λ_i = 2.0 default is not in its trade-off range, and the code default is 0.2 [Paper vs code].
• PSNR is measured against the original generator, not the original photo. The paper's PSNR compares watermarked and unwatermarked generations from the same latents §6.2. The paper says this makes the reported ~30 dB look low, because a real user would only ever see the watermarked image §6.2. The code logs PSNR the same way, between the watermarked and the original decoder's outputs finetune_ldm_decoder.py.
• Inpainting of masked regions is the weakest case. The paper says bit accuracy stays above 0.9 except when only the masked region is inpainted §6.3. The masked-only row in Table 1 gives 0.76 to 0.89 across the edits tab. 1.
• Collusion and purification results are empirical. The paper observes that the marking assumption holds for averaged models, although the process was not designed for it §7.2. For purification, artifacts appear before bit accuracy drops much §7.2 fig. 8.
• Cost. The fine-tuning step is cheap (about one minute on one GPU), but the paper's experiments used roughly 50,000 GPU-hours in total §8. Pre-training the extractor took about a day on eight GPUs app. A.1.
• Input normalization is unstated. The paper does not say how decoded images are normalized before the extractor. The code's choice is listed under Paper vs code, so a reimplementation may get different bit accuracy.

Lineage

• Builds on: HiDDeN, Zhu et al. (2018), arXiv:1807.09937 README.md. The paper keeps the encoder/decoder design but drops the adversarial and perceptual losses from the extractor's training, and adds tanh and scaling §4.1. The image perceptual loss is Watson-VGG from Czolbe et al. (2020) §4.2. The LDM is Rombach et al. (2022) §4.2.
• Compared against: DCT-DWT (Cox et al., 2007), SSL Watermark (Fernandez et al., 2022), FNNS (Kishore et al., 2021) and HiDDeN, all as post-generation methods §6.4. Removal attacks use the auto-encoders of Ballé et al. (2018), Cheng et al. (2020), Esser et al. (2021) and Rombach et al. (2022) app. A.5.
• Superseded by: not known from the wiki.