WMFormer++: Advanced Watermark Removal

• The paper introduces WMFormer++ as an advanced neural architecture that leverages a nested Transformer design to implicitly integrate watermark localization with image restoration.
• It employs a unified network with cross-channel attention, gated-Dconv FFN, and nested multi-scale decoding for effective feature fusion and detail preservation.
• Empirical results on LOGO and CLWD benchmarks demonstrate significant improvements in PSNR, SSIM, and watermark mask accuracy over previous methods.

WMFormer++ is an advanced neural architecture for visible watermark removal that introduces a nested Transformer with implicit joint learning to address copyright protection adversarially and inform robust watermarking strategies. Unlike prior approaches that rely on explicitly separated decoder branches for watermark localization and background restoration, WMFormer++ integrates these components holistically through implicit information routing, cross-channel attention, and multi-scale nested decoding (Huo et al., 2023).

1. Network Architecture and Nested Transformer Design

WMFormer++ operates on a watermarked RGB image $J \in \mathbb{R}^{H \times W \times 3}$, which is concatenated with an all-zero single-channel map as input. The embedding layer consists of a $3 \times 3$ convolution followed by a $1 \times 1$ convolution, projecting the input to a feature tensor $X_0$. The encoder comprises four stages of Transformer blocks, with each stage containing a Cross-Channel Multi-Head Attention (CC-MHA) module and a Gated-Dconv Feed-Forward Network (FFN). At each stage, spatial resolution halves and the number of feature channels doubles.

Latent codes $\{Z_1, Z_2, Z_3\}$ are extracted at the deepest three encoder layers. The nested decoder constructs a UNet-style pathway from each $Z_d$ of depth $d$. Pathways are densely interconnected at every spatial level to amalgamate features across resolutions, similar to a UNet++ structure. The final shared prediction head on each pathway produces both a restored image $\hat{I}^d$ and a watermark mask $\hat{M}^d$, but only $\hat{I}^3$ is used at test time for final output. This design enhances multi-scale information integration and end-to-end learnability.

2. Implicit Joint Learning Paradigm

Conventional separation between localization and restoration is replaced with implicit joint learning in WMFormer++. Each Transformer block employs the Gated-Dconv FFN, which splits projected features into a “gate” and “content.” The gate modulates the content through element-wise multiplication before recombination with the original features. This implicit routing enables the model to autonomously exchange signals between watermark localization and background restoration without explicit inter-branch connectors, reflecting the mutual dependence of these tasks.

3. Transformer Block: Gated-Dconv FFN and Cross-Channel Attention

The two principal components of each Transformer block are:

Gated-Dconv Feed-Forward Network (FFN):

$$U = X W_1, \quad [U_1, U_2] = \mathrm{split\_half}(U)$$

$$\mathbf{g} = \mathrm{GELU}(U_1), \quad Y = \left(\mathbf{g} \odot U_2\right) W_2 + X$$

where $W_1 \in \mathbb{R}^{C \times 2C}$, $W_2 \in \mathbb{R}^{C \times C}$, and $\odot$ denotes channel-wise multiplication. The gating vector $\mathbf{g}$ dynamically regulates channel content flow.

Cross-Channel Multi-Head Attention (CC-MHA):

Local context is aggregated via a $1 \times 1$ convolution and a $3 \times 3$ depth-wise convolution over $X$. Multi-head attention computes across channels: $$Q_i = XW^Q_i, \quad K_i = XW^K_i, \quad V_i = XW^V_i$$

$$A_i = \mathrm{softmax}(\sigma K_i^\top Q_i), \quad \mathrm{head}_i = V_i A_i$$

The outputs from all $h$ heads are concatenated and projected, with $\sigma$ a learnable scaling parameter. Channel attention extracts global semantic cues for localization, while the convolution focuses on finer details required for restoration.

4. Nested Multi-Scale Decoding and Feature Fusion

From $Z_d$ at each encoder stage $d \in \{1, 2, 3\}$, a dedicated decoder pathway unrolls with corresponding depth. At each spatial level $\ell$, decoder features $F^d_\ell$ from all eligible depths are fused using: $$F_\ell = \sum_{d:\,\ell\le d} \alpha^d_\ell\, \mathrm{UpSample}(F^d_{\ell+1}) + \sum_{d:\,\ell\le d} \beta^d_\ell\,(\mathrm{skip}(E_\ell))$$ where $\alpha^d_\ell$ and $\beta^d_\ell$ are learned weights and $E_\ell$ denotes encoder features. This nested architecture harmonizes textures from shallow decoders with global structures from deeper pathways, enhancing detail preservation and structural coherence in restored images.

5. Optimization Objectives and Deep Supervision

For each decoder depth $d \in \{1, 2, 3\}$, the model predicts both a reconstructed image $\hat{I}^d$ and a watermark mask $\hat{M}^d$. Supervision encompasses:

• Watermark mask loss (binary cross-entropy):

$$\mathcal{L}_{\text{mask}}^d = -\sum_{i, j} \left[(1 - M_{ij}) \log(1 - \hat{M}^d_{ij}) + M_{ij} \log(\hat{M}^d_{ij})\right]$$

• Image reconstruction loss ($L_1$):

$\mathcal{L}_{\text{image}}^d = \|I - \hat{I}^d\|_1$

• Perceptual loss (VGG):

$$\mathcal{L}_{\text{perc}}^d = \sum_{k=1}^3 \left\|\Phi^k_\text{VGG}(\hat{I}^d) - \Phi^k_\text{VGG}(I)\right\|_1$$

• Total objective:

$$\mathcal{L}_\text{all} = \sum_{d=1}^3 \lambda^d\left(\mathcal{L}_{\text{image}}^d + \lambda_{\text{mask}} \mathcal{L}_{\text{mask}}^d + \lambda_{\text{perc}} \mathcal{L}_{\text{perc}}^d\right)$$

Deep supervision at every decoder depth facilitates stable convergence and effective multi-scale learning.

6. Datasets, Benchmarking Protocols, and Metrics

Evaluation utilizes the following benchmarks:

Dataset	Watermark Type	Train/Test Sizes	Variation Parameters
LOGO‐L	Modest color	12,151 / 2,025	Standard visual watermarks
LOGO‐H	Large/transparent	≈12k / ≈2k	Larger, higher opacity
LOGO-Gray	Grayscale	≈12k / ≈2k	Grayscale watermarks
CLWD	Randomized logos	60k / 10k (PASCAL VOC base)	Size, position, rotation, $\alpha$

Metrics include PSNR (↑), SSIM (↑), LPIPS (↓) for LOGO series, and PSNR (↑), SSIM (↑), RMSE (↓), and RMSE$_w$ (↓) for CLWD.

7. Comparative Performance and Empirical Characterization

On all benchmarks, WMFormer++ achieves state-of-the-art results. Representative improvements:

• LOGO-H: PSNR 44.64 vs. 40.83, SSIM 0.9950 vs. 0.9919, LPIPS 0.50 vs. 0.89
• LOGO-L: PSNR 47.05 vs. 44.24, SSIM 0.9970 vs. 0.9954, LPIPS 0.31 vs. 0.54
• LOGO-Gray: PSNR 46.29 vs. 42.60, SSIM 0.9970 vs. 0.9944, LPIPS 0.21 vs. 0.53
• CLWD: PSNR 39.36 vs. 38.28, SSIM 0.9830 vs. 0.9814, RMSE 3.25 vs. 3.76, RMSE$_w$ 11.47 vs. 14.07

For CLWD watermark mask prediction: F1 0.8769, IoU 81.38% (best prior: F1 0.8234, IoU 74.63%).

Qualitative analysis identifies sharply defined watermark masks and background reconstructions without boundary artifacts or color shifts. The architecture’s reliance on channel attention, implicit joint learning, and nested decoders—without ad-hoc inter-branch modules—accounts for these gains (Huo et al., 2023).