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

2025-11-20

AI Summary

Medical image segmentation models exhibit insufficient robustness under distribution shifts and perturbations, while adversarial training (AT) suffers from the clean-robustness trade-off and high computational overhead. To address this, we propose a hierarchical noise-guided selective wavelet reconstruction framework: multi-level zero-mean noise injection models frequency-domain preference priors, enabling adaptive enhancement of directional structures and boundary responses in the wavelet domain. Our method improves robustness without requiring adversarial training, is plug-and-play, and compatible with mainstream segmentation architectures. It significantly mitigates performance degradation under strong attacks across diverse modalities—including CT and ultrasound—while simultaneously improving Dice and IoU scores on clean samples. When combined with AT, it further boosts robustness without compromising original accuracy. With low inference overhead, the approach demonstrates strong scalability and clinical deployment potential.

Technology Category

• Computer Vision: Adversarial Attacks & Robustness
• Machine Learning: Adversarial Learning & Robustness
• Natural Language Processing: Safety and Robustness

Application Category

• User Modeling, Personalization and Recommendation: Attacks and countermeasures in recommendation systems
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking
• Responsible Web: Machine-in-the-loop, human agency and autonomy

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.

Problem

• Improves medical image segmentation robustness against distribution shifts
• Reduces adversarial training costs while maintaining clean performance
• Enables frequency adaptation to suppress noise and enhance structures

Innovation

• Layer-wise noise injection for frequency-bias learning
• Selective wavelet reconstruction for frequency adaptation
• Backbone-agnostic framework with low inference overhead