AI Research Portal

This paper introduces MSAT-LDM, a novel framework for transferable high-fidelity watermarking in Latent Diffusion Models (LDM). The core idea is to integrate a modular watermark architecture with a Self-Augmented Training (SAT) strategy. SAT leverages an internally generated "free generation" distribution to train the watermark module, aligning training and testing phases without external data, while the modular architecture enables plug-and-play adaptation. The main empirical claim is that MSAT-LDM achieves robust watermarking, significantly improves watermarked image quality, and exhibits strong transfer performance without requiring external training data.

Existing training-based watermarking methods for AI-generated images face several challenges:

Generalization across diverse prompts: They often struggle to generalize effectively when applied to images generated from a wide range of prompts.
Image quality degradation: These methods can introduce visible artifacts, thereby degrading the quality of the watermarked images.

MSAT-LDM is proposed as a novel framework that integrates a modular watermark architecture with a Self-Augmented Training (SAT) strategy to achieve high fidelity, robustness, and efficient transferability watermarking for Latent Diffusion Models (LDM).
The paper demonstrates that self-augmented training significantly improves watermark fidelity and robustness by aligning training and testing distributions without external data, and provides theoretical insights into its generalization capability.
A pluggable watermark module is introduced that enables rapid few-shot adaptation to a variety of fine-tuned LDMs and LoRA-augmented models, significantly boosting practical applicability.
Extensive experiments using AI-generated prompts across diverse semantic categories show that MSAT-LDM not only achieves substantial improvements in image quality but also maintains competitive robustness and delivers superior transfer learning performance.

MSAT-LDM addresses the limitations of previous watermarking methods by integrating two key strategies: Self-Augmented Training (SAT) and a Modular Watermark Architecture.

Self-Augmented Training (SAT) The paper identifies an inconsistency between the training and testing phases of previous LDM watermarking methods. During testing, the pipeline is prompt → UNet → VAE decoder to generate watermarked images. Previous training methods, however, used image → VAE encoder → VAE decoder. MSAT-LDM aligns these phases by training the watermark module using an internally generated "free generation" distribution, which means generating images using an empty prompt. This eliminates the need for external datasets and reduces the distribution gap between training and testing.

Modular Watermark Architecture The watermark module is designed as a plug-and-play component, allowing it to be independently fine-tuned and adapted to different fine-tuned backbones or LoRA-enhanced variants with minimal overhead. This modularity is achieved by preserving the original VAE decoder parameters and adapting FSW's message injection strategy. Specific intermediate layers (input convolutional layer, mid block, and the first three upsampling blocks) are cloned and integrated into a message processor. A spatial transformer network is prepended to the message extractor (Tm) to enhance robustness against geometric distortions.

Objective The overall training pipeline involves a message processor plugged onto the VAE decoder. The objective is to minimize a unified loss function that combines message loss and image loss.

Key Equations The overall loss for MSAT-LDM is given by: E EE [lm (m, m', Am) + lı (Ιο, Ιω, λι)] xprompt ~ µp,ε~~µε,m~~µm

When using self-augmented training with an empty prompt, this simplifies to: EE [lm (m, m', Am) + lı (Ιο, Ιω, λι)] ε~~µε,m~~µm

The unified loss function l(h, z, m) is defined as: l(h, z, m) = lm(Tm(Dm(z, m)), m, Am) + lı (Dm (z, m), D(z), λι)

The image loss l₁(I₀, Iω, λ₁) is a combination of functions: l1 (Ιο, Ιω, λ1) = λI MSE(I₀, Iω) + λP LPIPS(I₀, Iω) + BAL(I₀, Iω)

The BAL (Balanced Attack Loss) term is defined as: BAL (I, I') = 1 / (c · h · w) Σ(i=1 to C) Σ(j=1 to h) Σ(k=1 to w) |I(i,j,k) - I'(i,j,k)| / (I(i,j,k) + 1)

The message loss lm(m, m', λm) is simply: lm (m, m', Am) = Am MSE(m, m')

Architecture:
- Builds upon FSW's framework for robust message embedding.
- Modular Watermark Architecture: Preserves original VAE decoder parameters. Clones specific intermediate layers (input convolutional layer, mid block, and the first three upsampling blocks) where FSW typically injects messages, integrating them into a message processor.
- Perspective Robustness: A spatial transformer network (Jaderberg et al. 2015) is prepended to the message extractor (Tm) to handle geometric distortions.
Loss/Objective:
- Unified loss function combines message loss (lm) and image loss (l₁).
- Image loss l₁ includes MSE, LPIPS perceptual loss, and BAL (Balanced Attack Loss).
- Message loss lm uses Mean Squared Error (MSE).
Optimization:
- Optimizer: AdamW optimizer (Loshchilov and Hutter 2019).
- Learning Rate: 2 × 10⁻⁵.
Augmentation:
- An attack layer with seven types of watermark attacks is used during training.
- At each training step, the watermarked image undergoes the attack layer with a random attack intensity.
- Attack intensity is scaled by a decay coefficient γ₀, which gradually increases from 0 to 1 to assist in convergence.
Data Handling:
- Self-Augmented Training: Utilizes an internally generated "free generation" distribution by using an empty prompt.
- Sample Size for Free Generation: 30k samples.
Implementation Details:
- SD Model: Stable Diffusion v1.5 (sd-legacy 2024) from Hugging Face.
- Image Size: 512 × 512 pixels.
- Latent Space Dimension: 4 × 64 × 64.
- Sampling: DDPM (Ho, Jain, and Abbeel 2020) sampling with 30 steps.
- Test Prompts: 1K AI-Generated Prompts, generated by GPT-4 across 10 broad categories (100 diverse prompts per category).
- Guidance Scale: 7.5 for SD image generation.
- Message Bits: 100 bits for main experiments, except for HiDDeN (48-bit) and Stable Signature (48-bit).
- Runs: All reported metrics are averaged over three independent runs, each initialized with a different random seed.

Datasets:
- AI-Generated Prompts (1K prompts generated by GPT-4).
Metrics:
- Image Quality: PSNR↑ (Peak Signal-to-Noise Ratio), SSIM↑ (Structural Similarity Index), FID↓ (Fréchet Inception Distance).
- Watermark Robustness: Bit accuracy↑ (percentage of correctly decoded bits).
Baselines:
- Post-hoc: DwtDctSvd, HiDDeN, StegaStamp.
- Diffusion-native: Stable Signature, FSW, Gaussian Shading.
Numeric Results (MSAT-LDM vs. Baselines, Table 1):
- PSNR: MSAT-LDM: 40.6 (best). Next best: DwtDctSvd: 37.9.
- SSIM: MSAT-LDM: 0.98 (best). Tied with DwtDctSvd: 0.98.
- FID: MSAT-LDM: 2.35 (best). Next best: DwtDctSvd: 7.04.
- Bit Accuracy (None attack): MSAT-LDM: 1.00 (best). Tied with StegaStamp, FSW, Gaussian Shading.
- Bit Accuracy (Adv. attack): MSAT-LDM: 0.99 (best). Next best: StegaStamp: 1.00 (tied), Gaussian Shading: 0.95.
- Bit Accuracy (Average across 7 attacks): MSAT-LDM consistently achieves high bit accuracy (e.g., 0.97 for Gaussian blur, 0.99 for Gaussian noise, 1.00 for brightness, 1.00 for contrast, 0.99 for desaturation, 0.98 for perspective warp, 0.96 for JPEG compression).
Transferability Results (Table 2):
- MSAT-LDM shows strong transferability across fine-tuned Stable Diffusion v1.5 models (GuoFeng, GuoFeng+ShuiMo, Realistic, Realistic+Detail) with high PSNR, SSIM, FID, and Adv. Acc. For example, on GuoFeng, MSAT-LDM achieves PSNR 36.6, SSIM 0.97, FID 4.88, and Adv. Acc 0.99.
Training Data Strategies (Figure 2):
- Training on real images (LI) consistently yields inferior performance compared to strategies using LDM-generated data.
- "Free" generation (using empty prompt) consistently achieves the highest image fidelity (FID) and its robustness improves rapidly with training size, matching or surpassing "LP" (LDM-generated images from LAION-400M prompts) at sufficient scale (~30k samples).
- Wasserstein distance (W₁) between "Free" training data and "AI-Gen Test" distribution is 504.4, which is smaller than W₁ between "LI" and "AI-Gen Test" (911.4).
Runtime/Training Cost: Not stated in the paper.

The prompt design used for generating test data may still harbor biases, necessitating future ablation studies to evaluate their impact comprehensively.
The assumption that the distribution of latent representations generated by all prompts via conditional sampling is equivalent to the one without a specific prompt (free generation) may not hold in practical scenarios due to model's parameter bias and training data limitations.

Method evidence: "MSAT-LDM integrates two key components: (1) Self-Augmented Training (SAT) leverages an internally generated "free generation" distribution to train the watermark module, aligning the training and testing phases without relying on external data." (section: Abstract)
Result evidence: "Extensive experiments show that MSAT-LDM achieves robust watermarking, significantly improves the quality of watermarked images across diverse prompts, and exhibits strong transfer performance—all without the need for external training data." (section: Abstract)
Limitation evidence: "Our prompt design may still harbor biases, necessitating future ablation studies to evaluate their impact comprehensively." (section: 6 Conclusion and Future Work)

MSAT-LDM: Toward Transferable High-Fidelity Watermarking for Latent Diffusion Model via Modular Self-Augmented Training

Summary

Problem

Key Contributions

Method

Technical Details

Results

Weaknesses / Limitations

Evidence

Additional Figures