MoE-GFN: Forensic Module in StableGuard
- The paper introduces MoE-GFN, an end-to-end trainable UNet that performs both watermark verification and tamper localization within the latent diffusion framework.
- MoE-GFN leverages a Mixture-of-Forensic-Experts block, fusing global, local, and frequency-domain cues through dynamic soft routing for precise forensic evidence integration.
- Empirically, MoE-GFN achieves near-perfect watermark extraction and superior tamper localization metrics under varied image degradations, demonstrating its robust performance.
Mixture-of-Experts Guided Forensic Network (MoE-GFN) is the forensic component of StableGuard, a unified framework for copyright protection and tamper localization in Latent Diffusion Models. In the StableGuard formulation, watermarked generation is written as , where is latent and is the bit watermark, while forensics is written as , where is a possibly tampered image, is the recovered watermark, and is the tampering mask. MoE-GFN is exactly : a single, end-to-end trainable network that performs reliable watermark verification / bit extraction and precise tampering localization from a potentially tampered image (Yang et al., 22 Sep 2025).
1. Position within StableGuard
StableGuard couples generation and forensics directly inside a latent diffusion model through two main modules. The first is MPW-VAE (Multiplexing Watermark VAE), a modified VAE decoder that can produce both watermarked and watermark-free images from the same latent. The second is MoE-GFN, which receives images produced by MPW-VAE, or synthetic tampered versions thereof, and jointly predicts a global watermark bit vector and a dense tampering mask . This design makes watermark extraction and tamper localization part of the same optimization target rather than post hoc add-ons (Yang et al., 22 Sep 2025).
The role of MoE-GFN is inseparable from StableGuard’s supervision strategy. MPW-VAE produces paired watermarked and watermark-free outputs, and these are fused with masks to create self-supervised forensic training samples. MoE-GFN then learns to exploit the watermark’s holistic perturbation patterns and to identify where those patterns are missing or inconsistent. In this arrangement, intact regions support watermark recovery, while altered regions provide the signal for localization.
This coupling is central to the framework’s forensic interpretation. Because the watermark is embedded during image generation rather than appended after generation, the forensic network is trained on watermark-bearing content that is native to the latent diffusion pipeline. A plausible implication is that the learned forensic cues are less dependent on brittle post-processing artifacts and more dependent on structured watermark consistency across space.
2. Architectural organization
MoE-GFN is a UNet-based forensic network. Its overall structure consists of a Stem module for dimensionality reduction and initial feature extraction, a UNet backbone with an encoder of multi-scale convolutional layers and a decoder with upsampling layers and skip connections, a Mixture-of-Forensic-Experts (MoFE) block placed in the decoder, and two forensic heads: one for watermark prediction and one for mask prediction (Yang et al., 22 Sep 2025).
The MoFE block is the defining architectural element. It contains three experts and a Dynamic Soft Router (DSR). The experts correspond to distinct forensic cue families: a Watermark Extraction Expert (WEE) for holistic watermark patterns, a Tampering Localization Expert (TLE) for local tampering artifacts, and a Boundary Enhancement Expert (BEE) for frequency-domain boundary cues. The router produces soft weights conditioned on the current features and fuses the expert outputs in an input-adaptive way.
| Component | Operating principle | Forensic role |
|---|---|---|
| WEE | Global, transformer-based expert | Captures holistic watermark patterns |
| TLE | Sub-patch transformer with 0 | Captures local, fine-grained tampering artifacts |
| BEE | Frequency-domain transformer with FFT/iFFT | Highlights high-frequency edges of tampered regions |
| DSR | Lightweight MLP with softmax weights | Fuses experts adaptively |
Formally, if 1 is the input feature map to MoFE, each expert defines a transformed feature map 2, and the router computes broadcastable soft weights 3. The unified representation is
4
This is a standard mixture-of-experts formulation with soft routing, but specialized to forensic evidence integration. WEE applies flattening, projection, and transformer processing over global spatial structure. TLE partitions the feature map into sub-patches of size 5, reshapes them, and applies transformer processing at sub-patch granularity. BEE applies FFT, processes the transformed representation with a transformer in the frequency domain, and returns to the spatial domain with iFFT.
The output heads are lightweight. The tampering head is a two-layer convolutional network that predicts a one-channel mask 6. The watermark head uses two convolutional layers and a fully connected layer to predict a bit vector 7, either as probabilities in 8 or as thresholded binary outputs in 9, where 0 is the watermark length.
3. Inputs, outputs, and self-supervised training signal
The direct input to MoE-GFN is a potentially tampered RGB image 1 of size 2. This image is synthesized from MPW-VAE outputs. MPW-VAE produces a watermark-free reconstruction 3 and a watermarked image 4, and then a random binary mask 5 is used to fuse watermarked and clean content according to
6
where 7 is the original real image and 8 is the VAE reconstruction. In the stated interpretation, 9 simulates human edits and 0 simulates AI edits (Yang et al., 22 Sep 2025).
MoE-GFN has two outputs. The first is watermark verification / extraction: an output vector 1 of length 2, trained with binary cross-entropy and used for watermark existence checking and bit-accurate extraction. The second is tampering localization: a probabilistic segmentation mask 3, which becomes a binary tamper mask after sigmoid and thresholding.
The training process is fully self-supervised in the sense stated for StableGuard: no human-annotated tampering masks are needed. The ground-truth tamper mask is exactly the random synthesis mask 4, and the watermark target is the embedded bit vector 5. The total objective is
6
The similarity term preserves visual quality between the watermark-free reconstruction and the watermarked image:
7
where PS is a perceptual similarity metric (LPIPS). The watermark term is binary cross-entropy over bits:
8
The tampering term combines weighted BCE and Dice loss:
9
The training algorithm proceeds by encoding 0 to latent 1, decoding 2 and 3, generating a random mask 4, forming 5, predicting 6, computing 7, and updating both the watermark adapter and MoE-GFN. This is the mechanism through which watermark embedding and forensic decoding co-adapt.
4. Guided evidence fusion and tampering-agnostic behavior
The term “Guided” in MoE-GFN refers to the Dynamic Soft Router’s control over expert fusion. The router computes weights 8 conditioned on the input features, and the unified representation is
9
The described visualization indicates that 0 is more activated in intact, watermarked regions, 1 is concentrated in corrupted or edited regions, and 2 highlights edges of tampered regions. Correspondingly, 3 dominates in untampered areas, 4 dominates around suspected manipulations, and 5 peaks at boundaries (Yang et al., 22 Sep 2025).
This expert partition matches the heterogeneity of forensic signals. Holistic watermark patterns are inherently distributed across the image. Tampering traces are local and often fine-grained. Boundary evidence is frequently stronger in the frequency domain than in raw spatial activations. The MoFE block therefore acts as a cue-type decomposition mechanism rather than a simple capacity multiplier.
MoE-GFN is also trained in a tampering-agnostic manner. The random mask 6 is not tied to a specific editing operator; it only defines watermarked versus non-watermarked regions. Two mask types are used: random binary masks and semantic masks from SAM, unpaired with explicit labels. Combined with 7 and 8, this yields a synthetic training regime that does not assume a particular tampering operator such as cropping, inpainting, or replacement.
The reported generalization targets reflect that strategy. MoE-GFN is stated to generalize to Stable Diffusion inpainting, SDXL, Kandinsky inpainting, ControlNet guiding edits, LaMa removal/inpainting, and to splicing, copy-paste, removal, and generic inpainting. This suggests that the model is trained to detect where watermark consistency disappears rather than how the modification was produced.
5. Empirical profile
The StableGuard experiments attribute strong watermark verification performance to the combined MPW-VAE and MoE-GFN design. For watermarking, bit accuracy on COCO and T2I is reported as up to 99.97–99.98% for a 32-bit watermark and >99.8% across different bit lengths. The comparison set includes HiDDeN, SepMark, WOUAF, Stable Signature, WaDiff, EditGuard, OmniGuard, and WAM (Yang et al., 22 Sep 2025).
For tampering localization on the AIGC tampering dataset, the reported results across five editing engines—SD inpainting, SDXL, Kandinsky, ControlNet, and LaMa—are F1: 9–0, AUC: 1–2, and IoU: 3–4. These are stated to surpass passive detectors such as MVSS-Net, IML-ViT, PSCC-Net, ObjectFormer, and HDF-Net, as well as proactive watermark-based detectors such as EditGuard, OmniGuard, and WAM.
The robustness profile is similarly explicit. Under Gaussian noise, JPEG compression, and Poisson noise, Bit Acc remains 5–6 and F1 stays high (>0.90). Under real-world degradations including JPEG/WebP/HEIC, severe downscaling, and brightness/contrast/saturation changes, StableGuard is reported to maintain high bit accuracy and reasonable F1. Under tampering ratios up to 90%, bit accuracy remains 7 at 90% tampering and >99.9% below 70%. Across tampering types—splicing, copy-paste, removal, and inpainting—the reported range is F1 8–9 and IoU 0–1.
Ablation evidence identifies the MoFE block and the Dynamic Soft Router as material contributors. Removing MoFE and using a plain UNet drops F1 from 0.980 to 0.931 and AUC from 0.992 to 0.920. Replacing DSR with naive summation (“w/o DSR”) reduces performance to F1 = 0.966 vs 0.980. The chosen decoder placement for MoFE (“Dec”) has 52.02M parameters and 78.51G FLOPs. For MoE-GFN specifically, the reported latency is 91.88 ms per image and VRAM usage is 428.32 MB.
6. Relation to broader mixture-of-experts research
Later mixture-of-experts work provides an interpretive framework for MoE-GFN’s routing and specialization, although these connections are not part of StableGuard’s original claim set. “Mixture of Experts with Soft Nearest Neighbor Loss: Resolving Expert Collapse via Representation Disentanglement” analyzes expert collapse as a condition in which one expert dominates routing or multiple experts learn redundant representations. It introduces Expert Specialization Entropy (ENT) and Pairwise Embedding Similarity (SIM) as diagnostics, and shows that SNNL-augmented MoE models can produce structurally diverse experts and more flexible routing on FashionMNIST, CIFAR10, and CIFAR100 (Agarap et al., 20 Mar 2026).
This suggests a direct diagnostic language for MoE-GFN. If a forensic MoE becomes dominated by one expert or if experts converge to similar filters, ENT and SIM would provide measurable indicators of collapse. The same paper further suggests that representation disentanglement before routing can improve specialization when the latent space is highly entangled. A plausible implication is that MoE-GFN could benefit from similar diagnostics or pre-conditioning when forensic evidence types become difficult to separate.
“Guided by the Experts: Provable Feature Learning Dynamic of Soft-Routed Mixture-of-Experts” provides a different kind of context: a theoretical analysis of jointly trained, soft-routed MoE models in a student-teacher setting. Its central conclusion is that the router’s learning process is “guided” by the experts, that sequential feature recovery occurs under moderate over-parameterization, and that redundant experts can be pruned and then fine-tuned with provable convergence (Liao et al., 8 Oct 2025).
For MoE-GFN, this does not constitute a direct guarantee, because StableGuard uses a deep UNet-based forensic network rather than the paper’s shallow student-teacher construction with Gaussian inputs and Hermite activations. Even so, the theory supplies a plausible conceptual reading of MoE-GFN’s Dynamic Soft Router: specialized expert features may be driving routing behavior rather than merely being selected by it. The same theoretical work also makes pruning and fine-tuning a natural extension for over-parameterized forensic MoE systems.
7. Limitations and extension paths
The StableGuard description identifies three direct limitations of MoE-GFN. First, localization accuracy is more sensitive than watermark accuracy to very aggressive compression or downscaling, because boundary cues are fragile. Second, the system is image-only and is not directly applied to video diffusion models. Third, expert capacity is fixed: WEE, TLE, and BEE are predefined expert types, with no online creation or removal of experts (Yang et al., 22 Sep 2025).
The extension paths follow from those constraints. The discussion explicitly notes video diffusion as a natural direction, since watermarking is applied through the VAE decoder; this would require temporal attention and a spatio-temporal MoFE. Other listed possibilities include adding experts for color inconsistencies, camera-model artifacts, or compression-specific traces, and using sparse gating or hierarchical MoE for larger capacity. Cross-modal forensics is also proposed, extending to multimodal latent diffusion models by leveraging text-conditional inconsistencies. For degradation robustness, the stated options are training with simulated compression and resampling in the MPW-VAE + MoE-GFN loop and designing degradation-aware experts.
Taken together, these limitations and extensions situate MoE-GFN as a specialized forensic module within latent-diffusion watermarking rather than a general forensic architecture. Its current form is defined by three interacting properties: end-to-end coupling to the generator, dynamic fusion of holistic, local, and frequency-domain forensic cues, and self-supervised tampering-agnostic training. Those properties determine both its reported strengths and the boundaries of its present scope.