Coefficient-Aware Global Assignment for Robust Latent-DCT Watermarking in Stable Diffusion

The rapid adoption of latent diffusion models [1] has re-duced the time and technical effort needed to create visual con-tent. Independent creators and small studios can now produce and distribute large image collections without maintaining a conventional production pipeline. Once an image leaves its original platform, however, prompts, account records, and metadata may be removed during conversion, redistribution, or editing. The creator may then have little evidence connecting the distributed file to its source. Invisible digital watermarking addresses this problem by embedding ownership information into the image itself [2], where it can remain recoverable after distribution and common image processing. Current diffusion watermarking methods do not fully ad-dress this setting. Methods that train or fine-tune water-marking components require data, computing resources, and deployment effort that may be impractical for smaller orga-nizations. Training-free methods remove this requirement but can introduce another restriction. Tree-Rings [3] modifies the initial diffusion noise in the Fourier domain, while Gaussian Shading [4] encodes the payload into a Gaussian latent before sampling. Both methods leave the model weights unchanged, but they require access to the generation process. They can protect images produced through a controlled generator but cannot directly watermark images that already exist. This leaves a practical need for a method that can protect existing image collections without retraining or modifying the diffusion model. Frozen-model latent watermarking provides a post-generation approach to this problem. DDIM inversion [5] maps an existing image into the latent space of Stable Diffusion. A block-based discrete cosine transform exposes latent-frequency coefficients in which a copyright-dependent payload can be embedded. Adam optimization [6] then refines the modified latent to recover visual fidelity, and statistical testing determines whether the corresponding watermark is present [7]–[9]. The model weights remain frozen through-out the process. However, the conventional implementation assigns every payload to the same mid-frequency coefficient template and applies a fixed 60-step optimization budget to every image. Latent coefficients vary with image content and inversion trajectories, so a fixed position may require a small modification in one image and a much larger modification in another. Optimization loss also stabilizes at different rates, yet the fixed budget continues processing even when further refinement produces little improvement. These choices can weaken watermark recovery after cumulative processing and increase the cost of protecting large image collections. This study examines whether both stages can be adapted to the latent content of each image. Embedding positions are selected from the available coefficients rather than imposed by a fixed template, and payload values are matched to coefficients that require smaller modifications. The resulting correspondence is retained with the per-image ownership record and reused during verification. A convergence-based early-stopping criterion monitors the optimization process and terminates refinement after the loss has stabilized. The evalu-ation compares the adaptive approach with conventional fixed embedding and measures image quality, watermark detection after individual and compound processing, and computation time. This design separates the effect of content-aware em-bedding from the effect of reducing the optimization budget without introducing model retraining. choose the set minimizing total embedding displacement for the image at hand. Let C̄(pᵢ) be the mean coefficient at candidate position pᵢ across the 64 blocks. With embedding strength δ, the assignment cost is (4) and the resulting objective is The experiments show that coefficient-payload assignment improves watermark recovery. (5) Under compound non-geometric processing, the fixed template is less reliable. The convergence-based criterion completes refinement in fewer iterations while preserving the robustness obtained through adaptive embedding. The results therefore indicate that processing time can be reduced without losing the watermark’s resistance to cumulative image operations. This extends the original frozen-model pipeline into a more efficient post-generation system for protecting existing images.

DCT Transform Figure 1 shows the overall framework. Section III summa-rizes the embedding and detection procedure of [8], [9], which is the baseline for the position-selection method proposed in this work. The payload occupies the fourth latent channel alone, which is partitioned into 64 disjoint tiles of 8 × 8 samples, each carried into the frequency domain by the two- dimensional DCT [10]: (1) where π is injective so no coordinate receives two entries. Selecting five of the 53 coordinates with minimum total cost forms a rectangular linear assignment problem with an exact polynomial-time solution. The Hungarian method [12] finds it at O(n3); Jonker–Volgenant [13] reaches the same optimum along shortest augmenting paths and scales better. Both are exact, so they return the same coordinates apart from ties; on the benchmark used here, the two solvers selected identical coordinates for every image. Choosing where to write leaves open what to write there. Direct embedding takes values in generated order, so a slot may receive one far from the resident coefficient, which forces the optimizer to induce substantial pixel-space displacement. Coefficient-aware matching instead scans the unused values and picks the closest, (6) for slot j, exchanging entries j and rj before writing. with for w > 0. All subsequent operations act on the 64 transformed tiles, never on pixels. The embedded payload derives from the owner’s copyright image. SHAKE256 [11] turns it into a byte stream that is deterministic in the image, so whoever later claims ownership can reproduce it. Each byte is rescaled onto a 256-level lattice, (2) and pushed through the inverse standard normal CDF, (3) Matching changes payload ordering but alters neither the selected coordinates nor the total capacity. Hungarian-Direct and JV-Coefficient-Matched solve the same assignment problem with exact solvers, and only JV-Coefficient-Matched applies payload matching. Any detection difference between them therefore cannot be attributed to the solver alone.

A. Watermark Embedding and Optimization Each block is written at the five selected coordinates by C′ᵦ(pₖ) = δw_idx(b,k), after which an inverse DCT returns the modified latent. Overwriting coefficients alone leaves visible traces, so this study uses Adam with three losses over 60 iterations. Mean squared error controls pixel difference. the structural similarity loss [14] maintains visual similarity and prevents color shift and blurring, (8) and the VGG loss keeps semantics aligned, focusing on high-level content including texture. where ϕl is the activation of layer l of a pretrained VGG network [15] and ClHlWl its size. These are combined. with sequentially adapted weights as Lₜ = λₘ(t)L_MSE + λₛ(t)L_SSIM + λᵥ(t)L_VGG, and the selected coefficients are projected back to their targets after every iteration. The stopping rule smooths this objective with an exponential moving average, (10) and monitors its relative decrease, (11) We refer to this convergence criterion as Relative-EMA. Optimization halts at the first t ≥ tmin for which ρτ < ϵrel across the window τ ∈ {t − W + 1, . . . , t}. We use W = 5, ϵrel = 5 × 10−3 and tmin = 15, and report the last iteration of the qualifying window. Smoothing distinguishes this rule from a threshold on the raw loss, which fires on single noisy iterations. The decoded result is then mixed with the original in pixel space by binary search on a mixing weight until the final SSIM reaches 0.92, fixing one quality budget across configurations so robustness cannot be traded against fidelity. B. Detection Verification is non-blind, as the right panel of Figure 1 shows: the test image is encoded, inverted, and transformed by block DCT, and the stored coordinates give observed coefficients cˆi. With s2 the estimated coefficient variance, the statistic is the squared distance to the targets, (12) A marked image has been pushed toward those targets, so T is small; an unmarked one holds coefficients unrelated to the key, so T is large. Formally T is approximately central χ2 with k = 320 degrees of freedom when the mark is present and noncentral χ² with λ = s⁻²Σᵢ(δwᵢ)² when it is absent, so the test is one-sided on the lower tail and a watermark is declared when p = Fχ′2 (T ; k, λ) ≤ 0.1. Since λ sums the realised payload rather than assuming unit variance, the quantization above shifts λ to its correct value instead of invalidating the reference. Under direct embedding the reference vector is reconstructible from the secret key alone. Under coefficient-aware matching it is the reordered payload, so the permutation must be stored and supplied at verification, an overhead of at most log2(320!) ≈ 2207 bits, about 276 bytes per image.

A. Experimental Setup Experiments use Stable Diffusion 2.1-base on an NVIDIA RTX 4090 (24 GB) under Ubuntu 24.04 and CUDA 13.0, over 100 WikiArt images at 512 × 512. Every configuration shares the payload length, embedding channel, 60-iteration baseline budget, threshold p ≤ 0.1, and final mixed SSIM of 0.92, and each runs at the full budget and again under the convergence- based early stopping (Relative-EMA) rule. Twelve attack conditions are applied: diffusion purification [16], two learned compressors [17], [18], JPEG at quality 50, brightness and contrast reduction, Gaussian noise (σn = 0.05), Gaussian blur, BM3D denoising [19], a 90◦ rotation, and two compound attacks, all norot applying the nine non-geometric operations in sequence and all adding rotation. Every result was re-scored from its raw chi-squared output at p ≤ 0.1, n = 100, and reported time is the full pipeline wall clock per image. B. Detection Rate and Robustness Table I reports the results. Any single non-geometric attack leaves every configuration detecting 98–100% of images, so single attacks separate nothing and the comparison rests on what happens when distortions accumulate. Under all norot the fixed template detects 69% of images and global assignment alone reaches 70%: solving (5) relo-cates the coordinates but leaves the displacement at each write unchanged, since its cost is a mean over 64 blocks and so optimizes a summary statistic rather than the coefficient each write meets. Adding coefficient-aware ordering on the same coordinates raises detection to 98%, 28 points above global assignment alone and 29 above the fixed baseline. Equation (6) explains why: the displacement is minimized at every write rather than fixed by sequence position, so the residual signal falls inside the frequency statistics the image already had, which cumulative processing perturbs less than a signal forced into an unfavorable block. The improvement belongs to the complete matching-based configuration, not to the JV solver. Across both budgets JV stays at 97–98% against 69–73% for the direct configurations, while average detection rises only from 82.0% to 83.9% because single attacks are already saturated, and pre-mixing PSNR is 23.38 dB against 23.10 dB. Origin false positives read 11% for Fixed-Direct and 6% for the assignment-based configurations, though negative-image calibration under the stored-permutation protocol would be needed before treating that as a deployment property. Rotation defeats all three, at 5–10% detection, because a quarter turn moves the block grid so the stored coordinates no longer index the coefficients carrying the payload. C. Effect of Early Stopping The lower block of Table I addresses the computational bottleneck. Figure 2 shows that the mean objective de-creases rapidly during the first ten iterations, begins to flatten near iteration 25, and improves only slightly thereafter. The convergence-based rule stops at an average of 49.4 iterations for Fixed, 50.1 for Hungarian, and 50.3 for JV-Coefficient-Matched. As shown in Figure 3, the overall runtime reduction is limited because early stopping affects only the Adam opti-mization stage, while the remaining pipeline components are Fig. 2. Mean total loss over 100 images. Dashed lines mark the mean stopping iteration chosen by (10)–(11): 49.4 Fixed, 50.1 Hungarian, 50.3 JV. Fig. 3. Execution time breakdown per image. The convergence-based early stopping (Relative-EMA) reduces the Adam optimization phase across all con-figurations, while components outside the optimization loop remain constant. End-to-end runtime falls by 10.6–11.6%. unchanged. Runtime decreases from 60.0 to 53.2 s for Fixed, from 59.4 to 53.1 s for Hungarian, and from 61.3 to 54.2 s for JV-Coefficient-Matched. These changes correspond to end-to-end reductions of 11.3%, 10.6%, and 11.6%, respectively. Robustness is maintained under early stopping. Compound non-geometric detection changes from 69% to 73% for Fixed, from 70% to 69% for Hungarian, and from 98% to 97% for JV-Coefficient-Matched. Average detection changes by no more than 0.8 percentage points across all configurations. Pre-mixing PSNR decreases by 0.57-0.71 dB, but the subsequent mixing search holds the final SSIM at 0.92 for every config-uration. Figure 4 shows representative final outputs after the mixing stage. Because the mixing search uses the same SSIM target regardless of the stopping iteration, the outputs under the two optimization budgets remain visually similar. Under the tested attacks and fixed-SSIM condition, the optimization can therefore stop earlier without reducing the robustness obtained from coefficient-aware payload matching.

This work addressed two limitations of an existing frozen-model latent-DCT watermarking pipeline: fixed embedding positions and a static optimization budget. The evaluation found that selecting positions based only on latent content offered little improvement over the fixed template under com-pound non-geometric processing. Detection increased from 69% to 70%. When payload values were matched to coef-ficients of similar magnitude, detection rose to 98%. Because the two exact assignment solvers selected identical positions across all 100 images, this result is associated with payload matching and the stored payload order rather than the position solver. Verification requires storing the reordered watermark sequence, with an upper-bound storage requirement of approx-imately 276 bytes per image. A convergence-based early-stopping criterion reduced aver-age end-to-end runtime by approximately 11%, while detection under compound non-geometric processing changed from 98% to 97%. The final mixing stage held SSIM at 0.92 for both optimization settings. Under the evaluated conditions, the optimization process can therefore terminate earlier while retaining the robustness obtained from coefficient-aware em-bedding. The method remains sensitive to geometric transformations, with detection falling to 5–10% under rotation. At the decision threshold of p ≤ 0.1, the recorded false-positive rates ranged from 6% to 11%, and the stored-payload procedure requires further calibration on unwatermarked images. Future work should investigate alignment methods for geometric attacks, repeat stochastic experiments across multiple runs, and evaluate the stored payload order on larger sets of unwatermarked images.