OptMark: Robust Multi-bit Diffusion Watermarking via Inference Time Optimization

TL;DR

OptMark introduces an inference time optimization method for embedding multi-bit watermarks into diffusion-generated images, enhancing robustness against various attacks while maintaining image quality and imperceptibility.

Contribution

The paper presents a novel optimization-based multi-bit watermarking technique that embeds watermarks into diffusion model latents with improved robustness and reduced memory usage.

Findings

Achieves invisible multi-bit watermarks with high robustness.

Resists generative, geometric, and editing attacks effectively.

Reduces memory consumption from linear to constant during optimization.

Abstract

Watermarking diffusion-generated images is crucial for copyright protection and user tracking. However, current diffusion watermarking methods face significant limitations: zero-bit watermarking systems lack the capacity for large-scale user tracking, while multi-bit methods are highly sensitive to certain image transformations or generative attacks, resulting in a lack of comprehensive robustness. In this paper, we propose OptMark, an optimization-based approach that embeds a robust multi-bit watermark into the intermediate latents of the diffusion denoising process. OptMark strategically inserts a structural watermark early to resist generative attacks and a detail watermark late to withstand image transformations, with tailored regularization terms to preserve image quality and ensure imperceptibility. To address the challenge of memory consumption growing linearly with the number of denoising steps during optimization, OptMark incorporates adjoint gradient methods, reducing memory usage from O(N) to O(1). Experimental results demonstrate that OptMark achieves invisible multi-bit watermarking while ensuring robust resilience against valuemetric transformations, geometric transformations, editing, and regeneration attacks.

Tables

Table 1. Table 1: Performance of multi-bit different watermarking methods under various attacks on DiffusionDB  (Gustavosta 2023 ) . “Average” indicates calculating the average score across cases under sixteen different attacks and the no-attack (“None”). “*” indicates that Gaussian Shading  (Yang et al. 2024b ) and RivaGAN  (Zhang et al. 2019 ) embed 64-bit and 32-bit hidden messages respectively, whereas all other methods are compared under the condition of embedding 48-bit messages. The underline indicates poor robust performance with Bit Acc. < 0.75 <0.75 and TPR < 0.5 <0.5 .

	Various Attack
	None		Geometric		Valuemetric		Editing		Regeneration		Average
Method	Bit Acc.	TPR	Bit Acc.	TPR	Bit Acc.	TPR	Bit Acc.	TPR	Bit Acc.	TPR	Bit Acc.	TPR
DwtDct (Cox et al. 2007)	0.828	0.576	0.501	0.000	0.509	0.363	0.719	0.256	0.494	0.000	0.573	0.125
DwtDctSvd (Cox et al. 2007)	1.000	1.000	0.468	0.000	0.701	0.405	0.837	0.671	0.605	0.022	0.679	0.340
RivaGAN* (Zhang et al. 2019)	0.994	0.994	0.742	0.492	0.974	0.966	0.914	0.775	0.570	0.003	0.835	0.641
SSL Watermark (Fernandez et al. 2022)	1.000	1.000	0.996	0.998	0.989	0.994	0.922	0.750	0.596	0.005	0.906	0.763
Stable Signature (Fernandez et al. 2023)	0.995	0.998	0.810	0.496	0.824	0.724	0.253	0.498	0.605	0.011	0.757	0.509
Gaussian Shading* (Yang et al. 2024b)	1.000	1.000	0.634	0.250	0.998	0.997	0.870	0.750	0.986	0.958	0.880	0.756
AquaLoRA (Feng et al. 2024)	0.963	0.979	0.690	0.271	0.954	0.973	0.858	0.702	0.930	0.955	0.866	0.741
OptMark (ours)	1.000	1.000	0.998	1.000	0.998	1.000	0.990	0.979	0.923	0.872	0.983	0.972

Table 2. Table 2: Performance of zero-bit different watermarking methods under various attacks.

Method	None	Geo.	Valu.	Edit.	Regen.	Avg.
Tree-Ring	1.000	0.773	0.970	0.765	0.953	0.874
RingID	1.000	0.750	0.999	0.717	0.814	0.841
WIND	1.000	0.985	0.976	0.748	1.000	0.930
OptMark	1.000	1.000	1.000	1.000	0.993	0.999

Table 3. Table 3: Quantitative analysis of the watermarked image quality. “w/o watermark” indicates the baseline using images generated by Stable Diffusion  (Rombach et al. 2022 ) without watermarks.

Method	$\text{FID}\downarrow$	$\text{CLIP Score}\uparrow$
w/o watermark	124.309	0.3686
SSL Watermark (Fernandez et al. 2022)	128.053	0.3555
Gaussian Shading (Yang et al. 2024b)	127.756	0.3646
OptMark (ours)	127.378	0.3630

Table 4. Table 4: Effect of different watermarks. “Structure” and “Detail” refer to the structure watermark and detail watermark, respectively. “Other Attacks” encompasses various attacks, including geometric, valuemetric, and editing attacks.

		Other Attacks		Regeneration Attack
Structure	Detail	Bit Acc.	TPR	Bit Acc.	TPR
$\text{usym}2713$	$\text{usym}2717$	0.961	0.935	0.834	0.567
$\text{usym}2717$	$\text{usym}2713$	0.984	0.990	0.794	0.407
$\text{usym}2713$	$\text{usym}2713$	0.993	1.000	0.923	0.872

Table 5. Table 5: Effect of watermarks’ different initialization. The robustness results here refer to the average scores calculated under four types of attacks and the non-attack scenario. Note that “Init. Cons.” denotes the Initial Constraint.

	Robustness		Image Quality
Method	Bit Acc.	TPR	${\Delta }_{\text{FID}}\downarrow$	${\Delta }_{\text{CLIP-Score}}\uparrow$
None	0.985	0.975	9.779	-0.0085
Init. Cons.	0.984	0.974	7.928	-0.0083
Init. Cons. + ${ℒ}_{low}$	0.984	0.973	4.434	-0.0059
Init. Cons. + ${ℒ}_{low}$+ ${ℒ}_{high}$	0.983	0.972	3.069	-0.0056

Table 6. Table 6: Effect of watermarks’ different initialization. The robustness results here refer to the average scores calculated under four types of attacks and the non-attack scenario.

	Robustness		Image Quality
Method	Bit Acc.	TPR	${\Delta }_{\text{FID}}\downarrow$	${\Delta }_{\text{CLIP-Score}}\uparrow$
$variance=0.001$	0.502	0.000	0.000	0.0000
$variance=0.1$	0.985	0.976	4.620	-0.0060
$variance=0.01$ (ours)	0.983	0.972	3.069	-0.0056

Table 7. Table 7: Effect of choosing different positions of detail watermarks. “Others” refers to the average bit accuracy of our OptMark under various attacks, including geometric, valuemetric, and editing attacks. “Regeneration” indicates the bit accuracy of our method under regeneration attacks.

	Robustness		Image Quality
Method	Others	Regeneration	$\text{FID}\downarrow$	$\text{CLIP Score}\uparrow$
$t_d=51$	0.998	0.893	251.540	0.3656
$t_d=151$	0.997	0.919	250.979	0.3696
$t_d=251$	0.996	0.931	248.978	0.3708
$t_d=351$	0.994	0.940	251.272	0.3632

Table 8. Table 8: Effect of different watermarks’ detector under 1200 1200 training steps. The robustness results here refer to the average scores calculated under four types of attacks and the non-attack scenario.

	Robustness		Image Quality
Method	Bit Acc.	TPR	$\text{FID}\downarrow$	$\text{CLIP Score}\uparrow$
DINO V1	0.984	0.974	250.782	0.3712
DINO V2	0.833	0.575	252.789	0.3634

Table 9. Table 9: Effect of watermarks’ inference steps T T . The robustness results here refer to the average scores calculated under four types of attacks and the non-attack scenario.

	Robustness		Image Quality
Method	Bit Acc.	TPR	$\text{FID}\downarrow$	$\text{CLIP Score}\uparrow$
$T=10$	0.983	0.972	250.088	0.3611
$T=20$	0.985	0.973	250.830	0.3704
$T=30$	0.984	0.972	243.949	0.3713
$T=50$	0.985	0.973	256.621	0.3628

Table 10. Table 10: Performance of different watermarking methods under the Imprint-Forgery  (Müller et al. 2025 ) , Imprint-Removal  (Müller et al. 2025 ) , and Averaging  (Yang et al. 2024a ) Attacks. Note that “G. Shad.” denotes Gaussian Shading  (Yang et al. 2024b ) .

		G. Shad.	Tree-Ring	OptMark
Attack	Step	Bit Acc.	TPR	Bit Acc.	TPR
Imprint-F	50	0.967	1.000	0.552	0.000
	100	0.978	1.000	0.557	0.000
	150	0.989	1.000	0.563	0.000
Imprint-R	50	0.183	0.164	0.937	0.998
	100	0.084	0.052	0.832	0.962
	150	0.026	0.033	0.791	0.917
Averaging	NA	0.245	0.142	0.996	1.000

Table 11. Table 11: Impact of Different Diffusion Samplers on OptMark’s Robustness.

	DDIM		DPM-Solver++
Robustness	Bit Acc.	TPR	Bit Acc.	TPR
Average	0.983	0.972	0.985	0.973

Table 12. Table 12: Full detecting results of different watermarking methods under various attacks on DiffusionDB  (Gustavosta 2023 ) . “*” indicates that Gaussian Shading  (Yang et al. 2024b ) and RivaGAN  (Zhang et al. 2019 ) can embed 64-bit and 32-bit hidden messages due to the method constraint, whereas all other methods are compared under the condition of embedding 48-bit messages. “Average-attack” indicates calculating the average score across cases under sixteen different attacks and the no-attack (“None”). The underline indicates poor robust performance with Bit Acc. < 0.75 <0.75 and TPR < 0.5 <0.5 . Note that “S. Sign.” and “G. Shad.” denote Stable Signature  (Fernandez et al. 2023 ) and Gaussian Shading  (Yang et al. 2024b ) , respectively.

		DwtDct		DwtDctSvd		RivaGAN*		SSL		S. Sign.		G. Shad.		AquaLoRA		OptMark
		Bit acc.	TPR	Bit acc.	TPR	Bit acc.	TPR	Bit acc.	TPR	Bit acc.	TPR	Bit acc.	TPR	Bit acc.	TPR	Bit acc.	TPR
	None	0.828	0.576	1.000	1.000	0.994	0.994	1.000	1.000	0.995	0.998	1.000	1.000	0.963	0.979	1.000	1.000
Geomtric	Horizontal Flip	0.474	0.000	0.438	0.000	0.506	0.000	1.000	1.000	0.676	0.000	0.553	0.000	0.651	0.000	1.000	1.000
	Rotation (40)	0.502	0.000	0.471	0.000	0.499	0.000	0.991	0.998	0.621	0.000	0.485	0.000	0.478	0.000	0.994	1.000
	Resize (0.6)	0.503	0.000	0.476	0.000	0.973	0.986	0.995	0.994	0.951	0.982	0.999	1.000	0.962	0.976	0.999	1.000
	Crop (0.6)	0.526	0.000	0.488	0.000	0.991	0.984	0.997	0.998	0.993	1.000	0.497	0.000	0.667	0.106	0.998	0.998
Valuemetric	Blur (11)	0.529	0.000	0.986	1.000	0.984	0.974	0.999	1.000	0.526	0.000	0.999	1.000	0.960	0.979	0.999	1.000
	Brightness (0.5)	0.489	0.000	0.635	0.032	0.976	0.965	0.999	1.000	0.992	0.998	0.999	1.000	0.955	0.975	0.999	1.000
	JPEG (50)	0.499	0.000	0.889	0.800	0.942	0.932	0.949	0.972	0.640	0.646	0.993	0.984	0.950	0.970	0.993	1.000
	Contrast (0.5)	0.488	0.000	0.415	0.032	0.974	0.972	0.999	1.000	0.970	0.980	0.999	1.000	0.951	0.972	1.000	1.000
	Saturation (1.5)	0.540	0.090	0.580	0.162	0.992	0.986	0.999	1.000	0.994	0.996	1.000	1.000	0.953	0.966	0.999	1.000
Editing	Meme Format	0.796	0.453	0.852	0.666	0.974	0.982	0.981	1.000	0.579	0.016	0.481	0.000	0.643	0.000	0.964	0.925
	Random Erase (0.1)	0.774	0.422	0.998	1.000	0.993	0.996	0.999	1.000	0.577	0.000	0.999	1.000	0.929	0.945	0.999	1.000
	Text Overlay	0.828	0.576	1.000	1.000	0.991	0.988	1.000	1.000	0.991	0.996	1.000	1.000	0.950	0.975	1.000	1.000
	InstructPix2Pix	0.478	0.000	0.496	0.016	0.699	0.123	0.708	0.000	0.542	0.000	0.997	1.000	0.911	0.886	0.995	0.991
Regen- eration	VAE-B (3)	0.493	0.000	0.612	0.002	0.567	0.002	0.626	0.008	0.639	0.016	0.980	0.937	0.936	0.964	0.896	0.812
	VAE-C (3)	0.493	0.000	0.602	0.016	0.553	0.000	0.579	0.004	0.651	0.018	0.981	0.953	0.940	0.970	0.904	0.820
	Diffusion (60)	0.495	0.000	0.602	0.048	0.590	0.004	0.582	0.002	0.527	0.000	0.997	0.984	0.915	0.930	0.968	0.984
	Average-attack	0.573	0.125	0.679	0.340	0.835	0.641	0.906	0.763	0.757	0.509	0.880	0.756	0.866	0.741	0.983	0.972

Equations

$$q(x_t|x_{t-1})=\mathcal{N}(x_t;\sqrt[]{1-{\beta }_t}{x}_{t-1},{\beta }_t\mathbf{I}),$$

$$x_t=\sqrt[]{{\bar{\alpha }}_t}{x}_0+\sqrt[]{1-{\bar{\alpha }}_t}ϵ,$$

$${\mathbb{E}}_{x_0,t\sim \text{Uniform}(1,T),ϵ\in \mathcal{N}\left(0,\text{I}\right)}\left[{\Vert ϵ-{ϵ}_{\theta }\left(x_t,t,\psi (p)\right)\Vert }_2^2\right],$$

$$\begin{array}{rl}{x}_{t-1}=& \sqrt{{\alpha }_{t-1}}\left(\frac{x_t-\sqrt{1-{\alpha }_t}{ϵ}_{\theta }\left(x_t\right)}{\sqrt{{\alpha }_t}}\right)\\ & +\sqrt{1-{\alpha }_{t-1}-{\sigma }_t^2}\cdot {ϵ}_{\theta }\left(x_t\right)+{\sigma }_t{ϵ}_t.\end{array}$$

$${\widehat{ϵ}}_t=\{\begin{array}{ll}{ϵ}_{\theta }\left({\mathcal{F}}_s\left(x_t,w_s\right),t,\psi \left(p\right)\right)& \text{ if }t=t_s,\\ {ϵ}_{\theta }\left({\mathcal{F}}_d\left(x_t,w_d\right),t,\psi (p)\right)& \text{ if }t=t_d,\\ {ϵ}_{\theta }(x_t,t,\psi (p))& \text{otherwise},\end{array}$$

$$\hat{m}=\left[sign\left(E_w\cdot a_1^{\top }\right),\cdots ,sign\left(E_w\cdot a_k^{\top }\right)\right],$$

$${\mathcal{L}}_{msg}=\frac{1}{k}\sum _{i=1}^{k}max(0,(\mu -(E_w\cdot a_i^{\top })\cdot m_i)).$$

$$x_T^w=w_s+\sqrt{\frac{\text{var}\left(x_T\right)-\text{var}\left(w_s\right)}{\text{var}\left(x_T\right)}}\cdot x_T,$$

$$x_T^w=\sqrt{\frac{\text{var}\left(x_T\right)}{\text{var}\left(x_T^w\right)}}\cdot x_T^w,$$

$${\mathcal{L}}_{init}={\mathcal{L}}_{mean}\left(x_T^w,x_T\right)={\left(\text{mean}\left(x_T^w\right)-\text{mean}\left(x_T\right)\right)}^2,$$

$$\begin{array}{r}{\mathcal{L}}_{low}={\mathcal{L}}_{mean}(w_s,w_s^{init})+{\mathcal{L}}_{var}(w_s,w_s^{init})\\ +{\mathcal{L}}_{mean}(w_d,w_d^{init})+{\mathcal{L}}_{var}(w_d,w_d^{init}),\end{array}$$

$${\mathcal{L}}_{high}={\mathcal{L}}_{kur}(w_s)+{\mathcal{L}}_{kur}(w_d)+{\mathcal{L}}_{ske}(w_s)+{\mathcal{L}}_{ske}(w_d),$$

$$\mathcal{L}={\lambda }_{msg}{\mathcal{L}}_{msg}+{\lambda }_{init}{\mathcal{L}}_{init}+{\lambda }_{low}{\mathcal{L}}_{low}+{\lambda }_{high}{\mathcal{L}}_{high},$$

$$\begin{array}{rl}\mathcal{L}\left(w\right)& =\mathcal{L}\left(x_T+{\int }_T^{0}f\left(x_t,t,c,w\right)dt\right)\\ & =\mathcal{L}\left(\text{ODESolve}\left(x_T,f,T,0,w\right)\right),\end{array}$$

$$\displaystyle\begin{split}\frac{dx_{t}}{dt}&=f\bigl{(}x_{t},t,c,w\bigr{)},\\ \frac{da_{t}}{dt}&=-a_{t}^{\top}\frac{\partial f(x_{t},t,c,w)}{\partial x_{t}},\\ \frac{\partial\mathcal{L}}{\partial w}&=\int_{0}^{T}a_{t}^{\top}\frac{\partial f\left(x_{t},t,c,w\right)}{\partial w}dt.\end{split}$$

$$[x_T,a_T,\frac{\partial \mathcal{L}}{\partial w}]=\text{ODESolve}\left(s_0,\text{dynamics},0,T,w\right)$$

$$x_T^w=w_s+\gamma \cdot x_T$$

$$\text{var}(x_T^w)$$

$$=\text{var}(w_s)+{\gamma }^2\text{var}(x_T)+2\text{cov}(w_s,\gamma \cdot x_T)$$

$${\gamma }^2\text{var}(x_T)+2\text{cov}(w_s,\gamma \cdot x_T)+\text{var}(w_s)-\text{var}(x_T)=0$$

$$\gamma =$$

$$\frac{\sqrt{{\left(\text{cov}\left(w_s,x_T\right)\right)}^2-\text{var}\left(x_T\right)\cdot \left(\text{var}\left(w_s\right)-\text{var}(x_T)\right)}}{\text{var}\left(x_T\right)}$$

$$\gamma =\pm \sqrt{\frac{\text{var}\left(x_T\right)-\text{var}\left(w_s\right)}{\text{var}\left(x_T\right)}}$$

$$x_T^w=w_s+\sqrt{\frac{\text{var}\left(x_T\right)-\text{var}\left(w_s\right)}{\text{var}\left(x_T\right)}}\cdot x_T$$

$$x_T^w=\sqrt{\frac{\text{var}\left(x_T\right)}{\text{var}\left(x_T^w\right)}}\cdot x_T^w$$

Full text

1 Introduction

In the AIGC era, diffusion models (Ho, Jain, and Abbeel 2020; Song, Meng, and Ermon 2020; Rombach et al. 2022) have become a cornerstone of digital content creation, enabling the generation of hyper-realistic images. This advancement revolutionizes visual content production while raising critical intellectual property and content safety challenges in the digital age. As a crucial copyright protection technology, invisible watermarking enables AIGC service providers to embed imperceptible identifiers into generated content, facilitating traceability and ownership verification. This paper explores multi-bit invisible watermarking for diffusion-generated content, focusing on copyright protection and traceability.

Current watermarking approaches fall into two camps: pixel-level and semantic-level. Pixel-level watermarking methods, such as HiDDeN (Zhu 2018), SSL (Fernandez et al. 2022), WAM (Sander et al. 2024), and Stable Signature (Fernandez et al. 2023), embed watermarks directly at the pixel level. While these methods are straightforward to implement, they exhibit limited robustness against regeneration attacks (Zhao et al. 2023). Semantic-level watermarking methods usually embed watermarks during the image generation process and alter the semantic layout of the generated images. A typical approach is to embed handcrafted watermark patterns in the diffusion noise. Compared with pixel-level methods, these approaches are more robust to regeneration attacks, yet they remain vulnerable to certain image transformations and often lack sufficient capacity to embed more bits. Specifically, Tree-Ring (Wen et al. 2023) is susceptible to cropping and scaling, while Gaussian Shading (Yang et al. 2024b) is vulnerable to geometric attacks that disrupt the order of patches, such as horizontal flipping. Furthermore, methods such as RingID (Ci et al. 2024b) and WIND (Arabi et al. 2024) lack sufficient capacity to embed adequate watermark bits, limiting their scalability. Overall, significant challenges remain in balancing robustness and capacity in existing approaches.

In this paper, we propose OptMark, a novel semantic-level multi-bit watermarking approach that ensures ample capacity while achieving comprehensive robustness against four common types of attacks: valuemetric, geometric, editing, and regeneration, as shown in Fig. 1. To achieve this, OptMark optimizes the watermarks in an end-to-end manner during the diffusion inference process. Unlike prior works that rely on handcrafted watermark patterns (Wen et al. 2023; Ci et al. 2024b; Yang et al. 2024b), our approach offers two key advantages through end-to-end learning:

1. Enhanced robustness: By seamlessly integrating with diverse training-time image augmentations, OptMark improves resilience against a wide range of attacks, whereas manually designed watermarks struggle to cover all possible scenarios.
2. Greater flexibility: End-to-end optimization allows for the efficient embedding of a larger number of bits, as the process is fully automated.

To establish this end-to-end optimization framework with comprehensive robustness, high image quality, and low GPU memory overhead, we introduce three key designs: 1) Comprehensive Robustness: We adopt a dual watermarking mechanism, optimizing a structure watermark in the initial diffusion noise to resist generative attacks and a detail watermark in one late denoising step to counter image transformations. 2) Minimal Impact on Image Quality: We develop specialized embedding strategies and constraints to regulate the shape and statistical properties of the learned watermarks, ensuring high image quality and imperceptibility. 3) Efficient GPU Memory Usage: To reduce GPU memory overhead, we introduce the adjoint method for computing gradients on learnable watermarks, lowering memory consumption from $O(N)$ to $O(1)$. Extensive experiments demonstrate that our method significantly outperforms state-of-the-art approaches in robustness, with sufficient bit capacity and high generated image quality.

2 Related Work

2.1 Pixel-Level Watermark

Pixel-level watermarking typically embeds invisible watermarks directly into the image pixel domain. Mainstream approaches can be categorized into two types: optimization-based methods and encoder-decoder methods. Representative optimization-based approaches, such as FNNS (Kishore et al. 2021) and SSL (Fernandez et al. 2022), iteratively optimize a small perturbation on the cover image so that the image features extracted by a pre-trained model can reliably recover the target watermark bits. In contrast, encoder-decoder methods (Zhu 2018; Tancik, Mildenhall, and Ng 2020; Fernandez et al. 2023; Ci et al. 2024a; Sander et al. 2024) train watermark encoders and decoders on a large set of images with different watermark bit sequences, enabling on-the-fly embedding of watermark bits into images. While pixel-level watermarking is imperceptible to the human eye, it has been shown to be inherently vulnerable to regeneration attacks (Zhao et al. 2023).

2.2 Semantic-Level Watermark

Semantic-level watermark approaches embed watermarks during the diffusion generation process, altering the semantic content and layout of the generated image, and improving robustness against regeneration attacks. Some methods train diffusion plugins (Feng et al. 2024; Min et al. 2024) for semantic watermarking, but they require expensive training and struggle to achieve optimal robustness. While Tree-Ring (Wen et al. 2023) pioneered another direction by injecting a handcrafted tree-ring pattern into the initial diffusion noise as a zero-bit watermark. Subsequent works (Ci et al. 2024b; Yang et al. 2024b; Zhang et al. 2024; Huang, Wu, and Wang 2024; Gunn, Zhao, and Song 2024) have improved its robustness or imperceptibility. However, they either remain vulnerable to geometric attacks (Yang et al. 2024b) or lack the capacity to embed sufficient multi-bit information (Ci et al. 2024b; Zhang et al. 2024; Huang, Wu, and Wang 2024). Our proposed method, OptMark, belongs to the semantic-level watermarking. It is the first approach to achieve both sufficient multi-bit capacity and comprehensive robustness against common image transformations and generative attacks.

3 Method

3.1 Preliminary

Diffusion Models. Diffusion models (Ho, Jain, and Abbeel 2020; Song, Meng, and Ermon 2020) progressively convert standard Gaussian noise $x_{T}\sim\mathcal{N}(0,\mathbf{I})$ into samples from the true data distribution $x_{0}\sim q(x)$ over $T$ reverse (denoising) steps. The forward (noising) process is defined as:

[TABLE]

where $\left\{\beta_{t}\right\}_{t=1}^{T}\in\left(0,1\right)$ is the scheduled variance, and $x_{t}$ can be sampled directly from $x_{0}$ as:

[TABLE]

where $\bar{\alpha}_{t}={\textstyle\prod_{i=0}^{t}(1-\beta_{t})}$ and $\epsilon\sim\mathcal{N}(0,\mathbf{I})$. Subsequently, a network $\epsilon_{\theta}$ is trained to predict the added noise at each step, with the following objective:

[TABLE]

where $x_{t}$ represents the noisy latent at timesteps $t$ and $\psi(p)$ denotes the embedding of the text input prompt $p$. The reverse (generation) process can be written as:

[TABLE]

When $\sigma_{t}=0$, it is a DDIM sampler (Song, Meng, and Ermon 2020). When $\sigma_{t}=\sqrt{\left(1-\alpha_{t-1}\right)/\left(1-\alpha_{t}\right)}\sqrt{1-\alpha_{t}/\alpha_{t-1}}$, it is a DDPM sampler (Ho, Jain, and Abbeel 2020).

Background and Task Definition. In the multi‑bit watermarking scenario for diffusion models, OptMark embeds a $k$-bit invisible watermark message $m$ into the generation process to produce a watermarked image $x_{0}^{*}$. When these images are disseminated online, they may undergo various attacks $\mathcal{T}$. For copyright verification or user identification, the model owner decodes the potentially distorted image $\mathcal{T}(x_{0}^{*})$ to recover $\hat{m}$ and compares it to the original watermark $m$.

3.2 Overview

Figure 1 illustrates the OptMark’s end-to-end pipeline, which comprises two stages: Watermark Encoding and Decoding. In the Watermark Encoding stage, learnable watermark vectors are injected into the diffusion latents during inference to produce a watermarked image $x_{0}^{*}$. An inference‑time optimization strategy balances watermark robustness against visual fidelity. In the Decoding stage, we employ a pre‑trained, self‑supervised image encoder (Caron et al. 2021) as the message decoder to extract the embedded watermark representation from versions of $x_{0}^{*}$ subjected to attacks $\mathcal{T}$. Finally, the $k$-bit message is recovered by computing the dot product between this representation and a pre-defined set of $k$ carrier vectors.

3.3 Dual-Watermark for Diffusion Models

Watermark Encoding

Compared with recent pixel‑level watermarking methods (Kishore et al. 2021; Fernandez et al. 2022), which exhibit poor robustness against regeneration attacks (Zhao et al. 2023), OptMark embeds messages directly into the diffusion denoising process and thus achieves significantly higher resistance to these attacks. The diffusion model’s denoising trajectory can be divided into two stages: structure formulation and detail refinement. We therefore propose injecting different watermarks at each stage, with each targeting a distinct semantic level, to enhance robustness against a wide range of attacks. However, since imprinting the watermark into the denoising process is an increasing entropy reaction, excessive introduction of the watermark can negatively impact the quality of image generation. To balance the watermark robustness and image quality, OptMark inserts exactly one watermark per stage: a structure watermark during the first stage, injected into high-level semantic features to leave a persistent mark that is difficult to erase through generative attacks; and a detail watermark during the second stage embedded at a finer, near-pixel level to withstand geometric and volumetric attacks while accelerating convergence.

We consider a standard diffusion framework using the DDIM sampler (Song, Meng, and Ermon 2020). Fig. 2 depicts the watermark embedding process in OptMark. Given standard Gaussian initial noise $x_{T}\sim\mathcal{N}\left(0,\mathbf{I}\right)$, the model predicts the noise $\epsilon_{\theta}$ at each denoising timestep $t$ via:

[TABLE]

where $w_{s}$ and $w_{d}$ represent the structure and detail watermark, respectively, both initialized with a Gaussian distribution. $\mathcal{F}_{s}$ and $\mathcal{F}_{d}$ specify the corresponding watermark‑embedding operator.

Choices of Watermark Position

We inject the structure watermark $w_{s}$ at the initial timestep $t_{s}=T$ for two reasons: (i) injecting at initialization enhances robustness against generative attacks; and (ii) the latents $x_{T}$ follow the standard normal distribution $\mathcal{N}\left(0,\mathbf{I}\right)$, which serves as a reference to constrain the post‑embedding distribution and thus minimize any degradation in generation quality.

For the detail watermark $w_{d}$, we need to select an appropriate timestep $t_{d}$ after the semantic generation process, ensuring that the introduction of $w_{d}$ does not distort the semantics of the generated image. At the same time, this step should not be too close to the pixel level, as pixel-level watermarks are more vulnerable to regeneration attacks and prone to introducing visible artifacts. Fig. 3 shows the evolution of the mean values of classifier-free guidance noise throughout the generation process: $s\cdot(Condition-Uncondition)$, where “Condition” and “Uncondition” represent the predicted noise with and without text conditioning, and $s$ is the guidance scale. We can observe that over timesteps 0 to 400, the variation in guidance noise decreases significantly, indicating that the fundamental semantics have been established. Based on the ablation study detailed in the Appendix Sec. B.2, we set $t_{d}\in[200,300]$ to balance watermark robustness and image quality.

Watermark Decoding

Following SSL (Fernandez et al. 2022), we employ a pre-trained image feature extractor $\mathcal{D}_{msg}$ (e.g., DINO (Caron et al. 2021)) as our message decoder. Given a watermarked image $x_{0}^{*}$, we compute its embedding $E_{w}=\mathcal{D}_{msg}\left(x_{0}^{*}\right)\in\mathbb{R}^{1\times D}$, and denote the secret $k$‑bit message as $m=\left(m_{1},\dots,m_{k}\right)\in\{-1,1\}^{k}$. We predefine a set of carrier vectors $\{a_{i}\}_{i=1}^{k}$, $a_{i}\in\mathbb{R}^{D}$, each initialized by whitening on a large natural‑image dataset to ensure that decoding on arbitrary (non‑watermarked) images yields i.i.d. Bernoulli(0.5) bits. The recovered message is then:

[TABLE]

During training, the watermark decoding loss is defined as the hinge loss with margin $\mu\geq 0$ on the projections:

[TABLE]

3.4 Balancing Robustness and Image Quality

Quality‑Preserving Components

To minimize watermarking’s impact on visual fidelity, we propose three complementary components: watermark initialization, embedding strategy, and regularization loss. Our optimization targets two criteria: (i) the latent distribution before and after watermark embedding remains as close as possible; and (ii) the embedded watermark follows a low‑variance Gaussian profile, as diffusion models are well trained to handle small Gaussian perturbations.

Based on the above design principles, we initialize both the structure watermark and detail watermark as $w_{s}^{init},w_{d}^{init}\sim\mathcal{N}\left(0,0.01\right)$. For the structure watermark, since it is embedded into the initial diffusion latent $x_{T}\sim\mathcal{N}\left(0,\mathbf{I}\right)$, we apply a two‑step normalization within the embedding operator $\mathcal{F}_{s}$ to preserve unit variance:

[TABLE]

[TABLE]

where $\text{var}(\cdot)$ indicates the variance of data. The derivation and proof can be found in the Appendix Sec. A. Additionally, we impose an L2 regularization to ensure that the mean of the watermarked initial diffusion latent remains nearly unchanged to its original value before watermarking:

[TABLE]

where $\text{mean}(\cdot)$ indicates the mean of data.

For the detail watermark, we also aim to minimize the impact of the embedding operator $\mathcal{F}_{d}$ on the DDIM sampling. By Eq. 4, the reverse process is robust to small Gaussian perturbations $\sigma_{t}\epsilon_{t}$. Thus, at $t=t_{d}$ we replace the term $\sigma_{t_{d}}\epsilon_{t_{d}}$ with the detail watermark $w_{d}\sim\mathcal{N}(0,0.01)$, initializing $\sigma_{t_{d}}=0.1$; for all other timesteps we use $\sigma_{t}=0$.

In addition, we further introduce losses to separately constrain the watermarks’ low-order statistics (mean and variance) and high-order statistics (kurtosis and skewness), ensuring they remain statistically similar to the small initial Gaussian noise, given by:

[TABLE]

[TABLE]

where $\mathcal{L}_{mean}(\cdot,\cdot)$ and $\mathcal{L}_{var}(\cdot,\cdot)$ indicates the L2 mean and variance loss. $\mathcal{L}_{kur}(x)=\left(\frac{1}{n}\sum_{i=1}^{n}{\left(\frac{x_{i}-mean\left(x\right)}{std\left(x\right)}\right)^{4}-3}\right)^{2}$ is the Kurtosis loss and $\mathcal{L}_{ske}=\left(\frac{1}{n}\sum_{i=1}^{n}{\left(\frac{x_{i}-mean\left(x\right)}{std\left(x\right)}\right)^{3}}\right)^{2}$ is the Skewness loss. These two high-order losses constrain the shape of the watermark distribution.

Final Objective

The final optimization objective is defined as a weighted combination of the watermark decoding loss and the image‑quality constraint terms:

[TABLE]

where $\lambda_{msg}$, $\lambda_{init}$, $\lambda_{low}$ and $\lambda_{high}$ are hyperparameters that balance the respective loss components.

3.5 Optimizing with Adjoint Sensitivity Method

The DDIM sampler (Song, Meng, and Ermon 2020) can be interpreted as an ordinary‑differential‑equation (ODE) solver. Our objective is to minimize $\mathcal{L}$ with respect to the watermark $w$. For simplicity, we merge $w_{s}$ and $w_{d}$ into a unified notation $w$. We optimize the watermark vector $w$ by minimizing:

[TABLE]

where $f$ predicts the denoising residuals, incorporating operations such as denoising noise prediction, classifier-free guidance, and scheduler scaling. A straightforward optimization approach is to back‑propagate through the DDIM solver. However, this requires storing the entire computation graph during DDIM inference, leading to GPU memory consumption proportional to the number of inference steps, $O\left(N\right)$. To address this, we adopt the Adjoint Sensitivity Method introduced in (Chen et al. 2018) to compute the gradient of $\mathcal{L}$ with respect to $w$, which reduces memory cost to $O\left(1\right)$. The key idea is to compute gradients by solving a second, adjoint ODE backward in time. First, we define three interdependent quantities: $x_{t}$ is the intermediate latents at timestep $t$; $a_{t}=\frac{\partial\mathcal{L}}{\partial x_{t}}$, is the gradient of $\mathcal{L}$ w.r.t $x_{t}$; $\frac{\partial\mathcal{L}}{\partial w}$ is the gradient of $\mathcal{L}$ w.r.t. $w$, which is also our target. The dynamics of these three quantities can be defined by the following equations:

[TABLE]

Subsequently, by making a single call to the ODE solver, we simultaneously perform backward integration along the diffusion path from timestep [math] to $T$ for all three quantities, ultimately obtaining the gradient of $\mathcal{L}$ with respect to $w$:

[TABLE]

where $s_{0}=[x_{0},a_{0},\mathbf{0}_{w}]$ is the initial state of the three quantities, dynamics are $\left[f,-a_{t}^{\top}\frac{\partial f}{\partial x_{t}},-a_{t}^{\top}\frac{\partial f}{\partial w}\right]$ defined in Eq. 15.

4 Experiments

4.1 Experimental Setup

Model and Dataset. We adopt widely-used StableDiffusion‑v2.1 (Rombach et al. 2022) as our generative model, and use the Stable‑Diffusion‑Prompts dataset (Gustavosta 2023) as the source of text prompts.

Evaluation Metrics. To evaluate robustness, we use bit accuracy as a metric and calculate the true positive rate (TPR) corresponding to a fixed false positive rate (FPR), which is set at $10^{-6}$, to assess the degradation of secret messages under various attacks. For image quality evaluation, we use the FID (Heusel et al. 2017) to assess the fidelity of the watermarked image distribution and the CLIP score (Radford et al. 2021) to measure the alignment between the generated images and their corresponding text prompts.

Implementation Details. For the diffusion model, we apply the DDIM (Song, Meng, and Ermon 2020) scheduler with 20 denoising steps to generate 1,000 images at $512\times 512$ resolution in the main experiments. We embed 48-bit secret messages ($k=48$) into each image, and the pre-defined key carrier’s dimension is 2048 ($D=2048$). The detail watermark is injected at step 251 ($t_{d}=251$, $15^{th}$ step). The loss weights $\lambda_{msg}$, $\lambda_{init}$, $\lambda_{low}$ and $\lambda_{high}$ are set to 0.1, 100, 1000, and 100, respectively. Inspired by SSL (Fernandez et al. 2022), DINO (Caron et al. 2021) is used as the pre-trained message decoder. We employ the Adam (Kingma 2014) optimizer with 1,200 optimization rounds, and the learning rate is 0.002.

4.2 Robustness of Watermark

The various attack methods that we implement can be divided into four categories: geometric attack (horizontal flip, random rotation of 40 degrees, resizing of 60%, and center cropping of 60%), valuemetric attack (color jitter with brightness 0.5, Gaussian blur with radius 11, contrast adjustment to 0.5, 50% JPEG compression, and saturation adjustment to 1.5), editing attack (Meme format, random erase with area ratio of 0.1, text overlay, and InstructPix2Pix (Brooks, Holynski, and Efros 2023)) and regeneration attack (two types of VAE regeneration attacks (Ballé et al. 2018; Cheng et al. 2020) from the CompressAI library (Bégaint et al. 2020) with a compression factor of 3, and a diffusion regeneration attack performed with 60 denoising steps (Zhao et al. 2023).) The processed samples after diverse attacks are shown in the Appendix Sec. B.5.

Multi-bit Methods Comparison

For multi‑bit watermarking, we evaluate our OptMark against seven baselines: DwtDct (Cox et al. 2007), DwtDctSvd (Cox et al. 2007), RivaGAN (Zhang et al. 2019), SSL Watermark (Fernandez et al. 2022), Stable Signature (Fernandez et al. 2023), Gaussian Shading (Yang et al. 2024b), and AquaLoRA (Feng et al. 2024). Except for Gaussian Shading and AquaLoRA, which embed watermarks in the diffusion latent space, all other methods operate in pixel space. Tab. 1 shows the watermark robustness comparison between other methods and our OptMark. We find that SSL Watermark (Fernandez et al. 2022) exhibits strong robustness against attacks, except for generative attacks, making it stand out among all pixel-space embedding methods. However, it is worth noting that all pixel space embedding methods exhibit little to no resistance to generative attacks. In contrast, the diffusion space embedding method Gaussian Shading (Yang et al. 2024b) and AquaLoRA (Feng et al. 2024) exhibits strong robustness against regeneration attacks but is rendered ineffective when facing geometric attacks. Unlike them, our OptMark is a highly comprehensive approach that demonstrates exceptional robustness against various attacks without evident weaknesses, achieving SOTA performance. A more detailed experiment on the performance of various methods against different attacks can be found in the Appendix Sec. B.6.

Zero-bit Methods Comparison

For zero-bit watermarking, we compare our OptMark with Tree-Rings (Wen et al. 2023), RingID (Ci et al. 2024b), and WIND (Arabi et al. 2024). All of these approaches embed semantic‑level watermarks within the diffusion latent space. Consistent with the standard evaluation for zero‑bit schemes, we report all results under TPR@FPR=$1\%$, with results shown in Tab. 2. Compared to alternative methods, our approach demonstrates superior robustness against all attack types, exhibiting no vulnerability to any specific attack and achieving the best overall robustness performance.

4.3 Quality of Watermarked Image

The qualitative image quality comparison is shown in Fig. 4. SSL Watermark (Fernandez et al. 2022) introduces noticeable artifacts due to the disturbance added in the pixel space. In contrast, Gaussian Shading (Yang et al. 2024b) only adds the watermark to the initial latent in the diffusion model without impacting the denoising process, resulting in image quality comparable to that of images without watermark. Although our OptMark injects two watermarks (structure and detail watermark) during the denoising process, the image quality remains unaffected compared to Gaussian Shading and images without watermark, and the semantic representation stays consistent with the corresponding text prompt, demonstrating the effectiveness of our method.

For a quantitative comparison of image quality, we compare the FID (Heusel et al. 2017) and CLIP Score (Radford et al. 2021). The FID is evaluated on the MS-COCO-2017 dataset (Lin et al. 2014). As shown in Table 3, our OptMark achieves the best performance in FID, indicating the closest alignment to the real data distribution. Furthermore, it demonstrates a CLIP Score comparable to Gaussian Shading (Yang et al. 2024b).

4.4 Ablation Study

In this section, to more clearly illustrate the changes in quality metrics, we introduce $\Delta_{\textbf{FID}}$ and $\Delta_{\textbf{CLIP-Score}}$, both of which are relative values compared to the baseline, i.e., “w/o watermark”. All training iterations in the following ablation studies are set to 1,200.

Effect of Dual Watermarks We conduct both quantitative and qualitative analyses to demonstrate the necessity of combining the structure watermark and the detail watermark. The quantitative results are shown in Tab. 4. From the table, it can be observed that the structure watermark, introduced during the structure formation stage, demonstrates stronger robustness against regeneration attacks compared to the detail watermark, which is introduced in the detail formulation stage. However, the structure watermark’s convergence is relatively slower, and optimization over 1200 iterations is insufficient for full convergence. As a result, its performance against conventional attacks is weaker than that of the detail watermark. The combination of both can accelerate convergence and result in a more robust performance under various attacks. For qualitative analysis, as shown in Fig. 5, the introduction of the detail watermark closer to the final image generation stage makes it prone to issues similar to those encountered in pixel-level watermarking methods (e.g., SSL (Fernandez et al. 2022)), such as the appearance of artifacts. In contrast, the structure watermark does not exhibit this problem. Since it is introduced at the semantic level, it also leads to some visual differences compared to the original image without watermarks. Furthermore, the combination of both watermarks helps mitigate artifacts.

Effect of Image Quality Constraints. To assess the effectiveness of the proposed image quality constraints, we perform an ablation study on each individual component. The results are presented in Fig. 6 and Tab. 5. Note that “Initial Constraint” includes the normalization step in $\mathcal{F}_{s}$ and $\mathcal{F}_{d}$, and the loss $\mathcal{L}_{init}$. From a qualitative perspective, as shown in Fig. 6, the realism and quality of the generated images progressively improve with the introduction of each quality-driven constraint. From a quantitative perspective, as shown in Tab. 5, our constraints achieve a significant improvement in FID with only a minimal loss in robustness.

Effect of Adjoint Method. In experiments, we find that under the DDIM setting with 20 inference steps, the GPU memory consumption of the naive optimization is about 52 GB, while for 30 steps, it increases to around 76 GB. Using the adjoint method, the memory consumption can be uniformly reduced to just 9 GB, making it feasible to scale to larger inference steps and more complex diffusion models.

5 Conclusion

This paper presents OptMark, a robust watermarking framework based on inference-time optimization. We propose a dual-watermark mechanism to enhance robustness, design a tailored objective and regularization scheme to preserve image fidelity, and integrate the adjoint sensitivity method for constant‑memory gradient computation. Extensive experiments show that OptMark delivers SOTA robustness across a diverse range of common attacks.

Appendix A Analytical Derivation of the Variance Constraint at Step $T$

The derivation of operation $\mathcal{F}_{s}$ for controlling the variance of output $x_{T}^{w}$, which is the combination of initial noise $x_{T}$ and structure watermark $w_{s}$, to be $\text{var}(x_{T}^{w})\approx 1$ in Sec. 3.4 of the main paper is as follows. Given that $x_{T}\sim\mathcal{N}\left(0,\mathbf{I}\right)$ and initial $w^{init}_{s}\sim\mathcal{N}\left(0,{0.01}\right)$, the combination of them $x_{T}^{w}$ can be defined as:

[TABLE]

where $\gamma\in(0,+\infty)$ is a variable control coefficient that ensures the variance of $x_{T}^{w}$ remains constant at $\text{var}(x_{T}^{w})\approx 1$. The reason we apply $m$ to $x_{T}$ instead of $w_{s}$ is that the variance of $w_{s}$ is very small. Scaling it does not effectively control the variance and would severely impact the robustness of the watermark. Our goal is to solve for $\gamma$, which has an exact solution. The derivation is as follows:

[TABLE]

where $\text{var}(\cdot)$ and $\text{cov}(\cdot,\cdot)$ indicates the variance and covariance calculations, respectively. To ensure that $\text{var}(x_{T}^{w})$ equals $\text{var}(x_{T})$, the linear equation can be written as follows:

[TABLE]

To solve this linear equation, we can obtain the exact solution for $\gamma$, as follows:

[TABLE]

The condition for real roots to exist is $\sqrt{\left(\text{cov}\left(w_{s},x_{T}\right)\right)^{2}-\text{var}\left(x_{T}\right)\cdot\left(\text{var}\left(w_{s}\right)-\text{var}(x_{T})\right)}>0$, but due to the small covariance $\text{cov}\left(w_{s},x_{T}\right)$, it is difficult for real roots to occur. Considering this situation, we assume that $w_{s}$ and $x_{T}$ are independent and identically distributed, such that their covariance $\text{cov}(w_{s},x_{T})$ is zero. Eq. 20 can be simplified to:

[TABLE]

Considering our scenario ($\gamma>0$), we only take the positive root, which leads to Eq. 5 in the main paper. The combination $x_{T}^{w}$ can be represented as :

[TABLE]

However, we observe that in the later stages of training, $w_{s}$ and $x_{T}$ does not follow the assumption of independent and identically distributed variables. As a result, Eq. 22 alone is insufficient to constrain the variance of $x_{T}^{w}$ to $\text{var}(x_{T})\approx 1$. To address this, we introduce an additional scaling step to ensure that $\text{var}(x_{T}^{w})$ matches $\text{var}(x_{T})$, as follows:

[TABLE]

We don’t initially scale the direct combination $x_{T}^{\prime w}=w_{s}+x_{T}$ as a whole because we don’t want to apply a large-scale transformation to $w_{s}$, which can affect the robustness of the watermark. Meanwhile, since adding a watermark is an entropy-increasing process, appropriately compressing $x_{T}$ (usually $\gamma<1$) to make room for $w_{s}$ can help with the learning of the watermark.

Appendix B More Experimental Results

B.1 Initialization of Watermarks

The watermarks are initialized following a Gaussian distribution, with a mean of 0 and various possible choices for the variance. Intuitively, a smaller initial variance results in a lesser impact on the generated image. As shown in Tab. 6, an excessively small initial variance, such as when $variance=0.001$, makes it difficult for the watermark to be optimized in any case, leading to a final outcome with almost no robustness, where the generated result is nearly identical to the baseline. A larger initial variance helps improve the robustness of the optimized images but also reduces the quality of the generated images. To find a balance, we set the variance to 0.01. Qualitative comparisons are shown in Fig. 7.

B.2 Discussion on Optimal Placement of Detailed Watermarks

In Sec. 3.3 of the main paper, we discuss the injection positions of the detail watermark. Here, we further validate our reasoning through experiments conducted on 100 randomly sampled cases. We test four different time injection points under the setting of a total of 20 inference steps: $t_{d}=51\ (0<t_{d}<100)$, $t_{d}=151\ (100<t_{d}<200)$, $t_{d}=251\ (200<t_{d}<300)$, and $t_{d}=351\ (300<t_{d}<400)$. The quantitative analysis is presented in Tab. 7. In terms of robustness performance, smaller values of $t_{d}$ lead to poorer performance under regeneration attacks but better performance under other attacks (geometric, valuemetric, and editing attacks) and vice versa. We believe this phenomenon occurs because smaller values of $t_{d}$ bring the watermark closer to the pixel level. While pixel-level watermarks exhibit weak robustness against regeneration attacks, they still maintain decent robustness against other attacks (e.g., SSL (Fernandez et al. 2022)). We find $t_{d}=251$ strikes a good balance between robustness and image quality. So we choose it as the default embedding timestep for detail watermark.

B.3 Results under Different Watermark Decoder

To assess the influence of the watermark decoder, we evaluate DINO V1-RN50 (Caron et al. 2021) and DINO V2-ViT-S (Oquab et al. 2023). Both are trained for 1200 steps. As shown in Tab. 8, DINO V1 demonstrates better robustness and image quality. Therefore, we chose DINO V1-RN50 as the default watermark decoder.

B.4 Results under Different Inference Steps

To evaluate the impact of different inference steps $T$ on our OptMark we randomly sample 100 cases for both quantitative and qualitative experiments. The quantitative results are shown in Tab. 9. From the perspective of watermark robustness, our OptMark exhibits stable performance without significant fluctuations across different inference steps. In terms of image quality, $T=30$ performs the best, followed by $T=20$. The qualitative results are illustrated in Fig. 8. From the figure, it can be observed that when $T=10$, the original images and watermarked images are slightly blurred and lack some high-frequency details. However, when $T\geq 20$, the generated details become more stable, and the images maintain stronger integrity. Although we apply the adjoint method, the memory usage of OptMark remains consistent across different inference step sizes. However, as $T$ increases, the training time becomes longer. To achieve a balance between efficiency, image quality, and watermark robustness, we finally set $T=20$.

B.5 Samples under Different Attacks

The visualization of samples under different attacks can be seen in Fig. 9. Specifically, the various attack methods we implement can be categorized into four types: geometric attack (horizontal flip, random rotation of 40 degrees, resizing of 60%, and center cropping of 60%), valuemetric attack (color jitter with brightness 0.5, Gaussian blur with radius 11, contrast adjustment to 0.5, 50% JPEG compression, and saturation adjustment to 1.5), editing attack (Meme format, random erase with a probability of 0.1, text overlay, and InstructPix2Pix (Brooks, Holynski, and Efros 2023) with the prompt: “turn it into an ink wash painting”) and regeneration attack (two types of VAE regeneration attacks (Ballé et al. 2018; Cheng et al. 2020) from the CompressAI library (Bégaint et al. 2020) with a compression factor of 3, and a diffusion regeneration attack performed with 60 denoising steps (Zhao et al. 2023).)

B.6 Detailed Quantitative Results

Tab. 12 presents the detailed detection results for all attacks. We observe that our OptMark is capable of handling all of them. In contrast, other methods exhibit at least one weakness in dealing with these attacks. For instance, nearly all pixel-level methods, including DwtDct (Cox et al. 2007), DwtDctSvd (Cox et al. 2007), RivaGAN*(Zhang et al. 2019), SSL(Fernandez et al. 2022), and Stable Signature (Fernandez et al. 2023), struggle with Generative attacks. Meanwhile, the semantic-level method Gaussian Shading (Yang et al. 2024b) and AquaLoRA (Feng et al. 2024) perform poorly under Geometric attacks.

B.7 Robustness to Forgery and Removal Attacks

We assess the impact of several attacks—previously shown to be highly disruptive to other watermarking methods—on our proposed OptMark framework, including the Imprint-Forgery and Imprint-Removal Attacks from  (Müller et al. 2025), as well as the Averaging Attack from  (Yang et al. 2024a). We evaluate robustness to all attacks on a set of 1,000 images, with results shown in Tab. 10. The results demonstrate that the Imprint-Forgery Attack fails to forge our OptMark even after running 150 steps, achieving only 0.563 average multi-bit accuracy. In terms of Imprint-Removal Attack at 50, 100, 150 steps, our method achieves bit accuracy of 0.937, 0.832, and 0.791, significantly outperforms Gaussian Shading, whose bit accuracy is consistently lower than 0.2. Under the Averaging Attack, which averages 1,000 watermarked images, OptMark maintains a bit accuracy of 0.996, demonstrating strong robustness. We attribute this resilience to our per-image optimization strategy: Averaging attacks are effective only when the watermark pattern is independent of image content, whereas OptMark generates image-specific watermarks, rendering such attacks ineffective.

B.8 Empty Prompt Example

We evaluate OptMark on images generated with an empty prompt. As shown in Fig. 10, this setting produces low‑quality, content‑diverse outputs. Despite these challenging conditions, OptMark maintains its full robustness.

B.9 The Impact of Training Steps on Robustness

Fig. 11 illustrates the relationship between training iterations and the watermark robustness of our OptMark. Larger training iterations can achieve better watermark robustness, and users can adjust the number of training iterations based on their requirements as a trade-off.

B.10 Generality of OptMark across different diffusion samplers.

To demonstrate the generality of our OptMark under different diffusion samplers, we conduct additional experiments using the DPM-Solver++ (Lu et al. 2025) sampler with 20 inference steps, with the results shown in Tab. 11. On a set of 1,000 images, the setting $t_{d}=251$ remains optimal, achieving an average multi-bit accuracy of 0.985 and a TPR of 0.973, thereby confirming the robustness and general applicability of OptMark.

B.11 More Qualitative Results

More qualitative results are shown in Fig. 12.

Appendix C Limitations

Although our OptMark demonstrates strong robustness against various types of attacks, its performance under regeneration attacks is slightly inferior to that of other semantic-level watermarking methods, such as Gaussian Shading. We believe the performance difference lies in the extraction methods: Gaussian Shading uses inversion to extract the watermark from the initial noise space, while OptMark uses the DINO network. We suspect that the watermark recovered from the inverted initial noise is inherently more robust to regeneration attacks than from the DINO latent space. To address this, we plan to explore alternative watermark extractors that offer more resilient extraction spaces—such as the denoising UNet used in diffusion models in the future.