PaperlayerView exact revision on arXiv

Source-linked AI summary

Masked Image Training for Generalizable Deep Image Denoising

Haoyu Chen, Jinjin Gu, Yihao Liu, Salma Abdel Magid, Chao Dong, Qiong Wang, Hanspeter Pfister, Lei Zhu

arXiv:2303.13132v1cs.CV

TL;DR

Deep denoisers may generalize poorly when real-world noise differs from the training distribution. The paper trains denoising networks to reconstruct randomly masked pixels and masks self-attention features; it reports improved generalization across multiple noise types despite Gaussian-noise training.

  • Problem

    Deep denoising models trained on one noise distribution can struggle when real-world noise differs from the training noise, limiting generalization.

    Source: §1 · pp. 1–2

  • Method

    The method randomly masks input pixels for reconstruction and masks features in Transformer self-attention during training.

    Source: §1 · p. 2

  • Results

    The method outperforms other models in generalization tests across multiple non-Gaussian noise types despite training on Gaussian noise.

    Source: §1 · p. 2, §4.1 · p. 5

  • Takeaways & Limitations

    The authors conclude that masked training is a promising approach for improving denoising generalization and developing more robust models for real-world applications.

    Source: §5 · p. 8

    The mask operation inevitably loses information, and preserving more image details remains future work.

    Source: §5 · p. 8

Abstract

from arXiv · show

When capturing and storing images, devices inevitably introduce noise. Reducing this noise is a critical task called image denoising. Deep learning has become the de facto method for image denoising, especially with the emergence of Transformer-based models that have achieved notable state-of-the-art results on various image tasks. However, deep learning-based methods often suffer from a lack of generalization ability. For example, deep models trained on Gaussian noise may perform poorly when tested on other noise distributions. To address this issue, we present a novel approach to enhance the generalization performance of denoising networks, known as masked training. Our method involves masking random pixels of the input image and reconstructing the missing information during training. We also mask out the features in the self-attention layers to avoid the impact of training-testing inconsistency. Our approach exhibits better generalization ability than other deep learning models and is directly applicable to real-world scenarios. Additionally, our interpretability analysis demonstrates the superiority of our method.

1. Introduction

Deep denoising models trained on a predefined noise distribution can struggle when real-world noise differs from training noise. The paper attributes this generalization problem to fitting noise patterns and proposes masked training to emphasize image-content reconstruction.

1. Introduction

  • Deep denoisers trained and evaluated on Gaussian noise can lose performance on other noise distributions, limiting their generalization to real-world settings.Prior approaches that model or collect target noise, or train directly on it, have not substantially improved generalization under noise mismatch.Source: §1 · p. 1
  • The authors argue that networks can reduce training loss by overfitting noise patterns instead of learning image content, which contributes to poor performance on unfamiliar noise.They describe fitting noise as easier than learning image content.Source: §1 · p. 2
  • Masked training randomly removes input pixels for reconstruction and masks features in Transformer self-attention to mitigate training-testing inconsistency.The method is intended to direct learning toward image-content reconstruction rather than training-noise overfitting.Source: §1 · p. 2
  • Trained on Gaussian noise, the proposed method reports improved performance across speckle, Poisson, salt-and-pepper, spatially correlated Gaussian, Monte Carlo-rendered, ISP, and mixed noise.The authors contrast this with existing methods that have yet to accurately remove all these diverse patterns.Source: §1 · p. 2

2. Related Works

Prior denoising work spans model-based methods, deep architectures, noise-specific training, and masked image modeling. The section contrasts their generalization limits with the proposed emphasis on reconstructing image content.

2. Related Works

  • Traditional model-based denoisers accommodate varied noise and can generalize well, but their image-content reconstruction is unsatisfactory.Source: §2 · p. 2
  • Deep denoising work has advanced through network architectures and Transformers, but stronger fitting ability does not necessarily yield better generalization.The cited discussion notes that a more efficient network can indicate worse generalization.Source: §2 · p. 2
  • Real-noise synthesis, collection, and blind denoisers can improve target-noise performance, yet still fail on noise outside their training distributions.Source: §2 · p. 2
  • SwinIR trained on immunohistochemistry images can denoise natural images, which the authors interpret as evidence that existing methods overfit training noise.Their approach instead emphasizes reconstructing natural-image textures and edges from the training set.Source: Figure 2 · p. 3
  • The proposed mask-and-complete strategy is illustrated as reconstructing the input to some extent even when many pixels are masked.Source: Figure 3 · p. 3
  • Masked image modeling predicts missing image content, but the authors note few successful applications to low-level vision despite its relevance to such tasks.Source: §2 · p. 3

3. Method

Masked training aims to make denoising networks reconstruct image content rather than overfit training noise. It combines input masking with attention masking to address the training–testing gap.

3. Method

  • A SwinIR model trained on Gaussian-noised immunohistochemistry images could still denoise natural images, supporting the authors’ noise-overfitting hypothesis.Source: §3 · p. 3
  • Masked training directs denoising networks toward reconstructing image textures and structures instead of relying on noise-pattern overfitting.The approach masks part of the input during training and trains the network to complete it.Source: §3 · p. 3
  • The method makes minimal changes to SwinIR, adding input and attention masks while retaining the Transformer architecture.A 1 × 1 convolution embeds pixels as feature tokens without allowing pixels to affect one another during embedding.Source: Figure 4 · p. 4, §3 · p. 4
  • The input mask randomly replaces embedded feature tokens with a mask token, training the network to reconstruct the image from limited information.The passage reports that reconstruction remains possible even when up to 90% of pixel information is destroyed.Source: §3 · p. 4
  • Training-noise performance does not drop significantly until a 75% masking ratio, while gains on noise outside training exceed training-set losses.Source: Figure 7 · p. 5
  • Attention masking applies masking within self-attention so the network can use unmasked tokens and work on unmasked test inputs.The passage says input masking alone creates a training–testing inconsistency that can increase output brightness.Source: §3 · pp. 4–5

4. Experiments

The experiments train models on Gaussian noise and test generalization across unseen synthetic and real noise settings. Results compare masked training with other models and analyze masking choices and learned representations.

4. Experiments

  • Evaluation uses Gaussian training noise at σ = 15 and tests six synthetic noise types, two real-noise settings, and PSNR, SSIM, and LPIPS.The real-noise settings include SIDD and Monte Carlo rendered images; test images come from CBSD68, Kodak24, McMaster, and Urban100.Source: §4 · p. 5

4.1. Resutls

  • Removing the input mask loses generalization to unseen noise, while removing the attention mask degrades quantitative performance and produces incorrect brightness.The ablation identifies input masking as the crucial factor for generalization; without the attention mask, most noise is still removed.Source: §4.1 · p. 5
  • Random input-mask ratios of 75%–85% give the best quantitative performance, while attention-mask ratios are optimal around 75%.Fixed input-mask ratios are less stable and perform worse; larger ratios improve generalization but can remove image detail.Source: §4.1 · p. 5
  • Masked training outperforms the compared models on non-Gaussian noise, with slower degradation as noise severity rises, despite training only on fixed-level Gaussian noise.Figure 11 compares PSNR and SSIM across noise types and levels; Table 3 describes tests on noise outside the training set.Source: §4.1 · pp. 5–6, Figure 11 · p. 7, Table 3 · p. 7
  • On synthesized ISP noise, the method beats the compared methods overall and improves SSIM and LPIPS over DnCNN despite matching its PSNR.The authors report a significant LPIPS lead and less visible residual noise than DnCNN.Source: §4.1 · p. 6, Figure 10 · p. 6
  • On Monte Carlo rendered images, the method outperforms all compared methods at both 128 spp and 64 spp, where existing methods fail visually.Figure 9 presents visual results; lower sample counts produce more severe noise.Source: Figure 9 · p. 6, §4.1 · p. 7

4.2. Generalization Analysis

  • With Gaussian-only training, the baseline training curve shows overfitting, whereas masked training continues to improve without overfitting.Figure 13 shows the training curves.Source: §4.2 · p. 7, Figure 13 · p. 8
  • CKA and feature-distribution analyses show more similar representations across noise types with masked training, while baseline deep-layer features differ across noise types.The authors also report that baseline deep layers correlate weakly with masked-training layers and that lower SRGA indicates better generalization.Source: §4.2 · pp. 7–8, Figure 12 · p. 8, Figure 14 · p. 8

5. Conclusion and Limitations

Masked training is presented as a promising way to improve denoising generalization, but masking inevitably loses information and preserving more detail remains future work.

5. Conclusion and Limitations

  • Masked training is presented as a promising approach to improving generalization in deep-learning image denoising.The authors frame it as a step toward more robust real-world models.Source: §5 · p. 8
  • The mask operation inevitably loses information, so preserving more detail remains a future-work challenge.Source: §5 · p. 8

A. Details of the Test Noise

The evaluation covers six synthetic noise types and two real noise types, with mixture noise designed to represent images affected by multiple degradations.

A. Details of the Test Noise

  • The synthetic test set covers speckle, Poisson, spatially correlated, salt-and-pepper, ISP, and mixture noise.The passage describes speckle as multiplicative, Poisson as signal-dependent, and spatially correlated noise as filtered additive Gaussian noise.Source: A. Details of the Test Noise · p. 13
  • Mixture noise combines Gaussian, Poisson, salt-and-pepper, and speckle noise at different levels to model multiple degradations.Speckle's multiplicative effect depends on its position in the sequence and compounds existing noise.Source: A. Details of the Test Noise · p. 13
  • The real-noise evaluation includes smartphone-image noise from SIDD and Monte Carlo rendered-image noise.Source: A. Details of the Test Noise · p. 13

B. Additional Comparisons

The method is reported to outperform other state-of-the-art models across four benchmarks and noise types, with a particularly large LPIPS lead. On limited SIDD data, clean-image masked pre-training performs best and preserves the most texture detail.

Additional Comparisons

  • The comparison includes DnCNN, RIDNet, RNAN, SwinIR, Restormer, and dropout, with dropout probability 0.7 before the output convolution.Dropout is included as a method proposed to improve generalization and relieve overfitting.Source: B. Additional Comparisons · p. 13
  • With only two SIDD training pairs, masked pre-training on clean images followed by target-data fine-tuning achieves the best results.It outperforms pre-training on Gaussian noise, and the visual comparison reports that it preserves the most texture detail.Source: B. Additional Comparisons · p. 14, Figure 17 · p. 14
  • On limited SIDD training data, masked-training models continue improving during training, while scratch training fluctuates and Gaussian pre-training without masking remains low.The clean-image-pretrained masked model performs better than the other masked-training variant.Source: B. Additional Comparisons · pp. 14–15
  • The method significantly outperforms other state-of-the-art models across four benchmarks and all evaluated noise types, with a significant lead in LPIPS.The authors associate the LPIPS lead with better human perceptual quality.Source: B. Additional Comparisons · p. 15
  • Visual comparisons report that the model without masked training is significantly limited across noise types, while the proposed model removes varied out-of-training noise effectively.Source: B. Additional Comparisons · p. 15

C. Additional Analyses of CKA

The CKA analysis compares network representations using Gram matrices and HSIC. It reports that masked training yields higher feature correlations across noise types than the baseline, whose deeper layers show low correlations across noise conditions and with the masked model.

Additional Analyses of CKA

  • CKA compares representations from two network layers computed on the same data points, using their Gram matrices and HSIC.The passage defines centered Gram matrices as part of the HSIC calculation.Source: C. Additional Analyses of CKA · p. 15
  • Deeper baseline-layer features have low correlation with all layers of the masked-training model, indicating different feature-learning patterns, especially in deeper layers.Source: C. Additional Analyses of CKA · p. 15
  • The baseline has low deep-layer feature correlation between Gaussian and Poisson noise, while masked training shows high correlations across the analyzed cases.The authors report a similar low-correlation pattern for other noise types in the baseline.Source: C. Additional Analyses of CKA · p. 15

Related papers

Loading 2303.13132v1…