MPW-VAE: Watermarking in Latent Diffusion Models

• The paper introduces MPW-VAE, integrating imperceptible watermarking within a frozen VAE decoder to produce dual outputs from a single latent code.
• It leverages latent residual-based adapters and a multiplexing design to generate both watermark-free and watermarked images for joint reconstruction and forensic training.
• Empirical results show high image fidelity (≈40 dB PSNR), near-perfect bit accuracy (~99.9%), and robust tamper localization, outperforming prior watermarking methods.

Searching arXiv for the cited paper and closely related watermarking methods mentioned in the provided data. The Multiplexing Watermark VAE (MPW-VAE) is the watermarking component of StableGuard, a diffusion-native framework for latent diffusion models that integrates an imperceptible holistic binary watermark into the image generation process while using the same watermark patterns as cues for tamper localization (Yang et al., 22 Sep 2025). In StableGuard, MPW-VAE replaces or wraps the standard VAE decoder at the second stage of the latent diffusion pipeline, after denoising has produced a clean latent $z_0$. Its defining property is multiplexing: from a shared latent code $z$, it can generate both a watermark-free reconstruction and a watermarked image, thereby coupling copyright protection with self-supervised forensic training in a single end-to-end design.

1. Position in the StableGuard framework

In a standard latent diffusion model such as Stable Diffusion, the pipeline first denoises sampled noise into a clean latent $z_0$ through the diffusion UNet and then decodes $z_0$ into an image $X$ through the VAE decoder. StableGuard inserts MPW-VAE at this second stage. The original decoder $\mathcal{D}_\nu$ is kept and frozen, while a watermark adapter is inserted after each decoder block to form the MPW-VAE decoder $\mathcal{D}_\phi$ (Yang et al., 22 Sep 2025).

This arrangement allows the adapters to be toggled on or off. From the same latent $z$, the system can therefore produce a watermark-free reconstruction

$\hat{X} = \mathcal{D}_\nu(z)$

and a watermarked image

$Y = \mathcal{D}_\phi(z).$

These paired outputs are central to the design. StableGuard uses the pairs $z$0, together with random fusion masks, to synthesize tampered images $z$1 and corresponding ground-truth tamper masks $z$2 for training the Mixture-of-Experts Guided Forensic Network (MoE-GFN). Within this overall system, MPW-VAE is both the embedding module, which ties watermarking to the latent diffusion model’s generative process, and the multiplexing source that makes self-supervised forensic learning possible.

A common misconception would be to treat MPW-VAE as a conventional post hoc watermarking module attached after generation. The StableGuard formulation explicitly differs from that pattern: the watermark is integrated into the decoder path of the latent diffusion model rather than appended as a separate image-space encoder (Yang et al., 22 Sep 2025).

2. Architectural formulation and watermark injection

StableGuard does not retrain the base VAE. The encoder $z$3 and decoder $z$4 are frozen, while MPW-VAE augments the decoder using latent residual-based adapters inserted after each decoder block (Yang et al., 22 Sep 2025). In the standard VAE setting, the encoder is $z$5, the decoder is $z$6, and the classical objective is

$z$7

with $z$8.

MPW-VAE preserves this pretrained backbone and injects the watermark through intermediate decoder features. Let the feature map at decoder block $z$9 be $z_0$0, and let the watermark be a binary bit vector

$z_0$1

The adapter $z_0$2 first encodes $z_0$3 through two fully connected layers,

$z_0$4

reshapes the result to match the feature map,

$z_0$5

concatenates it with the original block features, and applies two convolutional blocks followed by a residual addition: $z_0$6

$z_0$7

$z_0$8

The resulting watermark-conditioned decoder is

$z_0$9

whereas disabling the adapters recovers the vanilla decoder

$z_0$0

The paper characterizes this residual design as minimally perturbing the original decoder features while keeping image fidelity high and encoding watermark patterns that are subtle but traceable (Yang et al., 22 Sep 2025). It also states that the adapters are hierarchical, operating at different spatial scales after each decoder block, which is used to enhance watermark robustness.

3. Multiplexing, paired outputs, and synthetic tampering data

In StableGuard, “multiplexing” means that a single latent code $z_0$1 can be transformed into two consistent outputs with different forensic properties: a watermark-free reconstruction $z_0$2 and a watermarked version $z_0$3 (Yang et al., 22 Sep 2025). For each input image $z_0$4, the process is

$z_0$5

This yields the triplet $z_0$6, where $z_0$7 is the original real image, $z_0$8 is the reconstruction by the vanilla VAE, and $z_0$9 is the watermarked reconstruction.

StableGuard then synthesizes tampered examples by splicing watermarked and watermark-free regions. Let $X$0 be a random binary mask or a semantic mask from SAM. The fused image $X$1 is defined as

$X$2

Here, $X$3 carries watermark everywhere, whereas $X$4 or $X$5 supplies watermark-free source regions. Pixels with $X$6 are therefore “tampered” relative to the watermarked baseline, and the ground-truth tamper mask is exactly $X$7.

Mask generation uses two strategies. With 50% probability, StableGuard uses random binary masks composed of pure noise patterns. With 50% probability, it uses semantic masks from SAM segmentations, with masks kept non-paired across images to maximize shape diversity. The dataset construction is further used with different editing backbones—SD, SDXL, Kandinsky, ControlNet, and LaMa—in the AIGC tampering dataset (Yang et al., 22 Sep 2025).

This design produces a rich, tampering-agnostic dataset in which MoE-GFN is exposed to object-level and region-level splicing, to both “human-like” edits using original $X$8 and AI-generated content using $X$9, and to precise labels for tampered regions and watermark integrity without manual annotation. A plausible implication is that the term “multiplexing” refers not merely to dual decoding, but to the generation of matched supervision signals for both attribution and localization.

4. Watermark representation, objectives, and self-supervised optimization

StableGuard uses a binary bit watermark

$\mathcal{D}_\nu$0

with reported bit lengths

$\mathcal{D}_\nu$1

The watermark is embedded inside the VAE decoder through adapters; it is not added directly to the latent $\mathcal{D}_\nu$2 and not added directly to image pixels (Yang et al., 22 Sep 2025). The decoder may be written as

$\mathcal{D}_\nu$3

The data describe the induced perturbation as holistically spread over the image rather than taking the form of a local logo or visible mark.

Training is governed by the overall StableGuard objective. The similarity loss preserves fidelity between the watermarked image $\mathcal{D}_\nu$4 and the clean reconstruction $\mathcal{D}_\nu$5: $\mathcal{D}_\nu$6 where $\mathcal{D}_\nu$7 is a perceptual similarity term described as LPIPS-style feature distance. The watermark loss supervises extraction of the bit vector from images, including tampered ones: $\mathcal{D}_\nu$8 The tampering loss is

$\mathcal{D}_\nu$9

with

$\mathcal{D}_\phi$0

$\mathcal{D}_\phi$1

The total loss is

$\mathcal{D}_\phi$2

The training loop is explicitly self-supervised. Given a real image $\mathcal{D}_\phi$3, StableGuard encodes it into $\mathcal{D}_\phi$4, decodes $\mathcal{D}_\phi$5 and $\mathcal{D}_\phi$6, samples a mask $\mathcal{D}_\phi$7, constructs $\mathcal{D}_\phi$8, feeds $\mathcal{D}_\phi$9 to MoE-GFN, computes $z$0, and updates the adapter parameters of MPW-VAE and the forensic network parameters while keeping the base VAE encoder and decoder frozen (Yang et al., 22 Sep 2025). Supervision comes automatically from the known bit vector $z$1 and the constructed mask $z$2; no manual tamper annotation is required.

5. Interaction with MoE-GFN

MoE-GFN receives three types of images generated by MPW-VAE: pure watermarked images $z$3, tampered versions $z$4 built from $z$5 and $z$6, and watermarked reconstructions used for robustness evaluation. For each $z$7, it predicts both the watermark bits and the tamper mask: $z$8

Its internal MoFE block contains three experts (Yang et al., 22 Sep 2025). The Watermark Extraction Expert is a global transformer that captures global correlations of watermark patterns across the image: $z$9 The Tampering Localization Expert is a sub-patch transformer. It partitions the feature map into patches of size $\hat{X} = \mathcal{D}_\nu(z)$0, with $\hat{X} = \mathcal{D}_\nu(z)$1 in the implementation details, and is described as being sensitive to local inconsistencies indicative of tampering: $\hat{X} = \mathcal{D}_\nu(z)$2 The Boundary Enhancement Expert is a frequency-domain transformer: $\hat{X} = \mathcal{D}_\nu(z)$3

$\hat{X} = \mathcal{D}_\nu(z)$4

and is intended to highlight high-frequency edges where manipulations occur.

The outputs of the three experts are merged through a Dynamic Soft Router with router weights

$\hat{X} = \mathcal{D}_\nu(z)$5

where $\hat{X} = \mathcal{D}_\nu(z)$6 is a small MLP, producing a unified forensic feature

$\hat{X} = \mathcal{D}_\nu(z)$7

The interaction between MPW-VAE and MoE-GFN is not unidirectional. StableGuard describes a reciprocal relation in which MPW-VAE improves forensic accuracy by supplying paired, consistent watermarked and clean samples together with rich tampering variants, while MoE-GFN guides watermark embedding through backpropagation so that the embedded patterns become more useful for detection and localization (Yang et al., 22 Sep 2025). Ablation in Tab. 12 is reported to show that training MPW-VAE and MoE-GFN separately, without joint optimization, yields poorer localization than joint training.

6. Empirical characteristics, implementation, and relation to prior methods

The reported empirical profile of MPW-VAE is defined by three simultaneous targets: image fidelity, watermark robustness, and tamper localization (Yang et al., 22 Sep 2025). On COCO and T2I datasets, MPW-VAE is reported to achieve PSNR of approximately $\hat{X} = \mathcal{D}_\nu(z)$8–$\hat{X} = \mathcal{D}_\nu(z)$9 dB, SSIM of approximately $Y = \mathcal{D}_\phi(z).$0, LPIPS of approximately $Y = \mathcal{D}_\phi(z).$1–$Y = \mathcal{D}_\phi(z).$2, and FID of approximately $Y = \mathcal{D}_\phi(z).$3–$Y = \mathcal{D}_\phi(z).$4. The same section states that bit accuracy on clean images reaches up to 99.98%, and that compared to methods such as WAM, OmniGuard, WaDiff, and Stable Signature, MPW-VAE improves perceptual quality and FID while achieving near-perfect bit accuracy of about 99.9%.

Under noise, JPEG, Poisson, and many real-world degradations, the combination of MPW-VAE and MoE-GFN is reported to maintain bit accuracy above 99.6% in most noise and compression conditions, while F1 scores for tampering localization remain above 0.9 in many settings. Under varying tampering ratios, the paper reports bit accuracy of approximately 99.96–99.98% at 10–50% tampered area, and approximately 89.6% even at 90% tampering, outperforming post hoc methods. On the AIGC tampering dataset, the reported localization performance across SD-Inpainting, SD-XL, Kandinsky, ControlNet, and LaMa is $Y = \mathcal{D}_\phi(z).$5–$Y = \mathcal{D}_\phi(z).$6, $Y = \mathcal{D}_\phi(z).$7–$Y = \mathcal{D}_\phi(z).$8, and $Y = \mathcal{D}_\phi(z).$9–$z$00. On traditional datasets including CASIA, NIST, Columbia, Coverage, and IMD2020, StableGuard is reported often to exceed EditGuard, OmniGuard, and WAM, with $z$01–$z$02 and $z$03–$z$04.

The reported trade-offs are correspondingly explicit. Increasing the bit length from $z$05 to $z$06 slightly lowers PSNR and SSIM and increases LPIPS, while image quality remains high at approximately $z$07 dB PSNR. Stronger watermark embedding can marginally improve robustness but must be balanced with $z$08 to avoid artifacts. Joint optimization is described as crucial, because decoupling MPW-VAE from MoE-GFN reduces localization performance significantly. This suggests that MPW-VAE should be understood not simply as a watermark encoder, but as a component whose utility depends on co-adaptation with the forensic model.

Implementation details situate MPW-VAE within Stable Diffusion 2.1. The architecture uses the original VAE encoder $z$09 and decoder $z$10, inserts watermark adapters composed of two FC layers, two conv layers, and a residual connection after each decoder block, and trains with Adam at learning rate $z$11 for 10 epochs with batch size 16 on $z$12 RTX 4090D. The reported loss weights are $z$13, $z$14, and $z$15. In text-to-image generation, after denoising latent $z$16 with the UNet, decoding through $z$17 yields watermarked images, while the original decoder remains available for comparison and clean reconstruction.

The complexity overhead is described as small. Tab. 21 reports vanilla VAE latency of approximately 12.69 ms and MPW-VAE latency of approximately 14.99 ms; VRAM rises from 1682 MB to 2034 MB for the decoder, while MoE-GFN uses approximately 428 MB. The paper therefore characterizes the added overhead as negligible relative to the base VAE.

Relative to prior watermarking approaches, StableGuard distinguishes MPW-VAE from two families. First, post hoc image-space methods such as HiDDeN, SepMark, WAM, EditGuard, and OmniGuard attach an external encoder and decoder after generation, incur extra runtime and potential degradation, are decoupled from the diffusion process, and are often optimized separately from forensics. Second, diffusion-native latent watermarking methods such as Stable Signature, WOUAF, and WaDiff modify noise, UNet weights, or latent routing, but are described as focusing on source attribution or simple watermark detection rather than fine-grained tamper localization or self-supervised forensic training (Yang et al., 22 Sep 2025).

Within that comparison, MPW-VAE is defined by three features: a decoder-side latent residual adapter that directly modulates decoder features conditioned on a bit vector; a multiplexing design in which a single latent $z$18 yields both watermarked and watermark-free outputs; and joint optimization with MoE-GFN, through which embedding and forensics are trained together in a single end-to-end system. In the formulation of StableGuard, this shifts watermarking from post hoc, isolated watermarking to embedded, multiplexed, forensic-aware watermarking inside the latent diffusion pipeline.