Source-linked AI summary
MBRS : Enhancing Robustness of DNN-based Watermarking by Mini-Batch of Real and Simulated JPEG Compression
Zhaoyang Jia, Han Fang, Weiming Zhang
TL;DR
Existing DNN watermarking frameworks struggle with JPEG compression because JPEG is non-differential and simulated or two-stage training may fail to optimize robustness globally. MBRS randomly alternates real JPEG, simulated JPEG, and noise-free layers across mini-batches, augmented with SE blocks, a message processor, and an optional diffusion block; it achieves BER below 0.01% under JPEG compression while maintaining high image quality.
Problem
Existing DNN-based watermarking methods do not robustly handle JPEG compression, whose non-differential quantization prevents direct gradient propagation and limits training.
Method
MBRS randomly selects real JPEG, simulated JPEG, or a noise-free layer for each mini-batch, using SE blocks, a message processor, and optional diffusion blocks for crop robustness.
Results
The model achieves BER below 0.01% under JPEG compression and performs better in JPEG robustness and image quality.
Takeaways & Limitations
MBRS provides strong robustness against JPEG and other distortions, including Gaussian filtering, dropout, crop, and cropout attacks.
Takeaways & Limitations
During message expansion, the secret message length and cover-image shape must follow a specific relationship determined by the architecture.
Abstract
from arXiv · showhide
Based on the powerful feature extraction ability of deep learning architecture, recently, deep-learning based watermarking algorithms have been widely studied. The basic framework of such algorithm is the auto-encoder like end-to-end architecture with an encoder, a noise layer and a decoder. The key to guarantee robustness is the adversarial training with the differential noise layer. However, we found that none of the existing framework can well ensure the robustness against JPEG compression, which is non-differential but is an essential and important image processing operation. To address such limitations, we proposed a novel end-to-end training architecture, which utilizes Mini-Batch of Real and Simulated JPEG compression (MBRS) to enhance the JPEG robustness. Precisely, for different mini-batches, we randomly choose one of real JPEG, simulated JPEG and noise-free layer as the noise layer. Besides, we suggest to utilize the Squeeze-and-Excitation blocks which can learn better feature in embedding and extracting stage, and propose a "message processor" to expand the message in a more appreciate way. Meanwhile, to improve the robustness against crop attack, we propose an additive diffusion block into the network. The extensive experimental results have demonstrated the superior performance of the proposed scheme compared with the state-of-the-art algorithms. Under the JPEG compression with quality factor Q=50, our models achieve a bit error rate less than 0.01% for extracted messages, with PSNR larger than 36 for the encoded images, which shows the well-enhanced robustness against JPEG attack. Besides, under many other distortions such as Gaussian filter, crop, cropout and dropout, the proposed framework also obtains strong robustness. The code implemented by PyTorch \cite{2011torch7} is avaiable in https://github.com/jzyustc/MBRS.
1 INTRODUCTION
The paper targets weak JPEG robustness in DNN-based image watermarking by combining real and simulated compression during mini-batch training. It also augments the network with feature-learning, message-expansion, and diffusion components for image quality and robustness against crop attacks.
- Problem: JPEG compression is difficult for DNN-based watermarking because quantization is non-differential and simulated replacements do not reliably reproduce real-noise behavior.Existing approaches using OEDS or TSR remain insufficient under JPEG compression.
- MBRS method: MBRS randomly switches each mini-batch among real JPEG, simulated JPEG, and noise-free layers to train different parts of the watermarking task.Real JPEG trains decoder robustness, simulated JPEG supports joint encoder-decoder training, and the identity layer preserves decoding without compression.
- Network design: The proposed auto-encoder uses Squeeze-and-Excitation blocks and a message processor to improve feature learning and expand secret messages with redundancy.A strength factor adjusts the trade-off between encoded-image quality and robustness.
- Results: The experiments report decoding bit error rates below 0.01% under JPEG compression and high image quality, with robustness extending to Gaussian filter, dropout, crop, and cropout distortions.The contribution summary also emphasizes stronger JPEG robustness and the combined architecture components.
- Crop robustness: A diffusion block and inverse diffusion block spread the secret message across the image to improve robustness against crop and cropout attacks.The paper reports that decoding accuracy under other noises does not decrease too much.
2 RELATED WORKS
Related work established auto-encoder watermarking, adversarial training, JPEG simulation, and two-stage real-distortion training, but these approaches remained inadequate for JPEG robustness. The paper frames the limitation as a mismatch between differentiable simulations, non-differential JPEG, and stage-wise optimization.
- Deep Learning for Digital Watermarking: DNN-based watermarking commonly uses an auto-encoder-like encoder-noise-layer-decoder structure, sometimes with adversarial training for more realistic visual results.Prior work also explored domain-specific residual encoders and simulated print-shooting distortions.
- Deep Learning for Digital Watermarking: Two-stage separable training initializes the encoder and decoder without noise, then enhances the decoder alone using non-differential distortions.This framework was proposed to address the differential limitation of the noise layer.
- JPEG robustness gap: Existing DNN-based methods achieve image quality and some processing robustness but do not adequately handle JPEG compression distortion.The related-work discussion attributes this to JPEG’s non-differential process and insufficient feature preservation for decoding.
- JPEG simulations: JPEG-Mask simulates compression by zeroing selected high-frequency coefficients while retaining specified low-frequency regions of the image channels.Other work approximates JPEG quantization, including a piece-wise approximation near zero called JPEG-SS.
3 PROPOSED FRAMEWORK
The proposed framework combines a message processor, SE-based encoder and decoder, adversarial training, and MBRS noise-layer switching for robust blind watermarking. It also adds diffusion blocks for crop robustness and a strength factor to balance image quality with message recovery.
- Model Architecture: The architecture processes the secret message into a feature map, embeds it into the cover image, applies randomly selected noise, and decodes the message.The discriminator distinguishes encoded from cover images during adversarial training.
- Message Processor: The message processor expands the message and introduces redundancy before the message feature map is combined with image features.Its expansion uses convolution, batch normalization, ReLU activation, and transposed convolutions.
- Encoder and Decoder: SE blocks extract image features in the encoder and decoder, while the encoder minimizes image distortion and the decoder minimizes message reconstruction error.The encoder targets similarity between cover and encoded images; the decoder targets similarity between original and decoded messages.
- MBRS Noise Layer: MBRS randomly switches each mini-batch among real JPEG, simulated JPEG, and Identity noise layers to support robustness, joint training, and uncompressed decoding.Real JPEG trains decoding after compression, simulated JPEG propagates gradients to the encoder, and Identity preserves decoding without compression.
- MBRS Noise Layer: Unlike TSR’s two-stage switch, MBRS changes noise layers per mini-batch and is designed to search across the solution space rather than optimize stages separately.Momentum-based optimization helps retain update direction across mini-batches involving real JPEG.
- Strength Factor: The strength factor scales the residual added to the cover image, trading off encoded-image quality against message-recovery accuracy in a blind watermarking architecture.The cover image is not used during decoding.
4 EXPERIMENTS
The experiments use separate ImageNet training and COCO testing data, evaluate robustness with BER and image quality with PSNR and SSIM, and compare against three deep-learning watermarking baselines.
- Implementation Details: 10,000 ImageNet images train the model, while 5,000 COCO test images evaluate generalization.The implementation uses PyTorch and an NVIDIA RTX 2080ti.
- Metrics: Robustness is measured by decoded-message Bit Error Rate, while image quality is measured by PSNR and SSIM.These are the study’s two main evaluation dimensions.
- Baselines: The baselines are methods [27],, and [22], using an open-source pretrained model for [22] and published results for [27] and.The authors report using published results where reproducing the reported best performance was unsuccessful.
4.1 Visual Quality
Visual examples under JPEG compression show high encoded-image quality and zero measured bit error, with residuals concentrated more strongly in textured regions.
- Residual Analysis: Residual visualizations indicate that message information is embedded more heavily in texture-complex regions, supporting image transparency.The figure shows cover, encoded, residual, and normalized residual images from left to right.
4.2 Comparison with Previous Methods
The comparison evaluates MBRS against prior deep-learning watermarking methods under JPEG and other distortions using matched image sizes and message capacities. The framework reports low JPEG BER and improved robustness from an added diffusion block for crop-related distortions.
- Comparison Setup: The comparison uses 64-bit messages in 128 × 128 images and 625-bit messages in 400 × 400 images, matching or exceeding prior capacities.The authors adjust settings for a fair comparison across methods.
- JPEG Compression: Under real JPEG Q=50 testing, MBRS is evaluated against methods [27],, and [22] using JPEG-specific or published baseline results.Table 1 reports PSNR, SSIM, and BER-related comparisons, with SSIM unavailable for [27] and.
- Other Distortions: Adding an additive diffusion block substantially lowers BER under crop, cropout, and dropout while retaining strong JPEG and Gaussian-filter robustness.The comparison adjusts the strength factor at PSNR = 33.5.
- JPEG Compression: BER below 0.01% is reported for the proposed architecture under JPEG compression.This result is presented as evidence of enhanced JPEG robustness.
4.3 Ablation Study
The ablation study examines strength-factor trade-offs and the contributions of MBRS components, including real and simulated JPEG compression, identity layers, and diffusion blocks. Results show that combined noise layers improve JPEG robustness, while diffusion improves resistance to cropping distortions.
- 4.3.1 Strength Factor: Increasing the strength factor decreases PSNR and SSIM while improving extraction accuracy.The factor balances robustness against visual transparency and is adjusted for comparable image quality.
- 4.3.2 JPEG Training Quality: Training quality factor affects the distortion trade-off: Q=10 improves robustness to low-quality compression but performs worse at high quality than Q=50.The authors use JPEG-Mask and real JPEG with Q=50 because compression below 50 is considered visually excessive.
- 4.3.1 Strength Factor: BER can reach 0.00% at PSNR=34, while PSNR=42 yields BR=1.35%.These settings illustrate the application-dependent trade-off between image quality and decoding accuracy.
- 4.3.3 Real, Simulated, and Identity Combination: Removing any MBRS noise layer decreases robustness, while adding real JPEG substantially improves performance.Using only JPEG-Mask can overfit simulated noise, whereas real JPEG keeps decoding loss informative for real-compression robustness.
- 4.3.3 Real, Simulated, and Identity Combination: MBRS reduces BER to almost 0% at Q=50, outperforming OEDS, TSR, and TSR-S under comparable PSNR.The comparison attributes OEDS weakness to simulated-distortion overfitting and reports that TSR-S reduces BER by about 17% at Q=50.
- 4.3.4 Diffusion Block: Adding diffusion blocks substantially lowers BER under crop, cropout, and dropout distortions.The blocks address the message processor and decoder’s tendency to localize each bit in a small image region.
5 CONCLUSION
The paper proposes MBRS watermarking with randomly selected real JPEG, simulated JPEG, or noise-free layers during mini-batch training, supplemented by feature-processing and diffusion components. Experiments report improved JPEG robustness and image quality, while diffusion spreads each message bit across the image to improve crop robustness.
- 5 CONCLUSION: MBRS randomly selects simulated JPEG, real JPEG, or a noise-free layer for each mini-batch to enhance JPEG robustness.The method uses these layers for joint training, real-compression decoder robustness, and uncompressed decoding ability, respectively.
- 5 CONCLUSION: The network uses Squeeze-and-Excitation blocks, a message processor, and an optional diffusion block.The diffusion block is intended to improve robustness against crop attack.
- 5 CONCLUSION: Each message bit is otherwise embedded and extracted mainly within a nearby small image region.The additive diffusion block uses a fully connected layer to distribute each bit across the whole image.
- 5 CONCLUSION: Diffusion changes the spatial distribution of bit information between the original model and the diffusion-enhanced model.Figure 6 visualizes the normalized difference caused by changing one embedded bit.