Sample-Specific Noise Injection For Diffusion-Based Adversarial Purification

๐Ÿ“… 2025-06-06
๐Ÿ“ˆ Citations: 0
โœจ Influential: 0
๐Ÿ“„ PDF
๐Ÿค– AI Summary
To address the limited robustness of diffusion-based adversarial purification caused by fixed denoising steps $t^*$, this work identifies significant sample-wise variation in optimal denoising steps and proposes the first sample-adaptive noise injection mechanism. We introduce SSNI (Sample-Specific Noise Injection), a framework that estimates sample-specific noise scheduling via score normโ€”leveraging a pretrained score network to assess input cleanliness and a learnable reweighting function to dynamically modulate noise intensity during both forward and reverse diffusion processes. Evaluated on CIFAR-10 and ImageNet-1K, SSNI substantially improves post-purification classification accuracy and adversarial robustness: e.g., +4.2% robust accuracy against PGD attacks on ImageNet. The implementation is publicly available.

Technology Category

Application Category

๐Ÿ“ Abstract
Diffusion-based purification (DBP) methods aim to remove adversarial noise from the input sample by first injecting Gaussian noise through a forward diffusion process, and then recovering the clean example through a reverse generative process. In the above process, how much Gaussian noise is injected to the input sample is key to the success of DBP methods, which is controlled by a constant noise level $t^*$ for all samples in existing methods. In this paper, we discover that an optimal $t^*$ for each sample indeed could be different. Intuitively, the cleaner a sample is, the less the noise it should be injected, and vice versa. Motivated by this finding, we propose a new framework, called Sample-specific Score-aware Noise Injection (SSNI). Specifically, SSNI uses a pre-trained score network to estimate how much a data point deviates from the clean data distribution (i.e., score norms). Then, based on the magnitude of score norms, SSNI applies a reweighting function to adaptively adjust $t^*$ for each sample, achieving sample-specific noise injections. Empirically, incorporating our framework with existing DBP methods results in a notable improvement in both accuracy and robustness on CIFAR-10 and ImageNet-1K, highlighting the necessity to allocate distinct noise levels to different samples in DBP methods. Our code is available at: https://github.com/tmlr-group/SSNI.
Problem

Research questions and friction points this paper is trying to address.

Determining optimal sample-specific noise levels for adversarial purification
Adapting noise injection based on sample deviation from clean data
Improving diffusion-based purification accuracy and robustness via dynamic noise adjustment
Innovation

Methods, ideas, or system contributions that make the work stand out.

Sample-specific noise injection via score norms
Adaptive noise level adjustment per sample
Pre-trained score network estimates data deviation
Y
Yuhao Sun
School of Computing and Information Systems, The University of Melbourne
J
Jiacheng Zhang
School of Computing and Information Systems, The University of Melbourne
Zesheng Ye
Zesheng Ye
Postdoc Research Fellow at University of Melbourne
trustworthy machine learningmodel reprogramming
Chaowei Xiao
Chaowei Xiao
University of Wisconsin - Madison/NVIDIA
Trustworthy Machine LearningAdversarial Machine LearningAI SafetyRobust AISecurity
F
Feng Liu