Source-linked AI summary
RAIN: Region-Aware Inversion Network for Semantic Watermark Extraction
Zilai Li
TL;DR
Gaussian Shading extraction conventionally relies on multi-step inversion to recover noise, even though extraction depends on recovering the correct watermark regions. RAIN uses lightweight high-SNR endpoint regression and an anchor-residual estimator for one-step extraction; the authors report lower extraction cost than OSI and FARI. Its regression approach assumes that predicting a conditional mean or median remains sufficient for Gaussian Shading extraction.
Problem
Conventional Gaussian Shading extraction requires multi-step DDIM inversion to recover noise, although the watermark decision depends on the encoded Gaussian regions.
Source: §1 · p. 1, §3.1 · p. 5Method
RAIN uses high-SNR endpoint regression to predict recoverable noise statistics with a lightweight estimator decomposed into an image-like anchor and a noise-oriented residual.
Source: §1 · p. 2, §3.2 · p. 6, §3.4 · p. 6Results
The authors report that extraction computational cost is lower than OSI and FARI, and that very small networks can successfully extract the watermark using either endpoint function.
Source: §1 · p. 2, §3.3 · p. 6Takeaways & Limitations
Gaussian Shading extraction need not reconstruct the original terminal noise if the estimated latent preserves enough correct Gaussian decision regions.
Source: §3.1 · p. 5The method assumes that predicting the conditional mean or median of possible noises does not invalidate Gaussian Shading extraction.
Source: §3.1 · p. 5
Abstract
from arXiv · showhide
Semantic watermarks for diffusion models embed ownership information into the generative process while preserving perceptual quality, but Gaussian-Shading extraction conventionally requires multi-step diffusion inversion to recover the initial noise. Recent one-step methods show that this cost can be reduced substantially. We study this problem through extended flow matching and conditional regression. The key observation is that, near the high-SNR image endpoint, recovering a useful noise statistic given by the first-step output of the extended flow matching in the high-SNR regime is much simpler than reconstructing the full inverse trajectory, and Gaussian Shading only requires the recovered latent to remain in the correct watermark decision region. Based on this observation, we propose a lightweight, prompt-free extractor that decomposes endpoint recovery into an image-like anchor and a noise-oriented residual, which increases the capability of the model to utilize GPU parallel computation. The resulting method avoids iterative inversion and repeated evaluation of a diffusion-scale U-Net, providing an efficient one-step extraction pipeline with a concise theoretical interpretation. The computational cost of extracting noise is lower than that of both OSI and FARI. The github repo is there: https://github.com/TheLovesOfLadyPurple/RAIN-lightweight-NN-for-one-step-semantic-watermark-extraction
1 Introduction
Gaussian Shading (GS) embeds ownership information in the initial noise, but conventional extraction requires multi-step DDIM inversion. The paper proposes a lightweight one-step estimator based on high-SNR endpoint regression, where recovering a useful noise statistic is simpler than reconstructing the full inverse trajectory.
1 Introduction
- Conventional GS extraction uses multi-step DDIM inversion to recover initial noise, while recent one-step methods suggest that this cost may be reduced.Prior work also targets faster inversion or lower inference cost, but not specifically extraction.Source: §1 · pp. 1–2
- For GS, L2 regression may predict the mean of multiple noises compatible with an image, which can still suffice if the prediction preserves the encoded Gaussian regions.The authors assume this does not invalidate extraction, noting invalid extraction is rare in normal DDIM inversion.Source: §1 · p. 2
- In high-SNR Rectified Flow, predicting the conditional mean velocity for the remaining transport is equivalent to predicting a useful mean noise statistic.The authors hypothesize that this task can use a substantially smaller network than a full diffusion denoiser.Source: §1 · p. 2
- A lightweight estimator predicts recoverable noise statistics directly from an almost-clean image, avoiding iterative inversion and reducing extraction cost.The introduction reports better inference speed and parameter count than OSI and FARI.Source: §1 · p. 2
Background
The background frames diffusion and rectified flow as endpoint interpolation with learned velocity fields, then describes Gaussian Shading as watermarking Gaussian regions of the terminal latent. Extraction decodes those regions, so exact recovery of the original latent is unnecessary.
2.1 From DDPM Corruption to Extended Flow Matching and Rectified Flow
- A DDPM noisy state can be viewed as an interpolation between a clean image and a Gaussian-noise endpoint, with independently sampled endpoints in the standard construction.The noise schedule permits direct sampling of a noisy state at any timestep.Source: §2.1 · p. 3
- Flow matching learns a vector field from the observable state, but sample-wise interpolation velocity generally depends on hidden endpoints and is not a closed inference rule.Rectified flow permits more general, including deterministic, image-noise couplings; under squared regression its field is population-optimal.Source: §2.1 · pp. 3–4
2.2 Gaussian Shading
- Gaussian Shading repeats and encrypts payload bits, then uses them to select equal-probability regions for sampling each terminal Gaussian scalar.Uniform encrypted symbols and equal Gaussian probability regions preserve the marginal standard Gaussian prior.Source: §2.2 · p. 4
- Extraction maps the estimated terminal latent back to Gaussian regions, decrypts the symbols, and aggregates repeated bits by voting.For l = 1, decoding reduces to recovering each coordinate's sign, so exact reconstruction of the original noise is unnecessary.Source: §2.2 · p. 4
Approach
The approach reframes one-step extraction as conditional endpoint-statistic prediction, then uses an anchor–residual estimator to avoid full inverse-trajectory recovery and support parallel computation.
3.1 One-Step Recovery as Conditional Regression
- For Gaussian Shading, endpoint regression need only preserve the watermark decision region, not reproduce the generating noise exactly.Squared regression targets the conditional mean noise; L1 targets its conditional median.Source: §3.1 · p. 5
3.2 High-SNR Endpoint Regression
- Near the image endpoint, high-SNR flow matching makes predicting the mean of possible endpoint noises a simpler target than recovering the full inverse trajectory.This motivates a lightweight estimator rather than repeated evaluation of a diffusion-scale denoiser.Source: §3.2 · p. 6
3.3 Endpoint Parameterization and Perturbation
- Both DDPM-style and direct-noise parameterizations successfully extract the watermark with a very small neural network.A DDIM-based heuristic is presented as a more GPU-efficient alternative to the DDPM-based form.Source: §3.3 · p. 6
3.4 Anchor–Residual Reverse Estimator
- The estimator splits recovery into an image-like anchor and a residual noisy component, allowing F0 and F1 to be inferred in parallel.The anchor is trained against a clean representation from the original inference process.Source: §3.4 · pp. 6–7
- Extraction requires no text encoder, prompt embedding, cross-attention, or source diffusion U-Net, while adversarial training selects among common corruptions by reverse-recovery loss.The finite-set min–max procedure is reported to improve robustness without changing the endpoint objective.Source: §3.4 · pp. 7–8
4 Experiments
The experiments evaluate prompt-free extractors on matched watermark-image conditions, including distortion robustness, latency, and backbone cost. The ablation reports robustness gains from adversarial training, while the two-branch extractor has lower measured latency than single-network alternatives.
4.1.1 Current experimental setting
- Training uses 89,600 offline image-noise pairs and prompt-free width-64 NAFNet extractors with coordinatewise L1 endpoint supervision.The evaluation uses Stable Diffusion v1.5 at 512×512, with 15-step UniPC generation and a VAE decode-encode pathway before extraction.Source: §4.1.1 · pp. 8–9
4.1.2 Metrics
- Evaluation reports raw latent sign agreement, decrypted-and-voted GS bit accuracy, watermark detection rate, and traceability rate.Source: §4.1.2 · p. 9
4.1.3 Ablation results
- Adversarial training consistently improves robustness under severe distortions while preserving clean extraction accuracy across the six extractor combinations.The direct-xT and epsilon parameterizations remain closely matched; the two-branch model is more sensitive without distortion training and recovers most of the gap after it.Source: §4.1.3 · p. 9, Table 1 · p. 10
- The SD v1.5 ablation compares direct-xT, FARI-like, and DDIM-heuristic prediction, each with or without adversarial training on the same 1,000 cached COCO samples.All extractors receive the same attacked image and VAE-encoded latent.Source: §4.1.3 · p. 9, Table 1 · p. 10
- The two-branch allocation requires 2.993–3.085 ms per image, versus 4.847–4.977 ms for single-network alternatives.Measurements cover extractor-only inference on two RTX Pro 6000 GPUs; adversarial training has little effect on latency.Source: §4.1.3 · p. 9, Table 2 · p. 11
4.2.1 Robustness comparison protocol
- The matched SD v2.1 robustness comparison uses 1,000 COCO captions, shared generation and watermark settings, and ten FARI-style image distortions.It reports mean GS bit accuracy, detection rate, and traceability rate.Source: §4.2.1 · p. 10, Table 3 · p. 11
- The two-branch extractor uses 14.394 GFLOPs and 7.148 GMACs, about 47.4× fewer FLOPs than FARI and 47.2× fewer than OSI.It has 15.468M total parameters, compared with 868.469M for FARI and 865.911M for OSI.Source: §4.2.1 · p. 10
5 Conclusion
The conclusion frames Gaussian Shading extraction as decision-region recovery and proposes a prompt-free two-NAFNet reverse estimator. Its stated goal is one-step extraction, and the accompanying backbone-cost table measures a latent-to-noise pass while excluding image encoding and GS decoding.
5 Conclusion
- Gaussian Shading needs enough Gaussian decision-region labels to survive decryption and voting, not exact reconstruction of the terminal latent.Extended flow matching and DDPM-style epsilon regression learn conditional statistics of hidden endpoints rather than sample-specific hidden pairs.Source: §5 · p. 11
- The proposed prompt-free two-NAFNet reverse estimator uses intentional L1 loss and finite-set adversarial image distortions for decision-region recovery.Its two-branch anchor/interpolation parameterization is an optimization method, not a claim of a newly derived continuous trajectory.Source: §5 · p. 11
- Table 4 measures the backbone cost of one latent-to-noise extraction pass, excluding image encoding and GS decoding.Source: Table 4 · p. 12
6 Further Discussion
The discussion questions whether FARI’s trajectory explanation supports its one-step interpretation, distinguishing its regression objective from distillation objectives and pointing to DMD2’s removal of regression loss.
6 Further Discussion
- FARI’s simple regression loss differs from SOTA distillation losses [15] [19] [29] [30], which tend to generate a sample from the target distribution directly in one step.Source: §6 · p. 12
- The discussion says FARI’s objective matches extended RF-2 in the high-SNR stage [13], so one-step behavior depends on RF-2 rather than RF-1 trajectory characteristics.Source: §6 · p. 12
- DMD2 [28] removes regression loss because it violates the distribution-matching target, which the discussion offers as evidence against regression directly generating a real target-distribution sample.Source: §6 · p. 12