Secure and Efficient Watermarking for Latent Diffusion Models in Model Distribution Scenarios

1 Introduction

Substantial progress in latent diffusion models (LDMs) [Croitoru et al., 2023] have significantly enhanced the quality of image generation, which presents observable abilities in producing a wide scope of creative visuals, e.g., artistic works and realistic depictions. To safeguard the copyright of generative models [Gowal and Kohli, 2023] and prevent their misuse [Barrett, 2023], watermarking is one avenue for detecting generated content and tracing its source. Recently, there is a compelling trend for model producers to distribute LDMs to numberous model users by model sharing [Donahue and Kleinberg, 2021], disclosure [Azcoitia and Laoutaris, 2022], and trading [Pei et al., 2023]. Since a large amount of model users are granted with model architecture access and finetuning permission in these model distribution scenarios, effective watermark injection and robust watermark verification becomes more challenging compared with local model usage.

In order to support applications in model distribution scenarios, LDM watermarks need to accommodate serveral key constrains. (1) Since model networks and parameters will be distributed for personalized usage, it is possible for model users to bypass the watermark injection by model modifications. Therefore, security mechanism is indispensable to avoid watermark evading. (2) When the model is distributed to massive users, the watermark has to guarantee low injection time cost while spanning distinctive information for a large amount of user verification. (3) Due to the higher model access permission and larger user scale in model distribution scenarios, untrustworthy users pose a greater threat of model theft and leakage, making it essential for watermarking methods to ensure robustness against diverse adversaries.

A traditional watermarking solution is post-processing watermark that embeds watermarks after image generation (Figure 1(a)). However, untrustworthy users can remove posthoc watermark trivially. On the other hand, in-processing watermarks inject messages into the image generation process, which contain three category solutions based on modification ways. Whole model modifications [Zhao et al., 2023b; Feng et al., 2024] embed watermarks by training the entire generative models (Figure 1(b)), which require substantial training resources and thus inefficient in terms of model distribution senarios. Partial model modifications [Fernandez et al., 2023; Xiong et al., 2023] merely fine-tune the decoder of the LDMs (Figure 1(c)). However, these methods are vulnerable to multiple attacks [An et al., 2024] such as reconstructive attack [Zhao et al., 2023a] and adaptive adversarial sample attack [Jiang et al., 2023]. Random seed modifications [Wen et al., 2024; Yang et al., 2024; Ci et al., 2025] inject watermarks into the initial latent variable of LDMs which are time-efficient without model finetuning and robust against diverse attacks. But in model distribution scenarios, the untrustworthy user can easily change the initial latent vector to circumvent watermark injection.

In this work as shown in Figure 1(d), we extend the application of LDM watermarking to model distribution scenarios and propose a secure and efficient watermarking method, named as DistriMark. Considering the watermark injection efficiency, DistriMark is based on the random seed modification schema without any model fine-tuning. To avoid model user bypassing watermark injection, we propose a watermark-network controller module as a security mechanism, which establishes binding association between the VAE network in LDMs and the watermarked initial latent variable. In this way, LDMs can generate expected content only when the watermark is mandatory injected. To reduce the time cost of training the watermark-network controller module, we decouple the watermark injection and the security mechanism, ensuring that fine-tuning VAE only accomplishes the security mechanism without the burden of learning watermark patterns. Furthermore, we propose watermark generation module to transform the watermark into a watermarkspecific distribution and obtain a watermarked latent variable through sampling strategy. For watermark verification, the latent variable obtained by diffusion inversion is compared to the watermark distribution instead of a fixed watermarked variable. This watermark generation and verification strategies not only increases the security of plaintext watermarks, but also makes up the errors caused by diffusion inversion and enhance the robustness against various watermark attacks.

The main contributions are summarized as follows: (1) We propose new security mechanism to prevent watermark leakage and watermark escape in the model distribution scenarios by pseudo-random latent variable transformation and VAEbased fine-tuning strategy. We consider watermark randomness and watermark-model association as two constraints for enhancing watermarking security, which sheds new light on the real-world application of diffusion model watermarking. (2) We propose a novel model distribution scenario-oriented watermarking schema for LDMs. By injecting multi-bit watermarks into the initial latent variables and fixing the verification errors via watermark distribution verification and adversarial training strategy, our schema achieves both robustness and flexibility compared with existing fine-tuning and random seed-based watermarks. (3) DistriMark shows superior performance on effectiveness and robustness compared with existing six baselines over ten image processing attacks, challenging adaptive adversarial sample attacks and reconstructive attacks. DistriMark is more secure against watermark escape and leakage compared with existing random seed modification watermarks in the distribution scenarios.

2 RelatedWork

Diffusion Models has demonstrated prominent performance in image generation [Dhariwal and Nichol, 2021] with the support of methodologies [Song et al., 2020b] and sampling techniques [Song et al., 2020a]. Latent diffusion models optimize images in the latent space of pre-trained VAEs, further accelerating the practical applications of diffusion models while also raising concerns about potential abuse and intellectual property of models. The immense cost of training a diffusion model, which requires hundreds of GPU-days [Rombach et al., 2022], makes copyright protection for diffusion models crucial, especially when the model architecture and weights are distributed to users for deployment. We focus on the security and efficiency issues of model watermarking in distribution scenarios.

Watermarking for Latent Diffusion Models is primarily aimed at tracing the origins of generated images of the latent diffusion model. WDM [Zhao et al., 2023b] trains an autoencoder to stamp a watermark on all training data before re-training the generator from scratch. However, this approach suffers from inefficiencies in terms of computational resources and time. Stable Signature [Fernandez et al., 2023] and FSwatermark [Xiong et al., 2023] fine-tune VAE-Decoder to ensure that all generated images contain the watermark. However, these approaches are not resilient to diverse threats. Tree-ring [Wen et al., 2023] and ZoDiac [Zhang et al., 2024] propose random seed modification watermarks which show significant advantages in dealing with various processing attacks [An et al., 2024]. However, these methods lack secure mechanisms to guarantee watermark embedding in model distribution scenarios.

Model Watermarking Attacks on diffusion model watermarking primarily occur at two levels: image and model. At the image level, attacks such as image processing attacks, adaptive adversarial sample attacks [Jiang et al., 2023], and reconstruction attacks [Zhao et al., 2023a] are included. At the model level, attacks include techniques such as purification and model collison. Model purification will significantly reduce the detection accuracy of whole model modification watermark and partial model modification watermark. Model collison will deceive watermark detection. We propose watermark-network controller to avoid watermark verification issues related to model-level attacks and ensures image-level robustness by secure watermark distribution.

3 Methodology

3.1 Framework of DistriMark

In this work, we extend to achieve the security and embedding efficiency of watermarks in the model distribution scenarios. Our Distrimark embed the watermark into the latent variables of the diffusion model and enforce the mandatory embedding of the watermark whenever the model is utilized by leveraging the watermark-network controller. To guarantee the security of watermark distribution and maintain the unpredictability and fidelity of watermark, we propose a novel watermark distribution method Secure Latent Watermark Distribution. This method establishes a unified representation of latent variables and watermark information as shown in Figure 2. The watermark region follows a specific distribution, from which watermarked latent variables are sampled. The variability in latent variables across different outputs increases randomness and unpredictability, which ensures the security of watermark distribution. To safeguard the security of watermark embedding, we introduce Watermark-Network Controller, a security mechanism integrated into latent diffusion model components which binds the variational autoencoder with watermarked latent variables to prevent users from evading the watermark embedding process. This module binds the VAE-Decoder with watermarked latent variables through skip connections.

3.2 Watermark-Network Controller

To enforce the embedding of watermarks during model usage, the watermark-network controller directs image generation by using watermarked initial latent as control signals. watermark-network controller connects the watermarked initial latent variables and relevant components of the VAEDecoder through skip connections. Through fine-tuning the VAE-Decoder, images corresponding to the watermarked latent variables consistent with the original model, while corresponding to random latent variables are transformed to random noise. In the implementation of skip-connection, we design a network association to bind the watermarked initial variable to the intermediate layer variables.

The loss function employs the LPIPS loss and L2 distance between images, denoted as L1 and L2, respectively. L2(Do(z),Dv(z)) = ||Do(z) -Dv(z)||22 . D(·) denotes the decoding process of the variational autoencoder Do and Dv denote the original and the fine-tuned VAE-Decoder with skip connections respectively. When the VAE-Decoder is connected to the initial latent variable, the loss function is:

where Zr denotes the random noise. The factor λv is a constant. When the VAE-Decoder is connected to the random latent variable, the loss function is:

To prevent pixels with smaller values from being excessively altered, We calculate the difference across multiple channels between the output images of the original model and the fine-tuned model as the loss:

where θ indicates whether the initial latent variables are from the message encoder. ε, δ is the balancing weight. The overall loss for this step is as follows:

3.3 Secure Latent-Watermark Distribution

We assume a series of deterministic functions f(z; θ) parameterized by a vector θ. When θ is fixed and z ∼ N(1, 0), f(z; θ) can generate latent variables that conform to a specific distribution. Specially, the encoder outputs the mean vector and variance vector to simulate the deterministic function f(z; θ) and generates the initial latent variables through reparameterization.

The watermarked latent variables are put into the message decoder directly to train them in the self-supervision paradigm. The message loss is the Binary Cross Entropy (BCE) between m and the sigmoid σ(m′ ):

5

Figure 3: Generated image comparison under the security mechanism. Images in each sample from left to right are Watermarked Initial Latent Variables (WIL for short) without VAE-based fine-tuning (fine-tuning for short), WIL with fine-tuning, non-WIL without fine-tuning, and non-WIL with fine-tuning, representatively.

Since the training samples of the diffusion model are generated by progressively adding noise until they conform to a standard normal distribution, during the inference stage, the message encoder will output initial latent variables that follow the same distribution. The Kullback-Leibler (KL) divergence between the initial latent variables and the standard normal distribution is utilized as the loss function. The output follows a normal distribution, denoted as q(z) ∼ N(µ1, σ2 1 ) and the standard normal distribution is denoted as p(z) ∼ N(0, 1). The distribution loss is as follows:

3.4 Watermark Robustness Enhancement Module

Watermarking Verification. Watermark extraction involves diffusion inversion, an approximate process for obtaining initial hidden variables from generated images. Diffusion inversion [Dhariwal and Nichol, 2021] algorithmically retrieves the initial latent variables from images generated by a diffusion model. xt represents the image at the timestep t. Based on the assumption xt-1-xt ≈ xt+1-xt , diffusion inversion of the Denoising Diffusion Implicit Model (DDIM) [Song et al., 2020a] is formalized as follows:

where ¯α is the parameter of the diffusion model. t denotes the denoising timestep. ϵθ(xt) is the estimated noise for timestep t. ˆ x0 represents the prediction of the image at the current timestep and is defined as:

To mitigate the effects of diffusion inversion and raise the robustness of image processing, we introduce the watermark robustness enhancement module which employs adversarial training to raise performance of the message decoder. The loss function is binary cross entropy between the message m and the sigmoid δ(m′ ) which is the same as Equation 5.

Attack Simulation for Adversarial Training. Various attacks are common in practical image usage. Therefore, during the training process, we deploy an attack layer to watermarked images before employing a watermark extraction algorithm. This attack layer encompasses six common types of attacks: blur, Gaussian noise, brightness adjustment, contrast adjustment, saturation adjustment, and JPEG compression. To remain the differentiable of attack during training, we employ the differentiable simulation method to perform JPEG attack [Zhu et al., 2018].

4 Experiments

Implementation details. In this paper, we focus on textto-image generation, hence we utilized Stable Diffusion-v2 [Rombach et al., 2022]. The number of inference steps is 25 for both generation and detection process. Following the settings of existing works [Wen et al., 2024], we employ the prompts from StableDiffsionDB [Wang et al., 2022] with guidance scale of 5 during inference and an empty prompt during DDIM inversion. We utilize AdamW with a learning rate of 5 × 10-4 and weight decay of 0.01 during finetuning. All experiments are conducted on a single NVIDIA L40.

Watermarking baselines. We select six typical baselines: three official watermark of Stable Diffusion [Rombach et al., 2022] for cloud services called DwtDct [Cox et al., 2007], DwtDctSvd [Cox et al., 2007], and RivaGAN [Zhang et al., 2019], two multi-bit watermarking methods named FSwatermark [Xiong et al., 2023] and Stable Signature [Fernandez et al., 2023], and a fine-tuning-free semantic watermarking method called Tree-Ring [Wen et al., 2024].

Evaluation metrics. We measure the detection performance by the true positive rate (TPR) when the false positive rate (FPR) is at 1%. We measure the traceability performance by the bit accuracy. To measure the image generation quality, we compute the Peak Signalto-Noise Ratio (PSNR) [Hore and Ziou, 2010] and Structural Similarity score (SSIM) [Wang et al., 2004] for image distortion evaluation, the Fr´echet Inception Distance (FID) [Heusel et al., 2017] and CLIP score [Radford et al., 2021] for image diversity and semantic evaluation, and Natural Image Quality Evaluator (NIQE) [Mittal et al., 2012] and Perceptual Image Quality Evaluator (PIQE) [Venkatanath et al., 2015] for image quality evaluation.

4.1 Watermark Security against Escape

In order to make the watermark flexible for distributing the LDMs to a large number of model users with strong robustness, we leverage the semantic watermarking framework to inject the watermark message into the latent variable m. However, the model users can easily escape the watermark by replacing m with other random latent variables to obtain the non-watermarked generated images. To tackle this issue, we design the security mechanism to decrease the generated image quality when the model user escape the watermark. Representative non-watermarked images under the security mechanism are shown in Figure 3. Embedding watermark does not significantly affect the image quality metrics NIQE and PIQE, or semantic quality metric CLIP. As the quality of unauthorized images decreases further, the quality of watermarked images also slightly deteriorates.

4.2 Watermark Performance Comparison

We compare the performance of our method with existing six typical baselines over two tasks: (1) Detection. We consider all compared methods as single-bit watermarks with a unified watermark. We set the FPR to be 1% and test the TPR on 1, 000 watermarked images. (2) Traceability. Each compared method, excepting for the single-bit watermark TreeRing, serves as a multi-bit watermark. In our experiments, we assume that there are 1, 000 model users, each of them requires one watermark for model tracing. Each user generates 10 images, resulting in a dataset of 10, 000 watermarked images. During test, if an image contains a watermark, we then calculate the number of matched bits (Bit Accuracy) with the watermark of each user. The user with the highest Bit Accuracy is considered the traced user and verified. The comparison results are shown in Table 1. Our watermarking achieves strong robustness and significantly outperforms baselines in both tasks. In terms of bit accuracy, it surpasses the bestperforming baseline by approximately 6.7%. This can be attributed to the extensive diffusion of the watermark throughout the entire latent space, establishing a profound binding between the watermark and the image semantics.

4.3 Robustness against Image-Level Attacks

We examine watermark's robustness against three typical kinds of adversarial attacks, including image processing attack [Song et al., 2010] for transforming generated images, adaptive adversarial sample attacks [Jiang et al., 2023] for disturbing watermark verification, and reconstructive attacks [Ball´e et al., 2018; Cheng et al., 2020; Zhao et al., 2023a] for re-generating non-watermarked images.

Image Processing Attack. We select ten representative types of image-level noise shown in Table 2. Please refer to the Supplementary Materials for detailed parameter settings.

Adaptive Adversarial Sample Attack. To further enhance the attack ability, we assume that attackers can query a blackbox watermark verification interface and conduct querybased black-box attack [Jiang et al., 2023]. By iterative querying the verification interface this attack compute optimal perturbations that progressively bring the watermark-free initial image closer to the original image.

Reconstructive Attack. The core idea of the reconstructive attack is to add random noise to destroy the watermark and then reconstruct the image. We utilize the implementation of the paper [Zhao et al., 2023a] with denoising steps of 60.

Main Results. For each image processing attack, we report the average bit accuracy in Table 2. We see that our DistriMark watermark is indeed robust across all the transformations with the bit accuracy all above 0.9. DistriMark remarkably outperforms existing multi-bit watermarks on the average performance with more than 0.25 boost compared with the state-of-the-art (SoTA) results. Compared to the SoTA finetuning methods FSWatermark and Stable Signature, Distrimark has significant advantages in image resize, VAE-based compression algorithm, and reconstructive attack. Adaptive Samples Attack utilizes the same evaluation metric ℓ∞-norm as described in the paper [Jiang et al., 2023]. Figure 4 and Figure 5 display images and related parameters generated by Reconstructive Attack and Adaptive Adversarial Samples Attack under the same parameter conditions. Under both types of attacks, DistriMark demonstrates better robustness than the other two methods. The robustness stems from the watermark being embedded at the semantic level within the image, so more extensive attacks are needed at the pixel level to remove the watermark. Please refer to the Supplementary Materials for more results.

4.4 Robustness against Model-level Attacks

Consider the model-level attacks of the model for both single users and multiple users, we have included two types of model-level attacks: model purification and model collusion. Model purification. The adversary fine-tunes the Variational autoencoder to circumvent watermark embedding through the same training mode as Section 3. This involves removing the message loss Lm, and shifting the focus to the perceptual loss Lw between the original image and the one reconstructed by the LDM auto-encoder.

Model Collusion. We mainly considered two types of collusion attacks: Averaging weights and Finetuning. (1) Averaging weights. User(i) and User(j) can average their weights like Model Soup [Wortsman et al., 2022] to creat a new model to deceive identification. (2) Finetuning. Another form of collusion attack is when the user B generates a large number of watermarked latent variables and watermarked generated images, and fine-tunes the VAE of A so that A can use B's watermark latent variables to generate images.

Main results. Figure 6 shows the results of model purification attack. As for model purification, when the bit accuracy decreases, the image quality also declines. It is difficult to significantly reduce the bit accuracy without affecting the image quality. As for model collusion, because of the security mechanism, parameter averaging will cause a significant drop in image quality. As for finetuning, this could pose a threat of identity spoofing. However, this still does not break the security mechanism.

4.5 Watermarked Image Quality

Besides the qualitative examples of how the watermarked images are not sensitive for the human eyes to distinguish (see Figure 3), we further present quantitative evaluation of images generated by existing watermarking methods in Table 3. The results show that no matter the qualitative metrics, our DistriMark achieves comparable performance with existing works in the model distribution scenarios. For DistriMark, the initial watermark is only manifested in the selection of the initial latent variables, hence the image quality is consistent with the diffusion model inference. Although DistriMark shows comparable quality compared to the methods in the model distribution scenario, the method involves a trade-off between the quality of unauthorized images and watermarked images when fine-tuning the VAE-Decoder, which results in lower image quality compared to post-processing methods.

4.6 Ablation Study

Skip connection selection. The way latent variables are connected to the VAE-Decoder impacts how the VAE transforms images from the latent space to the pixel space. In Table 4, three methods were tested: no connection, single connection, and multi-level connection. Multi-level connections improve feature learning and enhance image quality.

Guidance scales. Larger guidance scales result in more faithful of the generated image adherence to prompts. Following existing works [Wen et al., 2024], we cover the range of 0 to 20. In Figure 7 (left), the watermark remains robust and reliable even at a guidance scale of 18.

Number of the inversion step. The inference step is often unknown in practice, which introduces a mismatch with the inversion step. From Figure 7 (right) we can see that the number of inference steps does not significantly affect the accuracy of inversion which is beneficial in practice.

5 Conclusion

We propose a novel distribution scenario-oriented watermarking schema for diffusion models and a new security mechanism to prevent watermark leakage and watermark escape in the model distribution scenarios. We separate the watermark injection from the security mechanism, ensuring that fine-tuning the VAE focuses solely on the security mechanism without the added task of learning watermark patterns. In the future, our research directions will include adversarial methods against forge attack [Lukas et al., 2023], security mechanism with higher image quality and more consideration about model distribution scenarios.

Acknowledgments

This work is supported by the National Natural Science Foundation of China (Grant No.s U24B20146, 62372044), and Beijing Municipal Science and Technology Commission Project (Z241100009124008).

Contribution Statement

This work was a collaborative effort by all contributing authors. Liangqi Lei and Keke Gai made equal contribution to this study and are designated as co-first authors. Jing Yu serving as the corresponding author, is responsible for all communications related to this manuscript.