Layer-wise Noise Guided Selective Wavelet Reconstruction for Robust Medical Image Segmentation

TL;DR

This paper introduces LNG-SWR, a novel wavelet-based method that enhances robustness of medical image segmentation models against distribution shifts and perturbations without significant trade-offs or high training costs.

Contribution

LNG-SWR is a new framework that injects noise at multiple layers and applies frequency-guided wavelet reconstruction, improving robustness and stability in medical image segmentation.

Findings

Consistent robustness gains on CT and ultrasound datasets.

Reduces performance drop under strong adversarial attacks.

Enhances robustness when combined with adversarial training.

Abstract

Clinical deployment requires segmentation models to stay stable under distribution shifts and perturbations. The mainstream solution is adversarial training (AT) to improve robustness; however, AT often brings a clean--robustness trade-off and high training/tuning cost, which limits scalability and maintainability in medical imaging. We propose \emph{Layer-wise Noise-Guided Selective Wavelet Reconstruction (LNG-SWR)}. During training, we inject small, zero-mean noise at multiple layers to learn a frequency-bias prior that steers representations away from noise-sensitive directions. We then apply prior-guided selective wavelet reconstruction on the input/feature branch to achieve frequency adaptation: suppress noise-sensitive bands, enhance directional structures and shape cues, and stabilize boundary responses while maintaining spectral consistency. The framework is backbone-agnostic and adds low additional inference overhead. It can serve as a plug-in enhancement to AT and also improves robustness without AT. On CT and ultrasound datasets, under a unified protocol with PGD-$L_{\infty}/L_{2}$ and SSAH, LNG-SWR delivers consistent gains on clean Dice/IoU and significantly reduces the performance drop under strong attacks; combining LNG-SWR with AT yields additive gains. When combined with adversarial training, robustness improves further without sacrificing clean accuracy, indicating an engineering-friendly and scalable path to robust segmentation. These results indicate that LNG-SWR provides a simple, effective, and engineering-friendly path to robust medical image segmentation in both adversarial and standard training regimes.